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) ### Core Types | Type | Description | Example | |------|-------------|---------| | `!fly.int_tuple` | Integer tuple (can be nested) | `(8, 16)`, `(8, (4, 2))` | | `!fly.layout` | (Shape, Stride) pair | `(8, 16):(1, 8)` (col-major) | | `!fly.memref` | Memory reference with layout | Typed pointer + layout info | ### Construction ```python import flydsl.expr as fx shape = fx.make_shape(8, 16) # IntTuple (8, 16) stride = fx.make_stride(1, 8) # IntTuple (1, 8) layout = fx.make_layout(shape, stride) # Layout (8,16):(1,8) # Shorthand with Python tuples layout = fx.make_layout((8, 16), (1, 8)) # Coordinates coord = fx.make_coord(i, j) # Nested shapes for hierarchical tiling shape_nested = fx.make_shape(9, (4, 8)) # (9, (4, 8)) # Identity layout identity = fx.make_identity_layout((M, N)) ``` ### Coordinate Mapping The fundamental operation maps logical coordinates to physical memory indices. **Formula**: `Index = sum(coord_i * stride_i)` ```python idx = fx.crd2idx(coord, layout) # Coordinate -> linear index coord = fx.idx2crd(idx, layout) # Linear index -> coordinate s = fx.size(layout) # Total element count (product of shape) ``` **Example**: For layout `(8, 16):(1, 8)` (8x16, column-major): - `crd2idx((3, 5), layout)` = `3*1 + 5*8` = 43 - `idx2crd(43, layout)` = `(43 % 8, 43 / 8)` = `(3, 5)` ### Query Operations ```python fx.size(layout) # Total element count fx.get_shape(layout) # Extract shape IntTuple fx.get_stride(layout) # Extract stride IntTuple fx.get(int_tuple, i) # Get i-th element fx.rank(int_tuple) # Number of top-level modes ``` ### Layout Algebra Operations #### Composition: `fx.composition(A, B)` Compose two layouts: `result(x) = A(B(x))`. Used to apply permutations or tile coordinate mappings. #### Complement: `fx.complement(tiler, target_size)` Compute remaining modes not covered by tiler, up to target_size. Internal building block for divides. #### Coalesce: `fx.coalesce(layout)` Simplify layout by merging adjacent modes. Preserves mapping but flattens structure. #### Right Inverse: `fx.right_inverse(layout)` Compute right inverse of layout mapping. #### Recast: `fx.recast_layout(layout, old_bits, new_bits)` Adjust layout for type width change (e.g., FP16->FP8). ### Product Operations (Combine Layouts) Products combine two layouts to create a larger layout: ```python fx.logical_product(layout, tiler) # Basic mode-wise concatenation fx.raked_product(thr, val) # Interleaved access pattern (see make_layout_tv for TV layouts) fx.blocked_product(layout, tiler) # Blocked access pattern fx.zipped_product(layout, tiler) # Zipped modes fx.tiled_product(layout, tiler) # Hierarchical tiled structure fx.flat_product(layout, tiler) # Flattened result ``` ### Divide Operations (Partition Layouts) Divides split a layout by a divisor, creating tile + rest dimensions: ```python fx.logical_divide(layout, divisor) # Basic partitioning (uses complement internally) fx.zipped_divide(layout, divisor) # Zipped division fx.tiled_divide(layout, divisor) # Hierarchical tiled division fx.flat_divide(layout, divisor) # Flattened division ``` ### Structural Operations ```python fx.select(int_tuple, indices=[0, 2]) # Pick specific modes fx.group(int_tuple, begin=1, end=3) # Group modes into nested tuple fx.append(base, elem) # Append mode fx.prepend(base, elem) # Prepend mode fx.slice(src, coord) # Slice at coordinate (None = keep mode) ``` --- ## 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): ... ```