前面兩天從論文理解 TileLoom 的 mapping、reuse 與 performance model。
接下來想要觀察編譯器產生實際的 IR，看這些概念如何對應。

我選矩陣相乘（matmul）當例子，因為 matmul 是 AI 模型裡常見的核心運算，而且數學關係單純。
輸入是 A[M, K] 和 B[K, N]，輸出是 C[M, N]。
我們可以把焦點留給 IR 如何逐步加入 tile、記憶體與硬體資訊。

版本範圍：今天的實作指令、圖片與 00～08 IR，來自舊 loom-dataflow standalone pipeline。可以從 examples/mm/IR 瀏覽這組 artifact，本文以 2026-07-31 的 loom-dataflow@39d2b858 固定實作版本。

今天會先複習 Loom 的完整編譯流程，再比較

00_from_helion_frontend.mlir
  → 01_tensor_canonicalized.mlir
  → 02_explicit_memory_access.mlir

本篇大綱

• 從 Loom 大架構定位這三份 IR
• 說明接下來會如何沿著同一個 matmul 追完整條 pipeline
• 比較 00 → 01 的 bufferization handoff 與 proxy copy
• 比較 01 → 02 新增的 L1 buffer、資料搬移與資源生命週期
• 回到 C++ pass 與 TableGen，確認畫面上的變化由哪段實作產生

先複習新版 Loom 的大架構

截至 2026-08-12 的新版 loom@f65d199b monorepo，把端到端流程分成四個子專案，由 root Python package 負責串接。這一節提供架構背景，後面的 00～08 IR 實作仍使用舊 pipeline。

簡單解釋這裡出現一些名詞

Helion 和 Triton 是什麼關係？

Helion 是 Python-embedded DSL。寫法接近帶有 tile 的 PyTorch，程式用 PyTorch operation 表達計算，再用 hl.tile 表達要切分的迭代空間。Helion 的一般編譯路徑會產生 Triton kernel，並自動探索 tile size、indexing、邊界 mask、program ID mapping 與其他 tuning choices。Program ID 是 Triton kernel 中工作執行個體的編號，mapping 會決定每個執行個體負責哪一塊資料。

Triton 位在較低的抽象位置。它讓開發者用 Python-like DSL 撰寫高效能 kernel，程式通常要更直接地處理 program ID、資料指標、mask 與 block shape。可以把兩者的關係整理成

一般 Helion 路徑

PyTorch-like Helion kernel
  → Helion frontend／autotuner
  → Triton kernel
  → GPU code

Loom 需要在 lowering 前期看見 symbolic tile size、structured loop、linalg.matmul 與 tensor SSA value，後面的 pass 才能列舉硬體 mapping、分析 reuse，並估計不同資料搬移方案。helion-mlir 因此重用 Helion frontend 產生的 Device IR FX graph，再改走 Affine＋Linalg-on-Tensors 的 MLIR 路徑

Loom 使用的路徑

PyTorch-like Helion kernel
  → Helion Device IR／FX graph
  → helion-mlir
  → affine + linalg-on-tensors MLIR
  → loom-dataflow

helion-mlir 會把 Helion Device IR 中的 _for_loop 轉成 affine.for／affine.parallel，把 addmm 等 ATen operation 經由 torch-mlir 降成 linalg。Triton lowering 會逐步消去上述高階資訊；Loom 需要在資訊仍可分析的 Helion frontend 階段接出另一條 MLIR 路徑。

今天選擇 Helion 還有三個實際原因

1. 在這份新版快照中，Loom root repository 的端到端使用方式，是用 helion.kernel() 定義 kernel，再交給 LoomKernel 執行 pipeline。
2. 我們閱讀的第一份 artifact 名為 00_from_helion_frontend.mlir，可以直接對照 Helion matmul 裡的 hl.tile、hl.dot 與輸出的 affine／linalg 結構。
3. Helion 把 indexing、mask 與 tuning 細節留給 frontend 與 autotuner，這幾篇可以把焦點放在 Loom 如何轉換 IR。

新版六個 stage 各自解決什麼問題？

架構文件也提到，使用者可以直接提供 assigned_block_size。此時 pipeline 會略過 Stage 3，將指定值送進 materialization。

