llvm : tutorial : jit

llvm sbcl

チュートリアル第四章『4. Kaleidoscope: Adding JIT and Optimizer Support — LLVM 3.4 documentation』の２。

タイトルにはJITとあるが、僕のこれまでの進め方ではチュートリアル通りのJITは実現できないので、似たような動作をする方法でごまかす。

実装

今回行うことは前回と似ている。
以下概要。

  1. ラッパー出力ストリームを作成。
  2. 関数定義の出力がなされたら、それをllvmのコマンド群およびgccに渡して、出力された関数に対応する共有ライブラリを作成する。
  3. 作成された共有ライブラリをロード。
  4. もし2の関数が無名関数(トップレベルの式)の場合は、その呼び出し結果を返し、通常の関数の場合は関数定義を行った旨を表示する。

残りはそのためのコード。

;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;; 前回までのコードの修正
(let ((fn-module (make-hash-table :test #'equal)))
  (defun register-fn (name fn) (setf (gethash name fn-module) fn))
  (defun deregister-fn (name)  (remhash name fn-module))
  (defun find-fn (name)        (gethash name fn-module))
  ;; 次の関数を追加: 定義/宣言済みの関数一覧を取得する
  (defun get-fn-list (&aux fns)
    (maphash (lambda (_ fn) (push fn fns)) fn-module)
    fns))

;; 一番最後に作成された関数(の名前)を保持する変数
(defvar *current-fn* nil)

;; 関数定義のパース: 定義された関数を保存するよう変更
(defun parse-definition ()
  (read-token)  
  (let ((def (make-define :proto (parse-prototype)
                          :body  (parse-expression))))
    (setf *current-fn* (prototype-name (define-proto def)))
    def))

;; トップレベルの式のパース: *current-fn*を"anonymous"にセットする
(defun parse-toplevel-exp ()
  (prog1 (make-define :proto (make-prototype :name "anonymous" :args '())
                      :body  (parse-expression))
    (setf *current-fn* "anonymous")))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;; llvmのアセンブリコード(関数定義)を受け取り、対応するlisp関数を作成するための関数群

;; llvmの関数定義コードをコンパイルし、[fn-name].so形式の共有ライブラリを作成する
(defun llvm-compile(input-stream fn-name)
  (zerop
   (sb-ext:process-exit-code
    (sb-ext:run-program
     "sh" `("-c" ,(format nil "llvm-as | ~
                               llc -f -o /tmp/~A.s; ~
                               gcc -shared -o /tmp/~A.so /tmp/~A.s"
                          fn-name fn-name fn-name))
     :search t :input input-stream :output t :error *error-output*))))

;; llvmの関数定義コードをコンパイル/ロードする
(defun llvm-compile-and-load (input-stream fn-name ret-type arg-types)
  (when (llvm-compile input-stream fn-name)
    (sb-alien:load-shared-object (format nil "/tmp/~A.so" fn-name))
    (eval `(sb-alien:define-alien-routine ,fn-name ,ret-type ,@arg-types))
    'done))

;;;;;;;;;;;;;;;
;;; 使用例　;;;
(with-input-from-string (in "
define i32 @plus(i32 %a, i32 %b) {
  %rlt = add i32 %a, %b
  ret i32 %rlt
}
")
  (llvm-compile-and-load in "plus" 'sb-alien:int '((a sb-alien:int) (b sb-alien:int))))
--> DONE

(plus 10 20)
--> 30

;;;;;;;;;;;;;;;;;;;;;;;
;;;; ラッパーストリーム
;; 以下のクラス/メソッドは、名前が異なる以外はほとんど同じ内容
(defclass llvm-jit-stream (sb-gray:fundamental-stream)
  ((stream :initarg :stream :reader stream-of)))

(defmethod stream-element-type ((stream llvm-jit-stream))
  (stream-element-type (stream-of stream)))

(defmethod close ((stream llvm-jit-stream) &key abort)
  (close (stream-of stream) :abort abort))

(defclass llvm-jit-character-output-stream
  (llvm-jit-stream sb-gray:fundamental-character-output-stream)
  ((col-index :initform 0   :accessor col-index-of)
   (buf       :initform '() :accessor buf-of)))

(defmethod stream-line-column ((stream llvm-jit-character-output-stream))
  (col-index-of stream))

(defmethod stream-write-char ((stream llvm-jit-character-output-stream) char)
  (with-accessors ((inner-stream stream-of) 
                   (cols col-index-of)
                   (buf  buf-of)) stream
    (push char buf)
    (when (char= char #\})
      ;; ここだけが前回と異なる
      ;;  - do-optimize関数の代わりに、llvm-eval関数を呼び出している
      ;;  - 結果の出力形式が若干変更されている
      (format inner-stream "~& ==> ~A~%"
              (llvm-eval (coerce (nreverse buf) 'string)))
      (setf buf '()))
    (if (char= char #\Newline)
        (setf cols 0)
      (incf cols))))

;;; llvmのコード評価用関数群
;; 補助関数: 既に定義/宣言されている関数宣言を関数定義の前に追加する
(defun append-declarations (llvm-program)
  (with-output-to-string (out)
    (dolist (fn (get-fn-list))
      (with-slots (ret-type name arg-types) fn
        (format out "~&declare ~A ~A (~{~A~^, ~})~%" ret-type name arg-types)))
    (format out llvm-program)))

;; 評価関数
(defun llvm-eval (llvm-program &aux (fn (find-fn *current-fn*)))
  (with-input-from-string (in (append-declarations llvm-program))
    (flet ((rettype ()  ; => sb-alien:double
             (intern (string-upcase (fn-ret-type fn)) :sb-alien))
           (argtypes () ; => ((name sb-alien:double) ...)
             (mapcar (lambda (name type)
                       `(,(intern (string-upcase name))
                         ,(intern (string-upcase type) :sb-alien)))
                     (fn-arg-names fn)
                     (fn-arg-types fn))))
      (if (llvm-compile-and-load in *current-fn* (rettype) (argtypes))
          (if (string= *current-fn* "anonymous")
              (prog1 (funcall 'anonymous) (unintern 'anonymous)) ; トップレベルの式の場合
            (format nil "define ~A" *current-fn*))               ; 関数定義の場合
        (error "Compilation failure")))))                        ; コンパイルに失敗した場合

実行。

> (defvar *jit-output-stream* 
    (make-instance 'llvm-jit-character-output-stream
                   :stream *standard-output*))
--> *JIT-OUTPUT-STREAM*

> (main-loop *llvm-pp-table* *jit-output-stream*)
ready> 1+2+3;
 ==> 6.0d0

ready> def foo(a b) a+b*3;
 ==> define foo

ready> foo(10,2);
 ==> 16.0d0

ready> extern cos(n);
ready> def bar(a b c) foo(a,b) - cos(c);
 ==> define bar

ready> bar(1,2,3);
 ==> 7.989992496600445d0

ready> ^D  ; 終了

> (bar 1.0d0 2.0d0 3.0d0)
--> 7.989992496600445d0

とりあえず四章も終了。