[DOC] Clean up MLA internal type usage and docs (#460)

Author: coderfeli

.claude/skills/debug-flydsl-kernel/SKILL.md

@@ -3,7 +3,7 @@ name: debug-flydsl-kernel
 description: >
   Debug FlyDSL GPU kernels that produce NaN, inf, wrong results, or crash.
   Covers cache invalidation, tracing pitfalls (runtime conditionals, range vs
-  range_constexpr), scf.for state packing, buffer_load addressing, MFMA operand
+  range_constexpr), loop-carried state packing, buffer_load addressing, MFMA operand
   layout verification, LDS bank conflict diagnosis, and systematic error
   isolation (all-1s test, single-partition test, host-side tensor inspection).
   Use when a FlyDSL kernel produces incorrect output or compilation errors.
@@ -34,7 +34,7 @@ compile_my_kernel.cache_clear()
 | All zeros output | Wrong output address, uninitialized temp buffer | Section 3 |
 | Partially wrong (>50% mismatch) | Wrong partition count, missing partitions, layout mismatch | Section 4 |
 | Small errors (1-5% mismatch) | FP8 quantization, scale factor, off-by-one masking | Section 5 |
-| Compilation error / crash | Type mismatch, scf.for state, range vs range_constexpr | Section 6 |
+| Compilation error / crash | Type mismatch, loop-carried state, range vs range_constexpr | Section 6 |
 | GPU hang | Infinite loop, deadlock in barrier, OOB memory access | Section 7 |
 
 ## 2. Debugging NaN
@@ -45,7 +45,7 @@ When ALL tokens in a partition are masked (out of context), `qk_max = -inf`. The
 
 **Fix**: Guard the exp calculation:
 ```python
-safe_diff = arith.select(qk_max > NEG_INF, diff, ZERO_F)
+safe_diff = (qk_max > NEG_INF).select(diff, ZERO_F)
 ```
 
 ### 2.2 Division by zero in normalization
@@ -54,8 +54,8 @@ When `exp_sum = 0` (all probs zero), `1/exp_sum = inf`.
 
 **Fix**:
 ```python
-safe_sum = arith.select(running_sum > ZERO_F, running_sum, arith.constant(1.0, type=T.f32))
-inv_sum = arith.constant(1.0, type=T.f32) / safe_sum
+safe_sum = (running_sum > ZERO_F).select(running_sum, fx.Float32(1.0))
+inv_sum = fx.Float32(1.0) / safe_sum
 ```
 
 ### 2.3 Host-side NaN check
@@ -143,7 +143,7 @@ Verify `_scale = softmax_scale * q_scale * k_scale` matches the reference. Commo
 
 ### 6.1 `range()` vs `range_constexpr()` inside @flyc.kernel
 
-FlyDSL's AST rewriter converts ALL `range()` to `scf.for` (runtime loops). Use `range_constexpr()` for compile-time unrolled loops:
+FlyDSL's AST rewriter converts runtime `range()` loops into MLIR loops. Use `range_constexpr()` for compile-time unrolled loops:
 ```python
 # WRONG: i becomes an ArithValue, can't index Python lists
 for i in range(4): result[i] = ...
@@ -160,23 +160,23 @@ FlyDSL tracing evaluates Python `if` at trace time. Runtime GPU values can't be
 if kv_tok < context_len:  # runtime comparison
     fx.printf(...)
 
-# CORRECT: use arith.select for runtime conditionals
-val = arith.select(kv_tok < context_len, good_val, bad_val)
+# CORRECT: use ArithValue.select for runtime value selection
+val = (kv_tok < context_len).select(good_val, bad_val)
 ```
 
 Python `if` is fine for COMPILE-TIME decisions (e.g., `if trans_v:` where trans_v is a Python bool).
 
-### 6.3 scf.for state packing
+### 6.3 Loop-carried state packing
 
-All loop-carried values must be raw SSA values (not Python wrappers):
+Prefer FlyDSL internal types (`fx.Int32`, `fx.Float32`, `Vector`, `ArithValue`) for loop-carried state. Unwrap only when a low-level helper explicitly requires raw `ir.Value`:
 ```python
 def _unwrap(v):
     return v.ir_value() if hasattr(v, 'ir_value') else v
 
 init_state = [_unwrap(v) for v in [val1, val2, vec_val]]
 ```
 
