rlx-ir

Tensor IR for the RLX ML compiler — standalone, serializable, optimizable.

The leaf of the workspace dep graph — adding deps here ripples through every backend. rlx-ir owns:

• DType — F32, F16, BF16, F64, I8, I16, I32, I64, U8, U32, Bool, C64. Promotion rules.
• Shape — static or dynamic dimensions, dtype-aware. No symbolic dims yet.
• Op — the IR vocabulary. JAX-shaped: MatMul, DotGeneral, Scan, SelectiveScan, DenseSolve, Fft, LayerNorm, Attention (with MaskKind), Quantize / Dequantize / FakeQuantize, If / While, Custom / CustomFn extension points, plus the full reverse-mode AD op set (*Backward variants). Adding an op means: this file, infer.rs, graph.rs builder, both backends' thunks, MPSGraph lowering, fusion patterns, cost model. ~6 files per new op.
• MaskKind — attention mask kind (None / Causal / SlidingWindow / Custom). Added Apr 2026 per plan #20.
• Graph — DAG builder + topo iterator.
• infer.rs — shape inference per op. Drives fusion and arena planning.
• verify.rs — graph-level invariant checks. Run after every optimizer pass in debug builds (plan #50).
• OpExtension / op_registry — the custom-op scaffold downstream crates register against (see [rlx-sparse], [rlx-linalg]).
• Logical kernels — one OpKind, native or common-IR lowering; see README-logical-kernels.md.
• Tick — CNTVCT_EL0-direct cycle counter for sub-millisecond measurement on Apple Silicon. Falls back to Instant. Use Tick::now() / tick.elapsed_ns(start) for any sub-ms measurement.

Install

[dependencies]
rlx-ir = "0.1"

Most users should depend on the rlx prelude crate instead, which re-exports rlx_ir as rlx::ir.

IR levels (HIR / MIR / LIR)

[GraphModule] tracks which stage you are in. Use [Graph::define] for fusion-first HIR builders; use [GraphModule::mir] for primitive MIR.

Pipeline: CompilePipeline::compile_module(module) in rlx_opt.

Quickstart

use rlx_ir::{DType, Graph, GraphModule, Shape};

// Fusion-first (recommended)
let module = Graph::define("layer", |m| {
    let x = m.input("x", Shape::new(&[2, 128], DType::F32));
    let w = m.param("w", Shape::new(&[128, 128], DType::F32));
    m.linear(x, w, None, None, Shape::new(&[2, 128], DType::F32))
});

// Primitive MIR (legacy)
let mut g = Graph::new("hello");
let x = g.input("x", Shape::new(&[1, 4], DType::F32));
let w = g.param("w", Shape::new(&[4, 2], DType::F32));
let y = g.matmul(x, w, Shape::new(&[1, 2], DType::F32));
g.set_outputs(vec![y]);
let module = g.module();

Build / test

cargo build -p rlx-ir
cargo test  -p rlx-ir       # 11 unit tests

Status

The Op enum is stable in shape but not in surface — minor 0.x bumps may add variants or refine attributes. Pin exact versions in production until 1.0.

QuantScheme and Op::DequantMatMul

rlx_ir::quant::QuantScheme describes how a packed weight tensor is laid out so the dequant kernel can decode it without a side lookup. Variants:

Op::DequantMatMul { scheme } takes 4 inputs for the legacy Int8 schemes (x, w_q, scale, zp) or 2 for the GGUF schemes (x, packed_w_bytes) — num_inputs() switches on scheme.is_gguf().

The CPU backend handles all listed schemes today; Metal lowering for the GGUF schemes is on the roadmap (Op::DequantMatMul falls through to the per-op thunk path, which dequants the weight to F32 scratch once before matmul — correct but doesn't keep packed bytes on the GPU).

Gotchas

• Op::Attention input count is variable: MaskKind::Custom → 4 inputs (Q, K, V, mask); other kinds → 3. num_inputs() knows.
• Op::DequantMatMul input count is variable too: 4 for legacy Int8, 2 for GGUF K-quant schemes (see table above).
• Op is not Copy (some variants own Vec/Box<Graph>). Match destructures must use .. for non-Copy fields.
• MaskKind IS Copy; *mask_kind is fine inside a match arm bound to &MaskKind.
• Don't add backend-flavored types here — keep rlx-ir portable.

License

GPL-3.0-only.