另外這個六階段架構流程和等一下要介紹 IR 教學 examples/mm/IR 裡的 00～08 檔名粒度不同。 examples/mm/IR 把 loom-dataflow 內部的 pass 結果逐步存下來。

參考：Loom Architecture（2026-08-12 快照）、Helion、helion-mlir

接下來沿著同一份 matmul 說明 2026-07-31 舊版 pipeline

跟著官方的例子來學習 examples/mm/IR

Helion frontend
  00_from_helion_frontend.mlir
        │ tensor canonicalization
        ▼
  01_tensor_canonicalized.mlir
        │ memory binding
        ▼
  02_explicit_memory_access.mlir
        │ hardware mapping
        ▼
  03_after_hardware_mapping.mlir
        │ reuse analysis
        ▼
  04_after_reuse_analyzation.mlir
        │ copy / broadcast enumeration
        ▼
  05_after_enumerate_broadcast.mlir
        │ block-size materialization + canonicalization
        ▼
  06_after_canonicalize.mlir
        │ one-shot bufferization
        ▼
  07_after_osb.mlir
        │ TT-oriented cleanup
        ▼
  08_tt-opt.mlir

repo 裡的 final.mlir 是另一份硬體模型輸入，不是 08_tt-opt.mlir 的下一階段輸出。

今天處理前面三份，先看計算語意如何整理成 Loom 能分析的形狀，再看資料搬移顯式化。下一步會把 logical tile grid 映射到硬體，接著分析 A、B tile 可以在哪些方向重用。最後再追 symbolic tile size materialize、one-shot bufferization 與 TT cleanup 後，IR 還留下什麼。

專案會持續更新，本文依據舊版 loom-dataflow@39d2b858 的例子與 pass 實作。

如何重現今天的兩個轉換

在 loom-dataflow repository build 完成後，可直接執行

這裡執行的是舊 standalone run_pipeline.sh。新版 monorepo 的 run_pipeline() 是另一條端到端路徑，兩者不混用。

./run_pipeline.sh mm "[1,2]"

也可以分開執行兩個 standalone driver

build/tool/loom-opt/single_stage/tensor_canonicalize \
  --input examples/mm/IR/00_from_helion_frontend.mlir \
  > /tmp/01_tensor_canonicalized.mlir

build/tool/loom-opt/single_stage/memory_binding \
  --input /tmp/01_tensor_canonicalized.mlir \
  > /tmp/02_explicit_memory_access.mlir

這裡刻意輸出到 /tmp，避免覆蓋 repository 已經保存的參考 artifact。

00：先讀懂 matmul 的計算骨架

函式輸入是

func.func @matmul(
  %x_arg:    memref<2048x256xf16>,
  %y_arg:    memref<256x256xf16>,
  %out__arg: memref<2048x256xf16>)

對應的矩陣形狀為

X   : M × K = 2048 × 256
Y   : K × N =  256 × 256
Out : M × N = 2048 × 256

tile_m、tile_n、tile_k 此時仍是帶有 upper bound 的符號：

%tm = loom.sym @tile_m {upper_bound = 2048 : index} : index
%tn = loom.sym @tile_n {upper_bound = 256 : index} : index
%tk = loom.sym @tile_k {upper_bound = 256 : index} : index

Frontend 先留下如何切 tile 的參數，後面的 solver 才會根據 ETG constraint、硬體模型與成本式決定數值。

外層 affine.parallel 走訪輸出矩陣的 M、N tile grid，內層 scf.for 沿 K 軸累加

affine.parallel (%mi, %ni) = (0, 0)
    to (symbol(%m_tiles), symbol(%n_tiles)) {
  %result = scf.for %ki = %c0 to %k_tiles step %c1
      iter_args(%acc = %zero_tile) -> (tensor<?x?xf16>) {
    %a = memref.subview %x_arg[%m0, %k0] [%tm, %tk] [1, 1]
    %b = memref.subview %y_arg[%k0, %n0] [%tk, %tn] [1, 1]
    %next = linalg.matmul ins(%a_tensor, %b_tensor) outs(%acc)
    scf.yield %next
  }
}

換成數學式就是

C_tile(mi, ni) = Σki A_tile(mi, ki) × B_tile(ki, ni)

memref.subview 只描述從完整矩陣取哪一塊 view，bufferization.to_tensor 讓 tensor-mode linalg.matmul 使用它。這份 IR 還看不到 A、B tile 要放在哪個 L1，也沒有 DRAM 到 L1 的搬移操作。

