Common Lisp で行列計算するブログを追試してみた (その3) #lisp

http://qiita.com/guicho271828/items/a095b106deec61477597 の続き。

アンローリングしてもここが遅い。

; 780: L2:   498B5011         MOV RDX, [R8+17]
; 784:       F30F10545A01     MOVSS XMM2, [RDX+RBX*2+1]
; 78A:       F30F59D1         MULSS XMM2, XMM1
; 78E:       488B5011         MOV RDX, [RAX+17]
; 792:       F30F105C4A01     MOVSS XMM3, [RDX+RCX*2+1]
; 798:       F30F58DA         ADDSS XMM3, XMM2
; 79C:       488B5011         MOV RDX, [RAX+17]
; 7A0:       F30F115C4A01     MOVSS [RDX+RCX*2+1], XMM3
; 7A6:       488D5102         LEA RDX, [RCX+2]            <<<<
; 7AA:       4C8D5302         LEA R10, [RBX+2]            <<<<
; 7AE:       4D8B6811         MOV R13, [R8+17]
; 7B2:       F3430F10545501   MOVSS XMM2, [R13+R10*2+1]   <<<<
; 7B9:       F30F59D1         MULSS XMM2, XMM1
; 7BD:       4C8B5011         MOV R10, [RAX+17]
; 7C1:       F3410F105C5201   MOVSS XMM3, [R10+RDX*2+1]   <<<<
; 7C8:       F30F58DA         ADDSS XMM3, XMM2
; 7CC:       4C8B5011         MOV R10, [RAX+17]
; 7D0:       F3410F115C5201   MOVSS [R10+RDX*2+1], XMM3
; 7D7:       4883C704         ADD RDI, 4
; 7DB:       4883C104         ADD RCX, 4
; 7DF:       4883C304         ADD RBX, 4

(row-major-aref array (+ i 2)) の時に、定数2 がオフセットとしてアセンブラに畳み込まれておらず、 LEAで明に計算されてしまっている(7A6)。

7B2は本当はこうなってほしい:

; 7B2:       F3430F10545501   MOVSS XMM2, [R13+R10*2+5]

そこで、SBCLの内部にダイブしてアセンブラをチューニングすることに。SBCLは好みの関数が特定の型に当てはまる時、手で書いた最速アセンブラに変換する機能が付いている。これは使う側のコード側で指定するのではなくて

;;;; 実際ではない例 : 使いにくい
(...
 ;; 普通バージョンを
 ;(my-function 0.1 0.2)
 ;; アセンブリバージョンにいちいち手で置き換える
 (my-function-in-super-fast-assembly/single-float/single-float 0.1 0.2)
 ...)

型がマッチした時には自動で置き換わる。マッチしない時は元の定義が使われる。

(define-VOP (my-function-in-super-fast-assembly/single-float/single-float)
  (:translate my-function)
  (arg-types single-float single-float)
  ...)

(...
 (my-function 0.1 0.2) ; コンパイラが自動でVOPのほうを使う
 ...)

さて、別に自分でアセンブリを書きたくはないので、SBCLにすでにある内部関数を使いたい。どうやらそれが sb-kernel:data-vector-ref-with-offset だ。

(define-vop (data-vector-ref-with-offset/simple-array-single-float-c)
  (:note "inline array access")
  (:translate data-vector-ref-with-offset)
  (:policy :fast-safe)
  (:args (object :scs (descriptor-reg)))
  (:info index offset)
  (:arg-types simple-array-single-float (:constant low-index)
              (:constant (constant-displacement other-pointer-lowtag
                                                4 vector-data-offset)))
  (:results (value :scs (single-reg)))
  (:result-types single-float)
  (:generator 4
   (inst movss value (make-ea-for-float-ref object index offset 4))))

しかし使うとコンパイル中にエラーになる。

