refactor(ops): migrate core/elementwise primitives into kernels/ops/ (completes wave 1)

crates/metaltile-cli/README.md

@@ -80,12 +80,18 @@ tile inspect --help
 
 | Command | Purpose |
 |---|---|
-| `tile bench` | Run the full benchmark suite. MetalTile kernels run against MLX Metal kernel reference. Use `--filter <op>` to narrow. Outputs per-op throughput ratio and correctness. |
-| `tile build` | Compile all registered kernels to MSL and report any errors. Use `--emit msl,metallib,swift,ir,all -o <dir>` to write artifacts: `.metal` sources, `kernels.metallib`, `MetalTileKernels.swift` wrappers, and `manifest.json`. |
+| `tile bench` | Run the benchmark suite (latency, GFLOP/s, %-peak, bottleneck). `--filter <op>` to narrow; `--backend metal\|cuda\|hip\|vulkan` to pick a device; an optional metal reference runs side-by-side where a bench defines one. |
+| `tile test` | Run the `#[test_kernel]` GPU correctness suite — each kernel's output vs its CPU oracle within tolerance. `--filter` / `--backend` as above. |
+| `tile build` | Compile all registered kernels and report errors. Use `--emit msl,metallib,swift,ir,all -o <dir>` to write artifacts: `.metal` sources, `kernels.metallib`, `MetalTileKernels.swift` wrappers, and `manifest.json`. |
 | `tile inspect --kernel <name>` | Print the IR (SSA-form) and/or generated MSL for one kernel. Use `--ir` for IR only, `--msl` for MSL only. |
-| `tile device` | Show GPU device info: name, Metal feature set, supported language version, max threadgroup size. |
-| `tile snap -o <file>` | Save current benchmark results as a JSON regression baseline file. |
-| `tile diff <file>` | Compare current benchmark results to a saved baseline. Reports regressions (throughput drops below threshold). |
+| `tile device` | Show GPU device info: name, feature set, supported language version, max threadgroup size. |
+| `tile snap -o <file>` | Save current bench results as a JSON regression baseline file. |
+| `tile diff <file>` | Compare current bench results to a saved baseline. Reports regressions. |
+| `tile clean` | Remove build artifacts and cached baselines. |
+| `tile config` | Print the effective merged config (defaults → `tile.toml` → `TILE_*` env → flags). |
+| `tile init` | Scaffold a new MetalTile kernel project. |
+| `tile update` | Self-update the `tile` binary. |
+| `tile completions <shell>` | Generate shell completion scripts (bash / zsh / fish). |
 
 ### Installation
 

crates/metaltile-codegen/README.md

@@ -1,14 +1,16 @@
 # metaltile-codegen
 
-Metal Shading Language (MSL) code generator for MetalTile kernels.
-Takes algorithm IR from `metaltile-core`, applies optimization passes
-and schedule lowering, and emits valid MSL source ready for the Metal
-compiler.
-
-This crate is the middle of the MetalTile compiler stack: it receives
-`Kernel` IR nodes, lowers tile-level ops into thread-mapped, vectorized
-MSL, and exposes `MslGenerator` for both programmatic use and the
-`tile inspect` / `tile build` CLI flows.
+Multi-backend GPU code generator for MetalTile kernels. Takes algorithm IR from
+`metaltile-core`, applies backend-independent optimization passes, then lowers
+through a `CodegenBackend` to the target source. **Metal (MSL) is the default and
+most mature backend**; `backend.rs` also defines the `Target` enum
+(`Metal` / `Cuda` / `Hip` / `Spirv`), the `TargetProfile` (lane width, MMA
+strategy), and the `cuda` / `hip` / `spirv` generators.
+
+This crate is the middle of the MetalTile compiler stack: it receives `Kernel`
+IR nodes, lowers tile-level ops into thread-mapped, vectorized target source, and
+exposes the per-backend generators (`msl::MslGenerator`, `cuda::CudaGenerator`, …)
+for both programmatic use and the `tile inspect` / `tile build` CLI flows.
 
 ## Position in the pipeline
 

crates/metaltile-codegen/src/spirv/mod.rs

