Kernel Code Cleanup

FlyDSL Kernel Code Cleanup

Maps legacy kernel constructs to the current fx.* surface. Companion to flydsl-kernel-authoring (API reference) and flydsl-tile-programming (authoring wizard).

Golden rule: in @flyc.kernel / @flyc.jit bodies, use fx.* and Python operators first. Drop to a raw dialect only at a hard boundary with no wrapper, and localize it.

Before applying a recipe, confirm the imported FlyDSL source path, version and native binding. A sibling aiter checkout can be pinned to a different release; validate ports separately against each repository's supported setup. Preserve FlyDSL-specific kernel interfaces and tuning choices when reusing aiter code. Kernel cleanup should use the existing API; do not expand it into expr/ or compiler changes unless the task calls for that. Honor the requested GPU scope.

Reuse kernels/common/act.py for shared activation components, kernels/common/tensor_shim.py for pointer/base extraction and cached dispatch, kernels/common/kernels_common.py for LOG2E, dtype and architecture lookup, and family common modules for specialized memory operations. Similar formulas can encode different rounding or scheduling contracts.

Cautions

• Surgical, behavior-preserving. Migration is a refactor: minimal diffs, match local style.
• Keep the requested scope. A broad kernel cleanup can include tuned paths; migrate them in bounded steps and retain their scheduling and ABI contracts.
• Verify. Compare before/after numerics and generated code with FLYDSL_RUNTIME_ENABLE_CACHE=0; isolate each specialization in a fresh dump directory so later shapes cannot overwrite earlier evidence.
• Raw boundaries are semantic. Preserve atomic scope/ordering, volatile and alias metadata, raw SSA contracts, and unsupported integer widths. Record why a boundary remains; moving it behind a new facade does not remove it.
• expr/ stays target-neutral: no rocdl/llvm/buffer imports in python/flydsl/expr/ top-level (guarded by test_expr_optional_rocdl.py).

1. ArithValue and index helpers (deprecated in expr/arith.py)

fx.Index maps to MLIR index. Prefer explicit-width fx.Int64/fx.Int32 for arithmetic, choosing width and signedness deliberately. Keep fx.Index where a launch, layout, loop or other API requires the index type; replacing it merely to remove the name can change the IR contract. Do not widen i32 counters or narrow an index without checking the consumer and supported bounds.

# Before
acc  = ArithValue(val) + peer
lane = ArithValue(tid) % fx.Index(64)
cond = arith.unwrap(idx >= limit)
off  = arith.index_cast(T.index, x)
# After
acc  = val + peer                    # val already fx.Float32 / fx.Vector
lane = tid % fx.Int64(64)
cond = (idx >= limit).ir_value()     # only if a raw scf.IfOp needs it
off  = fx.Index(x)                 # preserve this consumer's index contract

If an operand is a raw ir.Value, wrap it once at the source (fx.Float32(v)), not with ArithValue per use. Keep an explicit arith.*FOp only for non-default fastmath.

1b. Drop redundant fx.* wraps

Wrap only to introduce a type (Python literal / raw ir.Value) or change one. Re-wrapping an already-typed value is noise; double-wrapping is dead.

# Before
for i in range_constexpr(fx.Int32(N)):
    off = fx.Int64(fx.Int64(base) + fx.Int64(4))
tile = fx.make_layout(fx.Int32(BLOCK), fx.Int32(1))
idx  = fx.Int32(tx)                  # tx already fx.Int32
# After
for i in range_constexpr(N):
    off = base + fx.Int64(4)
tile = fx.make_layout(BLOCK, 1)      # builders take Python ints
idx  = tx

• Compile-time shapes/strides/bounds (make_layout, make_shape, range_constexpr, Constexpr) take plain Python ints.
• Wrap a runtime value once, at first typed use.
• A real cast (fx.Int64(i32) widen, fx.Int32(index) narrow) is not redundant — it replaces arith.index_cast.

2. buffer_ops → make_buffer_tensor + copy atoms

create_buffer_resource + manual offsets is legacy. Build a buffer-resource view with fx.rocdl.make_buffer_tensor(), then use layout ops + fx.copy (§7b); the OOB-checked V# descriptor is built for you.

# Before (manual offsets — see PA //4 offset bugs)
rsrc = buffer_ops.create_buffer_resource(A, max_size=True)
data = buffer_ops.buffer_load(rsrc, row * K + k, vec_width=4, dtype=fx.Float32)
buffer_ops.buffer_store(data, rsrc, row * N + col)
# After
bufA = fx.rocdl.make_buffer_tensor(A)
tA   = fx.make_view(fx.get_iter(bufA), fx.make_layout((M, K), (K, 1)))
copy = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32)
fx.copy(copy, fx.slice(tA, (None, tid)), rA)   # after partitioning tA (§7b: prefer fx.copy)

