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.
In this repository7 other files it gives its agents.
- aiter-config-shapeSkill
- aiter-op-testSkill
- flydsl-kernel-authoringSkill
- opus-kernel-best-practiceSkill
- opus-module-build-optimizationSkill
- review-prSkill
- validate-kernel-prSkill
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
- Very specific: line 19
Before applying the recipes, check requirements.txt and the imported FlyDSL version and module path.
- Very specific: line 89
If an operand is a raw ir.Value, wrap it once at the source (fx.Float32(v)), not with ArithValue per use.
- Very specific: line 190
Keep arith.cmpf / explicit *FOp only where no operator exists or fastmath is needed.
- Distinctive: line 200
Check the rewriter's loop contract in pinned FlyDSL 0.3.4.1:
- Very specific: line 232
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 …
- Very specific: line 287
Delete the now-unused per-kernel _encode_waitcnt / _s_waitcnt shims and magic *CNT_* constants once your changes make them dead.
- Very specific: line 289
Use the public fx.rocdl.sched_barrier / fx.rocdl.sched_group_barrier wrappers when exposed by the pinned version.
- Very specific: line 324
Default static=True leaves launch(smem=...) unset.
- Very specific: line 368
Before deleting a private AST-rewriting decorator, search all imports and decorator uses across aiter.
- Distinctive: line 491
Keep a broad comment-only cleanup separate from executable changes.
- Very specific: line 551
Triage: do mechanical swaps (operators, casts, vector.extract/bitcast) first.
- Distinctive: line 560
Use the existing test's actual CLI when it is a script rather than pytest.
- Distinctive: line 561
Check asserted comparisons and dispatch logs as well as the exit code.
≈ “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.
Files
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
| Location | Role |
aiter/ops/flydsl/kernels/ | @flyc.kernel device kernels and shared helpers (tensor_shim.py, kernels_common.py, …) |
aiter/ops/flydsl/*.py | Launch wrappers, compile helpers, public op entry points |
op_tests/test_flydsl_*.py | Top-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.pyuses_scf.IfOp/_rawpervasively). 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)
| Deprecated | Replacement |
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 counters | fx.Int64(...) or fx.Int32(...) when the consumer permits a fixed-width integer |
arith.index_cast(T.index, v) at an index-typed boundary | fx.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 contractIf 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 replacesarith.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 boundary2. 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)mirrorscreate_buffer_resource; passnum_records_bytes=for a const byte count, ormax_size=Falseto 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 onfx.rocdl.make_tdm_atomis one to migrate: itstensor_extents=andstrides=move into the global view's layout, and itsimm_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/storeoffsetis in elements (×sizeof(dtype)internally) — a classic bug. - aiter still ships
aiter/ops/flydsl/kernels/buffer_ops.pyfor legacy kernels; migrate off it when touching load/store paths.
3. Raw upstream dialects → fx.* and Python
arith
| Raw | Preferred |
arith.constant(42, index=True) | fx.Int64(42) |
arith.mulf/addf(a,b) | a * b / a + b |
arith.trunc_f(ty, v) / ext_f | v.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
| Raw | Preferred |
scf.ForOp | range_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_constexprrequests Python unrolling.range(..., init=[...])emits anscf.forwith explicit carried state and converts its bounds to index, including Python integer bounds. It does not discardinitmerely because the bounds are static.- Ordinary
rangewithoutinituses automatic carried-state dispatch and expectsi32bounds; index bounds are converted toi32, andi64is rejected.
Preserve the selected loop form, supported bounds and state types. See §5 for runtime branches inside helper functions.
vector
| Raw | Preferred |
vector.extract(v, static_position=[i]) | fx.Vector(v)[i] |
vector.bitcast(ty, v) | fx.Vector(v).bitcast(fx.Float32) |
vector.splat / const vector | fx.Vector.filled(width, val, fx.Float32) |
| build from scalars | fx.Vector.from_elements(...) |
| reg-memref load/store | fx.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,fxconstants. 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.*→fxmath helpers (expr/math.py); keepmath_dialect.fmaetc. 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.ptrwith publicfx.inttoptr, then convert it withfx.to_llvm_ptronly if the downstream operation truly needs!llvm.ptr:
Choose the correct address space, element type, and proven alignment. A raw byte address for LDS usesty = 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)fx.AddressSpace.Sharedinstead. mem_ops.get_llvm_ptr/element_ptralso fold in+ offset*dtype_bytesarithmetic; 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_waitcntshims and magic*CNT_*constants once your changes make them dead. - Use the public
fx.rocdl.sched_barrier/fx.rocdl.sched_group_barrierwrappers when exposed by the pinned version. The legacy wait form remains available as positionalfx.rocdl.s_waitcnt(bitfield)for a boundary the keyword form cannot express; localize it. - Scheduler-sensitive.
s_waitcntplacement 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=Trueleaveslaunch(smem=...)unset; onlystatic=Falseauto-inferssmemfromallocated_bytes. SmemPtr.get()caches its view — reusing it in an epilogue after ascf.forcauses a dominance error.SharedAllocatoravoids this (view taken per use); for legacy code, clearptr._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) inconst_expr.- For branches returning or updating values, check that the rewriter preserves
their structure and types. Keep manual
scf.IfOponly 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.jitand 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 rawT.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 themake_layout_tv+make_tiled_copypair above. Prefer it, or the explicit two-liner, over hand-building a TV layout inline inmake_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'stv_layout_{A,B,C}_tiled), then.get_slice(tid)+partition_S/.retile(frag). Seeexamples/03-tiledMma.pyin 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 viamake_fragment_{A,B,C}. - Layouts passed to
make_layout_tvmust be static (compile-time) — plain Python-int shapes/strides inmake_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 kwargsfx.copyfor partitioned tensors (partition_S/partition_D/ tiled divide);fx.gemmfor 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 tofx.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. Preferfx.copy/fx.gemmfor 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.compilecan 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. Seemla_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
@cachefactory returning a JitFunction is not a compiled cache; wrap its call in_run_compiled(nospecialization_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
- Find legacy usage (under
aiter/ops/flydsl/):
Resolve import aliases and inspect nested definitions/callers; text hits are candidates, not proof of duplication or dead code.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> - Triage: do mechanical swaps (operators, casts,
vector.extract/bitcast) first; structural ones (control flow,buffer_opsoffsets, MMA loops) next. - Migrate in small commits, one family at a time, matching local style.
- Verify:
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.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 - 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 --statandgit diff --check. - 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
| Legacy | Current |
ArithValue(x) + y | x + y (typed fx) |
arith.unwrap(v) / _to_raw(v) | v.ir_value() (boundary only) |
| index-typed arithmetic | explicit 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-div | fx.max / fx.min / fx.ceildiv when signedness and overflow behavior match |
vector.extract/bitcast/splat | fx.Vector(v)[i] / .bitcast(ty) / .filled(...) |
scf.ForOp / scf.IfOp | range_constexpr / range(..., init=) / Python if / const_expr |
buffer_ops.* + offsets | fx.rocdl.make_buffer_tensor + layout + fx.copy |
raw llvm/memref access | fx.make_view / fx.get_iter / SharedAllocator |
create_llvm_ptr(v, address_space=N) / manual IntToPtrOp | ptr.llvm_ptr / fx.to_llvm_ptr(ptr) (backend-resolved AS) |
rocdl.s_waitcnt(_encode_waitcnt(...)) / magic bitfield | fx.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 helper | local @flyc.jit with Python if and typed carried state |
fx.Int32(fx.Int32(x)) / wrapping const ints | plain Python int; wrap once |
rocdl.mfma_* raw intrinsic | fx.make_mma_atom(fx.rocdl.MFMA(...)) + fx.gemm for supported operands and semantics |
hand-built TV layout in make_tiled_copy | fx.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 runs | delete within scope; retain invariant and rationale comments |
per-call @flyc.jit on a hot path | _run_compiled(exe, *args) fast dispatch |
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
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.