00 → 01：把 handoff 整理成 Loom 能穩定辨認的形式

第一個明顯變化是 bufferization handoff operation 改用 Loom dialect

圖中有兩組變化

1. bufferization.to_tensor 變成 loom.bufferize_to_tensor。
2. bufferization.to_buffer 變成 loom.bufferize_to_memref。

loom.bufferize_to_tensor %subview_0[%0, %2] 的中括號明確帶入 logical tensor shape，也就是 A tile 的 tile_m × tile_k。這些 operation 都是 view／handoff，不會搬資料；DRAM→L1 的 loom.copy 會在 02 出現。

K loop 結果後面為何多了一次 copy

第二個變化發生在 K loop 結果寫回 output 之前

00 直接將 loop result %8 轉成 memref 後寫回。01 加入

%9 = tensor.empty(%0, %1) : tensor<?x?xf16>
%10 = linalg.copy ins(%8 : tensor<?x?xf16>)
                  outs(%9 : tensor<?x?xf16>)
                  -> tensor<?x?xf16>

後面的 writeback 改讀 %10。這個 proxy tensor 讓 memory analysis 看見一條獨立的 handoff SSA chain，同時保留 K loop 內 accumulator 的重用關係。

這個改寫來自 tensor_canonicalize driver 內的一串 passes。其中與本例畫面直接相關的是

LoopHandoffProxyCopyInsertionPass
  └─ 在 loop result 與 communication boundary 間加入
     tensor.empty + linalg.copy

CanonicalBufferizationToLoomPass
  ├─ bufferization.to_tensor → loom.bufferize_to_tensor
  └─ bufferization.to_buffer → loom.bufferize_to_memref

linalg destination specialization、elementwise fusion 與 extract-slice folding 也在同一個 driver 裡執行，但這份單純的 matmul dump 沒有呈現每個 pass 都改到 IR，閱讀 compiler pipeline 時，可以從實際 diff 判斷哪些 pass 在這個輸入上生效。

01 → 02：把記憶體與資料搬移寫進 IR

memory_binding 會先分析 tensor 的生命週期與 shape，建立 virtual buffer／physical buffer allocation plan，再依序 materialize allocation、管理 semaphore、將 tensor destination 綁到 buffer，最後改寫 load 與 writeback pattern。

先看 A tile 與 accumulator 的 before／after

紅框可以分成三組閱讀

1. C accumulator 有明確的 L1 buffer

01 以 tensor.empty 表示 tm × tn 的 destination；02 改成

%c_alloc = loom.alloc [%tm, %tn] on @L1 : memref<?x?xf16>
%c_buf = loom.semaphore_take %c_alloc
%c_tensor = loom.init_tensor %c_buf[%tm, %tn]
%zero = linalg.fill ins(%cst : f16) outs(%c_tensor)

• loom.alloc 宣告這個 shape 需要 L1 空間
• loom.semaphore_take 取得一個可追蹤生命週期的 buffer handle
• loom.init_tensor 將 buffer 綁回 tensor 語意，供 linalg.fill 與 linalg.matmul 使用

此時的 loom.alloc 是下游分析與 lowering 可見的配置，還不能當作實際 L1 位址

2. A、B tile 各自取得 L1 buffer

A tile shape 是 tm × tk，B tile shape 是 tk × tn

%a_alloc = loom.alloc [%tm, %tk] on @L1
%a_buf = loom.semaphore_take %a_alloc

%b_alloc = loom.alloc [%tk, %tn] on @L1
%b_buf = loom.semaphore_take %b_alloc

這些 allocation 位於 K loop 外。loop body 每一輪會依照 %ki 計算新的 subview offset，再把對應 tile 搬進 buffer

3. DRAM → L1 已成為明確操作

01 的 memref.subview + loom.bufferize_to_tensor 在 02 展開成

%a_src = loom.subview %arg0[%m0, %k0] [%tm, %tk] [1, 1],
  reuse : [seq = false, spat = false, temp = false]

loom.copy %a_src, %a_buf
  src_mem_space @mem_DRAM
  dst_mem_space @mem_L1,
  area : [1, 1]

%a_tensor = loom.bufferize_to_tensor %a_buf[%tm, %tk]