@@ -1057,7 +1057,7 @@ impl GlslGenerator {
                 // step on top of the hardware reciprocal gives correctly-
                 // rounded f32 divide assuming hardware fma is correctly
                 // rounded (true on AMD GFX1xxx V_FMAC_F32). Without this
-                // `test_logits_repetition_penalty` misses bit-exact tol=0
+                // `test_mt_logits_repetition_penalty` misses bit-exact tol=0
                 // by 1 ULP on 2 elements out of 8192.
                 if matches!(bop, BinOpKind::Div) && ty == "float" {
                     writeln!(out, "{pad}{ty} {v} = mt_fdiv({l}, {r});").ok();

crates/metaltile-core/README.md

@@ -67,7 +67,7 @@ kernel.body = block;
 | `shape` | `Shape`, `Dim`, `DimExpr`, `tile()` constructor |
 | `constexpr` | `ConstExpr` — symbolic constants resolved at kernel compile time |
 | `error` | `Error` enum and `Result<T>` alias |
-| `gpu_family` | `GpuFamily` — Apple GPU generation detection (M1/M2/M3+/M5) |
+| `protocol` | The `runner ↔ CLI` wire types (`ProtocolMessage`, `ProfileInfo`, bench/test/build results) for the `__tile_runner` subprocess, plus `GpuFamily` Apple-GPU generation detection (M1–M5) |
 | `kernel_registry` | `KernelEntry` — registry for kernel discovery via `inventory` |
 | `utils` | Internal helpers (bit manipulation, alignment) |
 
@@ -97,7 +97,8 @@ kernel.body = block;
 | `Dim` | A single dimension: `Known(usize)`, `ConstExpr(name)`, or `Any` | `src/shape.rs` | ✅ |
 | `DimExpr` | Symbolic dimension expression (Scale, Const, Var, Add, Range) | `src/shape.rs` | ✅ |
 | `ConstExpr` | Named compile-time constant used in shapes and kernel configs | `src/constexpr.rs` | ✅ |
-| `GpuFamily` | Apple GPU family level (7=M1, 8=M2, 9=M3/M4, 10=M5) | `src/gpu_family.rs` | ✅ |
+| `GpuFamily` | Apple GPU family level (7=M1, 8=M2, 9=M3/M4, 10=M5) | `src/protocol.rs` | ✅ |
+| `ProtocolMessage` | JSON-line wire format for the `__tile_runner` subprocess (bench/test/build/inspect results, `ProfileInfo`) | `src/protocol.rs` | ✅ |
 | `KernelEntry` | Inventory-registered kernel metadata | `src/kernel_registry.rs` | ✅ |
 | `Error` / `Result<T>` | Error enum and result alias | `src/error.rs` | ✅ |
 

crates/metaltile-core/src/protocol.rs

@@ -323,7 +323,7 @@ mod tests {
         assert_eq!(op_group("mt_softmax_f32"), "softmax");
         assert_eq!(op_group("mt_softmax_bf16"), "softmax");
         assert_eq!(op_group("softmax"), "softmax");
-        assert_eq!(op_group("ffai_vector_add_f16"), "ffai_vector_add");
+        assert_eq!(op_group("mt_vector_add_f16"), "mt_vector_add");
     }
 
     #[test]

crates/metaltile-macros/README.md

@@ -7,8 +7,9 @@ code.
 
 This crate is the front door of the MetalTile compiler: user-written
 `#[kernel]` functions enter here, and IR + dispatch surfaces exit. It
-also provides `shape!`/`tile!` constructors for shape annotations and
-`#[kernel(bench(...))]` for declarative benchmark registration.
+also provides `shape!`/`tile!` constructors for shape annotations,
+`#[kernel(variants(...))]` for compile-time specialisation, and the `#[bench]` /
+`#[test_kernel]` attributes for declarative bench / correctness-test registration.
 
 ## Position in the pipeline
 
@@ -66,10 +67,11 @@ pub fn scale<T>(a: Tensor<T>, factor: f32, out: Tensor<T>) {
 | Module | Purpose |
 |---|---|
 | `lib.rs` | All proc-macro entry points: `#[kernel]`, `#[constexpr]`, `#[scalar]`, `#[strided]`, `shape!`, `tile!`, `ValueRefs` / `OpFlags` derive macros |
-| `body_parser.rs` | `DslBodyParser` — walks `syn::Expr` trees and translates DSL calls into IR-building token streams |
-| `sig_parser.rs` | Signature parsing: `parse_kernel_params_generic`, `extract_constexprs_typed`, `extract_param_names` |
-| `bench_impl.rs` | `BenchArgs` parsing, `ClassKind` enum, `generate_submit` — used by `#[kernel(bench(...))]` |
-| `derive_op.rs` | Derive macros: `ValueRefs` (value-id traversal) and `OpFlags` (elementwise/side-effect/etc. predicates) |
+| `kernel/body.rs` | `DslBodyParser` — walks `syn::Expr` trees and translates DSL calls into IR-building token streams |
+| `kernel/sig.rs` | Signature parsing: `parse_kernel_params_generic`, `extract_constexprs_typed`, `extract_param_names` |
+| `kernel/variants.rs` | `#[kernel(variants(...))]` compile-time specialisation — stamps one kernel per parameter tuple (`VariantsSpec`, `substitute_fn`) |
+| `bench.rs` / `test.rs` | The `#[bench]` and `#[test_kernel]` attributes that register a bench/test setup into the `inventory` |
+| `derive/mod.rs` | Derive macros: `ValueRefs` (value-id traversal) and `OpFlags` (elementwise/side-effect/etc. predicates) |
 
 ## API reference
 
@@ -79,7 +81,8 @@ pub fn scale<T>(a: Tensor<T>, factor: f32, out: Tensor<T>) {
 |---|---|---|
 | `#[kernel]` | attribute | Parses a Rust function into IR + generates a module with `kernel_ir`, `kernel_ir_for`, `LaunchBuilder`, and `launch()` |
 | `#[autotune]` | attribute | Placed before `#[kernel]` to enable autotuning: `#[autotune(configs = [...], key = [M, N, K])]`. **Not yet implemented** — `AutotuneArgs` struct exists but parsing is a TODO in `expand_kernel`. |
-| `#[kernel(bench(...))]` | attribute | Registers a kernel for automatic benchmarking via `inventory::submit!`. Pass args inside `bench(...)` |
+| `#[kernel(variants(...))]` | attribute | Compile-time specialisation — stamps one kernel per parameter tuple (`variants(BITS = [2,4,8], suffix = "int{BITS}")`), constant-folding the values into the body. See the [Kernel Style Guide](../../docs/STYLE_GUIDE.md). |
+| `#[bench]` / `#[test_kernel]` | attribute | Register a kernel's throughput bench / GPU correctness test (and its setup callback) into the `inventory`. `#[bench(dtypes = [...])]`, `#[test_kernel(dtypes = [...], tol = [...])]`; both accept the same `variants(...)` syntax. |
 | `#[constexpr]` | attribute | Pass-through: marks a function parameter as a compile-time constant detected by `#[kernel]` |
 | `#[scalar]` | attribute | Pass-through: marks a `Tensor` parameter for `constant T&` lowering in MSL |
 | `#[strided]` | attribute | Pass-through: marks a `Tensor` parameter for strided lowering (shape + stride arrays emitted) |
@@ -137,29 +140,25 @@ Attributes placed on the function itself (before or alongside `#[kernel]`):
 | `#[scalar]` | Emits the parameter as `constant T& name` in MSL rather than `device T*`. Used for scalar values like `eps` or `scale`. |
 | `#[strided]` | Emits the parameter as `device T*` plus `constant uint* name_shape` and `constant uint* name_strides` in MSL. Used for non-contiguous tensor views. |
 
-### `#[kernel(bench(...))]` arguments
+### `#[bench]` / `#[test_kernel]`
 
-| Argument | Required | Purpose |
-|---|---|---|
-| `op` | yes | Bench table group, e.g. `"unary"`, `"binary"` |
-| `subop` | yes | Sub-operation label, e.g. `"exp"`, `"add"` |
-| `class` | yes | Dispatch class: `Unary`, `Binary`, `AllReduce`, `RowReduce`, `Arange`, `BinaryTwo`, `Select`, `RowNorm`, `Sort`, `Scan`, `ArgReduce`, `Random`, `FpQuantized`, `MatVec`, `MatVecMasked`, `QuantizedMatVec`, `Rope`, `Attention`, `StridedCopy` |
-| `tol` | yes | Maximum absolute correctness error, e.g. `1e-4` |
-| `input` | no | Input buffer init for unary: `Signed`, `Positive`, `Half`, `Unit` (default: `Half`) |
-| `input_a` / `input_b` | no | Input buffer init for binary (default: `Half`) |
-| `mlx` | no | MLX kernel name pattern; `{tn}` is replaced with the MLX type name |
-| `metal_file` | no | MLX reference .metal source path (loaded via `include_str!`) |
-| `dtypes` | no | `&'static [DType]` slice (default: `FLOAT_DTYPES`) |
-| `shapes` | no | Custom `ShapeSpec` array for complex dispatch shapes |
-| `start` / `step` | no | Arange start/step values (float literals) |
-| `reads` | no | Read count for bandwidth calculation (`RowNorm` class) |
-| `out_elements` | no | Output element count (`RowNorm` class; 1 = per-row scalar, >1 = full B×N) |
-| `tpg` | no | Threads per threadgroup override |
-| `pre_weight` / `pre_bias` / `post_eps` | no | RowNorm-specific: weight, bias, epsilon values |
-| `n` / `check_n` / `b` | no | Shape dimensions for complex dispatch (`RowNorm`, `Rope`) |
-| `h` / `l` / `d` / `n_per_group` | no | Rope-specific dimensions |
-| `group_size` | no | Quantization group size (`QuantizedMatVec` class) |
-| `m` / `pad` | no | StridedCopy dimensions |
+Benches and correctness tests are declared **next to the kernel** (in
+`kernel_benches` / `kernel_tests` modules) rather than via attribute arguments.
+Each attribute wraps a setup function that returns a builder:
+
+```rust,ignore
+#[bench(dtypes = [f32, f16, bf16])]
+fn bench_mt_scale(dt: DType) -> BenchSetup { /* BenchSetup::new(...).buffer(...).bytes_moved(...) */ }
+
+#[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-6, 1e-3, 1e-3])]
+fn test_mt_scale(dt: DType) -> TestSetup { /* TestSetup::new(...).input(...).expect(...) */ }
+```
+
+`dtypes` selects the float types to run; `tol` is the per-dtype absolute
+tolerance; both accept the same `variants(...)` syntax as `#[kernel]`. The bench
+name is the function name (there is no `name` key). A bench may attach an
+optional metal reference via `.with_reference(...)` on the `BenchSetup`. See the
+[Kernel Style Guide](../../docs/STYLE_GUIDE.md) for the full builder surface.
 
 ## Dependencies
 
@@ -185,7 +184,7 @@ Requires `[lib] proc-macro = true` in `Cargo.toml`.
 
 ## Extending
 
-- **New DSL intrinsic:** `src/body_parser.rs` — add a recognized function name to the
+- **New DSL intrinsic:** `src/kernel/body.rs` — add a recognized function name to the
   expression walker. Update the `Recognized call:` list in the module doc comment.
 
 - **New kernel parameter attribute:** `src/lib.rs` — add a new `#[proc_macro_attribute]`
@@ -196,10 +195,10 @@ Requires `[lib] proc-macro = true` in `Cargo.toml`.
   `#[proc_macro_attribute]` pass-through, parse its args in `expand_kernel`,
   and emit the corresponding token stream into the generated module.
 
-- **New bench class:** `src/bench_impl.rs` — add variant to `ClassKind` enum,
+- **New bench class:** `src/bench.rs` — add variant to `ClassKind` enum,
   add a match arm in `generate_submit` with its `ShapeSpec` and `BenchDispatch` variant.
 
-- **New bench argument:** `src/bench_impl.rs` — add field to `BenchArgs`, add parse
+- **New bench argument:** `src/bench.rs` — add field to `BenchArgs`, add parse
   arm in `BenchArgs::parse()`, consume in `generate_submit`.
 
 - **New shape/tile constructor syntax:** `src/lib.rs` — add a new `#[proc_macro]`

crates/metaltile-runtime/README.md

@@ -1,12 +1,16 @@
 # metaltile-runtime
 
-Apple Metal runtime dispatch for MetalTile GPU kernels.
-Manages Metal devices, compiles generated MSL into pipeline state objects,
-dispatches compute kernels, and returns output buffers to the host.
-
-This crate is the bottom of the MetalTile stack — it is the only crate
-that links against Apple's Metal framework, and all kernel execution
-ultimately flows through its `Context` type.
+GPU runtime dispatch for MetalTile kernels across backends. Compiles generated
+target source into pipeline state objects / modules, dispatches compute kernels,
+and returns output buffers to the host. **Metal (Apple) is the default**;
+`device/{cuda,hip,vulkan}/` add NVIDIA / AMD / portable devices behind the
+`cuda` / `hip` / `vulkan` Cargo features.
+
+This crate is the bottom of the MetalTile stack. It owns the device abstraction
+(`device/` — `metal_device.rs` plus the feature-gated cuda/hip/vulkan devices),
+the dispatch strategies (`dispatch/` — single, chained, buffer-plan, validate),
+and the compilation / PSO caches (`cache/`). Kernel execution flows through its
+per-backend `Context` / device types.
 
 ## Position in the pipeline