refactor(ops): migrate core/elementwise primitives into kernels/ops/ (completes wave 1) - #18
Conversation
d078f11 to
1fa8029
Compare
|
Wave 1 complete and the increment is the cleanest of the stack: One stack-management note for all four PRs: they all base on |
1fa8029 to
78fc51f
Compare
Agreed. We'll just merge in order. Was just rebasing and ensuring builds are all green. Working on wave 2 and 3 families. Will follow similar PR structure. |
|
@ekryski clean; |
Completes wave 1: moves the core/elementwise primitive kernels out of the
legacy mlx/ + ffai/ split into kernels/ops/ per the consolidation plan.
- mlx/{arange,arg_reduce,binary,binary_two,copy,gather_axis,hadamard,
hadamard_m,indexing,logsumexp,random,reduce,scan,scatter_axis,strided,
ternary,unary} -> kernels/ops/ (already mt_*; inventory-registered, so no
import sites to update beyond one snapshot test)
- ffai/{axpy_scalar_inplace,clamp_scalar,gather} + mt_vector_add.rs (renamed
vector_add.rs) -> kernels/ops/ with ffai_ -> mt_ prefixes
- delete ffai/arg_reduce.rs: ffai_argmax was a byte-for-byte duplicate of
mt_argmax (verified) -> real dedup
- fence.rs moves but stays intentionally undeclared (device/system fences have
no DSL representation — the '1 out of scope' op)
Updated the steel_msl_snapshots import, the op_group/name-filter test fixtures,
the plan (ops done, wave 1 complete) and KERNEL_AUDIT rows.
cargo build + workspace clippy clean; all ops kernel tests + steel_msl_snapshots
pass.
…probe
Finishes rehoming the last mlx/ stragglers:
- gated activations: fold mlx/{swiglu,clamped_swiglu,fused_gate_activation}.rs
(4 kernels) into one kernels/ops/gated_activation.rs. Names unchanged
(mt_swiglu, mt_clamped_swiglu, mt_fused_gate_gelu,
mt_fused_gate_clipped_swiglu); ops/ is the activation home (mt_silu, which
mt_swiglu cross-kernel-calls, lives in ops/unary.rs). 3 files -> 1.
- mlx/sgload_smoke.rs -> probe/simdgroup_load_probe.rs; mt_sgload_smoke ->
mt_simdgroup_load_probe. It's a HW-intrinsic validation kernel (sibling of
probe/mma_layout_probe + mpp_matmul_smoke), so probe/ is its real home, not
ops/. Renamed for what it does (validates the Op::SimdgroupLoad codegen path
via a byte-exact round-trip). Test file renamed to match.
Updated KERNEL_AUDIT (swiglu/fused rows) + the plan's ops listing. cargo build
+ clippy clean; gated-activation kernel tests + the probe integration test pass.
…t the probe convention Makes the probe/ module consistent: all three validation kernels now use the _probe suffix (mma_layout_probe, simdgroup_load_probe, mpp_matmul_probe). - probe/mpp_matmul_smoke.rs -> mpp_matmul_probe.rs; mt_mpp_matmul_smoke -> mt_mpp_matmul_probe; integration test tests/mpp_matmul_smoke.rs renamed to match. - STYLE_GUIDE §1: document the convention — a kernel that validates a codegen path / HW intrinsic is a 'probe' (mt_<thing>_probe, in probe/, outside kernels/<family>/). '_smoke' is reserved for *test* code (backend bring-up tests, codegen-sanity #[test]s), which stays test-side. - probe/mod.rs: add a module doc stating the same. No other kernels use the _smoke naming (the remaining *_smoke are legit tests: cuda/hip/vulkan backend bring-up + codegen-sanity test fns — left as-is). cargo build + clippy clean; probe integration + codegen unit tests pass.
fa62466 to
0c1a070
Compare
2 participants
What
Final wave-1 kernel-family port: moves the core / elementwise primitive
kernels out of the legacy
mlx/+ffai/split intokernels/ops/perdocs/specs/KERNEL_CONSOLIDATION_PLAN.md. With this, wave 1 is complete(
rope·norm·sampling·ops).Changes
mlx/{arange,arg_reduce,binary,binary_two,copy,gather_axis,hadamard,hadamard_m, indexing,logsumexp,random,reduce,scan,scatter_axis,strided,ternary,unary}→kernels/ops/(alreadymt_*; these are inventory-registered, so there wereno path-import sites to fix beyond one snapshot test).
ffai/{axpy_scalar_inplace,clamp_scalar,gather}+mt_vector_add.rs(→vector_add.rs) →kernels/ops/withffai_→mt_prefixes.ffai/arg_reduce.rs—ffai_argmaxwas a byte-for-byteduplicate of
mt_argmax(verified line-by-line), so it's dropped, not moved.fence.rsmoves but stays intentionally undeclared — device/system memoryfences and atomics have no
#[kernel]DSL representation (the "1 out of scope"op in the audit); the file is kept for its implementation notes.
steel_msl_snapshotsimport, theop_group/ name-filter unit-testfixtures, the plan (ops done, wave 1 complete), and the
KERNEL_AUDITrows.Generated MSL per kernel is unchanged — only module paths and inventory names.
Verification
cargo build -p metaltile-std+cargo clippy --workspace --all-targetsclean.tile test(binary, unary, ternary/select, copy,strided, arange, random, reduce, argmax/argmin, scan, indexing, gather(_axis),
scatter_axis, hadamard(_m), logsumexp, clamp, axpy, vector_add).
steel_msl_snapshotsintegration test passes.Note on stacking
Stacked on sampling (#17) → norm (#16) → conv+rope (#15) → docs (#14); the
diff includes those until they merge. Each family lands as its own PR.