FlyDSL Kernel Authoring Skill

Overview

FlyDSL is a Python DSL and MLIR-based compiler for writing AMD GPU kernels. It provides explicit layout algebra for controlling data movement, tiling, and memory access patterns.

Aiter dependency: check requirements.txt first (flydsl==0.3.4.1 as of this update). Use python -c 'import flydsl; print(flydsl.__file__)' to find the imported package; an editable checkout may differ from the pinned wheel. Targets in aiter: gfx942 (CDNA3), gfx950 (CDNA4), and gfx1250 (wave32). Check the chosen operation on the target architecture; not every fx.rocdl wrapper supports all three.

Scope (read this first): This skill is the authoring reference for aiter's pinned FlyDSL API. For migration of an existing kernel from raw dialects, ArithValue, buffer_ops, or SmemPtr, use the in-tree flydsl-kernel-code-cleanup skill. Keep fx.Index where an API requires MLIR index; it is not a blanket migration target.

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 - typed thread_idx/block_idx, raw thread_id/block_id, barrier, and SharedAllocator
• python/flydsl/expr/llvm.py and expr/rocdl/universal.py - typed memory operations and buffer-resource views
• python/flydsl/expr/rocdl/ - MFMA/WMMA and other ROCm intrinsics (package: cdna4, cluster, inline_asm, tdm_ops, universal)
• python/flydsl/utils/smem_allocator.py - legacy SmemAllocator (un-migrated kernels only)
• aiter/ops/flydsl/kernels/ - aiter's kernel examples and its legacy buffer_ops.py compatibility layer

2. Layout System (Core Abstraction)

Core Types

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

@flyc.kernel
def my_kernel(
    A: fx.Tensor,         # GPU tensor (memref via DLPack)
    B: fx.Tensor,
    N: fx.Constexpr[int], # Compile-time constant
):
    tid = fx.thread_idx.x  # fx.Int32
    bid = fx.block_idx.x   # fx.Int32
    # ... 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 public FlyDSL surface and aiter's aiter/ops/flydsl/kernels/preshuffle_gemm.py when writing new kernels:

Vec = fx.Vector

