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`) | 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. ```python # 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. ```python # 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. ```python # 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` | 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.minnumf(a,b)` | `fx.minnumf(a, b)` — different NaN semantics from `fx.min` | | `arith.ceildivsi/ceildivui(a,b)` | `fx.ceildiv(a, b)` | Keep `arith.cmpf` / explicit `*FOp` only where no operator exists or fastmath is needed. ### `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 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` | 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`, `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`. ```python # 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. ```python # 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**. ```python # 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()