flydsl-kernel-authoring

FlyDSL Kernel 编写技能

概述

FlyDSL 是一种基于 Python 的 DSL 和基于 MLIR 的编译器，用于在 AMD GPU（MI300X/MI350）上编写高性能 GPU 内核。它提供显式布局代数，用于控制数据移动、分块和内存访问模式。布局系统是 FlyDSL 区别于 Triton/Gluon 的核心抽象。

仓库：当前检出内容（以下示例假设可将其导入——见 12） 目标 GPU：gfx942（MI300X，CDNA3）、gfx950（MI350，CDNA4） Python：3.12，ROCm 7.2

范围（请先阅读）：此技能是参考手册——涵盖完整的布局代数 API、各操作的表格、MFMA/复制原子目录、环境变量以及详尽的故障排查清单。编写或审查内核代码时，可查阅此手册。如果你需要的是一套循序渐进的操作流程，将内核需求转化为完整且经过测试的内核（分类 -> 框架 -> 计算 -> 控制流 -> 测试），请使用 flydsl-tile-programming 技能，它是本参考手册的向导搭档。若要诊断已成功编译、但会产生 NaN/inf/错误结果的内核，请使用 debug-flydsl-kernel 技能。

---

1. 架构与编译

流水线

Python (@flyc.kernel/@flyc.jit)
  -> AST Rewriting (for/if -> scf.for/scf.if)
  -> MLIR Tracing (generates Fly dialect + gpu/arith/scf/memref/vector ops)
  -> MlirCompiler.compile() (Fly -> ROCDL -> LLVM -> HSACO binary)
  -> JITCFunction (ExecutionEngine wrapper)

关键 Pass

流水线由 RocmBackend._pipeline_parts() 构建，并分为三个阶段——各 Pass 的表格见 docs/architecture_guide.md §3。重点如下：

1. fly-rewrite-func-signature - 将函数 / SCF 边界处的 DSL 类型改写为打包的 LLVM 结构体
2. fly-layout-lowering - 将布局代数（fly.crd2idx、分区、除法）降为算术运算
3. fly-convert-atom-call-to-ssa-form + fly-promote-regmem-to-vectorssa - 将复制 / MMA 原子调用和寄存器内存提升为向量 SSA
4. convert-fly-to-rocdl - 将 Fly 操作转换为 ROCDL 内建函数
5. gpu-module-to-binary{format=fatbin} - 通过 LLVM AMDGPU 后端生成 HSACO 二进制文件

关键源码路径

• python/flydsl/compiler/ - JIT 编译（jit_function.py、kernel_function.py）
• python/flydsl/expr/ - DSL 表达式 API（primitive.py、derived.py、typing.py）
• python/flydsl/expr/primitive.py - 所有布局代数函数
• python/flydsl/expr/derived.py - CopyAtom、MmaAtom、TiledCopy、TiledMma 封装
• python/flydsl/expr/gpu.py - GPU 操作（thread_idx、block_idx、barrier）
• python/flydsl/expr/rocdl/ - MFMA/WMMA 和其他 ROCm 内建函数 （包：cdna3、cdna4、cdna5、rdna3、rdna4、cluster、inline_asm、tdm_ops、universal； 以及 utils.py / enum.py 辅助工具）
• python/flydsl/expr/gpu.py - 用于 LDS（共享内存）的 SharedAllocator、thread_id/block_id、barrier
• kernels/common/buffer_ops.py - 旧版原始 AMD 缓冲区加载 / 存储内建函数 （已在 #880 中从 flydsl.expr 移出；优先使用 fx.rocdl.make_buffer_tensor）
• kernels/ - 预构建内核，按子包组织：gemm/（preshuffle_gemm.py、mxfp4_preshuffle.py、...）、norm/（layernorm/softmax/rmsnorm）、attention/、moe/、mega_moe/、common/（包括 common/mma/）、comm/、conv/

---

2. 布局系统（核心抽象）

参见 references/layout-algebra.md：核心类型、构造、坐标映射、布局代数、乘积和除法。

---

3. 编写内核

基本模式

import flydsl.compiler as flyc
import flydsl.expr as fx
from flydsl.expr import const_expr, gpu, range_constexpr, rocdl

@flyc.kernel
def my_kernel(
    A: fx.Tensor,         # GPU tensor (memref via DLPack)
    B: fx.Tensor,
    N: fx.Constexpr[int], # Compile-time constant
):
    tid = gpu.thread_id("x")
    bid = gpu.block_id("x")
    # ... kernel body ...

