metaltile-std

MetalTile kernel standard library — benchmark metadata and type definitions. Provides the data types shared between #[bench_kernel]-annotated kernel definitions and the tile bench CLI runner. Contains no GPU runtime code.

Each #[bench_kernel] attribute (from metaltile-macros) generates an inventory::submit! { BenchSpec { ... } } alongside the kernel. The bench CLI collects all registered BenchSpec instances via inventory::iter, then runs each kernel against its MLX reference for throughput and correctness verification.

Supported Operations

See docs/KERNEL_AUDIT.md for the full per-op coverage table and docs/developing.md for how kernels are organised.

Position in the pipeline

metaltile-macros                         metaltile-cli
  (#[bench_kernel]             (tile bench collects
   generates BenchSpec)         inventory::iter::<BenchSpec>)
       │                                    │
       └────────── metaltile-std ───────────┘
                   (this crate)
                   BenchSpec · ShapeSpec · bench_types
                   runner · run_spec · stats

metaltile-std is the shared vocabulary between kernel definitions and the bench runner. It depends on the facade, core, codegen, and runtime crates to provide DType helpers, MSL generation utilities, and the inventory-based registration mechanism.

Quick start

Define a kernel with bench registration:

use metaltile::{bench_kernel, kernel};
use metaltile_std::bench_types::{FLOAT_DTYPES, OpBench};

#[bench_kernel(
    op    = "unary",
    subop = "exp",
    class = Unary,
    input = Signed,
    tol   = 1e-4,
    mlx   = "v_Exp{tn}{tn}",
    metal_file = "unary.metal",
)]
#[kernel]
pub fn mt_exp<T>(a: Tensor<T>, out: Tensor<T>) {
    let idx = program_id(0);
    store(out[idx], exp(load(a[idx])));
}

This single annotation registers the kernel for benchmarking under the "unary" group with sub-operation "exp". The tile bench CLI discovers it automatically — no manual registration needed.

Crate contents

API reference

Op catalog

Kernels live in two submodules depending on whether they have a side-by-side MLX reference:

mlx/ — MLX-compared kernels organized by category:

ffai/ — beyond-MLX kernels:

Benchmark spec reference

BenchSpec (in spec.rs) is the central registration type. Each #[bench_kernel(…)] annotation populates these fields:

ShapeSpec describes the benchmark setup:

BenchDispatch controls how the runner executes the kernel:

Dependencies

Internal

External

MSRV / platform

Rust: nightly (workspace-wide, for edition 2024). No platform gating — this crate's types compile everywhere. Benchmark execution requires macOS + Metal, but the types and BenchSpec registration compile on any host.

Feature flags

None — the crate has no Cargo features. NAX (Apple cooperative-tensor) kernels build by default; runtime gating happens via Context::chip_family().

Extending

• New MLX kernel: Create src/mlx/<name>.rs with #[bench_kernel(…)] + #[kernel] annotations. Add pub mod <name>; to src/mlx/mod.rs. The tile bench CLI discovers it automatically via inventory.

• New FFAI kernel (no MLX comparison): Create src/ffai/<name>.rs with #[bench_kernel(…)] + #[kernel] annotations. Add pub mod <name>; to src/ffai/mod.rs.

• New benchmark shape: src/spec.rs — add a ShapeSpec constant or update the relevant op file's #[bench_kernel] annotation. Common shapes use the constants at the top of spec.rs (ELEMENTWISE_N_BENCH, ROW_REDUCE_SHAPES, etc.).

• New BenchDispatch variant: src/spec.rs — add to the BenchDispatch enum. Add a match arm in src/run_spec.rs for the complex runner. Update the #[bench_kernel] proc-macro in metaltile-macros if a new ClassKind variant is needed.

• New dtype helper: src/bench_types.rs — add to dtype_label(), mlx_tname(), elem_bytes(), and dtype_tol() / dtype_tol_reduce().

• Tests to update: tile bench suite (macOS + Metal). Unit tests in src/bench_types.rs.

Related documentation

• Root README — project overview and architecture
• CONTRIBUTING — dev setup, PR process, CI
• metaltile-macros README — the #[bench_kernel] attribute that generates BenchSpec registration
• metaltile-cli README — the tile bench runner that consumes these specs

License

Licensed under the Apache License, Version 2.0.