agentleFS
ROCmROCm/aiter/.claude/skills/flydsl-kernel-code-cleanup/SKILL.md

flydsl-kernel-code-cleanup

Skill605 starsChanged 8 days agoMITView on GitHub

Clean up aiter FlyDSL kernels and shared helpers while preserving numerical behavior and performance. Use for raw-IR migrations, helper deduplication, dead-code removal, and reviews of those changes against the pinned FlyDSL API.

The file

Show its rules and text

We split this skill into the rules it gives an agent

13 rules. About 0 are common, 4 are distinctive, and 9 are very specific to this project.

  • ≈0

    Common · estimate

    Many files like this have a rule of this kind.

  • ≈4

    Distinctive · estimate

    This author's own choices. Worth a look.

  • 9

    Very specific

    Tied to their project's names, files or tools. Probably skip.

  • Common
  • Distinctive
  • Very specific
  1. Very specific:

    Before applying the recipes, check requirements.txt and the imported FlyDSL version and module path.

    line 19
  2. Very specific:

    If an operand is a raw ir.Value, wrap it once at the source (fx.Float32(v)), not with ArithValue per use.

    line 89
  3. Very specific:

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

    line 190
  4. Distinctive:

    Check the rewriter's loop contract in pinned FlyDSL 0.3.4.1:

    line 200
  5. Very specific:

    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 …

    line 232
  6. Very specific:

    Delete the now-unused per-kernel _encode_waitcnt / _s_waitcnt shims and magic *CNT_* constants once your changes make them dead.

    line 287
  7. Very specific:

    Use the public fx.rocdl.sched_barrier / fx.rocdl.sched_group_barrier wrappers when exposed by the pinned version.

    line 289
  8. Very specific:

    Default static=True leaves launch(smem=...) unset.

    line 324
  9. Very specific:

    Before deleting a private AST-rewriting decorator, search all imports and decorator uses across aiter.

    line 368
  10. Distinctive:

    Keep a broad comment-only cleanup separate from executable changes.

    line 491
  11. Very specific:

    Triage: do mechanical swaps (operators, casts, vector.extract/bitcast) first.

    line 551
  12. Distinctive:

    Use the existing test's actual CLI when it is a script rather than pytest.

    line 560
  13. Distinctive:

    Check asserted comparisons and dispatch logs as well as the exit code.

    line 561

≈ “Common” is an estimate: at least 1 in 5 of the 350,495 skill files we read have a rule of the same kind, not the same words. “Very specific” is our guess from the words: it names a file, a command, a setting or the project itself.

Keep what fits and drop the rest. Your version, ready for your agents.
Files
SKILL.md

FlyDSL Kernel Code Cleanup

Maps legacy kernel constructs to the fx.* surface available in aiter's pinned FlyDSL version. See the in-tree flydsl-kernel-authoring skill for API guidance.

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 the recipes, check requirements.txt and the imported FlyDSL version and module path. A sibling FlyDSL checkout can expose newer APIs than aiter supports. Kernel cleanup does not require a dependency bump or edits to FlyDSL itself; honor the task's explicit version, architecture and single-/multi- GPU scope.

Aiter layout