-Supported state types: `f32` (scalar), `f32x4` (vector), `i32`, `i64`, `index`.
+Supported state types: `f32` (scalar), vector values, `i32`, `i64`, `index`.
 
 ### 6.4 buffer_load type mismatch
 
@@ -186,21 +186,21 @@ k_addr_bytes = ...  # address in FP8 elements (= bytes for FP8)
 k_4xi32 = buffer_ops.buffer_load(k_rsrc, k_addr_bytes // 4, vec_width=4, dtype=T.i32)
 ```
 
-### 6.5 vector.store requires vector type
+### 6.5 Vector stores require vector values
 
-LDS `vector.store` requires the value to be a vector, not scalar:
+`Vector.store` requires the value to be a vector, not scalar:
 ```python
 # WRONG
-vector.store(scalar_i32, lds_ptr, [idx])
+Vec(scalar_i32).store(lds_ptr, [idx])
 
 # CORRECT
-vec = vector.from_elements(T.vec(1, T.i32), [scalar_i32])
-vector.store(vec, lds_ptr, [idx])
+vec = Vec.from_elements([scalar_i32], fx.Int32)
+vec.store(lds_ptr, [idx])
 ```
 
 ## 7. GPU Hang
 
-### 7.1 Infinite scf.for loop
+### 7.1 Infinite runtime loop
 
 If loop bounds are wrong (`stop < start` with unsigned comparison issues, or `step=0`), the GPU hangs. Verify bounds on host:
 ```python
@@ -239,8 +239,8 @@ sudo amdgpu-reset  # or reboot
 - [ ] `range_constexpr()` for all compile-time loops (not `range()`)
 - [ ] No Python `if` on runtime GPU values
 - [ ] `buffer_load` offset units match dtype (bytes/4 for i32)
-- [ ] `vector.store` uses vector type (not scalar)
-- [ ] `scf.for` state packed with `_unwrap()` (raw SSA values)
+- [ ] Vector stores use `Vector` values (not scalars)
+- [ ] `range(..., init=...)` state uses internal types, unwrapped only at hard boundaries
 - [ ] Output written to correct partition slot (`part_z`, not absolute index)
 - [ ] `exp_sums`/`max_logits` strides match actual tensor layout
 - [ ] Softmax guards against `-inf - (-inf) = NaN`

.claude/skills/flydsl-internal-types-cleanup/SKILL.md

@@ -0,0 +1,101 @@
+---
+name: flydsl-internal-types-cleanup
+description: >
+  Clean up FlyDSL kernel code by replacing direct scf/arith/vector/memref dialect calls
+  with FlyDSL internal types and helpers while preserving correctness, performance, and
+  generated ASM. Use when refactoring FlyDSL kernels, removing redundant wrappers, or
+  updating docs/skills to prefer fx.Int32/fx.Index/fx.Float32, ArithValue, Vector, and
+  range(..., init=...).
+---
+
+# FlyDSL Internal Types Cleanup
+
+## Default Rule
+
+Prefer FlyDSL internal types and high-level helpers in kernel code:
+
+- Constants and casts: `fx.Int32`, `fx.Int64`, `fx.Index`, `fx.Float32`
+- Arithmetic and comparisons: Python operators on `ArithValue` / `Numeric`
+- Runtime select: `cond.select(a, b)` or `arith.select(...)` only when a helper boundary requires it
+- Vectors: `Vector` (`Vec`) indexing, `Vec.from_elements`, `Vec.filled`, `.bitcast(...)`, `.to(...)`, `.store(...)`
+- Register memory: `fx.memref_alloca`, `fx.memref_load_vec`, `fx.memref_store_vec`
+- Runtime loops with carried state: `range(start, stop, step, init=[...])` using `fx.Index(...)` bounds
+- Compile-time loops: `range_constexpr(...)`
+
+Avoid new direct `scf.*`, `vector.*`, `memref.*`, `arith.index`, `arith.index_cast`, and `arith.trunc_f` in kernel bodies unless a lower-level boundary requires the exact op.
+
+## Replacement Map
+
+| Low-level form | Preferred form |
+|---|---|
+| `arith.constant(0, type=T.i32)` | `fx.Int32(0)` |
+| `arith.constant(0, index=True)` / `arith.index(0)` | `fx.Index(0)` |
+| `arith.constant(1.0, type=T.f32)` | `fx.Float32(1.0)` |
+| `arith.index_cast(T.i32, x)` | `fx.Int32(x)` |
+| `arith.index_cast(T.index, x)` | `fx.Index(x)` |
+| `vector.extract(v, static_position=[i], ...)` | `Vec(v)[i]` |
+| `vector.bitcast(T.vec(...), v)` | `Vec(v).bitcast(fx.Int32)` etc. |
+| `vector.from_elements(T.vec(n, T.i32), xs)` | `Vec.from_elements(xs, fx.Int32)` |
+| `vector.store(v, memref, [idx])` | `Vec(v).store(memref, [idx])` |
+| `arith.trunc_f(T.bf16x4, v)` | `Vec(v).to(fx.BFloat16)` |
+| `arith.addf/mulf` | `a + b`, `a * b` |
+| `arith.select(cond, a, b)` | `cond.select(a, b)` when `cond` is an `ArithValue` |
+
+## Important Exceptions
+
+Keep the exact lower-level op when it encodes semantics that internal types do not expose:
+
+- `llvm.InlineAsmOp` for hand-scheduled ISA snippets
+- `llvm.LoadOp` / `llvm.StoreOp` when `volatile`, `nontemporal`, address space, or alignment must be explicit
+- `arith.*FOp(..., fastmath=...)` when performance depends on fastmath flags
+- `arith.DivUIOp` / `arith.RemUIOp` for unsigned integer division/remainder
+- `rocdl.*` intrinsics and MFMA/WMMA/TDM ops
+- Backend dialect/C++ lowering docs and implementation code
+
+Do not hide these exceptions behind new helper wrappers just to remove the visible op. If exact semantics are required, keep the direct op at the boundary and document why.
+
+## Control Flow
+
+- Compile-time / constant conditions must be written as `if const_expr(condition): ...`. Do not rely on a plain Python `if` unless the condition is already a Python `bool`.
+- Use ordinary Python `if` on runtime values only when the AST rewriter keeps branch-local values and side effects correct.
+- For runtime branches inside nested helper functions, wrap the dispatch in a local `@flyc.jit` helper. This keeps branch side effects and loop-carried state in the right rewritten region.
+- For complex runtime branches with side effects, loop-carried state, or branch-local definitions, split branch bodies into local helper functions and dispatch through a local `@flyc.jit` helper. Verify correctness and ASM/perf.
+- Do not hand-write `scf.IfOp` in new kernel code unless the `@flyc.jit` helper pattern cannot express the required branch.
+- Use `range(..., init=[...])` for runtime loops with carried state; unwrap init values only if the API specifically requires raw `ir.Value`.
+
+Pattern:
+
+```python
+def _then_path():
+    ...
+
+def _else_path():
+    ...
+
+@flyc.jit
+def _dispatch():
+    if runtime_cond:
+        _then_path()
+    else:
+        _else_path()
+
+_dispatch()
+```
+
+## Verification Loop
+
+For performance-sensitive kernels:
+
+1. Record baseline shape coverage, timing, ASM hash, VGPR/SGPR counts, and spill counts.
+2. Apply one cleanup group at a time.
+3. Run correctness on small and large representative shapes.
+4. Compare performance; for strict cleanups, compare ASM hash.
+5. If performance drops or results change, revert that cleanup group and keep the lower-level op.
+
+Recommended checks:
+
+```bash
+PYTHONPATH=python:. FLYDSL_RUNTIME_ENABLE_CACHE=0 <kernel test command>
+FLYDSL_DUMP_IR=1 FLYDSL_RUNTIME_ENABLE_CACHE=0 PYTHONPATH=python:. <small compile command>
+sha256sum ~/.flydsl/debug/<kernel>/21_final_isa.s
+```