loom.subview 保留來源座標與 shape metadata，loom.copy 表示實體資料搬移，接著才建立 tensor view 給 linalg.matmul。B tile 也經過相同路徑。

目前所有 reuse flag 都是 false，copy area 也是 [1, 1]。硬體 mapping 與 reuse analysis 尚未執行，compiler 還沒有足夠資訊決定 row broadcast、column broadcast 或 temporal hoisting。

Output 的 L1 → DRAM writeback

K loop 結束後，02 也把輸出路徑改成 Loom operation

%out_l1 = loom.alloc [%tm, %tn] on @L1
%out_buf = loom.semaphore_take %out_l1
%out_tensor = loom.init_tensor %out_buf[%tm, %tn]
%copied = linalg.copy ins(%result) outs(%out_tensor)
%out_memref = loom.bufferize_to_memref %copied

loom.copy %out_memref, %out_subview
  src_mem_space @mem_L1
  dst_mem_space @mem_DRAM,
  area : [1, 1]

因此 02 已經能畫出完整的資料路徑

DRAM A ──copy──▶ L1 A ─┐
                       ├─▶ linalg.matmul ─▶ L1 C ──copy──▶ DRAM C
DRAM B ──copy──▶ L1 B ─┘

loom.semaphore_take/give 在這一層描述資源生命週期，也會阻止 compiler 將具有不同同步語意的 handle 合併。這裡尚未降成最終硬體 semaphore 指令。這份 dump 中，A、B 的 take 在 K loop 外、give 在 loop body 內，trip count 大於一時，take／give 如何配對值得在後續 materialization 與 backend lowering 繼續檢查，本文先保留這個觀察。

C++ pass 與 TableGen 如何對上這些 IR

這次需要看的實作集中在三個地方

memory_binding_pass.cpp 的兩個 rewrite pattern

讀取 A／B tile 的 ReadBlockLoadingLowering 會比對

memref.subview
  → loom.bufferize_to_tensor

它查詢 allocation plan 中對應的 L1 buffer，建立 loom.subview，插入 loom.copy，再讓新的 loom.bufferize_to_tensor 讀取 L1 semaphore buffer。

輸出端的 WriteBackLowering 則比對

loom.bufferize_to_memref
  → memref.copy 到 memref.subview

改寫後會得到 loom.bufferize_to_memref + loom.copy，copy 的 memory space 方向設為 @mem_L1 → @mem_DRAM。

LoomOps.td 定義 operation 長什麼樣

以 loom.copy 為例，TableGen 定義中包含

• source 與 destination memref
• src_mem_space、dst_mem_space
• copy area
• 選擇性的 mesh region 上下界
• reclaim attribute

因此後續 pass 能逐步補上 broadcast area 與 mesh coordinates，不必改用另一套 copy operation。LoomOps.cpp 再提供 type inference、canonicalization pattern 與 memory effect 等需要 C++ 實作的行為。

比較 IR 時要守住哪些不變條件

新增記憶體操作後，matmul 的數學關係仍應保持

圖上的 SSA 編號從 %14 變成 %16 沒有獨立語意，只是前面新增 operation 後重新編號。比較時應檢查 operation 類型、operand 來源、shape、index expression、memory space 與生命週期。

今天的結論

今天從 00 走到 02，可以看到三個具體變化

00：Helion frontend 產生 tiled matmul 計算骨架
01：建立 Loom handoff operation 與獨立 proxy tensor
02：加入 L1 buffer、DRAM/L1 copy 與資源生命週期

02 已經說清楚資料要在 DRAM 與 L1 之間移動，還沒決定 logical M／N tile 要放到哪個硬體座標。
明天會從 03_after_hardware_mapping.mlir 開始，觀察 affine.parallel 如何對應到 2D mesh，再沿著 subview index 分析 A、B tile 的 spatial、temporal 與 sequential reuse。

參考資料

• Loom Architecture
• PyTorch Helion
• helion-mlir：Helion Device IR 到 Affine＋Linalg-on-Tensors MLIR
• Triton language and compiler
• 舊版 loom-dataflow：matmul IR artifacts
• 本文固定的舊版 commit：loom-dataflow@39d2b858
• loom-dataflow：pipeline 與 standalone tools
• tensor canonicalization driver
• memory binding pass
• Loom operation TableGen definitions