• make_buffer_tensor(tensor, max_size=True) mirrors create_buffer_resource; pass num_records_bytes= for a const byte count, or max_size=False to derive from the layout.
• gfx1250 TDM uses a different atom — fx.rocdl.make_tdm_atom (raw VA, not a buffer resource).
• A scalar-base + per-thread-offset load with no layout form may stay on buffer_ops — note it. buffer_load/store offset is in elements (× sizeof(dtype) internally) — a classic bug.

3. Raw upstream dialects → fx.* and Python

arith

Keep arith.cmpf / explicit *FOp only where no operator exists or fastmath is needed.

scf

Check the rewriter's loop contract in the checkout version. In this checkout (FlyDSL 0.3.3):

• range_constexpr requests Python unrolling.
• range(..., init=[...]) emits an scf.for with explicit carried state and converts its bounds to index, including Python integer bounds. It does not discard init merely because the bounds are static.
• Ordinary range without init uses automatic carried-state dispatch and expects i32 bounds; index bounds are converted to i32, and i64 is rejected.

Preserve the selected loop form, supported bounds and state types. See §5 for runtime branches inside helper functions.

vector

llvm / memref / math

• llvm.* ptr math / load/store / const → layout views (fx.make_view, fx.get_iter), fx.Array + SharedAllocator, fx constants. Use existing intrinsic wrappers where equivalent; keep unsupported boundaries local to the kernel or shared helper. API extensions require their own task scope.
• memref.* → layout tensors/views + copy atoms.
• math.* → fx math helpers (expr/math.py); keep math_dialect.fma etc. only where no wrapper exists.

3b. fly.ptr → !llvm.ptr (backend-resolved address space)

When you hold an fx pointer (fly.ptr) and need a raw !llvm.ptr at a hard boundary, use the DSL primitive — it maps the pointer's semantic address space to the backend's LLVM address-space number for you. Don't hand-build one with a hardcoded <1> / <3> via IntToPtrOp.

# Before (hardcoded address space)
p = buffer_ops.create_llvm_ptr(lds_addr, address_space=3)
p = mem_ops._create_llvm_ptr(val, address_space=1)   # a.k.a. mem_ops.to_llvm_ptr
# After
p = ptr.llvm_ptr          # property on an fx pointer
p = fx.to_llvm_ptr(ptr)   # equivalent free function; backend resolves the AS

• Applies only when you already have a fly.ptr. A raw int/index address (e.g. an LDS byte offset with no pointer form) still needs manual construction — note it.
• mem_ops.get_llvm_ptr / element_ptr also fold in + offset*dtype_bytes arithmetic; keep the offset math (layout views / get_element_ptr) and only swap the final ptr cast for .llvm_ptr.
• Preserve byte versus element GEPs and alignment provenance. An equal numeric address alone does not guarantee equal memory instructions; compare the generated loads and stores when replacing an epilog pointer path.

3c. Manual s_waitcnt bitfields → fx.rocdl.s_waitcnt(vmcnt=/lgkmcnt=/expcnt=)

Hand-encoding a wait-counter bitfield (or calling rocdl.s_waitcnt(magic) with a raw number) is arch-fragile — the field widths differ per arch (CDNA3 lgkmcnt max 15 vs RDNA 63). The keyword form of fx.rocdl.s_waitcnt (expr/rocdl/universal.py) is arch-dispatched across gfx942/gfx950/gfx11xx/gfx120x and packs the correct bitfield for you.

# Before
rocdl.s_waitcnt(_encode_waitcnt(lgkmcnt=0))   # per-kernel encoder
rocdl.s_waitcnt(0)                             # raw "wait for everything"
_s_waitcnt(0xC07F)                             # magic LGKMCNT_0_ONLY bitfield
# After
fx.rocdl.s_waitcnt(lgkmcnt=0)                  # wait for LDS/SMEM only
fx.rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0, expcnt=0)  # matches raw s_waitcnt(0)
fx.rocdl.s_waitcnt(lgkmcnt=0)

