Parser 工具#
下面几项工具在 TVMScript 转换为 TIRx 的 parse time 生效。它们可以将 Python 计算得到的值直接写入 IR、提取可复用的代码片段,以及组织 parser 阶段使用的状态。
T.meta_var:保留 Python metadata#
T.meta_var(x) 原样返回 x,可以保留 parser 侧的对象句柄或 Python
metadata tuple。它不会绕过普通赋值规则:在 TVMScript 内赋值标量表达式仍会
创建原生的值绑定。Shape 参数、指令宽度和 TMA box 大小应在
@T.function 之前的 Python factory 中计算。
widths = T.meta_var((32, 64)) # tuple remains Python metadata
for j in T.unroll(widths[0]):
acc[0] = acc[0] + A[tx, j]
普通 range(n) 创建 serial TIRx loop;T.unroll(n) 标记该 loop 在
lowering 时展开。meta_var 不会让运行时值在输入变化后自动重新求值。
使用 persistent scheduler 时,应在每轮任务循环内,根据更新后的 cursor
重新计算 tile 坐标。
@T.inline:内联函数#
@T.inline 定义的函数会在 parse time 内联到每个调用位置,生成的代码中
不会保留函数调用。它遵循 Python 的 lexical(LEGB)scope 和 late binding;
函数参数会覆盖外层的同名变量:
@T.inline
def add_into(acc, x):
acc[0] = acc[0] + x
add_into(acc, A[tx, j]) # inlined -> acc[0] = acc[0] + A[tx, j]
@T.meta_class:组织 Parser 状态#
@T.meta_class 修饰普通的 Python class,使它的 instances 成为 parser meta
values。对象字段可以保存 buffers 和 scalars,因此可以把相关的 allocation
和状态放进同一个对象,再在 kernel body 中使用。
State 定义在 @T.function 外部;创建 s 和使用其字段的语句则放在
kernel 内部,并先分配 smem。
@T.meta_class
class State:
def __init__(self, smem):
self.acc = T.bind_(T.alloc_local([1], "float32"))
self.buf = T.bind_(T.decl_tensor(smem, ty_args=[
T.Tensor((64,), "float16", scope="shared.dyn")]))
s = State(smem.ptr_to([0]))
s.acc[0] = T.float32(0.0) # use its fields like ordinary buffers
# ... s.buf[i] ...
这种写法适合组织 kernel 的 pipeline state,例如 barriers、accumulators 和 临时 views,避免在函数中传递大量分散的局部变量。
T.constexpr#
以 T.constexpr 标注的参数由 @T.jit 的 .specialize(...)
固定。在函数体内,if T.constexpr(condition) 会在解析时选择 Python
分支,适用于 FA4 中的 is_causal 等已知 specialization flag。
普通 if 会构造 IR 控制流,并解析两个分支。