LocationRole
aiter/ops/flydsl/kernels/@flyc.kernel device kernels and shared helpers (tensor_shim.py, kernels_common.py, …)
aiter/ops/flydsl/*.pyLaunch wrappers, compile helpers, public op entry points
op_tests/test_flydsl_*.pyTop-level FlyDSL correctness / perf tests
op_tests/flydsl_tests/Additional FlyDSL kernel tests

Prefer from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled, ptr_arg. Reuse tensor_shim.py for compilation, cached dispatch and failure recovery. moe_kernels._run_compiled(exe, args) is a tuple-argument adapter for existing MoE/AOT callers; preserve that contract when consolidating launch code.

Reuse existing owners before adding helpers: act.py for activations; tensor_shim.py for pointer/base/dtype extraction; kernels_common.py for LOG2E and host or wrapping integer ceildiv; family common modules for reductions and specialized memory operations. Prefer typed fx.min/fx.max and fx.ceildiv where their signedness, NaN and overflow semantics match.

Cautions

  • Surgical, behavior-preserving. Migration is a refactor: minimal diffs, match local style.
  • Don't mass-rewrite heavily-legacy kernels unless asked (e.g. flash_attn_gfx950.py uses _scf.IfOp/_raw pervasively). Clean what the task touches.
  • Verify. Offset/type/SSA changes can shift results. Compare before/after numerics and generated code with FLYDSL_RUNTIME_ENABLE_CACHE=0; use a fresh dump directory per specialization so one shape cannot overwrite another.
  • Raw boundaries are semantic. Retain exact scope/ordering, volatile and alias metadata, raw SSA contracts, or integer widths unsupported by the pinned API. Record the specific reason; do not merely hide a dialect in a new facade.

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

DeprecatedReplacement
ArithValue(x) (wrap for operators)fx.Int32/Int64/Float32/Vector — already overload + - * / % << >> == < >
arith.unwrap(v) / arith._to_raw(v)v.ir_value(), only where a raw ir.Value is needed
index-typed arithmetic countersfx.Int64(...) or fx.Int32(...) when the consumer permits a fixed-width integer
arith.index_cast(T.index, v) at an index-typed boundaryfx.Index(v)

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.

1c. Thread and block IDs

In pinned FlyDSL 0.3.4.1, fx.thread_idx.x and fx.block_idx.x return fx.Int32 values. gpu.thread_idx.x and gpu.block_idx.x are the same typed accessors. Direct gpu.thread_id("x") / gpu.block_id("x") calls return raw MLIR values, so wrap those calls with fx.Int32(...) before typed arithmetic. Do not label the typed accessors legacy or remove a cast from a direct call merely because both spellings refer to the same GPU ID.

tid = fx.thread_idx.x                 # already fx.Int32
bid = fx.block_idx.x                  # already fx.Int32
raw_tid = gpu.thread_id("x")          # raw MLIR value
tid_from_raw = fx.Int32(raw_tid)      # typed boundary

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). Its default max_size=True leaves the record bound effectively unlimited; choose max_size=False or num_records_bytes= for hardware OOB checking.

# 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.cdna5.make_tiled_tdm_atom + fx.rocdl.cdna5.tdm_partition (raw VA, not a buffer resource). A site still on fx.rocdl.make_tdm_atom is one to migrate: its tensor_extents= and strides= move into the global view's layout, and its imm_offset= becomes the coordinate tensor's i32 tile index. Leave it alone only for the gather/scatter family (make_tensor_gather_descriptor / tensor_load_gather, per-row addresses) or a tile whose innermost dimension spans under 16 bytes per warp.
  • 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.
  • aiter still ships aiter/ops/flydsl/kernels/buffer_ops.py for legacy kernels; migrate off it when touching load/store paths.

3. Raw upstream dialects → fx.* and Python

arith

RawPreferred
arith.constant(42, index=True)fx.Int64(42)
arith.mulf/addf(a,b)a * b / a + b
arith.trunc_f(ty, v) / ext_fv.to(fx.BFloat16)
arith.index_cast(T.i32, v)fx.Int32(v)
arith.select(cond, t, f)cond.select(t, f)
arith.cmpi(slt, a, b)a < b
arith.maximumf/minimumf(a,b)fx.max(a, b) / fx.min(a, b)
arith.maxsi/maxui/minsi/minui(a,b)fx.max(a, b) / fx.min(a, b)
arith.maxnumf(a,b)fx.maxnumf(a, b) — different NaN semantics from fx.max
arith.ceildivsi/ceildivui(a,b)fx.ceildiv(a, b)
math.fma(a,b,c)fx.fma(a, b, c) — preserves fused rounding
math.absf(v)fx.absf(v) — check signed-zero/NaN behavior if replacing a select

Keep arith.cmpf / explicit *FOp only where no operator exists or fastmath is needed. Do not replace separate a * b + c operations with fx.fma unless fused rounding is intended.

scf

RawPreferred
scf.ForOprange_constexpr(N) (unrolled) or range(lo, hi, step, init=[...]) (runtime, loop-carried)
scf.IfOp(_raw(cond))Python if cond: (runtime) / if const_expr(flag): (compile-time)

Check the rewriter's loop contract in pinned FlyDSL 0.3.4.1:

  • 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

RawPreferred
vector.extract(v, static_position=[i])fx.Vector(v)[i]
vector.bitcast(ty, v)fx.Vector(v).bitcast(fx.Float32)
vector.splat / const vectorfx.Vector.filled(width, val, fx.Float32)
build from scalarsfx.Vector.from_elements(...)
reg-memref load/storefx.memref_load_vec(r) / fx.memref_store_vec(v, r)

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 aiter rather than extending FlyDSL as part of a cleanup.
  • 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
  • For a raw integer address, construct a semantic fly.ptr with public fx.inttoptr, then convert it with fx.to_llvm_ptr only if the downstream operation truly needs !llvm.ptr:
    ty = fx.PointerType.get(fx.Int32.ir_type, fx.AddressSpace.Global, 4)
    ptr = fx.inttoptr(ty, fx.Int64(address))
    llvm_ptr = fx.to_llvm_ptr(ptr)
    Choose the correct address space, element type, and proven alignment. A raw byte address for LDS uses fx.AddressSpace.Shared instead.
  • 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) packs target-specific bitfields for gfx942, gfx950, gfx11xx, and gfx120x. Pinned 0.3.4.1 rejects gfx1250 here; do not migrate a gfx1250 wait to this API. Use the target's supported wait ops (fx.rocdl.s_wait_loadcnt / s_wait_storecnt, or async wait where applicable) only after checking their ordering and generated ISA.

# 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 pinned 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.
  • Before deleting a private AST-rewriting decorator, search all imports and decorator uses across aiter. Decorators execute at module import time; update every caller to @flyc.jit and import the affected modules as part of validation.

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 in a FlyDSL source checkout.

# 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 in a FlyDSL source checkout.
  • 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 pinned 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 (aiter/ops/flydsl/kernels/tensor_shim.py): compile once, cache the CompiledFunction, fast-dispatch after.

from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled, ptr_arg

compiled = compile_my_kernel(...)          # {"launch": <exe>, ...}
_run_compiled(compiled["launch"],
              ptr_arg(out), ptr_arg(a), ptr_arg(b),
              a.stride(0), M, N, K, stream)
  • On the pinned runtime, flyc.compile can execute the launcher on the first call. Use the existing preload helper for no-dispatch materialization; never pass placeholder pointers to a compiling launch path without that guard.
  • Pass flat scalars/pointers (ptr_arg(t), data_ptr(), stride(i), sizes, stream) — bypasses DLPack. See mla_reduce_kernels.py, linear_attention_prefill_kernels.py, splitk_hgemm.py.
  • Reuse tensor_shim._run_compiled; don't add a second copy in wrapper modules.
  • Worth it for small kernels in tight loops, not cold one-shot launches. Arg order/types must match the compiled signature — verify.
  • A @cache factory returning a JitFunction is not a compiled cache; wrap its call in _run_compiled (no specialization_key — the factory fixes the config). Only a shared JitFunction with varying Constexpr args needs a key, and the key is just those Constexpr values.

10. Procedure

  1. Find legacy usage (under aiter/ops/flydsl/):
    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:
    black --check <changed-files> && ruff check <changed-files>
    FLYDSL_RUNTIME_ENABLE_CACHE=0 python3 -m pytest op_tests/test_flydsl_<kernel>.py -v
    # or: op_tests/flydsl_tests/test_flydsl_<kernel>.py
    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

LegacyCurrent
ArithValue(x) + yx + y (typed fx)
arith.unwrap(v) / _to_raw(v)v.ir_value() (boundary only)
index-typed arithmeticexplicit fx.Int64/Int32(...) where supported; retain fx.Index at index-typed boundaries
arith.mulf/addf/trunc_f/select*, +, .to(ty), .select(...)
raw integer min/max or ceil-divfx.max / fx.min / fx.ceildiv when signedness and overflow behavior match
vector.extract/bitcast/splatfx.Vector(v)[i] / .bitcast(ty) / .filled(...)
scf.ForOp / scf.IfOprange_constexpr / range(..., init=) / Python if / const_expr
buffer_ops.* + offsetsfx.rocdl.make_buffer_tensor + layout + fx.copy
raw llvm/memref accessfx.make_view / fx.get_iter / SharedAllocator
create_llvm_ptr(v, address_space=N) / manual IntToPtrOpptr.llvm_ptr / fx.to_llvm_ptr(ptr) (backend-resolved AS)
rocdl.s_waitcnt(_encode_waitcnt(...)) / magic bitfieldfx.rocdl.s_waitcnt(vmcnt=/lgkmcnt=/expcnt=) on supported targets; pinned 0.3.4.1 excludes gfx1250
SmemAllocator/SmemPtr + finalize()@fx.struct + fx.SharedAllocator().allocate(...).peek().view(...)
raw SCF or AST-rewriter calls in a plain helperlocal @flyc.jit with Python if and typed carried state
fx.Int32(fx.Int32(x)) / wrapping const intsplain Python int; wrap once
rocdl.mfma_* raw intrinsicfx.make_mma_atom(fx.rocdl.MFMA(...)) + fx.gemm for supported operands and semantics
hand-built TV layout in make_tiled_copyfx.make_layout_tv + fx.make_tiled_copy / make_tiled_copy_tv; _A/_B/_C for MMA operands
*_atom_call (loop or single atom)fx.copy / fx.gemm where equivalent; retain raw SSA and unsupported operand contracts
restated/dead/stale comments, blank runsdelete within scope; retain invariant and rationale comments
per-call @flyc.jit on a hot path_run_compiled(exe, *args) fast dispatch

Download now (ZIP)

Does it work?

Reviews from people and from AI agents that ran it: what they used it for, and whether it worked.

No reviews yet

People and AI agents can review whether this works in practice. Say what you used it for, and whether it worked.

  • It worked
  • Partly
  • It didn't
What did you use it for? e.g. Reviewing pull requests in a Rails app

Using an agent? It can post a review for you after it has run the file: the MCP tool registry_write, action report. How to connect one.

Your trail

Files you open will collect here, so you can find them again.

Open two or more, and you can remix the rules you like from each into one list. Kept in this browser only.

Now do this for your own.

Bring your CLAUDE.md, AGENTS.md or skills. agentleFS shows which of your rules are common, which are yours and which only fit one project, and keeps one version your agents read everywhere.