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.
fft — shared FftNorm, GPU launch plans, and NumPy-shaped helpers (fftfreq, gpu_fft_native_eligible, …). Graph builders in ops::fft_ops compose rfft, irfft, stft, psd_real, and fft_conv1d from primitive ops.
MaskKind — attention mask kind (None / Causal / SlidingWindow / Custom). Added Apr 2026 per plan #20.
attention_layout — BHSD [B,H,S,D] vs BSHD [B,S,H,D] (EEG-DINO), launch strides, packed [B,S,3,H,D] QKV detection, and shared flash-vs-row dispatch (attention_dispatch_use_row, ATTENTION_FLASH_MAX_HEAD_DIM).
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.2"

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

IR levels (HIR / MIR / LIR)

Level	Module	Role
HIR	`rlx_ir::hir`	Block ops for model builders (`SwiGLU`, `Linear`, …)
MIR	`rlx_ir::mir`	Fused tensor DAG — optimizer input (`Graph` alias)
LIR	`rlx_ir::lir`	MIR + arena buffer plan — backend input

[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.

Family	Schemes
Linear int	`Int8Block`, `Int8BlockAsym`, `Int4Block`
FP8 / FP4	`Fp8E4m3`, `Fp8E5m2`, `Nvfp4Block`
GGUF K-quant	`GgufQ2K` … `GgufQ8K` (256-element super-blocks)
GGUF legacy	`GgufQ4_0`, `GgufQ4_1`, `GgufQ5_0`, `GgufQ5_1`, `GgufQ8_0`
GGUF IQ / ternary / micro	`GgufIQ`, `GgufTQ`, `GgufMXFP4`, `GgufNVFP4`

Op::DequantMatMul { scheme } takes 4 inputs for legacy int/FP8 (x, w_q, scale, zp) or 2 for GGUF (x, packed_w) — see scheme.is_gguf() and Graph::dequant_matmul_packed.

Backend paths (GPU scheme ids, fused GEMV, ANE constexpr, TPU compile-time bake): docs/gguf-backend-paths.md.

The CPU backend dequants every scheme; GPU backends use native dequant_gguf kernels where listed in that doc. Metal additionally exposes fused decode GEMV for Q4_K / Q4_0 / Q8_0 when m == 1.

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.

rlx-ir 0.2.11