tx = fx.thread_idx.x
bx = fx.block_idx.x
by = fx.block_idx.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` when the operands are supported)

fx.thread_idx/fx.block_idx and gpu.thread_idx/gpu.block_idx are the same typed accessors. gpu.thread_id("x") and gpu.block_id("x") are available lower-level calls returning raw MLIR values; wrap them with fx.Int32(...) before typed arithmetic. fx.Index remains necessary at APIs whose contract is MLIR index. For migrating ArithValue, raw vector.*, and buffer_ops, see the flydsl-kernel-code-cleanup skill.

Parameter Types

Thread/Block Hierarchy

import flydsl.expr as fx
from flydsl.expr import gpu

tid_x = fx.thread_idx.x       # fx.Int32; preferred for typed arithmetic
bid_x = fx.block_idx.x        # fx.Int32
bid_y = fx.block_idx.y

# `gpu.thread_idx.x` / `gpu.block_idx.x` are equivalent typed accessors.
# Direct calls are lower-level and return raw MLIR values:
raw_tid_x = gpu.thread_id("x")
raw_bid_x = gpu.block_id("x")
tid_x = fx.Int32(raw_tid_x)   # if a direct call is needed

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 = fx.thread_idx.x
lane = tid % fx.Int32(64)
c_zero = fx.Int32(0)
c_limit = fx.Int32(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):
    ...

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.

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

# Correct
if lane == c_zero:
    ...

Frontend Semantic Restrictions

When writing or reviewing @flyc.kernel / @flyc.jit code, check these frontend contracts rather than applying ordinary Python control-flow rules.

1. Runtime branches can carry values out. Initialize a value before the branch when the rewriter must carry it through both paths, and keep the resulting types consistent. A branch-local value with no use outside needs no carried state.

   dst = a
   if cond:
       dst = b
   use(dst)

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.

   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.

   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. Plain helpers outside a rewritten body need @flyc.jit for runtime branches. A runtime if already inside a @flyc.kernel or @flyc.jit body is rewritten; do not add a second decorator just for that branch. For a separate helper that needs AST rewriting, use the public decorator:

   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):

# 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
# `init=` selects a runtime scf.for even with Python integer bounds.
_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:

Three critical pitfalls (all verified by debugging):

1. init= selects the explicit-carried-state scf.for path. Python-int bounds are accepted and converted to MLIR index; they do not make this loop unroll or discard init. Typed bounds also work here. For a loop without init, the automatic-carried-state path expects i32 bounds (index is cast to i32; i64 is rejected). Use range_constexpr to request Python unrolling.
2. Prefer internal types, but unwrap at hard boundaries. Most range(..., init=...) uses accept DSL numeric/vector values. If a lower-level helper explicitly expects raw ir.Value, unwrap with v.ir_value() / _raw(v) at that boundary only.
3. Clear SmemPtr._view_cache before epilogue. SmemPtr.get() caches the view it creates. If called inside the runtime loop body, the cached view is defined in the loop scope. Using it in the epilogue (outside the loop) causes an SSA dominance error. Fix:

   # After the runtime loop, before epilogue compute:
   my_smem_ptr._view_cache = None

Arithmetic Operations

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)

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

Internal Types: Vector and Numeric (PREFERRED)

Use FlyDSL's internal typed system instead of raw MLIR ops. The Vector class wraps vector<NxTy> with operator overloading and type-safe methods.

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.block_idx.x * tile_m

Prefer internal types over raw ops:

Use Vec.filled(...) for splats and Vec.from_elements(...) for vectors from scalars.

Arith Ops Availability Table

Printf Debugging

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

4. Data Movement Patterns

Layout-Based Copy (Preferred for Element-wise Kernels)

The standard pattern: divide tensor by tile size, slice by block/thread, copy via atoms.

@flyc.kernel
def my_kernel(A: fx.Tensor, B: fx.Tensor, BLOCK_DIM: fx.Constexpr[int]):
    bid = fx.block_idx.x
    tid = fx.thread_idx.x

    # 1. Divide tensor into blocks
    tA = fx.logical_divide(A, fx.make_layout(BLOCK_DIM, 1))
    tB = fx.logical_divide(B, fx.make_layout(BLOCK_DIM, 1))

    # 2. Select this block's tile
    tA = fx.slice(tA, (None, bid))
    tB = fx.slice(tB, (None, bid))

    # 3. Further divide for per-thread access
    tA = fx.logical_divide(tA, fx.make_layout(1, 1))  # 1 element per thread
    tB = fx.logical_divide(tB, fx.make_layout(1, 1))

    # 4. Allocate registers
    copyAtom = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32)
    rA = fx.make_rmem_tensor(1, fx.Float32)

    # 5. Copy: global -> register -> compute -> global
    fx.copy(copyAtom, fx.slice(tA, (None, tid)), rA)
    # ... compute on register values ...
    fx.copy(copyAtom, rA, fx.slice(tB, (None, tid)))

Vectorized Loads (Wide Copies)

VEC_WIDTH = 4
copy_bits = VEC_WIDTH * 32   # 128 bits
copyAtom = fx.make_copy_atom(fx.UniversalCopy(copy_bits), fx.Float32)

rA = fx.make_rmem_tensor(VEC_WIDTH, fx.Float32)

# Divide for VEC_WIDTH elements per thread
tA = fx.logical_divide(tA, fx.make_layout(VEC_WIDTH, 1))
fx.copy(copyAtom, fx.slice(tA, (None, tid)), rA)

# Load/store as vectors
vec = fx.memref_load_vec(rA)     # Load vector from register memref
fx.memref_store_vec(vec, rA)     # Store vector to register memref

TiledCopy Abstraction (for 2D Copies)

# Define thread and value layouts
thr_layout = fx.make_layout((4, 1), (1, 1))    # 4 threads
val_layout = fx.make_layout((1, 8), (1, 1))    # 8 values per thread

# Create copy atom
copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32)

# Build the thread-value (TV) layout and its tiler
tile_mn, tv_layout = fx.make_layout_tv(thr_layout, val_layout)
tiled_copy = fx.make_tiled_copy(copy_atom, tv_layout, tile_mn)

# Get this thread's slice and partition
thr_copy = tiled_copy.get_slice(tid)
partition_src = thr_copy.partition_S(src_tensor)
partition_dst = thr_copy.partition_D(dst_tensor)
frag = fx.make_fragment_like(partition_src)

# Execute copy: src -> fragment -> dst
fx.copy(copy_atom, partition_src, frag)
fx.copy(copy_atom, frag, partition_dst)

Buffer Load/Store

Build a buffer-resource view with make_buffer_tensor and move data through copy atoms. Its default max_size=True does not enforce a finite record bound; pass max_size=False or num_records_bytes= when hardware OOB checking is required.

buf = fx.rocdl.make_buffer_tensor(tensor)
tA  = fx.make_view(fx.get_iter(buf), fx.make_layout((M, N), (N, 1)))
copy = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32)
fx.copy(copy, fx.slice(tA, (None, tid)), rA)   # after partitioning tA

Legacy raw intrinsics (buffer_ops.create_buffer_resource / buffer_load / buffer_store, offset in elements) remain for un-migrated kernels; see the flydsl-kernel-code-cleanup skill to migrate them.

Copy Atom Types

gfx1250 TDM async copy — use fx.rocdl.cdna5.make_tiled_tdm_atom + fx.rocdl.cdna5.tdm_partition. The atom is built over the whole tensor, zipped_divide picks the block's tile, tdm_partition cuts one warp's share, and fx.copy moves it. Needs a raw VA (not make_buffer_tensor).

Three things that used to be explicit parameters now live in the tensors, and the hand math for them goes away:

• extents and strides — put them in the global view's layout, static or runtime; the lowering opens the dynamic leaves itself. There is no tensor_extents= or strides=.
• K advance — the coordinate tensor's surviving rest mode is an i32 tile number, not a byte imm_offset.
• LDS padding — row slack in the LDS box layout. A store cannot drain a padded tile, so for one make tdm_tile the padded width and give the view the real width; boundary_check drops the pad columns.

lds = fx.SharedAllocator().allocate(fx.Array[fx.Float16, M * N]).peek()
g_full = fx.Tensor(fx.make_view(fx.get_iter(A), fx.make_layout((M, K), (lda, 1))))
tile = (TM, TK)
atom, coord = fx.rocdl.cdna5.make_tiled_tdm_atom(
    fx.rocdl.TensorLoad(), g_full, fx.make_layout(tile, (LDS_ROW, 1)), tile, num_warps=4
)
blk = fx.zipped_divide(coord, tile)[None, (bid_x, None)]   # keep ONE rest for K
lds_box = fx.Tensor(fx.make_view(lds.ptr, fx.make_layout((tile, PIPE),
                                                         ((LDS_ROW, 1), PITCH))))
ts, tg = fx.rocdl.cdna5.tdm_partition(atom, wave, fx.make_layout(4, 1), lds_box, blk)
fx.copy(atom, tg[None, k_tile_i32], ts[None, stage])       # Global -> LDS
fx.rocdl.tdm_ops.tensor_wait(0)                                  # await async DMA

The older fx.rocdl.make_tdm_atom + fx.copy_atom_call + advance_tdm_atom carried the descriptor as atom state and is being retired — do not write new kernels against it. It is still required where the tiled atom cannot reach: the gather/scatter family (make_tensor_gather_descriptor / tensor_load_gather, for per-row addresses such as a KV page table), and a tile whose innermost dimension spans under 16 bytes per warp.

5. Shared Memory (LDS)

SharedAllocator Pattern (preferred for new kernels)

Declare the LDS storage as an @fx.struct of fx.Array fields, allocate it inside the kernel with fx.SharedAllocator, then .view() each field as a layout. In the default static=True mode the compiler sizes a per-leaf static LDS global, so launch(smem=...) is left unset.

import flydsl.expr as fx

@fx.struct
class SharedStorage:
    a: fx.Array[fx.Float16, 8192]
    b: fx.Array[fx.Float16, 8192]

@flyc.kernel
def my_kernel(A: fx.Tensor, ...):
    lds = fx.SharedAllocator().allocate(SharedStorage).peek()
    lds_a = lds.a.view(fx.make_layout((64, 128), (128, 1)))
    lds_b = lds.b.view(fx.make_layout((64, 128), (128, 1)))
    # ds_write / ds_read happen through copy atoms or view loads/stores

For the dynamic mode (static=False), the launch wrapper auto-infers smem from SharedAllocator.allocated_bytes when smem=None.

Legacy SmemAllocator

flydsl.utils.smem_allocator.SmemAllocator remains for un-migrated kernels. Its surface is __init__/finalize/get_base plus SmemPtr.get/load/store; prefer SharedAllocator for anything new.

LDS Capacity

6. MFMA Integration (Matrix Math)

MFMA atoms on gfx942/gfx950

mma = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 16, fx.Float16))
fx.gemm(mma, frag_C, frag_A, frag_B, frag_C)  # d, a, b, c

Pick the supported shape and dtype for the target. Keep a raw MFMA intrinsic only when the atom API cannot express its operand form or scheduling contract.

gfx1250 WMMA (wave32)

gfx1250 uses WMMA (not MFMA), M=N=16. Build the atom with rocdl.WMMA (arch dispatched: gfx11 v16 ABI, gfx12 / gfx1250 v8 ABI) and issue it via fx.make_mma_atom + fx.gemm / fx.mma_atom_call:

mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 128, fx.Float8E4M3FN))       # fp8 → f32
mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 32, T.i4, T.i32, sign_a=True, sign_b=True, clamp=True))

MX-scaled WMMA — rocdl.WMMAScale(m, n, k, elem_ty_a, ..., block_size=32) for the E8M0 block-scaled f8/f6/f4 format (16x16x128, or 32x16x128 fp4-only). Per-operand scales are atom state; block_size 32 → i32 scale (i64 for 16):

mma = fx.make_mma_atom(rocdl.WMMAScale(16, 16, 128, fx.Float8E4M3FN))
mma = fx.atom_set_value(mma, "scale_a", fx.Int32(scale_a))
mma = fx.atom_set_value(mma, "scale_b", fx.Int32(scale_b))
fx.gemm(mma, frag_C, frag_A, frag_B, frag_C)

GEMM Pattern (Preshuffle)

The preshuffle GEMM pattern in aiter/ops/flydsl/kernels/preshuffle_gemm.py:

1. B matrix is pre-shuffled to layout: (N/16, K/64, 4, 16, kpack_bytes)
2. A tiles loaded from global to LDS with XOR16 swizzle for bank-conflict avoidance
3. K64-byte micro-steps: each step issues 2x K32 MFMA operations
4. Ping-pong LDS (lds_stage=2) for overlapping loads with compute
5. Epilogue: either direct row-major store or CShuffle via LDS for packing

7. Reduction Patterns

Warp Reduction (AMD wave64)

XOR-shuffle-based intra-wave reduction:

val = fx.Float32(val)                                        # typed once, not per-iter
for sh in [32, 16, 8, 4, 2, 1]:
    peer = val.shuffle_xor(fx.Int32(sh), fx.Int32(64))       # typed numeric method
    val = val + peer                                         # or val.addf(peer, fastmath=...) for fastmath

Block Reduction

1. Intra-wave XOR shuffle (shifts: 32, 16, 8, 4, 2, 1)
2. Lane 0 writes per-wave partial to LDS
3. gpu.barrier()
4. Wave 0 reads and reduces NUM_WAVES partials from LDS

See the norm kernels in kernels/norm/ (e.g. rmsnorm_kernel.py, which defines local wave_reduce_add / block_reduce_add closures) for worked reductions.

8. Common Patterns and Recipes

Element-wise Kernel Template

@flyc.kernel
def elementwise_kernel(In: fx.Tensor, Out: fx.Tensor, BLOCK: fx.Constexpr[int], VEC: fx.Constexpr[int]):
    bid, tid = fx.block_idx.x, fx.thread_idx.x
    tile = BLOCK * VEC
    tIn = fx.logical_divide(In, fx.make_layout(tile, 1))
    tOut = fx.logical_divide(Out, fx.make_layout(tile, 1))
    tIn = fx.slice(tIn, (None, bid))
    tOut = fx.slice(tOut, (None, bid))
    tIn = fx.logical_divide(tIn, fx.make_layout(VEC, 1))
    tOut = fx.logical_divide(tOut, fx.make_layout(VEC, 1))
    copy = fx.make_copy_atom(fx.UniversalCopy(VEC * 32), fx.Float32)
    rIn = fx.make_rmem_tensor(VEC, fx.Float32)
    rOut = fx.make_rmem_tensor(VEC, fx.Float32)
    fx.copy(copy, fx.slice(tIn, (None, tid)), rIn)
    # Transform
    v = fx.Vector(fx.memref_load_vec(rIn))
    v = v * v  # example: square
    fx.memref_store_vec(v, rOut)
    fx.copy(copy, rOut, fx.slice(tOut, (None, tid)))

Element-wise Kernel Cookbook (GPU-Verified)

All recipes below follow the same vectorized copy_atom pattern (256 threads, vec_width=4, 128-bit loads). Only the compute section between memref_load_vec and memref_store_vec differs.

# --- Scale: C = A * scalar ---
vA = Vec(fx.memref_load_vec(rA))
scale = Vec.filled(vec_width, 2.0, fx.Float32)
vC = vA * scale

# --- Multiply: C = A * B ---
vC = Vec(fx.memref_load_vec(rA)) * Vec(fx.memref_load_vec(rB))

# --- Fused FMA: D = A * B + C with one rounding ---
vD = fx.fma(Vec(fx.memref_load_vec(rA)), Vec(fx.memref_load_vec(rB)),
            Vec(fx.memref_load_vec(rC)))
# Retain `a * b + c` when separate multiplication and addition are required.

# --- ReLU: C = max(A, 0) ---
vA = Vec(fx.memref_load_vec(rA))
zero_vec = Vec.filled(vec_width, 0.0, fx.Float32)
vC = vA.maximumf(zero_vec)

# --- Abs: C = |A| ---
vA = Vec(fx.memref_load_vec(rA))
vC = fx.absf(vA)

Naive GEMM Template (for understanding, not performance)

@flyc.kernel
def naive_gemm(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor,
               M: fx.Constexpr[int], N: fx.Constexpr[int], K: fx.Constexpr[int],
               BM: fx.Constexpr[int], BN: fx.Constexpr[int]):
    tid, bid = fx.thread_idx.x, fx.block_idx.x
    bm, bn = bid // (N // BN), bid % (N // BN)
    tm, tn = tid // BN, tid % BN
    row, col = bm * BM + tm, bn * BN + tn

    # Buffer-resource views with logical row-major layouts
    tA = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(A)), fx.make_layout((M, K), (K, 1)))
    tB = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(B)), fx.make_layout((K, N), (N, 1)))
    tC = fx.make_view(fx.get_iter(fx.rocdl.make_buffer_tensor(C)), fx.make_layout((M, N), (N, 1)))

    copy = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32)
    rA, rB, rC = (fx.make_rmem_tensor(1, fx.Float32) for _ in range(3))

    acc = fx.Float32(0.0)
    for k in range_constexpr(K):
        fx.copy(copy, fx.slice(tA, (row, k)), rA)
        fx.copy(copy, fx.slice(tB, (k, col)), rB)
        acc = acc + fx.Vector(fx.memref_load_vec(rA))[0] * fx.Vector(fx.memref_load_vec(rB))[0]

    fx.memref_store_vec(fx.Vector.from_elements([acc], fx.Float32), rC)
    fx.copy(copy, rC, fx.slice(tC, (row, col)))

9. Environment and Debugging

IR Dump

FLYDSL_DUMP_IR=1 FLYDSL_DUMP_DIR=./dumps python my_kernel.py

Produces numbered .mlir files per pipeline stage plus 22_final_isa.s for each compiled specialization in the pinned release.

Key Environment Variables

Disk Cache Invalidation

The JIT disk cache auto-invalidates when kernel source or closure values change. Set FLYDSL_RUNTIME_ENABLE_CACHE=0 only when modifying C++ passes or non-closure helper functions:

FLYDSL_RUNTIME_ENABLE_CACHE=0 python my_kernel.py  # or: rm -rf ~/.flydsl/cache

Source-to-Assembly Debug Info

FlyDSL supports source-to-assembly mapping for rocprofv3 ATT traces via the MLIR ensure-debug-info-scope-on-llvm-func pass (equivalent to Triton's add_di_scope).

How it works:

1. FlyDSL's FuncLocationTracker generates MLIR loc() metadata pointing to Python source lines
2. The ensure-debug-info-scope-on-llvm-func{emission-kind=LineTablesOnly} pass converts MLIR locations into LLVM DISubprogramAttr / DICompileUnitAttr metadata
3. The -g flag in gpu-module-to-binary preserves this metadata as .debug_line in the HSACO binary
4. rocprofv3 ATT reads .debug_line to produce code.json with "source_file:line" entries

Pipeline position: After reconcile-unrealized-casts, before gpu-module-to-binary:

... -> reconcile-unrealized-casts
    -> ensure-debug-info-scope-on-llvm-func{emission-kind=LineTablesOnly}  (conditional on enable_debug_info)
    -> gpu-module-to-binary{format=fatbin opts=-g}

Verification: With FLYDSL_DUMP_IR=1, check 22_final_isa.s for .file and .loc directives. The PA decode kernel achieves 99.9% coverage (1109/1110 ISA instructions mapped to source).

Key insight: Without this pass, MLIR loc() metadata is silently dropped during MLIR-to-LLVM-IR translation. The -g flag alone is useless — it preserves debug info, but there's none to preserve without the DI scope pass.

Autotune Module

FlyDSL includes a Triton-style flydsl.autotune module in its installed package:

from flydsl.autotune import autotune, Config, do_bench

@autotune(
    configs=[
        Config(block_dim=64, vec_width=4),
        Config(block_dim=128, vec_width=4),
        Config(block_dim=256, vec_width=4),
    ],
    key=['const_n'],     # re-tune when these arg values change
    warmup=5, rep=25,    # benchmark timing params
)
@flyc.jit
def myKernel(A, C, n: fx.Int32, const_n: fx.Constexpr[int],
             block_dim: fx.Constexpr[int], vec_width: fx.Constexpr[int],
             stream: fx.Stream = fx.Stream(None)):
    ...

• Config kwargs become Constexpr args injected into @jit call
• Config.num_warps, waves_per_eu, maxnreg are special compiler-level options
• First call benchmarks all configs; subsequent calls use cached best
• Disk cache at ~/.flydsl/autotune/{func_name}.json
• do_bench(fn, warmup=5, rep=25) benchmarks using CUDA/HIP events, returns median ms

IMPORTANT: waves_per_eu does NOT work via gpu-module-to-binary opts=. It needs to be set as an LLVM function attribute or through rocdl-attach-target. This is a known limitation.

DLTensorAdaptor bug: Do NOT use flyc.from_dlpack() with pre-wrapped tensors when calling a @jit function with varying Constexpr values. The DLTensorAdaptor caches MLIR types from the first ir.Context, which become invalid when a new context is created (causes segfault). Pass raw torch.Tensor objects instead.

10. Troubleshooting

Common Issues

1. Constants/casts: Prefer fx.Int32(...), fx.Int64(...), and fx.Float32(...) for fixed-width arithmetic. Keep fx.Index(...) where a layout, loop, or launch API requires MLIR index; it is not interchangeable with fx.Int64.
2. buffer_ops.buffer_load offset: The offset parameter is in ELEMENTS, not bytes.
3. Cache stale after code changes: The disk cache auto-invalidates on source/closure changes. Only set FLYDSL_RUNTIME_ENABLE_CACHE=0 or clear ~/.flydsl/cache/ if you changed C++ passes or non-closure helpers.
4. LDS overflow: Check capacity (64KB on gfx942, 160KB on gfx950). SharedAllocator sizes the LDS global for you; the compiler errors if a static allocation exceeds the arch limit.
5. Dynamic vs Constexpr: Constexpr[int] values are baked into IR -- different values produce different compiled kernels. Use Int32 for truly dynamic values.
6. Tensor layout marking: For dynamic shapes or alignment, use flyc.from_dlpack(tensor).mark_layout_dynamic(leading_dim=0, divisibility=4).
7. Legacy SmemAllocator finalize: only for the legacy SmemAllocator path — call allocator.finalize() inside the GPU module body (CompilationContext.get_current().gpu_module_body). SharedAllocator needs no finalize step.
8. AMD wavefront size: Always 64 on gfx9xx. Use shifts [32, 16, 8, 4, 2, 1] for full-wave reduction.
9. tile_k alignment for GEMM: Check the selected MMA atom and kernel's K-step constraints. The preshuffle GEMM's K64-byte path requires tile_k * elem_bytes divisible by 64; this is not a universal FlyDSL GEMM rule.
10. INT4 (W4A8): A matrix is int8, B matrix is packed int4 (2 values/byte), unpacked to int8 in-kernel.
11. Absolute value and FMA: fx.absf(v) and fx.fma(a, b, c) exist in pinned 0.3.4.1. FMA has fused rounding, unlike an assumed two-step multiply and add. Check signed-zero/NaN behavior before replacing a hand-written absolute-value select.
12. Scalar broadcast to vector: Use Vec.filled(width, value, fx.Float32) to create a splat constant vector. Do NOT use raw vector ops for ordinary arithmetic.
13. gfx1250 wait counters: In pinned FlyDSL 0.3.4.1, the keyword form of fx.rocdl.s_waitcnt supports gfx942/gfx950/gfx11xx/gfx120x but rejects gfx1250. Check the target's wait instructions and ordering before replacing a gfx1250 wait; an offline compile failure on both revisions does not prove identical ISA.

11. Comparison with Triton/Gluon

FlyDSL gives maximum control at the cost of verbosity. The layout algebra is the key differentiator -- it enables precise control over how data is arranged in registers, shared memory, and global memory, and how threads map to data.

12. Running Kernels

Run the aiter test for the kernel on a host with the target GPU. For an offline ISA comparison, use a separate dump directory per revision and set COMPILE_ONLY=1 with FLYDSL_GPU_ARCH=<target>. Offline compilation does not exercise device behavior; compare the same specialization, compiler version, and final ISA on both revisions. A pinned compiler limitation that occurs on both trees cannot establish ISA equivalence for that specialization.