AbstractTensor.lisp

Usage (WIP)

./roswell/caten.ros -i ./samples/gemm.toml --runtime ./runtimes/clang.lisp --debug 1 --test "M=100, N=100, K=100" --config OMP=1,MARCH\"native\"

Design (Memo)

• Eazy to read, eazy to develop, keep simplifiers hackable; 冗長なコード -> 拡張可能なSimplifierで最適化の方針

• Detach runtimes from ./source;

• CLIツールベースでやる，Coalton/Java/PythonあたりからBindingかけるようにする。

• etc...

Ideas

• Coalton-Based Implementation

• Lazy Graph Construction

• Only 27 principle operators. JIT Compiler

• Super eazy to add a new backend

• Keep the code simple and tiny.

• PyTorch-Level Speed

• Quantization Support

• Complex number support dedicated to HiFi-GAN

• Super fast to compile the model.

• Add: Eazy to add Python/Coalton binding.

• Eazy to add manual optimization techniques by users

• Documentation System

• Better Examples

• Support these ops and optimizations: Take, Slice

• Elegant Error

• Eazy to add graph-rewriting

• Every Ops are in-place, but for example add is defined as: (add (copy x) y).

• A Concept of Composite

• Super fast compiling time

• GraphView System

• Super eazy to debug/optimize/and understand the bottleneck.

• Kernel自動生成の考え方，Conv/ReLUなどは事前に自動生成しておく。

• Eazy (and fast) to rewriting a graph

  • CompileされたCompositeの一覧をどっかのLUTに割り振れるようにする。
  • Training/Inferenceの実行は，datenにGraphを記述することで実行できるようにする。

• randn/betaとかもコンパイルして生成

• winograd

• runtime template (ros create tekina)

• ./opsに色々最適化されたtomlを書いておく。arangeとかrandnも実装できるように

• コンパイラは，全部のTypeに対するtopi的なやつの自動生成を第一目標に，次にFusionしたカーネルの生成を目標にする。

  • だから，Polyhedral Compilerを実装しない方針にした代わりに
  • 計算中Indexを参照する
  • スカラー計算をやりやすく抽象化するとか
    • UOpsでスカラー計算表現した方がいい気がする。
  • ^とval_xxx使ってforでスカラー計算表現するとか
  • ができるべき。

• TODO: Beam Search, Auto Scheduler, Combine with cl-polyhedral to implement an optimizing and multiple-backend einsum!

• SIMDGroupでIfしたりIfしなかったりを生成できるようにしたい。

• we are too lazy to reimplement the backends. The smallest effort to write an extension!

• TinygradのScheduling, Loweringをもってくる。

• コンセプトはこう:

  • UOpsに対して難しい最適化Polyhedral Compilationはしない。(その代わり手動で各段階でループのオーダーだけチューニングしておく) -> after, simdify or grouping.
  • ^ メモリ帯域幅の事考えるとPolyhedral Compilationは欲しい。でもUOpのレベルじゃなくて，コンパイル元Lispコードの段階でできないか？(ループの順番入れ替えるだけだし，コードを難しくしたくない)
  • 27のPrinciple Operatorがある。(Including Take, If, etc...)
  • Composite = 27 Opsを組み合わせて複雑な命令(ReLU, Conv2D, Gemm...)を表現
  • In-Placeにこだわりすぎなくていい (LayerNormとか，ある程度でかいブロックでおk)
  • Composite Cache System. 自動でコンパイル結果をCache+割り当てする。 (~/tmpとかに配置)
  • Dynamicなグラフ書き換えに対応。
  • ^- defoptimizerみたいな感じで，書き換えのルールを作成。(OpenBLASのGemmに書き換える，など)
  • Composite call from compositeができるようにしたい
  • 今考える: Multiple BackendにCompileできるPrinciple Ops(Take)を実装. -> Scheduler -> Compile
  • [10, 10][10, 0]にしたらSIMD化できないけど，どう実装するのか気になる
  • ^- 上の最適化が必要になるのってどんな状況だろう？
    • Composite <-> Compositeの箇所でSliceは挟まないから，一旦Composite内部で考えればよさそう？
  • 要は早いConv2DのCompositeさえ入手できれば(PyTorchがそうなように) 十分な速度を期待できる。
  • Scalarに対するCompilerもつくる！

• POV: FlashAttention2が生成できる？ https://gist.github.com/xiabingquan/a4a9a743f97aadd531ed6218be20afd2

• POV: ~の記号があるとき -> ndarrayに対して処理できるようにしたい。

• Implement a pipeline for making a benchmark, for multiple runtimes, multiple composites.

• POV: Viewed Sin関数みたいなのを最適化？・・・一つのCompositeでは，Inputs, Outputsが全てBaseだと仮定する。

• 作成したRuntimeに対してのUnittestみたいなのを自動でできるようにする。(./aten-cli test-runtime --runtime xxx.lisp) tekina. (./Reports.mdみたいなの生成して対応してるOpSetの一覧の表を作る)

• ElementWiseな関数をFlattenする処理 ->

(defun sin (x)
if x is a base tensor -> x.flatten().sin().reshape_as_it_was()
if x is a view -> x.sin()
)

• daten.ros, abstracttensor/high-level/autograd, abstracttensor/high-level/quantization, nn, etc...

Usage (WIP)

• Lisp風DSLを記述 (他の言語から自動生成されること前提)
• gcc的に使えるようにしたい

Codegen

$ caten --runtime cuda
"
(composite ((A{T}[M, N], B{T}[M, N]) -> (A{T}[M, k]) where K = (* 2 M))
    (setf z1 (+ A B))
    (return z1))
"
↑ yamlでもいい気がする
conv2d:
    - in:
        - A{T}[M, N]<0 1>(where N = 1)
	- B{T} ...
    - out:
	- B # outはSymbolのみ
    - impl:
        - "(progn ...)"

Autodiff (WIP)

$ daten --runtime cuda --plugin cuda-backend.lisp
"
modules:
  function L1: A{Dense}[M, N] B{Dense}[M, N] -> A{T}[M, k]
  function L2: A{Dense}[M, N] B{Dense}[M, N] -> A{T}[M, k]

parameters:
  A{Grad}
  B{Grad}

main:
  output = (L1 (L2 A B) C)
"

• T = {Dense, Sparse, Complex}

設計的な

• AbstractTensor/IR 内部IR

  • Tensorの情報 (Name, Required-Grad, Dynamic Shape) を伝える
  • BodyのDAGを記述する部分
  • Iterations can be only expressed as permutations.

• AbstractTensor/Lang

  • DL Compiler専用プログラミング言語的な