Triton IR語法
Triton IR的語句遵從MLIR Dialect的語法定義規范，示例如下：
%3 = tt.splat %1 : i32 -> tensor<1024xi32> loc(#loc5)其中：
%0：右邊expression的結果值的名字（Value的name）
tt：表示Dialect名稱空間為tt（Triton）
splat：operation的名字
%1：operation的輸入
i32：%1的類型
tensor<1024*i32>：operation的結果類型（即3%的類型）
loc(%loc5)：對應源碼的行號，調試信息。
?如下是一個pytorch cat算子的Triton DSL（inductor產生）
@triton.jit
def triton_(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):xnumel = 3645440xoffset = tl.program_id(0) * XBLOCKxindex = xoffset + tl.arange(0, XBLOCK)[:]xmask = tl.full([XBLOCK], True, tl.int1)x0 = xindex % 890x1 = (xindex // 890)x2 = xindextmp0 = x0tmp1 = tl.full([1], 0, tl.int64)tmp2 = tmp0 >= tmp1tmp3 = tl.full([1], 390, tl.int64)tmp4 = tmp0 < tmp3tmp5 = tl.load(in_ptr0 + ((390*x1) + x0), tmp4, eviction_policy='evict_last', other=0.0)tmp6 = tmp0 >= tmp3tmp7 = tl.full([1], 890, tl.int64)tmp8 = tmp0 < tmp7tmp9 = tl.load(in_ptr1 + ((500*x1) + ((-390) + x0)), tmp6, eviction_policy='evict_last', other=0.0)tmp10 = tl.where(tmp4, tmp5, tmp9)tl.store(out_ptr0 + (x2), tmp10, None)
''', device_str='cuda')編譯生成的Triton IR如下：：
#loc = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":18:0)
module {tt.func public @triton_(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32} loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":18:0), %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32} loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":18:0), %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32} loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":18:0), %arg3: i32 {tt.divisibility = 16 : i32} loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":18:0)) attributes {noinline = false} {%cst = arith.constant dense<-390> : tensor<1024xi32> loc(#loc1)%cst_0 = arith.constant dense<500> : tensor<1024xi32> loc(#loc1)%cst_1 = arith.constant dense<0.000000e+00> : tensor<1024xf32> loc(#loc1)%cst_2 = arith.constant dense<390> : tensor<1024xi32> loc(#loc1)%cst_3 = arith.constant dense<390> : tensor<1024xi64> loc(#loc1)%cst_4 = arith.constant dense<890> : tensor<1024xi32> loc(#loc1)%c1024_i32 = arith.constant 1024 : i32 loc(#loc1)%0 = tt.get_program_id x : i32 loc(#loc2)%1 = arith.muli %0, %c1024_i32 : i32 loc(#loc3)%2 = tt.make_range {end = 1024 : i32, start = 0 : i32} : tensor<1024xi32> loc(#loc4)%3 = tt.splat %1 : i32 -> tensor<1024xi32> loc(#loc5)%4 = arith.addi %3, %2 : tensor<1024xi32> loc(#loc5)%5 = arith.remsi %4, %cst_4 : tensor<1024xi32> loc(#loc6)%6 = arith.divsi %4, %cst_4 : tensor<1024xi32> loc(#loc7)%7 = arith.extsi %5 : tensor<1024xi32> to tensor<1024xi64> loc(#loc8)%8 = arith.cmpi slt, %7, %cst_3 : tensor<1024xi64> loc(#loc8)%9 = arith.muli %6, %cst_2 : tensor<1024xi32> loc(#loc9)%10 = arith.addi %9, %5 : tensor<1024xi32> loc(#loc10)%11 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>> loc(#loc11)%12 = tt.addptr %11, %10 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32> loc(#loc11)%13 = tt.load %12, %8, %cst_1 evictionPolicy = evict_last : tensor<1024x!tt.ptr<f32>> loc(#loc12)%14 = arith.cmpi sge, %7, %cst_3 : tensor<1024xi64> loc(#loc13)%15 = arith.muli %6, %cst_0 : tensor<1024xi32> loc(#loc14)%16 = arith.addi %5, %cst : tensor<1024xi32> loc(#loc15)%17 = arith.addi %15, %16 : tensor<1024xi32> loc(#loc16)%18 = tt.splat %arg1 : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>> loc(#loc17)%19 = tt.addptr %18, %17 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32> loc(#loc17)%20 = tt.load %19, %14, %cst_1 evictionPolicy = evict_last : tensor<1024x!tt.ptr<f32>> loc(#loc18)%21 = arith.select %8, %13, %20 : tensor<1024xi1>, tensor<1024xf32> loc(#loc19)%22 = tt.splat %arg2 : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>> loc(#loc20)%23 = tt.addptr %22, %4 : tensor<1024x!tt.ptr<f32>>, tensor<1024xi32> loc(#loc20)tt.store %23, %21 : tensor<1024x!tt.ptr<f32>> loc(#loc21)tt.return loc(#loc22)} loc(#loc)
} loc(#loc)
#loc1 = loc(unknown)
#loc2 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":20:28)
#loc3 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":20:33)
#loc4 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":21:36)
#loc5 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":21:23)
#loc6 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":23:18)
#loc7 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":24:20)
#loc8 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":30:18)
#loc9 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":31:35)
#loc10 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":31:41)
#loc11 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":31:30)
#loc12 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":31:46)
#loc13 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":32:19)
#loc14 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":35:35)
#loc15 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":35:51)
#loc16 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":35:42)
#loc17 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":35:30)
#loc18 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":35:57)
#loc19 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":36:33)
#loc20 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":37:25)
#loc21 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":37:37)
#loc22 = loc("/tmp/torchinductor_vincent/yx/cyxlesomwjvogzqnvxmuj2p2z2gml7hud473ghmveu4fg6jbtlmz.py":37:4)Triton IR依賴的Dialects
編寫完triton程序后，導出的IR中，可以看到不止有triton IR，還包含其他的MLIR Dialects，其中包含：
-
Arith: addf, addi, andi, cmpf, cmpi, divf, fptosi, …
-
Math: exp, sin, cos, log, …
-
StructuredControlFlow（scf）: for, if, while, yield, condition
-
ControlFlow(cf): br, cond_br
Triton IR Operations
tt.call?(triton::CallOp)
語法:
operation ::= `tt.call` $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)
tt.call表示對同一個符號作用域內的函數的直接調用。
示例:
%2 = tt.call @my_add(%0, %1) : (f32, f32) -> f32
tt.func?(triton::FuncOp)
function聲明或定義，function是一個SSACFG region。
function內的Operation不能隱式地捕獲function外定義的值。所有外部引用都必須通過arguments或者attribute來傳遞。在MLIR中，function的arguments是通過第一個block的block arguments來表達的。
示例：
// External function definitions.
tt.func @abort()
tt.func @scribble(i32, i64, memref<? x 128 x f32, #layout_map0>) -> f64// A function that returns its argument twice:
tt.func @count(%x: i64) -> (i64, i64)attributes {fruit: "banana"} {return %x, %x: i64, i64
}// A function with an argument attribute
tt.func @example_fn_arg(%x: i32 {swift.self = unit})// A function with a result attribute
tt.func @example_fn_result() -> (f64 {dialectName.attrName = 0 : i64})// A function with an attribute
tt.func @example_fn_attr() attributes {dialectName.attrName = false}SSACFG region
SSACFG region內的語句滿足SSA形式，且不包含子Region（既不能包含scf.if/scf.for等），如下就是一個SSACFG region：
func.func @example(%a : i32) -> i32 {// 這是一個 SSACFG Region%cmp = arith.cmpi slt, %a, %c10 : i32cond_br %cmp, ^bb1, ^bb2^bb1:%x = arith.addi %a, %c1 : i32br ^exit(%x : i32)^bb2:%y = arith.subi %a, %c1 : i32br ^exit(%y : i32)^exit(%result : i32):return %result : i32
}如下不是一個SSACFG Region：
scf.if %cond {// 這里是一個新的 Region（嵌套）scf.yield
}Block Arguments
對如下函數：
func.func @foo(%arg0: i32, %arg1: f32) -> f32 {// 函數體直接使用 %arg0, %arg1%result = arith.addf %arg1, %arg1 : f32return %result : f32
}在MLIR的內部實現里，是把function的arguments作為function內第一個基本塊（entry block）的 block arguments 來存儲：
func.func @foo() -> f32 {
^bb0(%arg0: i32, %arg1: f32): // ← 參數實際屬于 entry block%result = arith.addf %arg1, %arg1 : f32return %result : f32
}這是因為MLIR的設計要求所有 SSA 值都由某個 block 或 op 產生，這樣做也解決了LLVM IR中的phi node的問題。
在LLVM IR中，通過phi node來匯聚不同前驅路徑的值，示例如下：
entry:br i1 %cond, label %left, label %rightleft:br label %mergeright:br label %mergemerge:%x = phi i32 [ %v1, %left ], [ %v2, %right ] ; ← φ 節點ret i32 %x在MLIR中，通過block arguments，可以實現同等的效果：
func.func @foo(%cond: i1, %v1: i32, %v2: i32) -> i32 {cf.cond_br %cond, ^left, ^right^left:cf.br ^merge(%v1 : i32) // 把 %v1 作為參數傳給 merge^right:cf.br ^merge(%v2 : i32) // 把 %v2 作為參數傳給 merge^merge(%x : i32): // ← block argument 取代 φreturn %x : i32tt.return?(triton::ReturnOp)
語法:
operation ::= `tt.return` attr-dict ($srcs^ `:` type($srcs))?
表達function的返回操作，擁有變長個數的操作數，操作數的個數和類型必須和function的簽名匹配。
示例：
tt.func @foo() : (i32, f8) {...tt.return %0, %1 : i32, f8
}tt.addptr?(triton::AddPtrOp)
語法:
operation ::= `tt.addptr` $ptr `,` $offset attr-dict `:` type($result) `,` type($offset)
張量或標量指針地址線性偏移計算。
示例：
%base = tt.splat %arg0 : !tt.ptr<f32> -> tensor<1024x!tt.ptr<f32>>
%idx = tt.make_range {start = 0, end = 1024} : tensor<1024xi32>// 生成偏移地址
%ptrs = tt.addptr %base, %idx: tensor<1024x!tt.ptr<f32>>, tensor<1024xi32>// 加載數據
%vals = tt.load %ptrs : tensor<1024xf32>tt.advance?(triton::AdvanceOp)
語法:
operation ::= `tt.advance` $ptr `,` `[` $offsets `]` attr-dict `:` type($result)
對 !tt.ptr<tensor<...>> 類型的指針按給定的 多維偏移量 進行偏移計算，返回一個新的張量指針。
示例：
scf.for %i = %c0 to %c128 step %c32iter_args(%tile_ptr = %base_ptr) -> (!tt.ptr<tensor<32x32xf16>>) {// 使用當前 tile%vals = tt.load %tile_ptr : !tt.ptr<tensor<32x32xf16>>// 推進到下一個 tile（第1個維度上推進 32，第2個維度保持不變）%next_ptr = tt.advance %tile_ptr, [%c32_i32, %c0_i32] : !tt.ptr<tensor<32x32xf16>>scf.yield %next_ptr : !tt.ptr<tensor<32x32xf16>>
}tt.assert?(triton::AssertOp)
語法:
operation ::= `tt.assert` $condition `,` $message attr-dict `:` type($condition)
tt.assert作用在device側，接收1個condition(i1 類型的標量或張量)和一個string. 如果condition為false，則打印message并終止程序。
示例：
%in_bounds = arith.cmpi slt, %idx, %size : i32
tt.assert %in_bounds, "index out of bounds"TODO
參考資料：
TritonOps — Triton documentation
