flydsl-kernel-authoring

name	flydsl-kernel-authoring
description	Comprehensive reference for authoring FlyDSL GPU kernels on AMD GPUs. Covers the layout algebra, tiled copy/MMA, buffer ops, loop-carried range loops, SharedAllocator (LDS), autotuning, and common patterns. Use when writing, reviewing, or understanding FlyDSL kernel code.
allowed-tools	Read Edit Bash Grep Glob Agent

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: /FlyDSL/ (installed in editable mode) 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. When modernizing or cleaning up an existing kernel — replacing legacy/deprecated constructs (ArithValue, fx.Index, buffer_ops, raw MLIR dialects, SmemPtr, per-tile *_atom_call, raw rocdl.mfma_*) with the current fx.* surface, trimming comments/dead code, or applying the _run_compiled fast launch path — use the flydsl-kernel-code-cleanup 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/buffer_ops.py - AMD buffer load/store intrinsics
• python/flydsl/expr/rocdl/ - MFMA/WMMA and other ROCm intrinsics (package: cdna4, cluster, inline_asm, tdm_ops, universal)
• python/flydsl/expr/gpu.py - SharedAllocator for LDS (shared memory), thread_id/block_id, barrier
• python/flydsl/utils/smem_allocator.py - legacy SmemAllocator (un-migrated kernels only)
• kernels/ - Pre-built kernels, organized into subpackages: gemm/ (preshuffle_gemm.py, mxfp4_preshuffle.py, ...), norm/ (layernorm/softmax/rmsnorm), attention/, moe/, mma/, common/, 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

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)

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

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:

fx.logical_product(layout, tiler)   # Basic mode-wise concatenation
fx.raked_product(thr, val)          # Interleaved access pattern (common for TiledCopy)
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:

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

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

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:

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_atom_call); raw buffer_ops is 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 flydsl-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

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

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.

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 = in_range.select(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:

if const_expr(trans_v):
    ...

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