flydsl-kernel-authoring

# FlyDSL Kernel Authoring Skill ## Overview FlyDSL is a Python DSL and MLIR-based compiler for writing high-performance GPU kernels on AMD GPUs (MI300X/MI350). It provides explicit layout algebra for controlling data movement, tiling, and memory access patterns. The layout system is the core abstraction that distinguishes FlyDSL from Triton/Gluon. **Repository**: this checkout (examples below assume it is importable — see 12) **Target GPU**: gfx942 (MI300X, CDNA3), gfx950 (MI350, CDNA4) **Python**: 3.12, ROCm 7.2 **Scope (read this first)**: This skill is the **reference** — the full layout-algebra API surface, per-op tables, MFMA/copy-atom catalogs, environment variables, and an exhaustive troubleshooting list. Reach for it to *look something up* while writing or reviewing kernel code. If instead you want a *guided, step-by-step procedure* that turns a kernel requirement into a finished, tested kernel (classify -> skeleton -> compute -> control flow -> test), use the **flydsl-tile-programming** skill, which is the wizard companion to this reference. For diagnosing a kernel that already compiles but produces NaN/inf/wrong results, use the **debug-flydsl-kernel** skill. --- ## 1. Architecture and Compilation ### Pipeline ``` 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) ``` ### Key Passes Pipeline is built by `RocmBackend._pipeline_parts()` and split into three stages — see `docs/architecture_guide.md` §3 for the per-pass table. Highlights: 1. `fly-rewrite-func-signature` - Rewrite DSL types at function / SCF boundaries to packed LLVM structs 2. `fly-layout-lowering` - Lower layout algebra (`fly.crd2idx`, partitions, divides) to arithmetic 3. `fly-convert-atom-call-to-ssa-form` + `fly-promote-regmem-to-vectorssa` - Lift copy/MMA atom calls and register memory to vector SSA 4. `convert-fly-to-rocdl` - Fly ops -> ROCDL intrinsics 5. `gpu-module-to-binary{format=fatbin}` - Emit HSACO binary via LLVM AMDGPU backend ### Key Source Paths - `python/flydsl/compiler/` - JIT compilation (jit_function.py, kernel_function.py) - `python/flydsl/expr/` - DSL expression API (primitive.py, derived.py, typing.py) - `python/flydsl/expr/primitive.py` - All layout algebra functions - `python/flydsl/expr/derived.py` - CopyAtom, MmaAtom, TiledCopy, TiledMma wrappers - `python/flydsl/expr/gpu.py` - GPU operations (thread_idx, block_idx, barrier) - `python/flydsl/expr/rocdl/` - MFMA/WMMA and other ROCm intrinsics (package: cdna3, cdna4, cdna5, rdna3, rdna4, cluster, inline_asm, tdm_ops, universal; plus `utils.py` / `enum.py` helpers) - `python/flydsl/expr/gpu.py` - `SharedAllocator` for LDS (shared memory), `thread_id`/`block_id`, `barrier` - `kernels/common/buffer_ops.py` - legacy raw AMD buffer load/store intrinsics (moved out of `flydsl.expr` in #880; prefer `fx.rocdl.make_buffer_tensor`) - `kernels/` - Pre-built kernels, organized into subpackages: `gemm/` (preshuffle_gemm.py, mxfp4_preshuffle.py, ...), `norm/` (layernorm/softmax/rmsnorm), `attention/`, `moe/`, `mega_moe/`, `common/` (incl. `common/mma/`), `comm/`, `conv/` --- ## 2. Layout System (Core Abstraction) See [references/layout-algebra.md](references/layout-algebra.md): core types, construction, coordinate mapping, layout algebra, products, and divides. --- ## 3. Writing Kernels ### Basic Pattern ```python 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) ``` ### Current Syntax Quick Reference Use the current public FlyDSL surface from `kernels/gemm/preshuffle_gemm.py` when writing new kernels: ```python 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 ``` Older code may use `gpu.thread_idx.x`, `gpu.block_idx.x`, `arith.constant(...)`, `T.i32`, raw `vector.*` helpers, the `ArithValue` wrapper, and `buffer_ops`. Keep those when editing existing code that already uses them heavily, but prefer `gpu.thread_id/block_id`, `fx.Int64`/`fx.Int32`/`fx.Float32`, `fx.Vector`, and `fx.rocdl.make_buffer_tensor` for new code. `ArithValue`, `fx.Index`, and `buffer_ops` are deprecated/legacy — for migrating an existing kernel see the **kernel-code-cleanup** skill. ### Parameter Types | Type | Description | At host boundary | |------|-------------|-----------------| | `fx.Tensor` | GPU tensor (memref) | Auto-converted from torch.Tensor via DLPack | | `fx.Constexpr[int]` | Compile-time constant | Different values -> different compiled kernels | | `fx.Int32` | Runtime i32 | Auto-converted from Python int | | `fx.Stream` | CUDA/HIP stream | `fx.Stream(None)` for default stream | ### Thread/Block Hierarchy ```python 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 ``` ### Control Flow ```python 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): ... ``` ### Runtime vs Compile-Time Conditions (Current Style) Use Python/DSL operators for runtime SSA comparisons. The AST rewriter lowers dynamic `if` conditions to `scf.IfOp`, and comparison operators like `==`, `<`, `>=` generate the needed MLIR predicates. ```python 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) ``` Use `const_expr(...)` only for values known at trace/compile time, such as Python booleans, constexpr arguments, loop-unroll choices, or type/layout branches: ```python if const_expr(trans_v): ... if const_expr(max_context_partition_num <= WARP_SIZE): ... ``` Do **not** wrap GPU runtime values in `const_expr`. Even with `@flyc.kernel(known_block_size=(256, 1, 1))`, `gpu.thread_id("x")`, `lane`, and `warp_id` are runtime SSA values; the compiler knows their range, not the current lane instance. ```python # Wrong: lane depends on gpu.thread_id("x") if const_expr(lane == c_zero): ... # Correct if lane == c_zero: ... ``` Keep explicit `arith.cmpi(...)` / `arith.unwrap(...)` for low-level manual MLIR construction, such as passing a raw condition to `scf.IfOp` directly: ```python cond = arith.unwrap(partition_idx >= visible_tile_count) if_op = scf.IfOp(cond, has_else=False) ``` ### Frontend Semantic Restrictions When writing or reviewing `@flyc.kernel` / `@flyc.jit` code, proactively avoid these patterns because they can conflict with MLIR construction even if they look valid in plain Python. 1. **Do not define values inside `if/else` and use them later outside the branch.** Keep a single explicit definition path. ```python if cond: dst = a else: dst = b use(dst) # avoid this pattern ``` 2. **Do not mutate captured outer variables inside nested helper functions.** Read-only closure capture is acceptable, but writes should go through explicit parameters and return values. ```python def kernel(): acc = fx.Float32(0.0) def helper(acc): acc = acc + fx.Float32(1.0) return acc acc = helper(acc) ``` 3. **Avoid early `return`, and do not place `return` / `yield` inside `if/else` branches.** Prefer a single explicit exit so the frontend can determine result types. ```python if cond: out = v0 else: out = v1 return out ``` 4. **Compile-time conditions must use `const_expr(...)`.** Use `if const_expr(flag): ...` for constexpr flags and other static decisions. A plain Python `if` is only safe when the condition is already a Python `bool`. 5. **Runtime branches inside helper functions should be dispatched via local `@flyc.jit`.** When a branch body has side effects, loop-carried values, or branch-local definitions, split the branch bodies into local helpers and wrap the `if` in a local JIT helper: ```python def then_path(): ... def else_path(): ... @flyc.jit def dispatch(): if runtime_cond: then_path() else: else_path() dispatch() ``` ### Runtime Loops with Loop-Carried Values (Software Pipelining) Use `init=` on `range()` to create a runtime loop with explicit SSA phi nodes for loop-carried state. This is required for software pipelining (prefetch patterns) where data must flow across iterations. **Pattern** (from `preshuffle_gemm.py`): ```python # 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) ``` **How it works in MLIR:** | Element | Meaning | |---|---| | `init=init_state` | List of SSA values that seed the runtime loop block arguments for iteration 0 | | `state` | The loop-carried block arguments (phi nodes) for THIS iteration | | `yield [...]` | Feeds values back as next iteration's `state` | | `results` | After loop exits, holds the last yielded values | **Three critical pitfalls (all verified by debugging):**