• Unset fields default to "no wait" (their per-arch max) — name only the counters you need.
• Delete the now-unused per-kernel _encode_waitcnt / _s_waitcnt shims and magic *CNT_* constants once your changes make them dead.
• Use the public fx.rocdl.sched_barrier / fx.rocdl.sched_group_barrier wrappers when exposed by the checkout version. The legacy wait form remains available as positional fx.rocdl.s_waitcnt(bitfield) for a boundary the keyword form cannot express; localize it.
• Scheduler-sensitive. s_waitcnt placement drives hot-loop pipelining in tuned attention/GEMM kernels — an op-identical swap can still shift the schedule. Verify perf (median-based), not just correctness, and don't mass-migrate pervasively-tuned kernels (e.g. flash_attn_gfx950.py, mla_fwd_decode_*).

4. SmemAllocator / SmemPtr → SharedAllocator

Legacy LDS path uses a manual base pointer, byte offsets, and finalize(). New kernels declare an @fx.struct of fx.Array fields and allocate via fx.SharedAllocator — the compiler sizes the LDS global; no finalize.

# Before
allocator = SmemAllocator(None, arch=GPU_ARCH, global_sym_name="smem")
base = allocator.get_base()
smem_a = SmemPtr(base, 0, dtype_, shape=(BLOCK_M * BLOCK_K,))
smem_b = SmemPtr(base, a_bytes, dtype_, shape=(BLOCK_K * BLOCK_N,))
allocator.finalize()
# After
@fx.struct
class SharedStorage:
    a: fx.Array[fx.Float16, BLOCK_M * BLOCK_K]
    b: fx.Array[fx.Float16, BLOCK_K * BLOCK_N]

lds   = fx.SharedAllocator().allocate(SharedStorage).peek()
lds_a = lds.a.view(fx.make_layout((BLOCK_M, BLOCK_K), (BLOCK_K, 1)))
lds_b = lds.b.view(fx.make_layout((BLOCK_K, BLOCK_N), (BLOCK_N, 1)))

• Default static=True leaves launch(smem=...) unset; only static=False auto-infers smem from allocated_bytes.
• SmemPtr.get() caches its view — reusing it in an epilogue after a scf.for causes a dominance error. SharedAllocator avoids this (view taken per use); for legacy code, clear ptr._view_cache = None.
• Structural change — migrate a kernel's whole LDS at once and re-run its test.

5. Runtime branches inside helper functions

Inside a rewritten @flyc.kernel or @flyc.jit function, ordinary Python if supports side effects and carried scalar/list/tuple state. Initialize carried values before the branch with matching types; None cannot become an SSA result. A branch producing values does not by itself require raw SCF.

A plain helper executing outside the rewriter's scope needs a local @flyc.jit boundary for its runtime if. Use that public decorator instead of calling ReplaceIfWithDispatch.scf_if_dispatch directly.

# Before
with _if_then(_scf.IfOp(_raw(ArithValue(q_start < seqlen_q)))):
    ...
# After
def then_path(): ...
def else_path(): ...

@flyc.jit
def dispatch():
    if q_start < seqlen_q:      # typed fx compare → scf.if
        then_path()
    else:
        else_path()

dispatch()

• A bare if cond: is fine for a simple guarded side effect — no helper needed.
• const_expr(flag) for compile-time branches; never wrap runtime SSA (gpu.thread_id, lane) in const_expr.
• For branches returning or updating values, check that the rewriter preserves their structure and types. Keep manual scf.IfOp only for a demonstrated unsupported control-flow contract; localize it.

6. Raw rocdl.mfma_* → MMA atom + fx.gemm

Raw intrinsics hardcode fragment types, the [a, b, c, 0, 0, 0] tuple, and the instruction. Build an atom and issue it; fragment layouts/packing are handled and you pick the atom family by target: MFMA for CDNA3/CDNA4, WMMA for gfx11/gfx1250.

# Before
c_frag = rocdl.mfma_f32_16x16x16f16(T.vec(4, T.f32), [a_frag, b_frag, c_frag, 0, 0, 0])
# After
mma = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 16, fx.Float16))   # → f32 acc
fx.gemm(mma, frag_C, frag_A, frag_B, frag_C)                    # d, a, b, c (prefer this)
fx.mma_atom_call(mma, frag_C, frag_A, frag_B, frag_C)           # single tile — prefer fx.gemm (§7b)

• fx.rocdl.MFMA(m, n, k, elem_ty_ab, elem_ty_acc=None) picks the intrinsic from shape+dtype. Scaled: fx.rocdl.cdna4.MFMA_Scale; gfx1250/gfx11: fx.rocdl.WMMA / WMMAScale.
• Build fragments with fx.make_fragment_like / make_fragment_{A,B,C}, not raw T.vec(...).
• Order is d, a, b, c (accumulator first).
• Structural — convert a complete supported MMA path and diff numerics. Retain raw calls for instructions or operand forms the builders cannot express, or when the alternative changes required semantics or scheduling.