@flyc.jit
def launch(
    A: fx.Tensor,
    B: fx.Tensor,
    N: fx.Constexpr[int],
    stream: fx.Stream = fx.Stream(None),
):
    my_kernel(A, B, N).launch(
        grid=(N // 256,), block=(256,), stream=stream
    )

# Usage:
import torch
A = torch.randn(1024, device="cuda", dtype=torch.float32)
B = torch.empty(1024, device="cuda", dtype=torch.float32)
launch(A, B, 1024)

当前语法速查

编写新内核时，请使用 kernels/gemm/preshuffle_gemm.py 中当前公开的 FlyDSL 接口：

Vec = fx.Vector

tx = gpu.thread_id("x")
bx = gpu.block_id("x")
by = gpu.block_id("y")

i32_m: fx.Int32
c_m = fx.Int64(i32_m)
c4 = fx.Int64(4)
zero_f = fx.Float32(0.0)

layout = fx.make_layout((4, 64), (64, 1))
coord = fx.idx2crd(tx, layout)
wave_id = fx.get(coord, 0)
lane_id = fx.get(coord, 1)

acc = Vec.filled(4, 0.0, fx.Float32)
v_i64 = Vec(raw_vec).bitcast(fx.Int64)
elem0 = v_i64[0]

buf = fx.rocdl.make_buffer_tensor(tensor)          # preferred buffer-resource view
tA  = fx.make_view(fx.get_iter(buf), fx.make_layout((M, N), (N, 1)))
# load/store tA via copy atoms (fx.copy); the raw buffer intrinsics are legacy

旧代码可能会使用 gpu.thread_idx.x、gpu.block_idx.x、arith.constant(...)、T.i32、原始 vector.* 辅助函数、ArithValue 封装和 buffer_ops。编辑大量使用这些内容的现有代码时，请保留它们；但编写新代码时，优先使用 gpu.thread_id/block_id、fx.Int64/fx.Int32/fx.Float32、fx.Vector 和 fx.rocdl.make_buffer_tensor。ArithValue、fx.Index 和 buffer_ops 已弃用 / 属于旧版接口——迁移现有内核时，请参阅 kernel-code-cleanup 技能。

参数类型

类型	描述	在宿主边界处
fx.Tensor	GPU 张量（memref）	通过 DLPack 从 torch.Tensor 自动转换
fx.Constexpr[int]	编译时常量	不同的值会生成不同的已编译内核
fx.Int32	运行时 i32	从 Python int 自动转换
fx.Stream	CUDA/HIP 流	默认流使用 fx.Stream(None)

线程/块层级

from flydsl.expr import gpu

tid_x = gpu.thread_id("x")   # Preferred current spelling
bid_x = gpu.block_id("x")
bid_y = gpu.block_id("y")

# Legacy spelling still appears in older kernels:
tid_x = gpu.thread_idx.x
bid_x = gpu.block_idx.x

gpu.barrier()                # Workgroup synchronization

控制流

from flydsl.expr import range_constexpr

# Compile-time unrolled loop (emitted inline in IR)
for i in range_constexpr(N):
    ...

# Runtime loop (lowered by AST rewriting)
for i in range(runtime_value):
    ...

运行时与编译时条件（当前写法）

运行时 SSA 比较应使用 Python/DSL 运算符。AST 重写器会将动态 if 条件降级为 scf.IfOp，而 ==、<、>= 等比较运算符会生成所需的 MLIR 谓词。

tid = gpu.thread_id("x")
lane = tid % fx.Int64(64)
c_zero = fx.Int64(0)
c_limit = fx.Int64(8)

# Preferred: readable DSL comparisons
if lane == c_zero:
    ...

in_range = lane < c_limit
val = fx.arith.select(in_range, good_val, zero_val)

# Avoid for simple integer comparisons
in_range = arith.cmpi(arith.CmpIPredicate.slt, lane, c_limit)

仅对跟踪/编译时已知的值使用 const_expr(...)，例如 Python 布尔值、constexpr 参数、循环展开选项或类型/布局分支：

if const_expr(trans_v):
    ...

if const_expr(max_context_partition_num <= WARP_SIZE):
    ...

不要将 GPU 运行时值包装在 const_expr 中。即使使用 @flyc.kernel(known_block_size=(256, 1, 1))，gpu.thread_id("x")、lane 和 warp_id 仍是运行时 SSA 值；编译器知道它们的取值范围，但不知道当前的通道实例。

# Wrong: lane depends on gpu.thread_id("x")
if const_expr(lane == c_zero):
    ...

# Correct
if lane == c_zero:
    ...

对于底层手动构建 MLIR 的场景，例如直接向 scf.IfOp 传递原始条件，应保留显式的 arith.cmpi(...) / arith.unwrap(...)：

cond = arith.unwrap(partition_idx >= visible_tile_count)
if_op = scf.IfOp(cond, has_else=False)

前端语义限制

编写或审查 @flyc.kernel / @flyc.jit 代码时，应主动避免以下模式，因为它们可能与 MLIR 构建冲突，即使在纯 Python 中看起来有效。

1. 不要在 if/else 中定义值并在分支外使用。 保持单一且显式的定义路径。

   if cond:
       dst = a
   else:
       dst = b
   use(dst)  # avoid this pattern
   

2. 不要在嵌套的辅助函数中修改捕获的外部变量。 可以只读捕获闭包变量，但写入操作应通过显式参数和返回值进行。

   def kernel():
       acc = fx.Float32(0.0)
   
       def helper(acc):
           acc = acc + fx.Float32(1.0)
           return acc
   
       acc = helper(acc)
   

3. 避免提前 return，也不要在 if/else 分支中放置 return / yield。 优先采用单一且显式的出口，以便前端确定结果类型。

   if cond:
       out = v0
   else:
       out = v1
   return out
   

4. 编译时条件必须使用 const_expr(...)。 对 constexpr 标志和其他静态决策使用 if const_expr(flag): ...。只有当条件已经是 Python bool 时，普通 Python if 才安全。

5. 辅助函数中的运行时分支应通过局部 @flyc.jit 进行分发。 如果分支主体包含副作用、循环携带值或分支局部定义，应将分支主体拆分为局部辅助函数，并将 if 包装在局部 JIT 辅助函数中：

   def then_path():
       ...
   
   def else_path():
       ...
   
   @flyc.jit
   def dispatch():
       if runtime_cond:
           then_path()
       else:
           else_path()
   
   dispatch()
   

使用循环携带值的运行时循环（软件流水线）

在 range() 上使用 init=，创建带有显式 SSA φ 节点的运行时循环，以携带循环状态。这对于软件流水线（预取模式）是必需的，因为数据必须跨迭代传递。

模式（来自 preshuffle_gemm.py）：

# Prologue: load first tile
tile_0 = prefetch(0)
init_state = [acc_init, tile_0_flat_val1, tile_0_flat_val2, ...]

# Runtime loop with loop-carried state
# Use typed DSL integer bounds (fx.Int64) so the AST rewriter does not treat this as a Python unrolled range.
_start = fx.Int64(0)
_stop = fx.Int64(N - 1)
_step = fx.Int64(1)
for iv, state in range(_start, _stop, _step, init=init_state):
    acc_in = state[0]
    tile_in = state[1:]

    next_tile = prefetch(iv + 1)      # load NEXT data
    acc_in = compute(acc_in, tile_in)  # compute CURRENT

    results = yield [acc_in] + next_tile  # carry to next iter

# Epilogue: process last tile from results
acc_final = results[0]
tile_final = results[1:]
compute(acc_final, tile_final)

MLIR 中的工作方式：

元素	含义
init=init_state	为运行时循环的第 0 次迭代的块参数提供初始值的 SSA 值列表
state	本次迭代的循环携带块参数（φ 节点）
yield [...]	将值传回，作为下一次迭代的 state
results	循环退出后，保存最后一次 yield 的值

三个关键陷阱（均已通过调试验证）：

1. 循环边界必须是 fx.Int64(...) 之类的有类型 DSL 整数，而不能是普通 Python 整数。 普通整数会使 AST 重写器展开循环，并悄悄忽略 init=。如果写成 range(0, 15, 1, init=...)，AST 重写器会将常量边界视为 Python range 并展开循环；只有普通 Python 整数边界会被展开，因此有类型的边界仍会生成运行时 scf.for。请改用 fx.Int64(0)、fx.Int64(15)、fx.Int64(1)。

2. 优先使用内部类型，但在硬性边界处解包。 大多数 range(..., init=...) 用法都接受 DSL 数值/向量值。如果低层辅助函数明确要求原始 ir.Value，只在该边界处使用 v.ir_value() / _raw(v) 解包。

3. 在内核顶部构建 LDS 视图，而不是在运行时循环内构建。 使用 fx.SharedAllocator().allocate(...).peek() 分配共享内存，并预先构建每个 .view()。在 scf.for 循环体内创建的视图定义于循环作用域中；在后置代码中（循环外）使用它会导致 SSA 支配错误。在顶部只构建一次，可确保它支配循环和后置代码。

算术运算

c42 = fx.Int64(42)                              # typed integer constant (preferred)
c3_14 = fx.Float32(3.14)                        # f32 constant (preferred)
mask = fx.Int32(0xFF)                            # i32 constant (preferred)

# Prefer operators / Numeric methods
result = a + b
result = a * scale
result = cond.select(true_val, false_val)
largest = fx.max(a, b)
smallest = fx.min(a, b)
tiles = fx.ceildiv(count, tile_size)  # signed/unsigned dispatch from typed operands

# Keep direct arith.*FOp only when explicit fastmath flags are required.

内部类型：Vector 和 Numeric（推荐）

请使用 FlyDSL 的内部类型化系统，而不是原始 MLIR 操作。Vector 类封装 vector<NxTy>，并提供运算符重载和类型安全的方法。

Vec = fx.Vector

# Wrap raw vector values
acc = Vec(frag_C.load())      # vector<Nxf32> → Vector with * / + operators

# Indexing (replaces vector.extract)
val = acc[idx]                 # returns Float32 scalar

# Bitcast (replaces vector.bitcast)
v_f32 = Vec(raw_vec).bitcast(fx.Float32)  # vector<Nxi32> → vector<Nxf32>

# Type conversion (replaces arith.trunc_f / arith.ext_f)
bf16_val = f32_val.to(fx.BFloat16)     # f32 → bf16

# Arithmetic — use Python operators, not arith.mulf/addf
result = (val * scale_a) * scale_b

# Splat constant vector
zeros = Vec.filled(N, 0.0, fx.Float32)

# Index cast — use fx.Int32 instead of arith.index_cast
idx = fx.Int32(gpu.block_id("x") * tile_m)

优先使用内部类型，而不是原始操作：

原始 MLIR 操作	对应的内部类型
vector.extract(v, static_position=[i], ...)	Vec(v)[i]
vector.bitcast(target_ty, v)	Vec(v).bitcast(Float32)
arith.trunc_f(ty, v)	v.to(BFloat16)
arith.mulf(a, b)	a * b
arith.addf(a, b)	a + b
arith.index_cast(T.i32, v)	fx.Int32(v)

向量广播请使用 Vec.filled(...)，由标量构造向量请使用 Vec.from_elements(...)。

Arith 操作可用性表

操作	函数	支持向量	备注
加法	a + b	是	只有需要显式 fastmath 时才使用直接 FOp
乘法	a * b	是	只有需要显式 fastmath 时才使用直接 FOp
取负	-a	是
最大值 / 最小值	fx.max(a, b) / fx.min(a, b)	是	浮点形式会传播 NaN；fx.maxnumf / fx.minnumf 不会
整数向上取整除法	fx.ceildiv(a, b)	是	直接的有符号/无符号操作；不同于布局 fx.ceil_div
比较	arith.cmpf(pred, a, b)	是	谓词在前；返回 i1/vec
选择	cond.select(t, f)	是
绝对值	无直接辅助函数	使用 -v、比较和 cond.select(...)
FMA	a * b + c	是	只有需要显式 fastmath 时才使用直接 FOp
常量广播	Vec.filled(width, val, dtype)	创建向量	用于标量广播

Printf 调试

fx.printf("tid={} bid={} val={}", tid, bid, value)

---

4. 数据移动模式

请参阅 references/data-movement.md：基于布局的复制、向量化加载、TiledCopy、缓冲区加载/存储、复制原子类型以及 gfx1250 TDM。

5. 共享内存（LDS）

请参阅 references/data-movement.md：SharedAllocator 模式以及各架构的 LDS 容量。

6. MFMA 集成（矩阵运算）

请参阅 references/mfma-and-reductions.md：MFMA 和 gfx1250 WMMA 原子操作、原始内建函数，以及预洗牌 GEMM 模式。

7. 归约模式

请参阅 references/mfma-and-reductions.md：warp 和 block 归约。

8. 常见模式和配方

请参阅 references/recipes.md：逐元素模板和操作指南，以及一个朴素 GEMM 模板。

9. 环境与调试

请参阅 references/environment-and-debugging.md：IR 转储、环境变量、磁盘缓存、源代码到汇编的调试信息，以及自动调优。

---

10. 故障排除

常见问题

1. 常量/类型转换：优先使用 fx.Int32(...)、fx.Int64(...) 和 fx.Float32(...)（fx.Index(...) 即将弃用，改用 fx.Int64(...)）。仅在底层边界处使用 arith.constant(...)。

2. 原始 buffer_load 偏移量：offset 参数的单位是元素，而非字节。 原始内建函数位于 kernels/common/buffer_ops.py，而非 flydsl.expr。

3. 代码更改后缓存过期：磁盘缓存会根据源代码/闭包更改自动失效。只有在更改了 C++ pass 或非闭包辅助函数时，才设置 FLYDSL_RUNTIME_ENABLE_CACHE=0 或清除 ~/.flydsl/cache/。

4. LDS 溢出：检查容量（gfx942 为 64KB，gfx950 为 160KB）。SharedAllocator 会为你设置 LDS 全局变量；如果静态分配超出架构限制，编译器会报错。

5. Dynamic 与 Constexpr：Constexpr[int] 值会固化到 IR 中——不同的值会生成不同的已编译内核。对于真正动态的值，请使用 Int32。

6. 张量布局标记：对于动态形状或对齐，请使用 flyc.from_dlpack(tensor).mark_layout_dynamic(leading_dim=0, divisibility=4)。

7. AMD wavefront 大小：在 gfx9xx 上始终为 64。完整 wave 归约时使用移位量 [32, 16, 8, 4, 2, 1]。

8. GEMM 的 tile_k 对齐：tile_k * elem_bytes 必须能被 64 整除（K64 字节微步长）。

9. INT4（W4A8）：A 矩阵为 int8，B 矩阵为打包的 int4（每字节 2 个值），并在内核中解包为 int8。

10. 绝对值：arith dialect 没有 absf，但 FlyDSL 导出了该函数——请使用 abs(v) 或 fx.absf(v)，而不是取反/比较/选择序列。

11. 标量广播到向量：使用 Vec.filled(width, value, fx.Float32) 创建广播常量向量。普通算术运算不要使用原始向量操作。

---

11. 与 Triton/Gluon 的比较

方面	FlyDSL	Triton	Gluon
布局控制	显式布局代数（Shape、Stride、Layout）	通过 block 指针隐式控制	隐式
分块	通过除法/乘积操作手动完成	使用 tl.program_id 自动分块	自动分块
内存访问	Copy 原子操作、缓冲区加载/存储、TiledCopy	tl.load/tl.store	gluon.load/gluon.store
MFMA	直接使用 rocdl.mfma_* 内建函数	tl.dot	gluon.dot
共享内存	SharedAllocator + @fx.struct 显式管理	隐式暂存区	隐式
抽象级别	低（接近硬件）	中	中高
编译	MLIR（Fly dialect -> LLVM -> HSACO）	MLIR（Triton dialect -> LLVM）	MLIR
控制能力	对数据布局和移动的控制能力最大	控制较少，自动化程度较高	控制能力最低

FlyDSL 以更冗长为代价，提供了最大的控制能力。布局代数是其关键优势所在——它能够精确控制数据在寄存器、共享内存和全局内存中的排列方式，以及线程如何映射到数据。

---

12. 运行内核

本地运行

使用可编辑安装（pip install -e .）时，只需仓库根目录即可。直接检出源码时，还需将构建好的绑定添加到路径中（CLAUDE.md）：

export PYTHONPATH="${PWD}/build-fly/python_packages:${PWD}:${PYTHONPATH}"
export LD_LIBRARY_PATH="${PWD}/build-fly/python_packages/flydsl/_mlir/_mlir_libs:${LD_LIBRARY_PATH}"

python3 my_kernel.py
python3 -m pytest tests/kernels/test_vec_add.py -v
bash scripts/run_benchmark.sh

在远程主机或容器中

询问用户主机、容器和检出路径；不要自行假设。 设置好 $HOST、$CONTAINER 和 $FLYDSL_ROOT 后：

ssh -o LogLevel=ERROR "$HOST" \
  "docker exec $CONTAINER bash -c 'cd $FLYDSL_ROOT && python3 my_kernel.py'"

如果检出目录就在主机本身上，则去掉 docker exec 封装。 若要先在此类主机上构建 FlyDSL，请使用 build-flydsl 技能。