(defun rm-gemm+static-size+unroll2-k (ma mb mc)
  (declare (optimize (speed 3) (debug 0) (safety 0) (space 0)))
  (declare (type (matrix 500 500) ma mb mc))
    ...
          (dotimes (col (/ cols 2))
            (sb-kernel:data-vector-set-with-offset mc mc-index 0
             (+ (sb-kernel:data-vector-ref-with-offset mc mc-index 0)
                (* cell (sb-kernel:data-vector-ref-with-offset mb mb-index 0))))

->
full call to SB-KERNEL:DATA-VECTOR-REF-WITH-OFFSET
This is probably a bug in SBCL itself. (Alternatively, SBCL ...
   [Condition of type SB-INT:BUG]

Restarts:
 0: [ABORT] Abort compilation.
 1: [*ABORT] Return to SLIME's top level.
 2: [ABORT] abort thread (#<THREAD "worker" RUNNING {10051DE033}>)

Backtrace:
  0: (SB-INT:BUG "full call to ~S" SB-KERNEL:DATA-VECTOR-REF-WITH-OFFSET)
  1: (SB-C::PONDER-FULL-CALL #<SB-C::COMBINATION :FUN #<SB-C::REF  :LEAF #<SB-C::GLOBAL-VAR :%SOURCE-NAME SB-KERNEL:DATA-VECTOR-REF-WITH-OFFSET :TYPE #1=#<SB-KERNEL:FUN-TYPE #> :DEFINED-TYPE #1# :WHERE-FRO..
  2: (SB-C::IR2-CONVERT-FULL-CALL #<SB-C::COMBINATION :FUN #<SB-C::REF  :LEAF #<SB-C::GLOBAL-VAR :%SOURCE-NAME SB-KERNEL:DATA-VECTOR-REF-WITH-OFFSET :TYPE #1=#<SB-KERNEL:FUN-TYPE #> :DEFINED-TYPE #1# :WHER..
  3: (SB-C::IR2-CONVERT-BLOCK #<SB-C::CBLOCK 11 :START c1 {1005211BF3}>)
  4: (IR2-CONVERT #<COMPONENT :NAME RM-GEMM+STATIC-SIZE+UNROLL2-K {1005223F73}>)

これは、以下の理由に依る。

例のVOPは 多次元配列に対してはマッチしない。
SB-KERNEL:DATA-VECTOR-REF-WITH-OFFSET は always-transferrable という fun-info bit がついている。これは、 VOPに変換されない full call としてコンパイルされるとエラーになることを示すビットである。 IR2 コンパイルの PONDER-FULL-CALL はこれをチェックしている。

解決策としては、多次元配列を1次元配列に変換するラッパ関数 (sb-kernel:%array-data-vector array) を用いる。この "変換" はコンパイル時のメタ情報が変わるだけなので、実行時のオーバヘッドはない。どっちにせよVOPにマッチして変換される。

結果、できたアンローリングは以下。

(defun rm-gemm+static-size+unroll2-k (ma mb mc)
  (declare (optimize (speed 3) (debug 0) (safety 0) (space 0)))
  (declare (type (matrix 500 500) ma mb mc))
  (let ((rows (array-dimension ma 0))
        (cols (array-dimension mb 1)))
    (declare (type fixnum rows cols))
    (dotimes (row rows)
      (dotimes (k cols)
        (let ((cell (aref ma row k))
              (mb-index (array-row-major-index mb k 0))
              (mc-index (array-row-major-index mc row 0)))
          (declare (fixnum mb-index mc-index))
          (dotimes (col (/ cols 2))
            (sb-kernel:data-vector-set-with-offset
             (sb-kernel:%array-data-vector mc) mc-index 0
             (+ (sb-kernel:data-vector-ref-with-offset
                 (sb-kernel:%array-data-vector mc) mc-index 0)
                (* cell (sb-kernel:data-vector-ref-with-offset
                         (sb-kernel:%array-data-vector mb) mb-index 0))))
            (sb-kernel:data-vector-set-with-offset
             (sb-kernel:%array-data-vector mc) mc-index 1
             (+ (sb-kernel:data-vector-ref-with-offset
                 (sb-kernel:%array-data-vector mc) mc-index 1)
                (* cell (sb-kernel:data-vector-ref-with-offset
                         (sb-kernel:%array-data-vector mb) mb-index 1))))))))
    mc))

(benchmark (10 t)
  ;; Evaluation took:
  ;;   2.262 seconds of real time
  ;;   2.260000 seconds of total run time (2.260000 user, 0.000000 system)
  ;;   99.91% CPU
  ;;   6,785,072,819 processor cycles
  ;;   32,784 bytes consed
  (rm-gemm+static-size+unroll2-k *ma* *mb* *mc*))
(benchmark (10 t)
  ;; Evaluation took:
  ;;   2.696 seconds of real time
  ;;   2.696000 seconds of total run time (2.696000 user, 0.000000 system)
  ;;   100.00% CPU
  ;;   8,088,876,210 processor cycles
  ;;   33,056 bytes consed
  (rm-gemm+static-size-k *ma* *mb* *mc*))

結構早くなった。以下ディスアセンブル。

; disassembly for RM-GEMM+STATIC-SIZE+UNROLL2-K
; Size: 208 bytes. Origin: #x1005962C88
; C88:       31FF             XOR EDI, EDI                    ; no-arg-parsing entry point
; C8A:       E9B3000000       JMP L5
; C8F:       90               NOP
; C90: L0:   31C9             XOR ECX, ECX
; C92:       E99A000000       JMP L4
; C97:       660F1F840000000000 NOP
; CA0: L1:   4869C7F4010000   IMUL RAX, RDI, 500
; CA7:       4801C8           ADD RAX, RCX
; CAA:       498B5511         MOV RDX, [R13+17]
; CAE:       F30F104C4201     MOVSS XMM1, [RDX+RAX*2+1]
; CB4:       4C69C1F4010000   IMUL R8, RCX, 500
; CBB:       4869C7F4010000   IMUL RAX, RDI, 500
; CC2:       31F6             XOR ESI, ESI
; CC4:       EB5E             JMP L3
; CC6:       660F1F840000000000 NOP
; CCF:       90               NOP
; CD0: L2:   488B5311         MOV RDX, [RBX+17]
; CD4:       4C8B5311         MOV R10, [RBX+17]
; CD8:       F3410F105C4201   MOVSS XMM3, [R10+RAX*2+1]
; CDF:       4D8B5111         MOV R10, [R9+17]
; CE3:       F3430F10544201   MOVSS XMM2, [R10+R8*2+1]
; CEA:       F30F59D1         MULSS XMM2, XMM1
; CEE:       F30F58D3         ADDSS XMM2, XMM3
; CF2:       F30F11544201     MOVSS [RDX+RAX*2+1], XMM2
; CF8:       488B5311         MOV RDX, [RBX+17]
; CFC:       4C8B5311         MOV R10, [RBX+17]
; D00:       F3410F105C4205   MOVSS XMM3, [R10+RAX*2+5] <<< Good!
; D07:       4D8B5111         MOV R10, [R9+17]
; D0B:       F3430F10544205   MOVSS XMM2, [R10+R8*2+5] <<< Good!
; D12:       F30F59D1         MULSS XMM2, XMM1
; D16:       F30F58D3         ADDSS XMM2, XMM3
; D1A:       F30F11544205     MOVSS [RDX+RAX*2+5], XMM2 <<< Good!
; D20:       4883C602         ADD RSI, 2
; D24: L3:   4881FEF4010000   CMP RSI, 500
; D2B:       7CA3             JL L2
; D2D:       4883C102         ADD RCX, 2
; D31: L4:   4881F9E8030000   CMP RCX, 1000
; D38:       0F8C62FFFFFF     JL L1
; D3E:       4883C702         ADD RDI, 2
; D42: L5:   4881FFE8030000   CMP RDI, 1000
; D49:       0F8C41FFFFFF     JL L0
; D4F:       488BD3           MOV RDX, RBX
; D52:       488BE5           MOV RSP, RBP
; D55:       F8               CLC
; D56:       5D               POP RBP
; D57:       C3               RET

自動アンローリングは次回。