7. Tiled copy/MMA: build from a TV layout, iterate with fx.copy / fx.gemm

7a. Build the tiled copy/MMA (TV layout)

A tiled copy is a copy atom laid over a thread-value (TV) layout plus a tiler. Build the TV layout from separate thread/value layouts with fx.make_layout_tv (returns (tile_mn, tv_layout)), pass both to fx.make_tiled_copy, slice per-thread with .get_slice(tid), then partition the tensor. See examples/02-tiledCopy.py.

# thread + value layouts -> (tile_mn, tv_layout) -> tiled copy
thr_layout = fx.make_layout((4, 1), (1, 1))
val_layout = fx.make_layout((1, 8), (1, 1))
copy_atom  = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32)
tile_mn, tv_layout = fx.make_layout_tv(thr_layout, val_layout)
tiled_copy = fx.make_tiled_copy(copy_atom, tv_layout, tile_mn)
thr_copy   = tiled_copy.get_slice(tid)
part_src   = thr_copy.partition_S(bA)   # bA = fx.slice(fx.zipped_divide(A, tile), (None, bid))
part_dst   = thr_copy.partition_D(bB)
frag       = fx.make_fragment_like(part_src)

• fx.make_tiled_copy_tv(atom, thr_layout, val_layout) is the one-call shortcut for the make_layout_tv + make_tiled_copy pair above. Prefer it, or the explicit two-liner, over hand-building a TV layout inline in make_tiled_copy.
• For copies matched to an MMA operand layout, do not hand-build a TV layout — use fx.make_tiled_copy_A/B/C(copy_atom, tiled_mma) (they read the atom's tv_layout_{A,B,C}_tiled), then .get_slice(tid) + partition_S / .retile(frag). See examples/03-tiledMma.py.
• Build the MMA with fx.make_tiled_mma(mma_atom, atom_layout); slice with .thr_slice(tid) / .get_slice(tid) and make fragments via make_fragment_{A,B,C}.
• Layouts passed to make_layout_tv must be static (compile-time) — plain Python-int shapes/strides in make_layout.

7b. Prefer fx.copy / fx.gemm over *_atom_call — even for a single atom

fx.copy / fx.gemm iterate the atom over a tiled/partitioned layout and take atom state as kwargs — no hand-written loop or atom_set_value. Prefer them over copy_atom_call / mma_atom_call not just for loops but for single-atom sites too, when supported by the checkout API: fx.copy(atom, src, dst) can issue the same atom over a single-tile partition. Verify the generated code; API-level equivalence alone does not prove identical ISA or performance.

# Before — loop
for k in range_constexpr(K_TILES):
    fx.copy_atom_call(copy_atom, part_src[k], frag[k])
for k in range_constexpr(K_TILES):
    fx.mma_atom_call(mma, frag_C, frag_A[k], frag_B[k], frag_C)
# Before — single atom (helpers, one tile)
fx.copy_atom_call(copy_atom, fx.slice(tiles, (None, idx)), r)
fx.mma_atom_call(mma, frag_C, frag_A, frag_B, frag_C)
# After — same in both cases
fx.copy(copy_atom, part_src, frag)                                     # loop or single
fx.copy(copy_atom, fx.slice(tiles, (None, idx)), r)                    # single-atom swap
fx.gemm(mma, frag_C, frag_A, frag_B, frag_C)
fx.gemm(mma, frag_C, frag_A, frag_B, frag_C, scale_a=sa, scale_b=sb)   # atom state as kwargs

• fx.copy for partitioned tensors (partition_S/partition_D / tiled divide); fx.gemm for the MMA loop (accumulator-first order).
• A supported single-atom swap can be one-for-one (fx.copy_atom_call(a, s, d) → fx.copy(a, s, d)); no new TV layout is needed. Don't manufacture a TV layout for a degenerate single-tile load whose thread→data mapping is a mandatory swizzle — just pass the existing single-tile slice to fx.copy.
• Keep copy_atom_call_ssa / mma_atom_call_ssa (the SSA-returning variants are a different primitive) and any raw atom call whose operands have no tensor/ partition form to pass. Prefer fx.copy / fx.gemm for supported tensor forms.
• Diff numerics and ISA; for scheduler-sensitive hot loops compare repeated, paired graph timings. Unchanged resources alone do not prove unchanged time.

8. Trim comments and dead code

Cut low-value comments and dead code within the requested scope. A line-count reduction is useful context, not evidence of correctness or complete cleanup.

Remove: comments that restate code; commented-out / dead blocks; per-line step narration; ASCII banners (keep one concise header per section); stale comments that contradict the code; unused locals/imports/helpers you made redundant; runs of 2+ blank lines.

Keep: the why — non-obvious layout/stride math, swizzle rationale, ABI quirks, offset-unit gotchas, invariants, spec/ISA references.

• Keep a broad comment-only cleanup separate from executable changes. Compare ASTs (ignoring docstrings when appropriate) to verify that claim.
• Confirm callers, re-exports, generated/JIT lookup and import side effects before deleting a helper or argument. Remove pre-existing dead code when it is within the requested cleanup scope; retain comments explaining a real invariant.
• Similar formulas are not necessarily interchangeable: preserve activation rounding and batch scheduling, signed/unsigned extrema and packed bit widths. Extract the common operation and keep meaningful variants explicit.

9. Cut launch overhead with _run_compiled

Calling a @flyc.jit wrapper directly re-runs per-call dispatch (DLPack, arg marshalling, cache lookup). On hot paths use _run_compiled (kernels/common/tensor_shim.py): compile once, cache the CompiledFunction, fast-dispatch after.

from kernels.common.tensor_shim import _run_compiled

compiled = compile_my_kernel(...)          # {"launch": <exe>, ...}
_run_compiled(compiled["launch"], out.data_ptr(), a.data_ptr(), b.data_ptr(),
              a.stride(0), M, N, K, stream)

• Pass flat scalars/pointers (data_ptr(), stride(i), sizes, stream) — it bypasses DLPack. See pa_decode_fp8.py.
• Reuse the shim, including its failed-compile context cleanup; don't duplicate it.
• The cold flyc.compile call can execute the launcher. Pass real arguments; compiling with placeholder pointers is not automatically a no-dispatch preload.
• Worth it for small kernels in tight loops, not cold one-shot launches. Arg order/types must match the compiled signature — verify.

10. Procedure

1. Find legacy usage (under kernels/):

   rg -n 'ArithValue|_to_raw|arith\.(unwrap|index|index_cast)|fx\.Index\(' <file>
   rg -n 'buffer_ops\.(create_buffer_resource|buffer_load|buffer_store)' <file>
   rg -n '_mlir\.dialects|from flydsl\.expr import' <file>
   rg -n '\b(scf\.(For|If)Op|vector\.(extract|bitcast|splat)|llvm\.(load|store|mlir))' <file>
   rg -n 'SmemPtr|SmemAllocator|\.finalize\(\)' <file>
   rg -n 'fx\.(Int32|Int64|Float32)\(fx\.(Int32|Int64|Float32)\(' <file>
   rg -n 'rocdl\.mfma_|\bmfma_(f32|i32)_|copy_atom_call|mma_atom_call' <file>
   rg -n 'create_llvm_ptr|_create_llvm_ptr|get_llvm_ptr|IntToPtrOp' <file>
   rg -n 's_waitcnt\(|_encode_waitcnt|_s_waitcnt|CNT_[0-9A-Z_]*=|0x[Cc]07[Ff]' <file>
   rg -n 'LOG2E|log2e|def .*sigmoid|def .*tanh|def .*ceildiv|def .*ptr' <file>
   

   Resolve import aliases and inspect nested definitions/callers; text hits are candidates, not proof of duplication or dead code.
2. Triage: do mechanical swaps (operators, casts, vector.extract/bitcast) first; structural ones (control flow, buffer_ops offsets, MMA loops) next.
3. Migrate in small commits, one family at a time, matching local style.
4. Verify:

   bash scripts/check_python_style.sh --include-local
   FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 -m pytest tests/kernels/test_<kernel>.py -v
   

   Use the existing test's actual CLI when it is a script rather than pytest. Check asserted comparisons and dispatch logs as well as the exit code; compare numerical results, ISA and performance for the changed paths.
5. Review the actual merge-base diff and all remaining candidates. Check callers of shared helpers, excluded paths importing them, and the final tree after any upstream merge. Inspect git diff --stat and git diff --check.
6. Report coverage and residuals. Separate native correctness/performance, compile/ISA-only cases, skips and baseline failures. Record each retained low-level boundary's reason and the tested commit; inspect CI failures before declaring completion. Rerun checks affected by subsequent code changes.

Quick reference