Common Lispのマクロで型変数もどきのことをする
Introduction
こんにちは！
自分は趣味として、Common Lispで深層学習兼行列演算ライブラリを書いています。
この記事は、そんなcl-waffeの開発途中に起きたトラブルをマクロで解決したという内容です。
cl-waffeの開発中に以下のような状況に出くわします:
- 行列に対する演算（例えばベータ分布の生成）を定義する。このコードはSBCLに頑張って最適化して欲しいから、至る所に型宣言を設ける。
- 例えば最初single-floatに対して定義したベータ分布のサンプリング関数だが、自分のライブラリではdouble-floatも扱えるため、double-floatに対しても定義したいし、将来的に半精度にも対応したい。
- そうなると型宣言のみが違う瓜二つの関数を手動で何回も書き直さないといけないのでは？これを手作業は骨が折れるよ・・・
そもそも扱う型は全てnumberのsubtypeなんだから自動で定義できるよね？→型変数が欲しくなった。
型変数とは
自分はあくまで純粋関数型やろうと思ってHaskellやって折れた型理論にトラウマがある系Common Lisperなので正確なことは言えないが、自分がやりたい枠組みの中では：
リストの要素数を求める時、関数length(x: T)は引数T(=リスト,ベクトル, 文字列...)の型の名前が異なっていても、抽象的に見て、要素数を持つ型であったら定義できる。
つまり、xの型自体をパラメーターにすることができる。
- xの型
Tを引数に取る -> Tに対する関数length(x: T)の一対一対応をとる関数 - 型Tに属するxに対してxの長さを返す関数
length(x: T)
このような二層構造にできる。
1.の関数が取る引数Tのことを型変数と呼ぶ。
具体例で説明すると:
-
F(x: T) = xを定義する(xは型T)
T=single-floatを代入すれば
-> F(x: single-float) = x
T=booleanを代入すれば
-> F(x: boolean) = x
T=fixnumを代入すれば
-> F(x: fixnum) = x
こういうことができれば良い。
ジェネリック関数と違いコンパイラはF(x: fixnum)=xのxはfixnumだと認識してくれる。よって共通部分の多いコードを削減することに役立つ。
自分が欲しいのは何か
このように動くものを想定する。
(define-with-typevar example-function U (x y)
(declare (optimize (speed 3))
(type U x y))
(+ x y))
このマクロは名前の通り、example-functionを定義し、Uを型変数に、xとyを引数として展開する。
body部分のコードではUを型の名前として使っているが、このシンボルは展開時に自動で置き換えることにする。
CommonLispの関数+はnumberに対する関数として定義されているが、実行時に引数がsingle-floatやdouble-float等の型宣言がされてないと、SBCLがコンパイル時にうまくインライン化してくれない(i.e.: 速度がものすごく低下する)。
自分の行列演算ライブラリでは、現在扱えるデータ型は以下の二通り:
- single-float (float32)
- double-float (float64)
define-with-typevarマクロには型がsingle-floatとdouble-floatそれぞれの場合の関数のプログラム本文を吐いて欲しいし、SBCLにはそれぞれの場合をうまく最適化したアセンブリを吐いて欲しい。
データ型がsingle-float double-floatの時の場合分けはどうしようか？ここでsingle-floatとdouble-floatの場合についてジェネリック関数を定義すればいいと思ったが、実はCommonLispのジェネリックは遅いし最適化がすごくしにくい。
こういう場合は普通にcaseで分岐して、それぞれの型のコードを呼び出すようにすればいい。
ちなみにこのような方針でジェネリック関数をインライン化しているライブラリも複数存在する。
(詳細まではソースコード読んでないけど、define-method-combinationとcaseで同じようなことをしてたはず)
話題が逸れたが、とにかくこのような振る舞いをするものが欲しくなる。
-
define-with-typevarマクロに関数の振る舞いを記述する。 - 現在ライブラリが対応しているデータ型に応じて、single-floatに対して演算する
example-function-fやdouble-floatに対して演算をするexample-function-dを定義する。 -
example-function関数を定義する。これはライブラリが現在使っているデータ型をもとにexample-function-fやexample-function-dを適切に割り当てる役割を持つ。
実装
関数の本文を各データ型に対応するように書き直せばいいのだが、要はこれをリストに対してすればいいのだ。
"F(x: T) = x".replace("T", "single-float")
"F(x: T) = x".replace("T", "double-float")
該当するソースコードはこちら
定数値は以下の通り
(defparameter *dtypes* `(:float :double))
(defparameter *dtype-prefixes* `(-f -d))
(defparameter *dtype-cl-names* `(single-float double-float))
(defconstantを使うとREPLでのデバック時に手間が増えてしまうので、普段はdefparameterを使っています)
-
*dtypes*: 対応するデータ型のKeyword -
*dtype-prefixes*: 後述のマクロによって自動定義される関数名のsuffix(変数の命名ミスっったぁ〜〜〜😡) -
*dtype-cl-names*: 各Keywordに対応するsymbolの名前
さて、本題のマクロの全文はこんな感じ：
(defmacro define-with-typevar
(function-name
type-specifier
(&rest args)
&body body
&aux (fnames (map
'list
#'(lambda (p)
(symb function-name p)) ; (symb 'a 'b) -> 'ab
*dtype-prefixes*))
(params (get-params args))) ; これから解説
(multiple-value-bind (body declarations doc) (alexandria:parse-body `,body
:documentation t) ; Bodyをパースしてくれる便利なやつ
`(progn
; それぞれのdtypeに対して、受け取ったbodyを定義する。
,@(loop for i fixnum upfrom 0 below (length fnames)
collect (define-lisp-code-with
params
(nth i fnames)
(nth i *dtype-cl-names*)
body
declarations
type-specifier))
(defun ,function-name (,@args)
,doc
(case mgl-mat:*DEFAULT-MAT-CTYPE* ; 自作のライブラリでは、dtypeはmgl-matに依存する
(:half
(error "No implementation"))
(:float
(,(car fnames) ,@params))
(:double
(,(second fnames) ,@params))
(T
(error "No dtype")))))))
続いて、Body部分をパースして、型変数を置き換える部分
map-tree関数はS式をパースして
- listならmapcarして再帰
- list以外ならlambdaを適用
する。
(defun map-tree (fn tree)
(let ((tree (funcall fn tree)))
(if (listp tree)
(mapcar (lambda (subtree)
(map-tree fn subtree))
tree)
tree)))
(defun replace-lisp-code-with-dtype (body type-var dtype)
(map-tree #'(lambda (code)
(typecase code
(symbol
(if (equal (symbol-name code)
(symbol-name type-var))
dtype
code))
(single-float (coerce code dtype))
(double-float (coerce code dtype))
(T
code)))
body))
(defun define-lisp-code-with (args
fname-with-prefix
dtype
body
declarations
type-specifier)
`(defun ,fname-with-prefix (,@args)
(locally
,@(unless (eql dtype 'double-float)
(replace-lisp-code-with-dtype
declarations
type-specifier
dtype))
,@(replace-lisp-code-with-dtype
body
type-specifier
dtype))))
replace-lisp-code-with-dtype関数は受け取ったシンボルが型変数として使われているものと同じならsingle-floatやdouble-floatに置き換える。
ついでにコード内部で登場した小数値は自動でcoerceしてくれる。(e.g.: 1.0e0 -> 1.0d0 そうじゃないとエラー吐くので)。もちろんコンパイル時のみなのでcoerceのoverheadはなし。
define-lisp-code-with関数は読んで字の如くといった感じ。ちなみに、dtypeがdouble-floatの時declare部分を無視するようにしたが、理由は後ほど。
Keyword引数等の処理
Common Lispは関数定義に使う引数と、funcallで使える引数のフォーマットが違うから、define-with-typevarのargsをそのままfuncallの引数に使うとエラーが起こる。
;こういう関数があったとして
(defun a (x y &key (a 1))
~
)
;そのままfuncallできない
(funcall #'a x y &key (a 1))
example-functionの例で言えば、(defun example-function (&rest inputs))として定義すればいいのだけど、これはSLIMEの補完やdocumentの生成時にわかりにくくなるから好ましくない。
だからこうした：
以下の関数を処理することを考える。
(define-with-typevar complicated-argument-func U (x y &key (a 1))
(declare (optimize (speed 3))
(type U x y a))
(+ x y a))
-
complicated-argument-funcは一旦(complicated-argument-func x y &key (a 1))として定義 -
(x y &key (a 1))の引数のシンボルだけ取り出す、この場合(x y a) -
complicated-argument-func-fcomplicated-argument-func-dはこのシンボルに対しての関数にする。
2.の振る舞いを実現するための関数がget-paramsです。
寝る前に脳死で書いたカスみたいなコードをご容赦ください・・・
(defun get-params (list)
(reverse
(delete-duplicates
(flatten
(loop for i fixnum upfrom 0 below (length list)
collect (let ((sym (nth i list)))
(typecase sym
(symbol
(if (find sym `(&optional &rest &key &aux))
nil
sym))
(list
(if (= (length sym) 2)
(car sym)
(get-params sym))))))))))
振る舞い
実際に展開してみよう。
(print
(macroexpand
`(define-with-typevar complicated-argument-func U (x y &key (a 1))
(declare (optimize (speed 3))
(type U x y a))
(+ x y a))))
(PROGN
(DEFUN COMPLICATED-ARGUMENT-FUNC-F (A Y X)
(LOCALLY
(DECLARE (OPTIMIZE (SPEED 3))
(TYPE SINGLE-FLOAT X Y A))
(+ X Y A)))
(DEFUN COMPLICATED-ARGUMENT-FUNC-D (A Y X) (LOCALLY (+ X Y A)))
(DEFUN COMPLICATED-ARGUMENT-FUNC (X Y &KEY (A 1))
NIL
(CASE *DEFAULT-MAT-CTYPE*
(:HALF (ERROR "No implementation"))
(:FLOAT (COMPLICATED-ARGUMENT-FUNC-F A Y X))
(:DOUBLE (COMPLICATED-ARGUMENT-FUNC-D A Y X))
(T (ERROR "No dtype")))))
このような展開式が得られる。
ちなみにマクロを介して関数を定義して、body部分でdeclare formを使う場合の注意点だが、ソースコード内でそのマクロが適切に読み込めていないと
; in: DEFINE-WITH-TYPEVAR BACKWARD1
; (DECLARE (OPTIMIZE (SPEED 3) (SAFETY 0))
; (TYPE CL-WAFFE:WAFFETENSOR CL-WAFFE:TENSOR))
;
; caught ERROR:
; There is no function named DECLARE. References to DECLARE in some contexts
; (like starts of blocks) are unevaluated expressions, but here the expression is
; being evaluated, which invokes undefined behaviour.
Unhandled TYPE-ERROR in thread #<SB-THREAD:THREAD "main thread" RUNNING
{100DF18393}>:
The value
NIL
is not of type
SB-C::ENVIRONMENT
こんな感じの非自明なエラーに遭遇する。
.asdファイルの依存関係を書き直してもマクロが一番最初に読み込まれず(i.e.: コンパイラはdefine-with-typevarを未定義の関数と見做している)関数の引数として(declare)関数の引数を使っているみたいな感じで処理されるからこうなる。
同一のコードを別の環境に持っていけばなんの修正もなく動くのだが、これはQuicklisp/ASDFが変更されたマクロを自動で読み直してくれないことに由来するバグ（？）らしく、これに出くわした時どう修正すればいいのか何もわからない。今回は適当に.fasl消したら動いた。とりあえずプロジェクトの根幹となる重要なマクロを含むファイルはasdfの依存関係の一番上に持ってこよう。
最適化に関する諸問題
このマクロを使ってdouble-floatに対する関数を定義して、(declare (optimize (speed 3)))宣言をするとしょっちゅうSBCLがこの警告を出してくる。
note: doing float to pointer coercion (cost 13) to "<return value>"
single-floatを最適化する時は発生しないが、double-floatにsetqを含むコード(incfマクロ等も)があるとこうなる。
調べた結論としては、single-floatとdouble-floatは確かに同一のLispコードを動かすことができるが、最適化時にSBCLの振る舞いが異なるためこうなるらしい。なのでこれは今後の課題として、double-floatの場合だけは最適化宣言を外すことにした。(まぁどうせ深層学習でdouble-floatはあんま使わないだろうし・・・）
Disassemble
定義したマクロを用いてbeta関数のサンプリングアルゴリズムを定義した。
以下のように自動で生成された関数は手動で型宣言を加えている。（そのままでは返り値がうまく型推論されないので）
(declaim (ftype (function
(single-float
single-float
single-float)
single-float)
!beta-bb-f
!beta-bc-f))
(declaim (ftype (function
(double-float
double-float
double-float)
double-float)
!beta-bb-d
!beta-bc-d))
(disassemble #'!beta)
; disassembly for !BETA
; Size: 179 bytes. Origin: #x5407A29C ; !BETA
; 29C: 498B4510 MOV RAX, [R13+16] ; thread.binding-stack-pointer
; 2A0: 488945F8 MOV [RBP-8], RAX
; 2A4: 4D8B85C8110000 MOV R8, [R13+4552] ; tls: *DEFAULT-MAT-CTYPE*
; 2AB: 488B057EFFFFFF MOV RAX, [RIP-130] ; '*DEFAULT-MAT-CTYPE*
; 2B2: 4983F8FF CMP R8, -1
; 2B6: 4C0F444001 CMOVEQ R8, [RAX+1]
; 2BB: 4180F809 CMP R8B, 9
; 2BF: 0F8487000000 JEQ L3
; 2C5: 4C3B056CFFFFFF CMP R8, [RIP-148] ; :HALF
; 2CC: 7460 JEQ L2
; 2CE: 4981F81F851250 CMP R8, #x5012851F ; ':FLOAT
; 2D5: 743E JEQ L1
; 2D7: 4981F88F841250 CMP R8, #x5012848F ; ':DOUBLE
; 2DE: 7519 JNE L0
; 2E0: 488B55E0 MOV RDX, [RBP-32]
; 2E4: 488B7DE8 MOV RDI, [RBP-24]
; 2E8: 488B75F0 MOV RSI, [RBP-16]
; 2EC: B906000000 MOV ECX, 6
; 2F1: FF7508 PUSH QWORD PTR [RBP+8]
; 2F4: E929CF30FC JMP #x50387222 ; #<FDEFN !BETA-D>
; 2F9: L0: 4883EC10 SUB RSP, 16
; 2FD: 488B154CFFFFFF MOV RDX, [RIP-180] ; "No dtype"
; 304: B902000000 MOV ECX, 2
; 309: 48892C24 MOV [RSP], RBP
; 30D: 488BEC MOV RBP, RSP
; 310: E8CB1E99FE CALL #x52A0C1E0 ; ERROR
; 315: L1: 488B55E0 MOV RDX, [RBP-32]
; 319: 488B7DE8 MOV RDI, [RBP-24]
; 31D: 488B75F0 MOV RSI, [RBP-16]
; 321: B906000000 MOV ECX, 6
; 326: FF7508 PUSH QWORD PTR [RBP+8]
; 329: E914CF30FC JMP #x50387242 ; #<FDEFN !BETA-F>
; 32E: L2: 4883EC10 SUB RSP, 16
; 332: 488B151FFFFFFF MOV RDX, [RIP-225] ; "No implementation"
; 339: B902000000 MOV ECX, 2
; 33E: 48892C24 MOV [RSP], RBP
; 342: 488BEC MOV RBP, RSP
; 345: E8961E99FE CALL #x52A0C1E0 ; ERROR
; 34A: CC10 INT3 16 ; Invalid argument count trap
; 34C: L3: CC1A INT3 26 ; UNBOUND-SYMBOL-ERROR
; 34E: 00 BYTE #X00 ; RAX(d)
!beta -> !beta-f !beta-dの呼び出しに無駄なオーバーヘッドは無さそう。(Inline化してもいいかもね)
!beta-fのdisassembleも余計な型比較とかは増えてないので大丈夫だと思う。
ベンチマーク
ここに、
- 型変数を使って定義した
!beta - 型変数を介さないで定義した
!beta-test(single-float版)
を用意した。どちらもsingle-floatに対して実行する
(time (dotimes (i 100)
(!beta `(100 100) 2.0 1.0)))
Evaluation took:
0.114 seconds of real time
0.113210 seconds of total run time (0.112310 user, 0.000900 system)
99.12% CPU
264,092,798 processor cycles
4,127,168 bytes consed
(time (dotimes (i 100)
(!beta-test `(100 100) 2.0 1.0)))
Evaluation took:
0.112 seconds of real time
0.109619 seconds of total run time (0.108438 user, 0.001181 system)
98.21% CPU
259,233,172 processor cycles
4,140,544 bytes consed
このように性能低下は見られなかったので、single-float演算に対しての影響は一切ないと思う。
最後に
このマクロは自分のプロジェクトの中だけで使おうと思ってるので厳密性に欠けると思います。
将来的に自分のライブラリを半精度に対応させたいと考えているんので、その時にまたこのマクロが活躍することを祈る・・・
CommonLispは言語仕様の古さや設計時の時代背景をマクロで無理やり拡張できるのでいいですね、しかも拡張した構文は処理の透明度が高いので最適化との両立がしやすいです。
参考文献
Discussion