rlx_cpu/thunk/types.rs
1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4#![allow(unsafe_op_in_unsafe_fn)]
5use crate::thunk::*;
6
7/// Instrumentation: count of `FusedNomicLayer` thunks emitted by the full-layer
8/// fusion in this process. The fusion is otherwise invisible through the
9/// `Session` API, so out-of-tree model tests (rlx-models' Nomic parity test)
10/// load this to confirm the fusion actually fired (non-vacuous parity). Reset
11/// by storing 0 before a compile.
12pub static FUSED_NOMIC_LAYER_COUNT: std::sync::atomic::AtomicU64 =
13 std::sync::atomic::AtomicU64::new(0);
14
15/// A pre-compiled kernel call with all args resolved to arena offsets.
16///
17/// `IntoStaticStr` gives every variant a zero-alloc `&'static str` name (its
18/// identifier) via `<&'static str>::from(&thunk)`, used by the profiler and
19/// tracer (`thunk_kind_name`). Deriving it means a newly added variant is
20/// labeled automatically instead of silently landing in an `"Other"` bucket.
21#[derive(Clone, strum_macros::IntoStaticStr)]
22pub enum Thunk {
23 /// Skip (Input/Param already in arena)
24 Nop,
25 /// Fused element-wise region: ONE pass applies the whole `chain` of scalar
26 /// steps per output element (bias-add → relu → … collapse into a single
27 /// kernel, no intermediate tensors). This is RLX's XLA-style auto-fusion
28 /// path on CPU — the GPU backends already run the same `ElementwiseRegion`
29 /// IR via a generated kernel; this is the scalar interpreter twin so any
30 /// fused chain runs without a hand-written per-op kernel.
31 /// `dst`/`input_offs` are byte offsets; `input_modulus[i]` is input i's
32 /// broadcast period (0 ⇒ read element `gid`; scalar inputs read element 0).
33 ElementwiseRegion {
34 dst: usize,
35 len: u32,
36 input_offs: Vec<usize>,
37 chain: Vec<rlx_ir::op::ChainStep>,
38 scalar_input_mask: u32,
39 input_modulus: [u32; 16],
40 },
41 /// C = A @ B (BLAS sgemm)
42 Sgemm {
43 a: usize,
44 b: usize,
45 c: usize,
46 m: u32,
47 k: u32,
48 n: u32,
49 },
50 /// `C[m,n] = A[m,k](f32) @ B[k,n](BF16)` — dequant-on-the-fly GEMM with a
51 /// BF16 right-hand (half the weight memory traffic; used for a bf16-resident
52 /// LM head). `b` is a BF16 arena buffer (`k*n` `u16`), `a`/`c` are f32.
53 SgemmBf16 {
54 a: usize,
55 b: usize,
56 c: usize,
57 m: u32,
58 k: u32,
59 n: u32,
60 },
61 /// `C[m,n] = A[m,k](f32) @ B[k,n](F16)` — IEEE-half twin of [`Self::SgemmBf16`].
62 /// `b` is an F16 arena buffer (`k*n` `u16`), `a`/`c` are f32. Without this a
63 /// generic `Op::MatMul` with an F16 weight read the half bytes as f32 garbage.
64 SgemmF16 {
65 a: usize,
66 b: usize,
67 c: usize,
68 m: u32,
69 k: u32,
70 n: u32,
71 },
72 /// `C[m,n] = op(A) @ op(B)` with optional transpose of each operand,
73 /// done via cblas trans flags (no materialized transpose). Emitted when
74 /// the thunk compiler folds a last-two-axis `Op::Transpose` feeding a
75 /// matmul (the dominant cost in matmul backprop). `a`/`b` are the
76 /// *pre-transpose* operands. lda/ldb are derived from the trans flags.
77 SgemmT {
78 a: usize,
79 b: usize,
80 c: usize,
81 m: u32,
82 k: u32,
83 n: u32,
84 ta: bool,
85 tb: bool,
86 },
87 /// Fused SGD-with-momentum parameter update, one pass per param.
88 /// Computes `v' = mom·v + g ; p' = p − lr·v'` reading `param`/`vel`/
89 /// `grad` and writing `p_out`/`v_out`. Emitted when the thunk compiler
90 /// folds the in-graph SGD chain `Sub(p, Mul(Add(Mul(v,mom),g), lr))`
91 /// (4 `BinaryFull` ops + 2 full-size constants per param) into a single
92 /// kernel — the dominant cost of the fully-fused training step. All
93 /// fields are arena byte offsets except the scalar hyperparameters.
94 SgdMomentum {
95 param: usize,
96 vel: usize,
97 grad: usize,
98 p_out: usize,
99 v_out: usize,
100 lr: f32,
101 mom: f32,
102 len: u32,
103 },
104 /// Complex (C64) dense GEMM `C = A·B`. Operands are interleaved
105 /// `[re, im]` f32; `a`/`b`/`c` are byte offsets, `m`/`k`/`n` are
106 /// complex-element matrix dims (`A` `[m,k]`, `B` `[k,n]`, `C` `[m,n]`).
107 CgemmC64 {
108 a: usize,
109 b: usize,
110 c: usize,
111 m: u32,
112 k: u32,
113 n: u32,
114 },
115 /// f64 dense solve `x = A⁻¹·b` via LAPACK dgesv.
116 /// `a`, `b`, `x` are byte-offsets into the arena. `n` is the matrix
117 /// dimension; `nrhs` is 1 for a vector RHS or >1 for multi-RHS.
118 /// The kernel materializes scratch copies of A and b internally
119 /// (LAPACK overwrites both with LU factors and solution).
120 DenseSolveF64 {
121 a: usize,
122 b: usize,
123 x: usize,
124 n: u32,
125 nrhs: u32,
126 },
127 /// f32 twin of `DenseSolveF64`. Calls LAPACK `sgesv` (or the
128 /// no-blas Rust fallback). Same arena byte-offset contract.
129 DenseSolveF32 {
130 a: usize,
131 b: usize,
132 x: usize,
133 n: u32,
134 nrhs: u32,
135 },
136 /// Cholesky factorization `A = L·Lᵀ`. `a` input `[n,n]`, `l` output `[n,n]`
137 /// lower-triangular (upper zeroed). f32 arena; promoted to f64 for LAPACK
138 /// `dpotrf`.
139 CholeskyF32 {
140 a: usize,
141 l: usize,
142 n: u32,
143 },
144 /// Triangular solve `op(A)·X = B` (f32 arena; promoted to f64 for BLAS
145 /// `trsm`). `a` `[n,n]`, `b`/`x` `[n,nrhs]`.
146 TriangularSolveF32 {
147 a: usize,
148 b: usize,
149 x: usize,
150 n: u32,
151 nrhs: u32,
152 lower: bool,
153 transpose: bool,
154 },
155 /// Determinant (`log_abs=false`) or `log|det|` (`log_abs=true`) of `a`
156 /// `[n,n]` → scalar `out`. f32 arena; promoted to f64 for LAPACK LU.
157 DetF32 {
158 a: usize,
159 out: usize,
160 n: u32,
161 log_abs: bool,
162 },
163 /// Sort (`arg=false`, output values) or argsort (`arg=true`, output f32
164 /// indices) along an axis, viewed as `outer × axis_dim × inner`.
165 SortF32 {
166 src: usize,
167 dst: usize,
168 outer: u32,
169 axis_dim: u32,
170 inner: u32,
171 descending: bool,
172 arg: bool,
173 },
174 /// One factor of the thin SVD of `a` `[m,n]`. `part`: 0=U `[m,k]`, 1=S `[k]`,
175 /// 2=Vt `[k,n]` (`k=min(m,n)`). f32 arena; promoted to f64 for LAPACK `gesdd`.
176 SvdF32 {
177 a: usize,
178 out: usize,
179 m: u32,
180 n: u32,
181 part: u8,
182 },
183 /// One factor of the thin QR of `a` `[m,n]`. `part`: 0=Q `[m,k]`, 1=R `[k,n]`
184 /// (`k=min(m,n)`). f32 arena; promoted to f64 for LAPACK `geqrf`/`orgqr`.
185 QrF32 {
186 a: usize,
187 out: usize,
188 m: u32,
189 n: u32,
190 part: u8,
191 },
192 /// Batched f64 dense solve. `a`, `b`, `x` are byte-offsets to
193 /// the leading slice; `batch` is the number of independent
194 /// systems. Per slice the kernel calls `dgesv(A_i, b_i, n, nrhs)`
195 /// — LAPACK has no batched dgesv on Accelerate, so we loop.
196 BatchedDenseSolveF64 {
197 a: usize,
198 b: usize,
199 x: usize,
200 batch: u32,
201 n: u32,
202 nrhs: u32,
203 },
204 /// Batched f32 dense solve — loop of `sgesv` per batch slice.
205 BatchedDenseSolveF32 {
206 a: usize,
207 b: usize,
208 x: usize,
209 batch: u32,
210 n: u32,
211 nrhs: u32,
212 },
213 /// Batched f64 matmul. Both inputs and output have a leading
214 /// batch axis of size `batch`. Per-batch independent dgemm:
215 /// `C[i] = A[i] @ B[i]` for `i in 0..batch`. Used by VJP rules
216 /// that emit per-batch outer products (e.g., BatchedDenseSolve
217 /// VJP). The unbatched `Dgemm` thunk handles the rank-2 case.
218 BatchedDgemmF64 {
219 a: usize,
220 b: usize,
221 c: usize,
222 batch: u32,
223 m: u32,
224 k: u32,
225 n: u32,
226 },
227 /// Batched f32 matmul — same loop-per-batch shape as
228 /// `BatchedDgemmF64` but calling `sgemm`. Needed for attention
229 /// patterns where both operands carry a batch dim (e.g. q@k^T
230 /// and attn@v in decomposed self-attention). The 2-D `Sgemm`
231 /// flatten trick is wrong in that case because it treats `b` as
232 /// a single shared RHS across every batch.
233 BatchedSgemm {
234 a: usize,
235 b: usize,
236 c: usize,
237 batch: u32,
238 m: u32,
239 k: u32,
240 n: u32,
241 /// `true` when operand `a`/`b` has batch dim 1 and is BROADCAST across
242 /// every output batch (its per-matrix stride is 0 — reuse matrix 0).
243 /// Without this the loop over-reads a batch-1 operand for `bi>0`.
244 a_bcast: bool,
245 b_bcast: bool,
246 },
247 /// C = A @ B via Accelerate cblas_dgemm. Mirror of `Sgemm` at f64.
248 Dgemm {
249 a: usize,
250 b: usize,
251 c: usize,
252 m: u32,
253 k: u32,
254 n: u32,
255 },
256 /// f64 N-D index walk used for both `Op::Transpose` and `Op::Expand`.
257 /// `in_strides` carries 0s on broadcast axes (Expand) or permuted
258 /// strides (Transpose). Mirror of `Thunk::Transpose` at f64.
259 TransposeF64 {
260 src: usize,
261 dst: usize,
262 in_total: u32,
263 out_dims: Vec<u32>,
264 in_strides: Vec<u32>,
265 },
266 /// f64 element-wise activation. Single-input, single-output. The
267 /// kernel always reads from `src` and writes to `dst`, so it works
268 /// whether or not the planner aliased the two slots.
269 ActivationF64 {
270 src: usize,
271 dst: usize,
272 len: u32,
273 kind: Activation,
274 },
275 /// Element-wise complex squared-magnitude: `|z|² = re² + im²`.
276 /// Reads the C64 input at `src` as `2·len` f32 ([re,im] pairs),
277 /// writes `len` f32 to `dst`.
278 ComplexNormSqF32 {
279 src: usize,
280 dst: usize,
281 /// Logical element count (number of complex values).
282 len: u32,
283 },
284 /// Wirtinger backward for [`ComplexNormSqF32`]: `dz = g · z` as
285 /// C64. Reads `z` at `2·len` f32 + `g` at `len` f32; writes
286 /// `2·len` f32 to `dz`.
287 ComplexNormSqBackwardF32 {
288 z: usize,
289 g: usize,
290 dz: usize,
291 len: u32,
292 },
293 /// Element-wise C64 conjugate: writes `[re_i, -im_i]` per element.
294 /// Layout matches the rest of C64 here ([re,im] interleaved f32).
295 ConjugateC64 {
296 src: usize,
297 dst: usize,
298 len: u32,
299 },
300 /// C64 element-wise activation. Only kinds with well-defined
301 /// complex extensions are supported: Neg, Exp, Log, Sqrt.
302 /// Everything else (Sigmoid, Tanh, Relu, Abs, Sin/Cos/Tan/Atan,
303 /// Round, GeLU family) is rejected at lowering — those don't have
304 /// single natural complex definitions. `len` is the **complex
305 /// element count** (the f32 buffer holds `2·len` floats).
306 ActivationC64 {
307 src: usize,
308 dst: usize,
309 len: u32,
310 kind: Activation,
311 },
312 /// f64 contiguous reduction along a single axis range. Layout
313 /// `[outer, reduced, inner]` in memory; output is `[outer, inner]`.
314 /// Sum only for now (Mean composes via 1/N multiply post-pass).
315 ReduceSumF64 {
316 src: usize,
317 dst: usize,
318 outer: u32,
319 reduced: u32,
320 inner: u32,
321 },
322 /// f64 plain copy (Reshape / Cast at the same dtype). Mirrors `Copy`
323 /// but at 8 bytes per element.
324 CopyF64 {
325 src: usize,
326 dst: usize,
327 len: u32,
328 },
329 /// i64 element copy (Reshape/Cast on i64 tensors).
330 CopyI64 {
331 src: usize,
332 dst: usize,
333 len: u32,
334 },
335 /// Truncating f32 → i64 (ONNX Cast; toward zero).
336 CastF32ToI64 {
337 src: usize,
338 dst: usize,
339 len: u32,
340 },
341 CastF32ToF64 {
342 src: usize,
343 dst: usize,
344 len: u32,
345 },
346 /// Narrowing f64 → f32 (`x as f32`, round-to-nearest). Without this the
347 /// generic Copy fallback byte-copies `len` f32-width elements out of the
348 /// 8-byte-per-element f64 buffer, so it reads each f64's two 4-byte halves
349 /// as two garbage f32s (round-trip F32→F64→F32 returned garbage).
350 CastF64ToF32 {
351 src: usize,
352 dst: usize,
353 len: u32,
354 },
355 /// Truncating f64 → i64 (ONNX Cast; toward zero). The generic 4-byte Copy
356 /// fallback reads/writes the wrong lane width against the 8-byte f64/i64
357 /// buffers (and would be a bit-copy, not a numeric convert).
358 CastF64ToI64 {
359 src: usize,
360 dst: usize,
361 len: u32,
362 },
363 /// Truncating f64 → i32 (ONNX Cast; toward zero). f64 is 8 bytes, i32 is 4;
364 /// the generic Copy reads only the low 4 bytes of each f64 as garbage.
365 CastF64ToI32 {
366 src: usize,
367 dst: usize,
368 len: u32,
369 },
370 /// f64 → Bool (`x != 0.0`). f64 is 8 bytes, Bool is 1; the generic Copy
371 /// misreads widths entirely.
372 CastF64ToBool {
373 src: usize,
374 dst: usize,
375 len: u32,
376 },
377 /// Narrowing i64 → i32 (`x as i32`, wrapping). 8-byte → 4-byte; the generic
378 /// Copy would keep only the low 4 bytes of each i64 (correct only for the
379 /// bit-truncation case, wrong as a Copy of `len` 4-byte lanes off an 8-byte
380 /// buffer). Mirror of the existing `CastI32ToI64` widening.
381 CastI64ToI32 {
382 src: usize,
383 dst: usize,
384 len: u32,
385 },
386 /// i64 → f64 (`x as f64`). Both are 8 bytes, but the generic 4-byte Copy
387 /// would move only half of each element AND bit-reinterpret (not a numeric
388 /// convert). Uses the existing i64/f64 accessors.
389 CastI64ToF64 {
390 src: usize,
391 dst: usize,
392 len: u32,
393 },
394 /// i32 → f64 (`x as f64`). 4-byte src → 8-byte dst; the generic Copy reads
395 /// too few source bytes and writes the wrong width.
396 CastI32ToF64 {
397 src: usize,
398 dst: usize,
399 len: u32,
400 },
401 /// f32 → f16 (`half::f16::from_f32`, round-to-nearest-even). Output is
402 /// 2-byte f16; there is no `sl_f16` accessor so the kernel writes the
403 /// little-endian f16 bytes straight into the arena.
404 CastF32ToF16 {
405 src: usize,
406 dst: usize,
407 len: u32,
408 },
409 /// f16 → f32 (`half::f16::to_f32`, exact). Input is 2-byte f16 read from the
410 /// arena bytes; output is native f32.
411 CastF16ToF32 {
412 src: usize,
413 dst: usize,
414 len: u32,
415 },
416 /// f32 → bf16 (`half::bf16::from_f32`, round-to-nearest-even). 2-byte bf16
417 /// output written as little-endian arena bytes.
418 CastF32ToBf16 {
419 src: usize,
420 dst: usize,
421 len: u32,
422 },
423 /// bf16 → f32 (`half::bf16::to_f32`, exact). 2-byte bf16 input read from the
424 /// arena bytes; output is native f32.
425 CastBf16ToF32 {
426 src: usize,
427 dst: usize,
428 len: u32,
429 },
430 /// Truncating f32 → i32 (ONNX Cast; toward zero).
431 CastF32ToI32 {
432 src: usize,
433 dst: usize,
434 len: u32,
435 },
436 /// i64 → f32 (ONNX Cast on shape scalars, e.g. Albert head-dim).
437 CastI64ToF32 {
438 src: usize,
439 dst: usize,
440 len: u32,
441 },
442 /// bool → i32 (BERT attention mask grid).
443 CastBoolToI32 {
444 src: usize,
445 dst: usize,
446 len: u32,
447 },
448 CastBoolToF32 {
449 src: usize,
450 dst: usize,
451 len: u32,
452 },
453 /// f32 → Bool (`x != 0.0`). Lets logical And/Or lower their {0,1} f32
454 /// arithmetic result straight back to a 1-byte Bool without a scalar
455 /// `Compare` (whose thunk has no broadcast and would read a scalar rhs
456 /// out of bounds).
457 CastF32ToBool {
458 src: usize,
459 dst: usize,
460 len: u32,
461 },
462 /// i32 → f32 (BERT attention mask cast before subtract).
463 CastI32ToF32 {
464 src: usize,
465 dst: usize,
466 len: u32,
467 },
468 /// i32 → i64 (widen integer indices/ids; e.g. an ONNX graph with I32
469 /// `context_tokens` feeding a codebook Gather whose kernel wants i64).
470 CastI32ToI64 {
471 src: usize,
472 dst: usize,
473 len: u32,
474 },
475 /// i32 → bool: `out[i] = (in[i] != 0) as u8`. Element-wise; the generic Copy
476 /// fallback byte-copies (4-byte i32 → 1-byte bool) and reads garbage. Needed by
477 /// MOSS-TTS's `Cast(attention_mask i32 → bool)` in the RoPE-position cumsum.
478 CastI32ToBool {
479 src: usize,
480 dst: usize,
481 len: u32,
482 },
483 /// i64 → bool: `out[i] = (in[i] != 0) as u8`. Element-wise; Copy would treat
484 /// 8-byte lanes as 1-byte bools. Soprano KV-backbone casts `attention_mask`
485 /// (I64 ones) to Bool before the causal Where(-inf) mask.
486 CastI64ToBool {
487 src: usize,
488 dst: usize,
489 len: u32,
490 },
491 /// bool → i64: `out[i] = in[i] as i64` (0/1). Widen a bool mask for CumSum/arith.
492 CastBoolToI64 {
493 src: usize,
494 dst: usize,
495 len: u32,
496 },
497 /// Generic scalar cast for any numeric dtype pair not covered by a
498 /// hand-written fast-path arm above (exotic widths: I8/I16/U8/U32, and
499 /// F16/BF16 against any non-F32 dtype). Converts element-by-element with
500 /// correct numeric semantics (float→int saturates, int→int wraps,
501 /// round-to-nearest for F16/BF16) instead of the corrupting generic 4-byte
502 /// `Copy` that silently mis-strides / bit-reinterprets. C64 is never routed
503 /// here (no scalar conversion is defined for complex).
504 CastGeneric {
505 src: usize,
506 dst: usize,
507 len: u32,
508 src_dtype: rlx_ir::DType,
509 dst_dtype: rlx_ir::DType,
510 },
511 /// f64 element-wise binary with broadcast. `len`/`lhs_len`/`rhs_len`
512 /// are element counts; kernel does `out[i] = lhs[i % lhs_len] OP rhs[i % rhs_len]`.
513 /// Mirror of `BinaryFull` at 8 bytes per element.
514 BinaryFullF64 {
515 lhs: usize,
516 rhs: usize,
517 dst: usize,
518 len: u32,
519 lhs_len: u32,
520 rhs_len: u32,
521 op: BinaryOp,
522 /// Output shape dims (row-major). Empty in the fast path. See
523 /// `BinaryFull` doc for the broadcast convention.
524 out_dims_bcast: Vec<u32>,
525 bcast_lhs_strides: Vec<u32>,
526 bcast_rhs_strides: Vec<u32>,
527 },
528 /// f64 concat — byte-for-byte mirror of `Concat` but copies
529 /// 8 bytes per element. Element-counted offsets/strides match
530 /// the f32 variant; the executor scales by elem_size internally.
531 ConcatF64 {
532 dst: usize,
533 outer: u32,
534 inner: u32,
535 total_axis: u32,
536 inputs: Vec<(usize, u32, u32)>,
537 },
538 /// C64 element-wise binary with broadcast. Same `len` /
539 /// `lhs_len` / `rhs_len` semantics as `BinaryFull` but each
540 /// "element" is one complex value (8 bytes = `[re, im]` as two
541 /// f32s). The executor reads the underlying f32 buffer at
542 /// `2·len` floats and walks element pairs. Supports Add / Sub /
543 /// Mul / Div; Max / Min / Pow have no single natural complex
544 /// definition and panic at lowering.
545 BinaryFullC64 {
546 lhs: usize,
547 rhs: usize,
548 dst: usize,
549 /// Complex element count (NOT f32 count). f32 buffer length
550 /// is `2·len`.
551 len: u32,
552 lhs_len: u32,
553 rhs_len: u32,
554 op: BinaryOp,
555 out_dims_bcast: Vec<u32>,
556 bcast_lhs_strides: Vec<u32>,
557 bcast_rhs_strides: Vec<u32>,
558 },
559 /// Bounded scan. Holds a recursively-compiled body schedule + a
560 /// pre-initialized body arena snapshot (constants filled). Each
561 /// outer execution clones the snapshot, copies the carry-in slot
562 /// from the outer arena, runs the body schedule `length` times,
563 /// then writes the final carry to the outer arena.
564 ///
565 /// Single-carry MVP — body has exactly one Input and one output,
566 /// both same shape and dtype.
567 Scan {
568 body: Arc<ThunkSchedule>,
569 body_init: Arc<Vec<u8>>, // pristine body arena bytes
570 body_input_off: usize, // byte offset of the body's carry-Input slot
571 body_output_off: usize, // byte offset of the body's output slot
572 outer_init_off: usize, // outer-arena offset of the initial carry
573 outer_final_off: usize, // outer-arena offset of the final carry / trajectory base
574 length: u32,
575 carry_bytes: u32, // carry size in bytes
576 /// When true, write each step's carry to the outer arena at
577 /// offset `outer_final_off + t * carry_bytes`, producing a
578 /// `[length, *carry]` stacked trajectory. When false, only the
579 /// final carry lands at `outer_final_off`.
580 save_trajectory: bool,
581 /// Per-step `xs` inputs. For each: (body_x_input_off,
582 /// outer_xs_base_off, per_step_bytes). Per iteration `t`, the
583 /// executor copies `outer_xs_base_off + t * per_step_bytes`
584 /// into `body_x_input_off`. Empty when the scan has no xs.
585 xs_inputs: Arc<Vec<(usize, usize, u32)>>,
586 /// Broadcast inputs — values constant across iterations. For
587 /// each: (body_bcast_input_off, outer_bcast_off, total_bytes).
588 /// Filled into `body_buf` ONCE before the scan loop starts
589 /// (xs in contrast are re-filled every iteration). Empty when
590 /// the scan has no bcasts.
591 bcast_inputs: Arc<Vec<(usize, usize, u32)>>,
592 /// Number of trajectory checkpoints (when `save_trajectory`).
593 /// `0` or `length` ⇒ save every iteration. Otherwise save only
594 /// `K` rows at indices `floor((k+1) * length / K) - 1` for
595 /// `k in 0..K`. Last index is always `length-1` so the final
596 /// carry is always cached.
597 num_checkpoints: u32,
598 },
599
600 /// Reverse-mode AD companion to `Thunk::Scan`. Walks `t = length-1
601 /// .. 0`, threading `dcarry` through the body's VJP. Per iteration:
602 /// writes `carry_t` (from outer init or trajectory), each `xs_i[t]`
603 /// slice, and the current `dcarry` into the body_vjp's Input
604 /// slots, runs body_vjp, reads new `dcarry` from its single output.
605 /// f64 carry only — the upstream-accumulation step in trajectory
606 /// mode does an element-wise f64 add.
607 ScanBackward {
608 body_vjp: Arc<ThunkSchedule>,
609 body_init: Arc<Vec<u8>>,
610 body_carry_in_off: usize, // body_vjp's mirrored body-carry-input slot
611 body_x_offs: Arc<Vec<usize>>, // body_vjp's mirrored x_t_i Input slots, in xs order
612 body_d_output_off: usize, // body_vjp's "d_output" Input slot
613 body_dcarry_out_off: usize, // body_vjp's gradient output
614 outer_init_off: usize, // original init carry
615 outer_traj_off: usize, // [length-or-K, *carry] trajectory base
616 outer_upstream_off: usize, // upstream gradient (carry shape, or [length, *carry])
617 /// Per-xs entries: (outer_xs_base_off, per_step_bytes). Read
618 /// `xs_i[t]` from `outer_xs_base_off + t * per_step_bytes`.
619 outer_xs_offs: Arc<Vec<(usize, u32)>>,
620 outer_dinit_off: usize, // output: dinit
621 length: u32,
622 carry_bytes: u32,
623 /// Bytes per element in the carry tensor: 4 for f32, 8 for f64.
624 /// Used to dispatch the trajectory-mode upstream accumulation
625 /// kernel (the dcarry += upstream\[t\] add must use the right
626 /// floating-point type — a hard-coded f64 add silently does
627 /// nothing for an f32 carry whose `cb` isn't divisible by 8).
628 carry_elem_size: u32,
629 save_trajectory: bool, // true → upstream is per-step; false → just final
630 /// Recursive checkpointing config. `0` or `length` ⇒ full
631 /// trajectory cached, no recompute (existing behavior).
632 /// `0 < K < length` ⇒ trajectory has only K rows; the executor
633 /// recomputes intermediate carries via `forward_body` between
634 /// checkpoints. Memory: O(K · carry_bytes); time: O(length).
635 num_checkpoints: u32,
636 /// Forward body schedule (same compiled body as the forward
637 /// Op::Scan), used for recompute when `num_checkpoints` is
638 /// active. `None` for the All strategy.
639 forward_body: Option<Arc<ThunkSchedule>>,
640 /// Pristine forward body arena bytes (constants filled).
641 forward_body_init: Option<Arc<Vec<u8>>>,
642 /// Forward body's carry-Input and output slot offsets — needed
643 /// to seed/read the body during recompute.
644 forward_body_carry_in_off: usize,
645 forward_body_output_off: usize,
646 /// Forward body's per-step xs Input slots (one per outer xs).
647 /// Same indexing convention as `body_x_offs`.
648 forward_body_x_offs: Arc<Vec<usize>>,
649 },
650
651 /// Companion to `ScanBackward` that materializes one stacked
652 /// `dxs_i`. Same backward loop; per iteration, after running
653 /// body_vjp, copies its `body_dxs_out_off` slot into the outer
654 /// arena at `outer_dxs_off + t * per_step_bytes`. dcarry threading
655 /// is identical — we still need it for the body_vjp recurrence
656 /// even though we don't write it back to the outer arena.
657 ScanBackwardXs {
658 body_vjp: Arc<ThunkSchedule>,
659 body_init: Arc<Vec<u8>>,
660 body_carry_in_off: usize,
661 body_x_offs: Arc<Vec<usize>>,
662 body_d_output_off: usize,
663 body_dcarry_out_off: usize,
664 body_dxs_out_off: usize, // the body_vjp output we extract per step
665 outer_init_off: usize,
666 outer_traj_off: usize,
667 outer_upstream_off: usize,
668 outer_xs_offs: Arc<Vec<(usize, u32)>>,
669 outer_dxs_off: usize, // base of the stacked [length, *per_step] output
670 length: u32,
671 carry_bytes: u32,
672 /// Same role as `Thunk::ScanBackward::carry_elem_size`.
673 carry_elem_size: u32,
674 per_step_bytes: u32, // bytes per row of the dxs output
675 save_trajectory: bool,
676 /// Recursive checkpointing config. Same semantics as
677 /// `Thunk::ScanBackward::num_checkpoints` — `0` or `length`
678 /// means "save every step's carry"; `0 < K < length` means
679 /// the trajectory has only K rows and the executor recomputes
680 /// intermediate carries via `forward_body` (which must be
681 /// `Some`). Implemented via segment-cached recompute,
682 /// mirroring the `ScanBackward` path.
683 num_checkpoints: u32,
684 forward_body: Option<Arc<ThunkSchedule>>,
685 forward_body_init: Option<Arc<Vec<u8>>>,
686 forward_body_carry_in_off: usize,
687 forward_body_output_off: usize,
688 forward_body_x_offs: Arc<Vec<usize>>,
689 },
690 /// User-defined sub-graph (`Op::CustomFn`) — runs `fwd_body` once.
691 /// Per execution: clone `body_init`, copy each primal input from the
692 /// outer arena into its body Input slot, run the body schedule,
693 /// copy the body's single output back to the outer arena.
694 CustomFn {
695 body: Arc<ThunkSchedule>,
696 body_init: Arc<Vec<u8>>,
697 /// Per primal input: (body_input_off, outer_input_off, bytes).
698 inputs: Arc<Vec<(usize, usize, u32)>>,
699 body_output_off: usize,
700 outer_output_off: usize,
701 out_bytes: u32,
702 },
703 /// C = A @ B; C += bias; C = act(C)
704 FusedMmBiasAct {
705 a: usize,
706 w: usize,
707 bias: usize,
708 c: usize,
709 m: u32,
710 k: u32,
711 n: u32,
712 act: Option<Activation>,
713 },
714 /// out = LN(x + residual + bias, gamma, beta)
715 FusedResidualLN {
716 x: usize,
717 res: usize,
718 bias: usize,
719 g: usize,
720 b: usize,
721 out: usize,
722 rows: u32,
723 h: u32,
724 eps: f32,
725 has_bias: bool,
726 },
727 /// out = RmsNorm(x + residual + bias, gamma, beta)
728 FusedResidualRmsNorm {
729 x: usize,
730 res: usize,
731 bias: usize,
732 g: usize,
733 b: usize,
734 out: usize,
735 rows: u32,
736 h: u32,
737 eps: f32,
738 has_bias: bool,
739 },
740 /// out = bias_add(data, bias, m, n) for Binary::Add with broadcast
741 BiasAdd {
742 src: usize,
743 bias: usize,
744 dst: usize,
745 m: u32,
746 n: u32,
747 },
748 /// Element-wise binary op with NumPy-style broadcast.
749 ///
750 /// Fast path (`lhs_len == rhs_len == len`): plain element-wise loop,
751 /// SIMD-vectorized on aarch64 for `Add`/`Mul`. `bcast_*` fields
752 /// are unused.
753 ///
754 /// Broadcast path: uses `out_dims_bcast` + `bcast_lhs_strides` +
755 /// `bcast_rhs_strides` to compute per-cell indices into each
756 /// operand. The strides are precomputed at thunk-construction
757 /// time from the operands' true shapes (with stride 0 on any axis
758 /// where the operand has size 1). This is the only correct way
759 /// to handle bidirectional broadcasts like `[N, 1] op [1, S]
760 /// → [N, S]`, which simple `i % lhs_len` modulo indexing maps to
761 /// wrong cells.
762 BinaryFull {
763 lhs: usize,
764 rhs: usize,
765 dst: usize,
766 len: u32,
767 lhs_len: u32,
768 rhs_len: u32,
769 op: BinaryOp,
770 /// Output shape dims (row-major). Empty in the fast path.
771 out_dims_bcast: Vec<u32>,
772 /// Per-dim stride into `lhs` (0 where lhs broadcasts).
773 bcast_lhs_strides: Vec<u32>,
774 /// Per-dim stride into `rhs`.
775 bcast_rhs_strides: Vec<u32>,
776 /// Element size (4 = F32, 8 = I64).
777 elem_bytes: u8,
778 },
779 /// Activation in-place
780 ActivationInPlace {
781 data: usize,
782 len: u32,
783 act: Activation,
784 },
785 /// Gather axis=0: table\[idx\] → out
786 Gather {
787 table: usize,
788 table_len: u32,
789 idx: usize,
790 dst: usize,
791 num_idx: u32,
792 trailing: u32,
793 /// 1 when the index tensor is i64 (ONNX Gather indices).
794 idx_i64: u8,
795 /// Element size of table/output (4 = f32, 8 = i64).
796 table_bytes: u8,
797 },
798 /// Narrow: copy slice (`elem_bytes` = source element size: 4 for f32, 8 for f64).
799 Narrow {
800 src: usize,
801 dst: usize,
802 outer: u32,
803 src_stride: u32,
804 dst_stride: u32,
805 inner: u32,
806 elem_bytes: u8,
807 },
808 /// Reverse (flip) element order along the axes in `rev_mask`. `dims` is the
809 /// row-major input shape; output shape is identical. Dtype-agnostic
810 /// (byte-copy of `elem_bytes` per element).
811 Reverse {
812 src: usize,
813 dst: usize,
814 dims: Vec<u32>,
815 rev_mask: Vec<bool>,
816 elem_bytes: u8,
817 },
818 /// Copy (reshape, expand)
819 Copy {
820 src: usize,
821 dst: usize,
822 len: u32,
823 },
824 /// LayerNorm standalone
825 LayerNorm {
826 src: usize,
827 g: usize,
828 b: usize,
829 dst: usize,
830 rows: u32,
831 h: u32,
832 eps: f32,
833 },
834 /// GroupNorm on NCHW `[N,C,H,W]`.
835 GroupNorm {
836 src: usize,
837 g: usize,
838 b: usize,
839 dst: usize,
840 n: u32,
841 c: u32,
842 h: u32,
843 w: u32,
844 num_groups: u32,
845 eps: f32,
846 },
847 /// BatchNorm inference: frozen mean/var, feature axis last.
848 BatchNormInference {
849 src: usize,
850 g: usize,
851 b: usize,
852 mean: usize,
853 var: usize,
854 dst: usize,
855 count: u32,
856 channels: u32,
857 eps: f32,
858 },
859 BatchNormInferenceBackwardInput {
860 x: usize,
861 gamma: usize,
862 mean: usize,
863 var: usize,
864 dy: usize,
865 dx: usize,
866 count: u32,
867 channels: u32,
868 eps: f32,
869 },
870 BatchNormInferenceBackwardGamma {
871 x: usize,
872 mean: usize,
873 var: usize,
874 dy: usize,
875 dgamma: usize,
876 count: u32,
877 channels: u32,
878 eps: f32,
879 },
880 BatchNormInferenceBackwardBeta {
881 dy: usize,
882 dbeta: usize,
883 count: u32,
884 channels: u32,
885 },
886 /// LayerNorm2d on NCHW (SAM / candle semantics).
887 LayerNorm2d {
888 src: usize,
889 g: usize,
890 b: usize,
891 dst: usize,
892 n: u32,
893 c: u32,
894 h: u32,
895 w: u32,
896 eps: f32,
897 },
898 /// ConvTranspose2d on NCHW.
899 ConvTranspose2d {
900 src: usize,
901 weight: usize,
902 dst: usize,
903 n: u32,
904 c_in: u32,
905 h: u32,
906 w_in: u32,
907 c_out: u32,
908 h_out: u32,
909 w_out: u32,
910 kh: u32,
911 kw: u32,
912 sh: u32,
913 sw: u32,
914 ph: u32,
915 pw: u32,
916 dh: u32,
917 dw: u32,
918 groups: u32,
919 },
920 /// 3D convolution on NCDHW. Input [N, C_in, D, H, W], weight
921 /// [C_out, C_in/groups, kD, kH, kW], output [N, C_out, D_out, H_out, W_out].
922 /// Naive direct convolution; the depth-axis analogue of `Thunk::Conv2D`.
923 Conv3d {
924 src: usize,
925 weight: usize,
926 dst: usize,
927 n: u32,
928 c_in: u32,
929 d: u32,
930 h: u32,
931 w: u32,
932 c_out: u32,
933 d_out: u32,
934 h_out: u32,
935 w_out: u32,
936 kd: u32,
937 kh: u32,
938 kw: u32,
939 sd: u32,
940 sh: u32,
941 sw: u32,
942 pd: u32,
943 ph: u32,
944 pw: u32,
945 dd: u32,
946 dh: u32,
947 dw: u32,
948 groups: u32,
949 },
950 /// ConvTranspose3d on NCDHW. Weight [C_in, C_out/groups, kD, kH, kW].
951 /// `output_padding` is already folded into `d_out/h_out/w_out`.
952 ConvTranspose3d {
953 src: usize,
954 weight: usize,
955 dst: usize,
956 n: u32,
957 c_in: u32,
958 d: u32,
959 h: u32,
960 w_in: u32,
961 c_out: u32,
962 d_out: u32,
963 h_out: u32,
964 w_out: u32,
965 kd: u32,
966 kh: u32,
967 kw: u32,
968 sd: u32,
969 sh: u32,
970 sw: u32,
971 pd: u32,
972 ph: u32,
973 pw: u32,
974 dd: u32,
975 dh: u32,
976 dw: u32,
977 groups: u32,
978 },
979 /// Nearest 2× upsample on NCHW (per-batch slice).
980 ResizeNearest2x {
981 src: usize,
982 dst: usize,
983 n: u32,
984 c: u32,
985 h: u32,
986 w: u32,
987 },
988 /// Nearest NCDHW resample to an explicit spatial size.
989 Interpolate3d {
990 src: usize,
991 dst: usize,
992 n: u32,
993 c: u32,
994 d_in: u32,
995 h_in: u32,
996 w_in: u32,
997 d_out: u32,
998 h_out: u32,
999 w_out: u32,
1000 },
1001 /// SAM2 axial 2-D RoPE on `[batch, seq, num_heads * head_dim]`.
1002 AxialRope2d {
1003 src: usize,
1004 dst: usize,
1005 batch: u32,
1006 seq: u32,
1007 hidden: u32,
1008 end_x: u32,
1009 end_y: u32,
1010 head_dim: u32,
1011 num_heads: u32,
1012 theta: f32,
1013 repeat_factor: u32,
1014 },
1015 /// RMSNorm: out = (x / sqrt(mean(x^2) + eps)) * gamma + beta. No mean
1016 /// subtraction, hence cheaper than LayerNorm. Used by Llama-class models.
1017 RmsNorm {
1018 src: usize,
1019 g: usize,
1020 b: usize,
1021 dst: usize,
1022 rows: u32,
1023 h: u32,
1024 eps: f32,
1025 },
1026 /// DiT adaLN-Zero modulation: `out = norm(x)·(1+scale)+shift`, one pass —
1027 /// mean/var (or RMS) and the per-conditioning affine fused, the pre-norm
1028 /// activation never materialized. `scale`/`shift` share the modulation
1029 /// shape (`[B,1,D]`); `mod_off[row]` is the element offset into each for
1030 /// output `row` (the built-in broadcast that avoids an `Op::Expand`).
1031 AdaLayerNorm {
1032 src: usize,
1033 scale: usize,
1034 shift: usize,
1035 dst: usize,
1036 rows: u32,
1037 h: u32,
1038 eps: f32,
1039 /// `true` ⇒ LayerNorm (mean+var); `false` ⇒ RMSNorm (rms only).
1040 layer_norm: bool,
1041 /// Element count of the `scale`/`shift` buffers (they share a shape).
1042 mod_len: u32,
1043 /// Per-row element offset into `scale`/`shift` (broadcast map).
1044 mod_off: Vec<u32>,
1045 },
1046 /// DiT gated residual: `out = x + gate·y`. `gate` carries the modulation
1047 /// shape (`[B,1,D]`); `gate_off[row]` is the element offset into `gate` for
1048 /// output `row`.
1049 GatedResidual {
1050 x: usize,
1051 y: usize,
1052 gate: usize,
1053 dst: usize,
1054 rows: u32,
1055 h: u32,
1056 /// Element count of the `gate` buffer.
1057 gate_len: u32,
1058 /// Per-row element offset into `gate` (broadcast map).
1059 gate_off: Vec<u32>,
1060 },
1061 /// Packed `[dx ∥ dscale ∥ dshift]` for [`Op::AdaLayerNormBackward`].
1062 AdaLayerNormBackward {
1063 x: usize,
1064 scale: usize,
1065 shift: usize,
1066 dy: usize,
1067 dst: usize,
1068 rows: u32,
1069 h: u32,
1070 eps: f32,
1071 layer_norm: bool,
1072 mod_len: u32,
1073 mod_off: Vec<u32>,
1074 },
1075 /// Packed `[dx ∥ dy ∥ dgate]` for [`Op::GatedResidualBackward`].
1076 GatedResidualBackward {
1077 x: usize,
1078 y: usize,
1079 gate: usize,
1080 dy: usize,
1081 dst: usize,
1082 rows: u32,
1083 h: u32,
1084 gate_len: u32,
1085 gate_off: Vec<u32>,
1086 },
1087 /// Softmax
1088 Softmax {
1089 data: usize,
1090 rows: u32,
1091 cols: u32,
1092 },
1093 /// Inclusive (or exclusive) cumulative sum along the last axis
1094 /// (callers pre-flatten higher-dim cumsums via reshape views).
1095 Cumsum {
1096 src: usize,
1097 dst: usize,
1098 rows: u32,
1099 cols: u32,
1100 exclusive: bool,
1101 /// Element dtype — the scan runs in f32 or i64 (integer position ids /
1102 /// ragged offsets, e.g. MOSS-TTS's `cumsum(attention_mask)-1` RoPE positions
1103 /// must NOT go through the f32 path, which reads the i64 buffer as garbage).
1104 dtype: rlx_ir::DType,
1105 },
1106 /// Inclusive (or exclusive) cumulative product / maximum along the last
1107 /// axis (f32). `is_max` selects max over product. Mirrors [`Thunk::Cumsum`]'s
1108 /// row/col flattening; the reference O(L) scan behind `Op::{CumProd, CumMax}`.
1109 CumScan {
1110 src: usize,
1111 dst: usize,
1112 rows: u32,
1113 cols: u32,
1114 exclusive: bool,
1115 is_max: bool,
1116 },
1117 /// Mamba-style selective scan (plan #15).
1118 /// Inputs: x, delta \[b,s,h\], a \[h,n\], b \[b,s,n\], c \[b,s,n\].
1119 /// Output: y \[b,s,h\]. State h carries through the seq.
1120 SelectiveScan {
1121 x: usize,
1122 delta: usize,
1123 a: usize,
1124 b: usize,
1125 c: usize,
1126 dst: usize,
1127 batch: u32,
1128 seq: u32,
1129 hidden: u32,
1130 state_size: u32,
1131 },
1132
1133 /// Gated DeltaNet linear-attention scan (Qwen3.5/3.6 trunk).
1134 /// Inputs: q, k, v `[b, s, h, n]`; g, beta `[b, s, h]`. Output:
1135 /// `[b, s, h, n]`. See `Op::GatedDeltaNet` for math.
1136 GatedDeltaNet {
1137 q: usize,
1138 k: usize,
1139 v: usize,
1140 g: usize,
1141 beta: usize,
1142 /// Arena offset of the initial `[b, h, n, n]` state (valid only when
1143 /// `carry_state`; may legitimately be 0). When carrying, the initial
1144 /// state is loaded from here and the final state written back in place.
1145 state: usize,
1146 dst: usize,
1147 batch: u32,
1148 seq: u32,
1149 heads: u32,
1150 state_size: u32,
1151 /// When true (Kimi-K3 KDA), `g` is `[b, s, h, n]` (per key channel).
1152 gate_per_channel: bool,
1153 /// When true, resume from / write back the external `state` (decode).
1154 /// An explicit flag — NOT `state != 0`, since offset 0 is a valid slot.
1155 carry_state: bool,
1156 },
1157
1158 /// Multi-layer (optionally bidirectional, optional carry) LSTM with
1159 /// packed weights. See `Op::Lstm`. `h0`/`c0` are valid only when
1160 /// `carry`; `dst` is `[b, s, D*h]`.
1161 Lstm {
1162 x: usize,
1163 w_ih: usize,
1164 w_hh: usize,
1165 bias: usize,
1166 h0: usize,
1167 c0: usize,
1168 dst: usize,
1169 batch: u32,
1170 seq: u32,
1171 input_size: u32,
1172 hidden: u32,
1173 num_layers: u32,
1174 bidirectional: bool,
1175 carry: bool,
1176 },
1177 /// GRU (gate order r, z, n) with separate `b_ih`/`b_hh`. See `Op::Gru`.
1178 Gru {
1179 x: usize,
1180 w_ih: usize,
1181 w_hh: usize,
1182 b_ih: usize,
1183 b_hh: usize,
1184 h0: usize,
1185 dst: usize,
1186 batch: u32,
1187 seq: u32,
1188 input_size: u32,
1189 hidden: u32,
1190 num_layers: u32,
1191 bidirectional: bool,
1192 carry: bool,
1193 },
1194 /// Elman RNN (`tanh`/`relu`), single merged bias. See `Op::Rnn`.
1195 Rnn {
1196 x: usize,
1197 w_ih: usize,
1198 w_hh: usize,
1199 bias: usize,
1200 h0: usize,
1201 dst: usize,
1202 batch: u32,
1203 seq: u32,
1204 input_size: u32,
1205 hidden: u32,
1206 num_layers: u32,
1207 bidirectional: bool,
1208 carry: bool,
1209 relu: bool,
1210 },
1211 /// Mamba-2 / SSD scalar-decay SSM scan. See `Op::Mamba2`.
1212 Mamba2 {
1213 x: usize,
1214 dt: usize,
1215 a: usize,
1216 b: usize,
1217 c: usize,
1218 dst: usize,
1219 batch: u32,
1220 seq: u32,
1221 heads: u32,
1222 head_dim: u32,
1223 state_size: u32,
1224 },
1225
1226 /// 1×1 conv fast path (plan #26). The general Conv2D thunk
1227 /// runs the textbook 7-deep loop; a 1×1 stride-1 padding-0
1228 /// groups-1 conv is mathematically a per-batch matmul, and
1229 /// dispatching it through BLAS is 3-10× faster than the
1230 /// scalar nest. Common case: ViT patch-projection follow-on,
1231 /// transformer "expert" reductions in some MoE designs.
1232 ///
1233 /// Per batch: weight `[c_out, c_in]` × input `[c_in, h*w]`
1234 /// = output `[c_out, h*w]`.
1235 Conv2D1x1 {
1236 src: usize,
1237 weight: usize,
1238 dst: usize,
1239 n: u32,
1240 c_in: u32,
1241 c_out: u32,
1242 hw: u32,
1243 },
1244
1245 /// Fused dequant + matmul (plan #5). Today supports
1246 /// `QuantScheme::Int8Block` (symmetric); other schemes panic
1247 /// at lowering time with a clear message until kernels are added.
1248 DequantMatMul {
1249 x: usize,
1250 w_q: usize, // packed i8 bytes for Int8 schemes
1251 scale: usize, // [k/block, n] f32 scale
1252 zp: usize, // [k/block, n] f32 zero-point (0 for sym)
1253 dst: usize,
1254 m: u32,
1255 k: u32,
1256 n: u32,
1257 block_size: u32,
1258 is_asymmetric: bool,
1259 },
1260
1261 /// Codebook weight-synthesis matmul (single-level vector quant).
1262 /// Weight is stored transposed `[n, k]`; each contiguous
1263 /// `entry_dim`-length sub-vector along `k` is replaced by a learned
1264 /// codebook centroid `codebook[index]`. The centroid is reconstructed
1265 /// inside the matmul inner loop — the weight is never materialized.
1266 SynthMatMul {
1267 x: usize, // [m, k] f32 activations
1268 indices: usize, // [n, k/entry_dim] u8 codebook indices
1269 codebook: usize, // [num_entries, entry_dim] f32 centroids
1270 dst: usize, // [m, n] f32 output
1271 m: u32,
1272 k: u32,
1273 n: u32,
1274 entry_dim: u32,
1275 num_entries: u32,
1276 },
1277
1278 /// KAN Gaussian-RBF spline activation (per-channel learned basis).
1279 /// Shape-preserving: `y[.., c] = Σ_g coeff[c,g]·exp(-((x−center_g)·inv_h)²)`.
1280 SplineActivation {
1281 x: usize, // [rows, channels] f32
1282 coeff: usize, // [channels, num_basis] f32
1283 dst: usize, // [rows, channels] f32
1284 rows: u32,
1285 channels: u32,
1286 num_basis: u32,
1287 grid_min: f32,
1288 grid_max: f32,
1289 },
1290
1291 /// GGUF-format dequant + matmul. Weight is a packed byte tensor
1292 /// in one of the K-quant super-block layouts (Q4_K, Q5_K, Q6_K,
1293 /// Q8_K). Scales / mins live inside the packed bytes — no
1294 /// side-channel scale tensor.
1295 ///
1296 /// Today this is a "dequant-to-scratch then sgemm" kernel — it
1297 /// keeps the *arena* memory footprint down (weights stay packed)
1298 /// but the dequant itself happens per matmul. A future fully
1299 /// fused tile-streaming kernel would close the compute gap.
1300 DequantMatMulGguf {
1301 x: usize, // f32 activations [m, k]
1302 w_q: usize, // packed weight bytes (k*n elements packed)
1303 dst: usize, // f32 output [m, n]
1304 m: u32,
1305 k: u32,
1306 n: u32,
1307 scheme: rlx_ir::quant::QuantScheme,
1308 },
1309
1310 /// Int4 block dequant + matmul (packed nibbles, side scale/zp).
1311 DequantMatMulInt4 {
1312 x: usize,
1313 w_q: usize,
1314 scale: usize,
1315 zp: usize,
1316 dst: usize,
1317 m: u32,
1318 k: u32,
1319 n: u32,
1320 block_size: u32,
1321 is_asymmetric: bool,
1322 },
1323
1324 /// FP8 dequant + matmul (per-tensor or per-column scale).
1325 DequantMatMulFp8 {
1326 x: usize,
1327 w_q: usize,
1328 scale: usize,
1329 dst: usize,
1330 m: u32,
1331 k: u32,
1332 n: u32,
1333 e5m2: bool,
1334 },
1335
1336 /// NVFP4 (E2M1) block dequant + matmul — 16-wide groups, FP8 scales.
1337 DequantMatMulNvfp4 {
1338 x: usize,
1339 w_q: usize,
1340 scale: usize,
1341 global_scale: usize,
1342 dst: usize,
1343 m: u32,
1344 k: u32,
1345 n: u32,
1346 },
1347
1348 /// MxFp4x2 two-level residual E2M1 DequantMatMul. `w_q` = `[plane0|plane1]`
1349 /// nibbles, `scale` = `[s0|s1]` f32 per (k/`group`, n) block.
1350 DequantMatMulMxFp4x2 {
1351 x: usize,
1352 w_q: usize,
1353 scale: usize,
1354 dst: usize,
1355 m: u32,
1356 k: u32,
1357 n: u32,
1358 group: u32,
1359 },
1360
1361 /// MLX affine / mxfp DequantMatMul — host dequant via `rlx-mlx-io`.
1362 DequantMatMulMlx {
1363 x: usize,
1364 w_q: usize,
1365 scale: usize,
1366 zp: usize,
1367 dst: usize,
1368 m: u32,
1369 k: u32,
1370 n: u32,
1371 scheme: rlx_ir::quant::QuantScheme,
1372 },
1373
1374 /// Native low-precision scaled GEMM (FP8/FP6/FP4) — CPU reference oracle.
1375 /// TN layout: lhs [m,k], rhs [n,k] codes, out [m,n] f32.
1376 ScaledMatMul {
1377 lhs: usize,
1378 rhs: usize,
1379 lhs_scale: usize,
1380 rhs_scale: usize,
1381 bias: usize, // valid iff has_bias
1382 dst: usize,
1383 m: u32,
1384 k: u32,
1385 n: u32,
1386 lhs_fmt: rlx_ir::ScaledFormat,
1387 rhs_fmt: rlx_ir::ScaledFormat,
1388 layout: rlx_ir::ScaleLayout,
1389 has_bias: bool,
1390 },
1391
1392 /// Native low-precision *grouped* (MoE) TN GEMM — expert-indexed
1393 /// `ScaledMatMul`. `weight [E,n,k]` codes, per-expert scales/bias.
1394 ScaledGroupedMatMul {
1395 input: usize,
1396 weight: usize,
1397 input_scale: usize,
1398 weight_scale: usize,
1399 expert_idx: usize,
1400 bias: usize, // valid iff has_bias
1401 dst: usize,
1402 m: u32,
1403 k: u32,
1404 n: u32,
1405 num_experts: u32,
1406 lhs_fmt: rlx_ir::ScaledFormat,
1407 rhs_fmt: rlx_ir::ScaledFormat,
1408 layout: rlx_ir::ScaleLayout,
1409 has_bias: bool,
1410 },
1411
1412 /// Quantize f32 → packed low-precision codes (reads the scale tensor).
1413 ScaledQuantize {
1414 x: usize,
1415 scale: usize,
1416 dst: usize,
1417 rows: u32,
1418 cols: u32,
1419 fmt: rlx_ir::ScaledFormat,
1420 layout: rlx_ir::ScaleLayout,
1421 },
1422
1423 /// Compute the per-tensor / per-block scale tensor for `Op::ScaledQuantize`.
1424 ScaledQuantScale {
1425 x: usize,
1426 dst: usize,
1427 rows: u32,
1428 cols: u32,
1429 fmt: rlx_ir::ScaledFormat,
1430 layout: rlx_ir::ScaleLayout,
1431 },
1432
1433 /// Reconstruct f32 from packed codes (`Op::ScaledDequantize`).
1434 ScaledDequantize {
1435 codes: usize,
1436 scale: usize,
1437 dst: usize,
1438 rows: u32,
1439 cols: u32,
1440 fmt: rlx_ir::ScaledFormat,
1441 layout: rlx_ir::ScaleLayout,
1442 },
1443
1444 /// Fused LoRA matmul (plan #9): out = x·W + scale * (x·A)·B.
1445 /// `r` is the LoRA rank (typically 4-64) — the rank-r
1446 /// intermediate `x·A` lives in scratch, never on the arena.
1447 LoraMatMul {
1448 x: usize,
1449 w: usize,
1450 a: usize,
1451 b: usize,
1452 dst: usize,
1453 m: u32,
1454 k: u32,
1455 n: u32,
1456 r: u32,
1457 scale: f32,
1458 },
1459 /// Fused sample: logits [batch, vocab] → token ids \[batch\].
1460 /// See Op::Sample. Output values are f32-encoded usize indices
1461 /// (matches the rest of the IR's "ids as f32" convention).
1462 Sample {
1463 logits: usize,
1464 dst: usize,
1465 batch: u32,
1466 vocab: u32,
1467 top_k: u32, // 0 = disabled
1468 top_p: f32, // 1.0 = disabled
1469 temperature: f32, // 1.0 = neutral
1470 seed: u64,
1471 },
1472 /// ONNX `RandomNormalLike` fill.
1473 RngNormal {
1474 dst: usize,
1475 len: u32,
1476 mean: f32,
1477 scale: f32,
1478 key: u64,
1479 op_seed: Option<f32>,
1480 },
1481 /// ONNX `RandomUniformLike` fill.
1482 RngUniform {
1483 dst: usize,
1484 len: u32,
1485 low: f32,
1486 high: f32,
1487 key: u64,
1488 op_seed: Option<f32>,
1489 },
1490 /// Attention SDPA. `mask` is the offset of the optional mask tensor
1491 /// (only meaningful when `mask_kind == MaskKind::Custom`); other
1492 /// kinds synthesize the mask in-kernel.
1493 ///
1494 /// Q/K/V each carry a `_row_stride` (elements per source row).
1495 /// Defaults to `heads * head_dim` — matches the standalone
1496 /// "Q/K/V are their own contiguous buffers" case. The Narrow→
1497 /// Attention fusion below rewrites these to the parent QKV stride
1498 /// (typically `3 * heads * head_dim`) so the kernel reads QKV
1499 /// directly without materializing the per-head buffers (plan #46).
1500 Attention {
1501 q: usize,
1502 k: usize,
1503 v: usize,
1504 mask: usize,
1505 out: usize,
1506 batch: u32,
1507 /// Query sequence length.
1508 seq: u32,
1509 /// Key/value sequence length. Differs from `seq` during cached decode.
1510 kv_seq: u32,
1511 heads: u32,
1512 /// GQA/MQA: number of key/value heads that the query heads share.
1513 /// Equals `heads` for plain MHA (the common case); only `< heads` on
1514 /// the non-fused standalone `Op::Attention` path.
1515 kv_heads: u32,
1516 head_dim: u32,
1517 /// V/output per-head width. Equals `head_dim` for symmetric SDPA; only
1518 /// `!= head_dim` for MLA-style asymmetric attention (V read + output
1519 /// write use this; Q/K scores still use `head_dim`).
1520 v_head_dim: u32,
1521 mask_kind: rlx_ir::op::MaskKind,
1522 /// Softmax score scale (`Op::Attention::score_scale`). `head_dim^-0.5`
1523 /// when the op left it unset. Must be honored — Gemma 4 uses `1.0`
1524 /// (Q/K are per-head RMS-normed, so no `1/sqrt(d)` pre-scale).
1525 scale: f32,
1526 /// Attention logit soft-cap (`Op::Attention::attn_logit_softcap`).
1527 /// `cap*tanh(score/cap)` applied pre-softmax; 0 = disabled. Gemma 2
1528 /// uses 50.0 — must be honored to match HF / rlx-metal / rlx-cuda.
1529 softcap: f32,
1530 q_row_stride: u32,
1531 k_row_stride: u32,
1532 v_row_stride: u32,
1533 /// Memory layout flag. `false` (the historical default) →
1534 /// `[B, S, H, D]` row-major: per-head offset is
1535 /// `bi*S*H*D + si*H*D + hi*D`. `true` → `[B, H, S, D]`
1536 /// (head-major), matching the convention used by rlx-cuda /
1537 /// rlx-rocm / rlx-tpu: per-head offset is
1538 /// `bi*H*S*D + hi*S*D + si*D`. Detected at lowering time
1539 /// from the input shape vs `num_heads` / `head_dim`.
1540 bhsd: bool,
1541 },
1542 /// [`Op::AttentionBackward`] — emits dQ, dK, or dV (see `wrt`).
1543 AttentionBackward {
1544 q: usize,
1545 k: usize,
1546 v: usize,
1547 dy: usize,
1548 mask: usize,
1549 out: usize,
1550 batch: u32,
1551 seq: u32,
1552 kv_seq: u32,
1553 heads: u32,
1554 head_dim: u32,
1555 mask_kind: rlx_ir::op::MaskKind,
1556 wrt: rlx_ir::op::AttentionBwdWrt,
1557 bhsd: bool,
1558 },
1559 /// [`rlx_ir::Op::AttentionBackwardAll`] — computes dQ, dK and dV from ONE
1560 /// score+softmax recompute, writing them into three disjoint sub-ranges of
1561 /// this node's own packed `[3B,S,H,D]` output slot (`out_q`/`out_k`/`out_v`
1562 /// are byte offsets into the same slot, so the planner manages it as one
1563 /// buffer). Downstream axis-0 `Narrow`s (view ops) recover dQ/dK/dV.
1564 AttentionBackwardAll {
1565 q: usize,
1566 k: usize,
1567 v: usize,
1568 dy: usize,
1569 mask: usize,
1570 out_q: usize,
1571 out_k: usize,
1572 out_v: usize,
1573 batch: u32,
1574 seq: u32,
1575 kv_seq: u32,
1576 heads: u32,
1577 head_dim: u32,
1578 mask_kind: rlx_ir::op::MaskKind,
1579 bhsd: bool,
1580 },
1581 /// RoPE (rotary position embeddings).
1582 /// `src_row_stride` is elements per source row (defaults to `hidden`
1583 /// for the standalone case; set to `qkv_axis * inner` when the
1584 /// thunk fusion pass below rewires Rope to read directly from the
1585 /// fused QKV buffer — plan #45).
1586 Rope {
1587 src: usize,
1588 cos: usize,
1589 sin: usize,
1590 dst: usize,
1591 batch: u32,
1592 seq: u32,
1593 hidden: u32,
1594 head_dim: u32,
1595 n_rot: u32,
1596 cos_len: u32,
1597 src_row_stride: u32,
1598 /// `true` = GPT-J / llama.cpp-NORM interleaved pairs `(2i, 2i+1)`;
1599 /// `false` = HF / NeoX rotate-half pairs `(i, i+n_rot/2)`.
1600 interleaved: bool,
1601 },
1602 /// Fused attention block: QKV proj → split → \[RoPE\] → SDPA → output proj.
1603 /// All intermediates stay in L1 cache. Zero arena writes between ops.
1604 FusedAttnBlock {
1605 hidden: usize,
1606 qkv_w: usize,
1607 out_w: usize,
1608 mask: usize,
1609 /// How to mask attention scores. `Custom` reads the per-key `mask`
1610 /// buffer (BERT-style padding); `Causal` / `SlidingWindow` are
1611 /// synthesized in-kernel with no buffer; `None` applies no mask.
1612 /// Mirrors [`Thunk::Attention`]'s `mask_kind` — the attention-block
1613 /// fusion MUST carry it through, otherwise a causal decoder silently
1614 /// attends to future tokens (the fused kernel would only ever apply
1615 /// the per-key padding mask).
1616 mask_kind: rlx_ir::op::MaskKind,
1617 out: usize,
1618 qkv_b: usize,
1619 out_b: usize, // 0 = no bias
1620 cos: usize,
1621 sin: usize,
1622 cos_len: u32, // 0 = no RoPE
1623 batch: u32,
1624 seq: u32,
1625 hs: u32,
1626 nh: u32,
1627 dh: u32,
1628 /// QKV-projection input width (`in_dim`) and output-projection output
1629 /// width (`out_dim`). Both equal `hs` in the BERT case; MoonViT-style
1630 /// attention has `in_dim != hs` (input hidden ≠ nh·dh) and/or
1631 /// `out_dim != hs`. The kernel reads the input as `m·in_dim`, runs the
1632 /// QKV GEMM with K=in_dim, the out GEMM with N=out_dim, and writes
1633 /// `m·out_dim`; the attention core still works on the `hs`-wide qkv.
1634 in_dim: u32,
1635 out_dim: u32,
1636 has_bias: bool,
1637 has_rope: bool,
1638 /// RoPE pairing: `true` = GPT-J interleaved `(2i,2i+1)`, `false` = NeoX
1639 /// rotate-half `(i,i+d/2)`. Captured from the fused `Op::Rope` so the
1640 /// inline rope matches the standalone kernel (a GptJ model must not be
1641 /// silently rotated NeoX-style by the fused path).
1642 interleaved: bool,
1643 },
1644 /// Fused ENTIRE transformer layer: attention + residual + LN + FFN + residual + LN.
1645 /// Combines ~10 thunks into 1. All intermediates on stack. Zero arena traffic.
1646 FusedBertLayer {
1647 // attention
1648 hidden: usize,
1649 qkv_w: usize,
1650 qkv_b: usize,
1651 out_w: usize,
1652 out_b: usize,
1653 mask: usize,
1654 // LN1
1655 ln1_g: usize,
1656 ln1_b: usize,
1657 eps1: f32,
1658 // FFN (GELU)
1659 fc1_w: usize,
1660 fc1_b: usize,
1661 fc2_w: usize,
1662 fc2_b: usize,
1663 // LN2
1664 ln2_g: usize,
1665 ln2_b: usize,
1666 eps2: f32,
1667 // output
1668 out: usize,
1669 // dims
1670 batch: u32,
1671 seq: u32,
1672 hs: u32,
1673 nh: u32,
1674 dh: u32,
1675 int_dim: u32,
1676 },
1677 /// Fused Nomic transformer layer: attention+RoPE + residual + LN + SwiGLU FFN + residual + LN.
1678 FusedNomicLayer {
1679 hidden: usize,
1680 qkv_w: usize,
1681 out_w: usize,
1682 mask: usize,
1683 cos: usize,
1684 sin: usize,
1685 cos_len: u32,
1686 ln1_g: usize,
1687 ln1_b: usize,
1688 eps1: f32,
1689 fc11_w: usize,
1690 fc12_w: usize,
1691 fc2_w: usize,
1692 ln2_g: usize,
1693 ln2_b: usize,
1694 eps2: f32,
1695 out: usize,
1696 batch: u32,
1697 seq: u32,
1698 hs: u32,
1699 nh: u32,
1700 dh: u32,
1701 int_dim: u32,
1702 /// RoPE pairing, threaded from the consumed `FusedAttnBlock`
1703 /// (`true` = GPT-J interleaved, `false` = NeoX rotate-half).
1704 interleaved: bool,
1705 },
1706 /// Fused SwiGLU: out\[r,i\] = x\[r,i\] * silu(x[r, n_half+i]).
1707 /// Input: [outer, 2*n_half] — concatenated up||gate per row.
1708 /// Output: [outer, n_half].
1709 FusedSwiGLU {
1710 src: usize,
1711 dst: usize,
1712 n_half: u32,
1713 total: u32,
1714 gate_first: bool,
1715 },
1716 /// Concat along an axis: output[outer, axis, inner] = inputs concatenated.
1717 /// Each entry of `inputs` is (src_offset, axis_len_for_that_input) in u32
1718 /// elements. `outer`, `inner`, and `total_axis_len` are pre-computed
1719 /// at compile time to avoid per-run shape work.
1720 Concat {
1721 dst: usize,
1722 outer: u32,
1723 inner: u32,
1724 total_axis: u32,
1725 /// `(src_offset, axis_extent, input_numel)` — `input_numel` enables
1726 /// outer-dim broadcast when rank-deficient inputs are concatenated.
1727 inputs: Vec<(usize, u32, u32)>,
1728 },
1729 /// Element-wise comparison: out = (lhs CMP rhs) ? 1 : 0 (Bool u8 or F32 0/1).
1730 /// `lhs_scalar` / `rhs_scalar` broadcast a 1-element operand across `len`
1731 /// (ONNX Less/Greater with a scalar Constant).
1732 Compare {
1733 lhs: usize,
1734 rhs: usize,
1735 dst: usize,
1736 len: u32,
1737 op: CmpOp,
1738 /// Nonzero when lhs/rhs are i64 (mask/range ops).
1739 inputs_i64: u8,
1740 /// Input element size (1 = Bool, 4 = F32, 8 = I64).
1741 inputs_elem_bytes: u8,
1742 /// Output element size (1 = Bool, 4 = F32).
1743 dst_elem_bytes: u8,
1744 lhs_scalar: bool,
1745 rhs_scalar: bool,
1746 /// TRUE element counts of each operand. Needed when an operand is a
1747 /// *partial* broadcast (`[.,L,1]` into `[.,L,L]`): its numel is neither
1748 /// 1 (scalar) nor `len` (full), so the flat fast path would read past
1749 /// it into adjacent arena memory.
1750 lhs_len: u32,
1751 rhs_len: u32,
1752 /// Shape-aware broadcast metadata (mirrors `BinaryFull`). Empty
1753 /// `out_dims_bcast` selects the flat scalar/full fast path; when
1754 /// non-empty, each output element indexes its operands via row-major
1755 /// coords · strides — the ONLY correct path for a trailing/partial
1756 /// size-1 broadcast, which the flat path silently miscomputes.
1757 out_dims_bcast: Vec<u32>,
1758 bcast_lhs_strides: Vec<u32>,
1759 bcast_rhs_strides: Vec<u32>,
1760 },
1761 /// Reduction along a contiguous range of axes. Input layout (after
1762 /// shape decomposition) is `[outer, reduced, inner]`; output is
1763 /// `[outer, inner]`. The single-axis cases (axis=0 → outer=1;
1764 /// axis=last → inner=1) and contiguous multi-axis (e.g. reduce over
1765 /// [0, 1] of an [N, C, H, W] tensor → outer=1, reduced=N*C, inner=H*W)
1766 /// all map onto this triplet. Non-contiguous axes are not supported
1767 /// and bail to Nop in the compile pass.
1768 Reduce {
1769 src: usize,
1770 dst: usize,
1771 outer: u32,
1772 reduced: u32,
1773 inner: u32,
1774 op: ReduceOp,
1775 },
1776 /// Histogram: bucket `n` input elements into `bins` equal-width buckets over
1777 /// `[min, max]` and write f32 counts into `dst[bins]`. Out-of-range elements
1778 /// are dropped; `x == max` lands in the last bin. See `Op::Histogram`.
1779 Histogram {
1780 src: usize,
1781 dst: usize,
1782 n: u32,
1783 bins: u32,
1784 min: f32,
1785 max: f32,
1786 },
1787 /// Index of the max (`is_max`) or min along the reduced axis; writes the
1788 /// winning index as an f32 into `dst`.
1789 ArgReduce {
1790 src: usize,
1791 dst: usize,
1792 outer: u32,
1793 reduced: u32,
1794 inner: u32,
1795 is_max: bool,
1796 },
1797 /// Top-K **indices** along the last axis. Input shape `[outer, axis_dim]`,
1798 /// output `[outer, k]` (f32 or i64 per `indices_i64`). Ties broken by
1799 /// smaller index. Used by MoE gating + beam search.
1800 TopK {
1801 src: usize,
1802 dst: usize,
1803 outer: u32,
1804 axis_dim: u32,
1805 k: u32,
1806 indices_i64: u8,
1807 },
1808 /// Indexed batched matmul: out\[i\] = input\[i\] @ weight[expert_idx\[i\]].
1809 /// Naive impl per token; for real MoE workloads, sort-by-expert + run
1810 /// segmented GEMM would amortize. Done when there's a workload.
1811 GroupedMatMul {
1812 input: usize,
1813 weight: usize,
1814 expert_idx: usize,
1815 dst: usize,
1816 m: u32,
1817 k_dim: u32,
1818 n: u32,
1819 num_experts: u32,
1820 },
1821 /// GGUF K-quant packed expert stack + grouped matmul (MoE FFN).
1822 DequantGroupedMatMulGguf {
1823 input: usize,
1824 w_q: usize,
1825 expert_idx: usize,
1826 dst: usize,
1827 m: u32,
1828 k_dim: u32,
1829 n: u32,
1830 num_experts: u32,
1831 scheme: rlx_ir::quant::QuantScheme,
1832 },
1833 /// MLX-affine grouped expert matmul: per row, dequant expert `idx[row]`'s
1834 /// affine slab (codes+scales+biases) and matmul. `[M,K]×dequant([N,K])→[M,N]`.
1835 DequantGroupedMatMulMlx {
1836 input: usize,
1837 w_q: usize,
1838 scale: usize,
1839 zp: usize,
1840 expert_idx: usize,
1841 dst: usize,
1842 m: u32,
1843 k_dim: u32,
1844 n: u32,
1845 num_experts: u32,
1846 /// Packed-code bytes per expert (`total_w_q_bytes / num_experts`).
1847 slab_bytes: u32,
1848 scheme: rlx_ir::quant::QuantScheme,
1849 /// Scales/biases stored as BF16 (2 bytes) instead of F32 — halves the
1850 /// resident scale memory (the MoE scale slabs rival the packed codes at
1851 /// 2-bit). Decoded to f32 per-expert on the fly.
1852 scale_bf16: bool,
1853 },
1854 /// Materialize packed MoE weights to F32 `[E, K, N]` (autodiff helper).
1855 DequantMoEWeightsGguf {
1856 w_q: usize,
1857 dst: usize,
1858 k_dim: u32,
1859 n: u32,
1860 num_experts: u32,
1861 scheme: rlx_ir::quant::QuantScheme,
1862 },
1863 /// Scatter-add: dst[indices\[i\] * trailing + j] += updates[i * trailing + j].
1864 /// Output is zeroed first; multiple updates to the same row accumulate.
1865 ScatterAdd {
1866 updates: usize,
1867 indices: usize,
1868 dst: usize,
1869 num_updates: u32,
1870 out_dim: u32,
1871 trailing: u32,
1872 },
1873 /// ONNX ScatterND: copy `data`, write `updates` at multi-index locations.
1874 /// Inputs `[data, indices, updates]`; `indices_i64` selects I64 vs F32 index storage.
1875 ScatterNd {
1876 data: usize,
1877 indices: usize,
1878 updates: usize,
1879 dst: usize,
1880 data_shape: Vec<u32>,
1881 indices_shape: Vec<u32>,
1882 data_len: u32,
1883 updates_len: u32,
1884 indices_len: u32,
1885 indices_i64: u8,
1886 reduction: rlx_ir::ScatterNdReduction,
1887 },
1888 /// ONNX ScatterElements along `axis`.
1889 ScatterElements {
1890 data: usize,
1891 indices: usize,
1892 updates: usize,
1893 dst: usize,
1894 data_shape: Vec<u32>,
1895 data_len: u32,
1896 updates_len: u32,
1897 indices_len: u32,
1898 indices_i64: u8,
1899 axis: i32,
1900 reduction: rlx_ir::ScatterNdReduction,
1901 },
1902 /// ONNX GatherND.
1903 GatherNd {
1904 data: usize,
1905 indices: usize,
1906 dst: usize,
1907 data_shape: Vec<u32>,
1908 indices_shape: Vec<u32>,
1909 data_len: u32,
1910 indices_len: u32,
1911 out_len: u32,
1912 indices_i64: u8,
1913 batch_dims: i32,
1914 },
1915 /// ONNX GatherElements / take_along_axis.
1916 GatherElements {
1917 data: usize,
1918 indices: usize,
1919 dst: usize,
1920 data_shape: Vec<u32>,
1921 indices_shape: Vec<u32>,
1922 data_len: u32,
1923 indices_len: u32,
1924 out_len: u32,
1925 indices_i64: u8,
1926 /// Byte width of a data/output element (GatherElements preserves dtype).
1927 /// Non-f32 data (e.g. int64 token ids) must be gathered at its true width.
1928 data_elem_bytes: u8,
1929 axis: i32,
1930 },
1931 /// Ternary select: out = cond != 0 ? on_true : on_false.
1932 /// `*_scalar` broadcast a 1-element operand across `len`.
1933 Where {
1934 cond: usize,
1935 on_true: usize,
1936 on_false: usize,
1937 dst: usize,
1938 len: u32,
1939 elem_bytes: u8,
1940 /// Element size for cond (1 = Bool mask, 4 = F32 0/1).
1941 cond_elem_bytes: u8,
1942 cond_scalar: bool,
1943 true_scalar: bool,
1944 false_scalar: bool,
1945 },
1946 /// Single-rounded fused multiply-add: `dst = a*b + c` (one rounding —
1947 /// enables error-free transforms / compensated arithmetic).
1948 Fma {
1949 a: usize,
1950 b: usize,
1951 c: usize,
1952 dst: usize,
1953 len: u32,
1954 elem_bytes: u8,
1955 },
1956 /// General N-D transpose / broadcast. `out_dims[i]` is the output's dim
1957 /// i length; `in_strides[i]` is the input stride (in elements) used to
1958 /// index that dim — 0 for broadcast dims (Expand). `in_total` is the
1959 /// total element count in the source buffer (≤ output total when
1960 /// broadcasting). Strides are pre-computed at compile time.
1961 Transpose {
1962 src: usize,
1963 dst: usize,
1964 in_total: u32,
1965 out_dims: Vec<u32>,
1966 in_strides: Vec<u32>,
1967 elem_bytes: u8,
1968 },
1969 /// Gather along an arbitrary axis. `outer = product(dims[..axis])`,
1970 /// `trailing = product(dims[axis+1..])`, `axis_dim` = the dimension
1971 /// being indexed into. Output: outer × num_idx × trailing.
1972 /// (axis=0 still routes to the simpler Thunk::Gather fast path.)
1973 GatherAxis {
1974 table: usize,
1975 idx: usize,
1976 dst: usize,
1977 outer: u32,
1978 axis_dim: u32,
1979 num_idx: u32,
1980 trailing: u32,
1981 idx_i64: u8,
1982 table_bytes: u8,
1983 },
1984 /// 2D pooling (Max or Mean). Input layout [N, C, H, W], output
1985 /// [N, C, H_out, W_out]. Padding is implicit-zero; Mean divides by
1986 /// the full kernel area (matches torch's `count_include_pad=True`).
1987 Pool2D {
1988 src: usize,
1989 dst: usize,
1990 n: u32,
1991 c: u32,
1992 h: u32,
1993 w: u32,
1994 h_out: u32,
1995 w_out: u32,
1996 kh: u32,
1997 kw: u32,
1998 sh: u32,
1999 sw: u32,
2000 ph: u32,
2001 pw: u32,
2002 kind: ReduceOp,
2003 },
2004 /// 3D pooling (Max or Mean). Input `[N, C, D, H, W]`, output
2005 /// `[N, C, D_out, H_out, W_out]`. Mean divides by the full kernel volume
2006 /// (`count_include_pad=True`).
2007 Pool3D {
2008 src: usize,
2009 dst: usize,
2010 n: u32,
2011 c: u32,
2012 d: u32,
2013 h: u32,
2014 w: u32,
2015 d_out: u32,
2016 h_out: u32,
2017 w_out: u32,
2018 kd: u32,
2019 kh: u32,
2020 kw: u32,
2021 sd: u32,
2022 sh: u32,
2023 sw: u32,
2024 pd: u32,
2025 ph: u32,
2026 pw: u32,
2027 kind: ReduceOp,
2028 },
2029 /// 2D convolution. Input [N, C_in, H, W], weight [C_out, C_in_per_group, kH, kW],
2030 /// output [N, C_out, H_out, W_out]. Bias is a separate Op::Binary::Add
2031 /// after the conv (matching the IR's input layout — Op::Conv has 2 inputs).
2032 /// Naive direct convolution; sufficient for correctness, not optimised.
2033 Conv2D {
2034 src: usize,
2035 weight: usize,
2036 dst: usize,
2037 n: u32,
2038 c_in: u32,
2039 h: u32,
2040 w: u32,
2041 c_out: u32,
2042 h_out: u32,
2043 w_out: u32,
2044 kh: u32,
2045 kw: u32,
2046 sh: u32,
2047 sw: u32,
2048 ph: u32,
2049 pw: u32,
2050 dh: u32,
2051 dw: u32,
2052 groups: u32,
2053 },
2054
2055 // ── Backward / training kernels ─────────────────────────────
2056 /// Real INT8 matmul with i32 accumulation.
2057 /// `out[m, n] = requantize(bias[n] + Σₖ (x[m,k]-x_zp)·(w[k,n]-w_zp), mult, out_zp)`
2058 /// Reads `x` and `w` as i8, `bias` as i32; writes `out` as i8.
2059 /// Same kernel shape as `rlx_cortexm::dense::dense_i8` — promoted
2060 /// to a desktop thunk so a quantized graph compiled here doesn't
2061 /// have to round-trip through fake-quant.
2062 QMatMul {
2063 x: usize,
2064 w: usize,
2065 bias: usize,
2066 out: usize,
2067 m: u32,
2068 k: u32,
2069 n: u32,
2070 x_zp: i32,
2071 w_zp: i32,
2072 out_zp: i32,
2073 mult: f32,
2074 },
2075
2076 /// Real INT8 conv2d, NCHW layout. Same loop shape as `Thunk::Conv2D`
2077 /// but with i8 reads, i32 accumulation, and per-output requantize
2078 /// to i8. Bias is i32 in the accumulator scale.
2079 QConv2d {
2080 x: usize,
2081 w: usize,
2082 bias: usize,
2083 out: usize,
2084 n: u32,
2085 c_in: u32,
2086 h: u32,
2087 w_in: u32,
2088 c_out: u32,
2089 h_out: u32,
2090 w_out: u32,
2091 kh: u32,
2092 kw: u32,
2093 sh: u32,
2094 sw: u32,
2095 ph: u32,
2096 pw: u32,
2097 dh: u32,
2098 dw: u32,
2099 groups: u32,
2100 x_zp: i32,
2101 w_zp: i32,
2102 out_zp: i32,
2103 mult: f32,
2104 },
2105
2106 /// INT8 quantize. Reads `x` as f32, writes `q` as i8.
2107 /// `chan = (i / inner) % chan_dim` selects the per-channel
2108 /// scale/zp; `chan_axis` is informational only (the kernel uses
2109 /// `chan_dim` and `inner` directly).
2110 /// For per-tensor, `chan_dim = 1` and `inner = len` so `chan` is
2111 /// always 0.
2112 Quantize {
2113 x: usize,
2114 q: usize,
2115 len: u32,
2116 chan_axis: u32,
2117 chan_dim: u32,
2118 inner: u32,
2119 scales: Vec<f32>,
2120 zero_points: Vec<i32>,
2121 },
2122
2123 /// INT8 dequantize — inverse of `Thunk::Quantize`.
2124 Dequantize {
2125 q: usize,
2126 x: usize,
2127 len: u32,
2128 chan_axis: u32,
2129 chan_dim: u32,
2130 inner: u32,
2131 scales: Vec<f32>,
2132 zero_points: Vec<i32>,
2133 },
2134
2135 /// QAT fake-quantize. Per-channel (or per-tensor) symmetric
2136 /// quantize-then-dequantize on the fly. Computes
2137 /// `s[c] = max(|x[..., c, ...]|) / q_max`
2138 /// then
2139 /// `out[i] = clamp(round(x[i]/s[c]), -q_max, q_max) * s[c]`
2140 /// with `q_max = {127, 7, 1}` for `bits = {8, 4, 2}`. Same
2141 /// channel-layout convention as `Thunk::Quantize`: every
2142 /// element's channel is `(i / inner) % chan_dim`. The kernel
2143 /// does two passes — one to scan max-abs per channel, one to
2144 /// quant-dequant per element.
2145 FakeQuantize {
2146 x: usize,
2147 out: usize,
2148 len: u32,
2149 chan_axis: u32,
2150 chan_dim: u32,
2151 inner: u32,
2152 bits: u8,
2153 /// STE variant — informational on the forward side (output is
2154 /// the same regardless), kernel-relevant in the matching
2155 /// `FakeQuantizeBackward` thunk.
2156 ste: rlx_ir::op::SteKind,
2157 /// Scale-tracking strategy. `PerBatch` recomputes
2158 /// `max_abs/q_max` every call (the original path). `EMA{decay}`
2159 /// blends per-batch max-abs into the `state_off` buffer; `Fixed`
2160 /// reads `state_off` and never updates it.
2161 scale_mode: rlx_ir::op::ScaleMode,
2162 /// `Some(off)` for `EMA` and `Fixed`; `None` for `PerBatch`.
2163 /// Points at a `[chan_dim]` f32 buffer holding the running scale
2164 /// per channel.
2165 state_off: Option<usize>,
2166 },
2167
2168 /// Backward pass for `Op::FakeQuantize` under one of four STE
2169 /// variants. Computes `dx[i]` from the f32 forward input `x` and
2170 /// the upstream gradient `dy`, using the same per-channel scale
2171 /// scheme as the forward.
2172 FakeQuantizeBackward {
2173 x: usize,
2174 dy: usize,
2175 dx: usize,
2176 len: u32,
2177 chan_axis: u32,
2178 chan_dim: u32,
2179 inner: u32,
2180 bits: u8,
2181 ste: rlx_ir::op::SteKind,
2182 },
2183
2184 /// LSQ forward — same kernel shape as `FakeQuantize` Fixed mode.
2185 /// Reads scale from `scale_off` (a `[chan_dim]` Param tensor).
2186 FakeQuantizeLSQ {
2187 x: usize,
2188 scale_off: usize,
2189 out: usize,
2190 len: u32,
2191 chan_axis: u32,
2192 chan_dim: u32,
2193 inner: u32,
2194 bits: u8,
2195 },
2196
2197 /// LSQ backward, x-gradient. STE-clipped: passes upstream
2198 /// through inside the quantization range, zeros outside.
2199 FakeQuantizeLSQBackwardX {
2200 x: usize,
2201 scale_off: usize,
2202 dy: usize,
2203 dx: usize,
2204 len: u32,
2205 chan_axis: u32,
2206 chan_dim: u32,
2207 inner: u32,
2208 bits: u8,
2209 },
2210
2211 /// LSQ backward, scale-gradient. Per-channel:
2212 /// `dscale[c] = sum_i ψ(x[i]/s[c]) · upstream[i]`
2213 /// where `ψ(z) = -z + round(z)` if `|z| ≤ q_max` else
2214 /// `sign(z) · q_max`. Output shape: `[chan_dim]`.
2215 FakeQuantizeLSQBackwardScale {
2216 x: usize,
2217 scale_off: usize,
2218 dy: usize,
2219 dscale: usize,
2220 len: u32,
2221 chan_axis: u32,
2222 chan_dim: u32,
2223 inner: u32,
2224 bits: u8,
2225 },
2226
2227 /// ReLU backward: `dx[i] = dy[i] if x[i] > 0 else 0`.
2228 ReluBackward {
2229 x: usize,
2230 dy: usize,
2231 dx: usize,
2232 len: u32,
2233 },
2234 /// f64 sibling of `ReluBackward` — same shape as the f32 variant
2235 /// but reads/writes 8 bytes per element. Required because
2236 /// `ReluBackward`'s `&[f32]` slot view returns half of every f64
2237 /// otherwise → backward silently produces 0 gradients on an f64
2238 /// graph. Mirrors the `ActivationBackwardF64` split.
2239 ReluBackwardF64 {
2240 x: usize,
2241 dy: usize,
2242 dx: usize,
2243 len: u32,
2244 },
2245
2246 /// Generic element-wise activation backward.
2247 /// `dx[i] = (d/dx act(x))[i] · dy[i]`. The closure dispatch is
2248 /// per-element; expensive activations (Gelu) recompute internals
2249 /// inline rather than threading an extra "saved y" tensor through.
2250 ActivationBackward {
2251 x: usize,
2252 dy: usize,
2253 dx: usize,
2254 len: u32,
2255 kind: Activation,
2256 },
2257 /// f64 sibling of `ActivationBackward` — slot offsets, len in
2258 /// elements; kernel reads/writes 8 bytes per element. Required
2259 /// because `ActivationBackward`'s `&[f32]` slot view silently
2260 /// returns garbage on an f64 graph (cb % 4 still works but every
2261 /// loaded value is half of an f64 → wrong gradient).
2262 ActivationBackwardF64 {
2263 x: usize,
2264 dy: usize,
2265 dx: usize,
2266 len: u32,
2267 kind: Activation,
2268 },
2269
2270 /// LayerNorm backward — input gradient. Recomputes mean/var/x̂ from
2271 /// `x` and emits the closed-form `d_x` per row.
2272 LayerNormBackwardInput {
2273 x: usize,
2274 gamma: usize,
2275 dy: usize,
2276 dx: usize,
2277 rows: u32,
2278 h: u32,
2279 eps: f32,
2280 },
2281
2282 /// LayerNorm backward — gamma gradient. `d_gamma[d] = Σ_row dy·x̂`.
2283 LayerNormBackwardGamma {
2284 x: usize,
2285 dy: usize,
2286 dgamma: usize,
2287 rows: u32,
2288 h: u32,
2289 eps: f32,
2290 },
2291
2292 RmsNormBackwardInput {
2293 x: usize,
2294 gamma: usize,
2295 beta: usize,
2296 dy: usize,
2297 dx: usize,
2298 rows: u32,
2299 h: u32,
2300 eps: f32,
2301 },
2302 RmsNormBackwardGamma {
2303 x: usize,
2304 gamma: usize,
2305 beta: usize,
2306 dy: usize,
2307 dgamma: usize,
2308 rows: u32,
2309 h: u32,
2310 eps: f32,
2311 },
2312 RmsNormBackwardBeta {
2313 x: usize,
2314 gamma: usize,
2315 beta: usize,
2316 dy: usize,
2317 dbeta: usize,
2318 rows: u32,
2319 h: u32,
2320 eps: f32,
2321 },
2322 RopeBackward {
2323 dy: usize,
2324 cos: usize,
2325 sin: usize,
2326 dx: usize,
2327 batch: u32,
2328 seq: u32,
2329 hidden: u32,
2330 head_dim: u32,
2331 n_rot: u32,
2332 cos_len: u32,
2333 },
2334 CumsumBackward {
2335 dy: usize,
2336 dx: usize,
2337 rows: u32,
2338 cols: u32,
2339 exclusive: bool,
2340 },
2341 GatherBackward {
2342 dy: usize,
2343 indices: usize,
2344 dst: usize,
2345 outer: u32,
2346 axis_dim: u32,
2347 num_idx: u32,
2348 trailing: u32,
2349 },
2350
2351 GroupNormBackwardInput {
2352 x: usize,
2353 gamma: usize,
2354 beta: usize,
2355 dy: usize,
2356 dx: usize,
2357 n: u32,
2358 c: u32,
2359 h: u32,
2360 w: u32,
2361 num_groups: u32,
2362 eps: f32,
2363 },
2364 GroupNormBackwardGamma {
2365 x: usize,
2366 dy: usize,
2367 dgamma: usize,
2368 n: u32,
2369 c: u32,
2370 h: u32,
2371 w: u32,
2372 num_groups: u32,
2373 eps: f32,
2374 },
2375 GroupNormBackwardBeta {
2376 dy: usize,
2377 dbeta: usize,
2378 n: u32,
2379 c: u32,
2380 h: u32,
2381 w: u32,
2382 },
2383
2384 /// 2D max-pool backward (NCHW). Recomputes the argmax position
2385 /// inside each window and accumulates `dy` into `dx` at that
2386 /// position. Output is zeroed first; ties resolve to the first
2387 /// hit (lowest (kh,kw) index), matching what the forward kernel
2388 /// does with `acc.max(v)`.
2389 MaxPool2dBackward {
2390 x: usize,
2391 dy: usize,
2392 dx: usize,
2393 n: u32,
2394 c: u32,
2395 h: u32,
2396 w: u32,
2397 h_out: u32,
2398 w_out: u32,
2399 kh: u32,
2400 kw: u32,
2401 sh: u32,
2402 sw: u32,
2403 ph: u32,
2404 pw: u32,
2405 },
2406
2407 /// 2D conv backward w.r.t. input (`dx = conv_transpose(dy, w)`).
2408 /// `dy [N, C_out, H_out, W_out]`, `w [C_out, C_in_per_group, kH, kW]`,
2409 /// `dx [N, C_in, H, W]`.
2410 Conv2dBackwardInput {
2411 dy: usize,
2412 w: usize,
2413 dx: usize,
2414 n: u32,
2415 c_in: u32,
2416 h: u32,
2417 w_in: u32,
2418 c_out: u32,
2419 h_out: u32,
2420 w_out: u32,
2421 kh: u32,
2422 kw: u32,
2423 sh: u32,
2424 sw: u32,
2425 ph: u32,
2426 pw: u32,
2427 dh: u32,
2428 dw: u32,
2429 groups: u32,
2430 },
2431
2432 /// 2D conv backward w.r.t. weight. `x [N, C_in, H, W]`,
2433 /// `dy [N, C_out, H_out, W_out]`, `dw [C_out, C_in_per_group, kH, kW]`.
2434 /// `dw` is zeroed before accumulation.
2435 Conv2dBackwardWeight {
2436 x: usize,
2437 dy: usize,
2438 dw: usize,
2439 n: u32,
2440 c_in: u32,
2441 h: u32,
2442 w: u32,
2443 c_out: u32,
2444 h_out: u32,
2445 w_out: u32,
2446 kh: u32,
2447 kw: u32,
2448 sh: u32,
2449 sw: u32,
2450 ph: u32,
2451 pw: u32,
2452 dh: u32,
2453 dw_dil: u32,
2454 groups: u32,
2455 },
2456
2457 MaxPool3dBackward {
2458 x: usize,
2459 dy: usize,
2460 dx: usize,
2461 n: u32,
2462 c: u32,
2463 d: u32,
2464 h: u32,
2465 w: u32,
2466 d_out: u32,
2467 h_out: u32,
2468 w_out: u32,
2469 kd: u32,
2470 kh: u32,
2471 kw: u32,
2472 sd: u32,
2473 sh: u32,
2474 sw: u32,
2475 pd: u32,
2476 ph: u32,
2477 pw: u32,
2478 },
2479
2480 Conv3dBackwardInput {
2481 dy: usize,
2482 w: usize,
2483 dx: usize,
2484 n: u32,
2485 c_in: u32,
2486 d: u32,
2487 h: u32,
2488 w_in: u32,
2489 c_out: u32,
2490 d_out: u32,
2491 h_out: u32,
2492 w_out: u32,
2493 kd: u32,
2494 kh: u32,
2495 kw: u32,
2496 sd: u32,
2497 sh: u32,
2498 sw: u32,
2499 pd: u32,
2500 ph: u32,
2501 pw: u32,
2502 dd: u32,
2503 dh: u32,
2504 dw: u32,
2505 groups: u32,
2506 },
2507
2508 Conv3dBackwardWeight {
2509 x: usize,
2510 dy: usize,
2511 dw: usize,
2512 n: u32,
2513 c_in: u32,
2514 d: u32,
2515 h: u32,
2516 w: u32,
2517 c_out: u32,
2518 d_out: u32,
2519 h_out: u32,
2520 w_out: u32,
2521 kd: u32,
2522 kh: u32,
2523 kw: u32,
2524 sd: u32,
2525 sh: u32,
2526 sw: u32,
2527 pd: u32,
2528 ph: u32,
2529 pw: u32,
2530 dd: u32,
2531 dh: u32,
2532 dw_dil: u32,
2533 groups: u32,
2534 },
2535
2536 /// NCHW im2col for conv backward-weight matmul. Output `[M, C·kH·kW]`
2537 /// with `M = N · H_out · W_out`. `n == 0` means infer batch from `x`.
2538 Im2Col {
2539 x: usize,
2540 col: usize,
2541 n: u32,
2542 c_in: u32,
2543 h: u32,
2544 w: u32,
2545 h_out: u32,
2546 w_out: u32,
2547 kh: u32,
2548 kw: u32,
2549 sh: u32,
2550 sw: u32,
2551 ph: u32,
2552 pw: u32,
2553 dh: u32,
2554 dw_dil: u32,
2555 },
2556
2557 /// Fused softmax + cross-entropy loss against a dense target
2558 /// distribution. `logits [N, C]`, `targets [N, C]`, output `[N]`
2559 /// per-row loss `lse(logits[n]) - Σ_c targets[n,c]·logits[n,c]`.
2560 /// Numerically stable (max-subtract before exp).
2561 SoftmaxCrossEntropyDense {
2562 logits: usize,
2563 targets: usize,
2564 dst: usize,
2565 n: u32,
2566 c: u32,
2567 },
2568
2569 /// Fused softmax + cross-entropy loss with f32-encoded integer
2570 /// labels. `logits [N, C]`, `labels [N]`, output `[N]` per-row loss.
2571 /// Numerically stable (max-subtract before exp).
2572 SoftmaxCrossEntropy {
2573 logits: usize,
2574 labels: usize,
2575 dst: usize,
2576 n: u32,
2577 c: u32,
2578 },
2579
2580 /// Backward of the fused loss above.
2581 /// `dlogits[n, k] = (softmax(logits[n])[k] - one_hot(labels[n])[k]) * d_loss[n]`.
2582 SoftmaxCrossEntropyBackward {
2583 logits: usize,
2584 labels: usize,
2585 d_loss: usize,
2586 dlogits: usize,
2587 n: u32,
2588 c: u32,
2589 },
2590
2591 /// User-registered custom op (CPU side). Lowered from `Op::Custom`.
2592 /// `kernel` is resolved against the global CPU kernel registry at
2593 /// compile time and stored as `Arc<dyn CpuKernel>` so execution
2594 /// avoids per-call lookups. v1: f32 contiguous only — see
2595 /// `op_registry::CpuKernel::execute_f32`.
2596 CustomOp {
2597 kernel: Arc<dyn CpuKernel>,
2598 inputs: Vec<(usize, u32, Shape)>, // (offset, len_elements, shape)
2599 output: (usize, u32, Shape), // (offset, len_elements, shape)
2600 attrs: Vec<u8>,
2601 },
2602
2603 /// 1D FFT along the last axis. Input/output are `[..., 2N]`
2604 /// real-block layout (first N real, second N imag along the
2605 /// transformed axis). `outer` is the product of all leading axes;
2606 /// `n_complex` is N (the number of complex points). Both halves
2607 /// of the real-block layout are read together by the kernel.
2608 /// `dtype` selects the f32 or f64 path; the two share structure
2609 /// but not buffers, so a flag at compile time avoids per-row
2610 /// dispatch.
2611 /// CPU reference 3D Gaussian splat render ([`rlx_ir::Op::GaussianSplatRender`]).
2612 GaussianSplatRender {
2613 positions_off: usize,
2614 positions_len: usize,
2615 scales_off: usize,
2616 scales_len: usize,
2617 rotations_off: usize,
2618 rotations_len: usize,
2619 opacities_off: usize,
2620 opacities_len: usize,
2621 colors_off: usize,
2622 colors_len: usize,
2623 sh_coeffs_off: usize,
2624 sh_coeffs_len: usize,
2625 meta_off: usize,
2626 dst_off: usize,
2627 dst_len: usize,
2628 width: u32,
2629 height: u32,
2630 tile_size: u32,
2631 radius_scale: f32,
2632 alpha_cutoff: f32,
2633 max_splat_steps: u32,
2634 transmittance_threshold: f32,
2635 max_list_entries: u32,
2636 },
2637 GaussianSplatRenderBackward {
2638 positions_off: usize,
2639 positions_len: usize,
2640 scales_off: usize,
2641 scales_len: usize,
2642 rotations_off: usize,
2643 rotations_len: usize,
2644 opacities_off: usize,
2645 opacities_len: usize,
2646 colors_off: usize,
2647 colors_len: usize,
2648 sh_coeffs_off: usize,
2649 sh_coeffs_len: usize,
2650 meta_off: usize,
2651 d_loss_off: usize,
2652 d_loss_len: usize,
2653 packed_off: usize,
2654 packed_len: usize,
2655 width: u32,
2656 height: u32,
2657 tile_size: u32,
2658 radius_scale: f32,
2659 alpha_cutoff: f32,
2660 max_splat_steps: u32,
2661 transmittance_threshold: f32,
2662 max_list_entries: u32,
2663 loss_grad_clip: f32,
2664 sh_band: u32,
2665 max_anisotropy: f32,
2666 },
2667 /// Strict IR stage 1 — project + bin + sort + rays ([`Op::GaussianSplatPrepare`]).
2668 GaussianSplatPrepare {
2669 positions_off: usize,
2670 positions_len: usize,
2671 scales_off: usize,
2672 scales_len: usize,
2673 rotations_off: usize,
2674 rotations_len: usize,
2675 opacities_off: usize,
2676 opacities_len: usize,
2677 colors_off: usize,
2678 colors_len: usize,
2679 sh_coeffs_off: usize,
2680 sh_coeffs_len: usize,
2681 meta_off: usize,
2682 meta_len: usize,
2683 prep_off: usize,
2684 prep_len: usize,
2685 width: u32,
2686 height: u32,
2687 tile_size: u32,
2688 radius_scale: f32,
2689 alpha_cutoff: f32,
2690 max_splat_steps: u32,
2691 transmittance_threshold: f32,
2692 max_list_entries: u32,
2693 },
2694 /// Strict IR stage 2 — tile raster from prepare buffer ([`Op::GaussianSplatRasterize`]).
2695 GaussianSplatRasterize {
2696 prep_off: usize,
2697 prep_len: usize,
2698 meta_off: usize,
2699 meta_len: usize,
2700 dst_off: usize,
2701 dst_len: usize,
2702 count: usize,
2703 width: u32,
2704 height: u32,
2705 tile_size: u32,
2706 alpha_cutoff: f32,
2707 max_splat_steps: u32,
2708 transmittance_threshold: f32,
2709 max_list_entries: u32,
2710 },
2711 Fft1d {
2712 src: usize,
2713 dst: usize,
2714 outer: u32,
2715 n_complex: u32,
2716 inverse: bool,
2717 norm_tag: u32,
2718 dtype: rlx_ir::DType,
2719 },
2720 FftButterflyStage {
2721 state_src: usize,
2722 state_dst: usize,
2723 gate_src: usize,
2724 rev_src: usize,
2725 tw_re_src: usize,
2726 tw_im_src: usize,
2727 batch: u32,
2728 n_fft: u32,
2729 stage: u32,
2730 },
2731 LogMel {
2732 spec: usize,
2733 filters: usize,
2734 dst: usize,
2735 outer: u32,
2736 n_fft: u32,
2737 n_bins: u32,
2738 n_mels: u32,
2739 },
2740 LogMelBackward {
2741 spec: usize,
2742 filters: usize,
2743 dy: usize,
2744 dst: usize,
2745 outer: u32,
2746 n_fft: u32,
2747 n_bins: u32,
2748 n_mels: u32,
2749 },
2750 WelchPeaks {
2751 spec: usize,
2752 dst: usize,
2753 welch_batch: u32,
2754 n_fft: u32,
2755 n_segments: u32,
2756 k: u32,
2757 },
2758}
2759
2760/// Compiled thunk schedule — the runtime hot path.
2761/// Nop thunks are filtered out at compile time for zero iteration overhead.
2762#[derive(Clone)]
2763pub struct ThunkSchedule {
2764 pub thunks: Vec<Thunk>,
2765 /// TIDE merged placement mask (union across layers).
2766 pub moe_resident: Option<std::sync::Arc<[bool]>>,
2767 /// Per MoE layer placement (`layer[e]`); preferred when set.
2768 pub moe_resident_layers: Option<std::sync::Arc<Vec<std::sync::Arc<[bool]>>>>,
2769 /// MoE router TopK capture (per-layer refresh).
2770 pub moe_topk_capture: Option<std::sync::Arc<crate::moe_topk_capture::MoeTopkCapture>>,
2771 /// Cached config values.
2772 pub mask_threshold: f32,
2773 pub mask_neg_inf: f32,
2774 pub score_skip: f32,
2775 /// Pre-compiled closure dispatch (zero match overhead). `Arc` (not
2776 /// `Box`) so the schedule can be `Clone` — multiple parallel
2777 /// executors share the same compiled closures (they're read-only
2778 /// `Fn(*mut u8)` so concurrent dispatch is safe; the arena pointer
2779 /// they receive is the only mutable state and is per-executor).
2780 pub compiled_fns: Vec<Arc<dyn Fn(*mut u8) + Send + Sync>>,
2781 /// Runtime-mutable RNG policy for [`Thunk::RngNormal`] / [`Thunk::RngUniform`].
2782 pub rng: Arc<std::sync::RwLock<rlx_ir::RngOptions>>,
2783}
2784
2785impl ThunkSchedule {
2786 pub fn strip_nops(&mut self) {
2787 self.thunks.retain(|t| !matches!(t, Thunk::Nop));
2788 // compiled_fns must be rebuilt after stripping — caller should
2789 // call strip_nops() before compile_closures().
2790 self.compiled_fns.clear();
2791 }
2792}
2793
2794/// Get the arena byte offset for a node.
2795pub(crate) fn node_offset(arena: &Arena, id: NodeId) -> usize {
2796 if arena.has_buffer(id) {
2797 arena.byte_offset(id)
2798 } else {
2799 usize::MAX
2800 }
2801}
2802
2803/// Every byte-offset that a thunk reads from. Used by the Narrow→Rope
2804/// fusion (#45) to verify a Narrow's dst has exactly one consumer
2805/// before eliding it. Conservative: when in doubt about reads (an op
2806/// not yet listed here), the fusion will skip — correctness over
2807/// completeness.
2808pub(crate) fn thunk_read_offsets(t: &Thunk) -> Vec<usize> {
2809 match t {
2810 Thunk::Sgemm { a, b, .. } => vec![*a, *b],
2811 Thunk::SgemmBf16 { a, b, .. } => vec![*a, *b],
2812 Thunk::SgemmF16 { a, b, .. } => vec![*a, *b],
2813 Thunk::SgemmT { a, b, .. } => vec![*a, *b],
2814 Thunk::SgdMomentum {
2815 param, vel, grad, ..
2816 } => vec![*param, *vel, *grad],
2817 Thunk::DenseSolveF64 { a, b, .. } => vec![*a, *b],
2818 Thunk::DenseSolveF32 { a, b, .. } => vec![*a, *b],
2819 Thunk::BatchedDenseSolveF64 { a, b, .. } => vec![*a, *b],
2820 Thunk::BatchedDgemmF64 { a, b, .. } => vec![*a, *b],
2821 Thunk::BatchedSgemm { a, b, .. } => vec![*a, *b],
2822 Thunk::FusedMmBiasAct { a, w, bias, .. } => vec![*a, *w, *bias],
2823 Thunk::ElementwiseRegion { input_offs, .. } => input_offs.clone(),
2824 Thunk::BiasAdd { src, bias, .. } => vec![*src, *bias],
2825 Thunk::BinaryFull { lhs, rhs, .. } => vec![*lhs, *rhs],
2826 Thunk::BinaryFullF64 { lhs, rhs, .. } => vec![*lhs, *rhs],
2827 Thunk::BinaryFullC64 { lhs, rhs, .. } => vec![*lhs, *rhs],
2828 Thunk::ComplexNormSqF32 { src, .. } => vec![*src],
2829 Thunk::ComplexNormSqBackwardF32 { z, g, .. } => vec![*z, *g],
2830 Thunk::ConjugateC64 { src, .. } => vec![*src],
2831 Thunk::Scan {
2832 outer_init_off,
2833 xs_inputs,
2834 ..
2835 } => {
2836 let mut v = vec![*outer_init_off];
2837 for (_, outer_xs_off, _) in xs_inputs.iter() {
2838 v.push(*outer_xs_off);
2839 }
2840 v
2841 }
2842 Thunk::ScanBackward {
2843 outer_init_off,
2844 outer_traj_off,
2845 outer_upstream_off,
2846 outer_xs_offs,
2847 ..
2848 } => {
2849 let mut v = vec![*outer_init_off, *outer_traj_off, *outer_upstream_off];
2850 for (off, _) in outer_xs_offs.iter() {
2851 v.push(*off);
2852 }
2853 v
2854 }
2855 Thunk::ScanBackwardXs {
2856 outer_init_off,
2857 outer_traj_off,
2858 outer_upstream_off,
2859 outer_xs_offs,
2860 ..
2861 } => {
2862 let mut v = vec![*outer_init_off, *outer_traj_off, *outer_upstream_off];
2863 for (off, _) in outer_xs_offs.iter() {
2864 v.push(*off);
2865 }
2866 v
2867 }
2868 Thunk::CustomFn { inputs, .. } => {
2869 inputs.iter().map(|(_, outer_off, _)| *outer_off).collect()
2870 }
2871 Thunk::ActivationInPlace { data, .. } => vec![*data],
2872 Thunk::LayerNorm { src, g, b, .. } | Thunk::GroupNorm { src, g, b, .. } => {
2873 vec![*src, *g, *b]
2874 }
2875 Thunk::BatchNormInference {
2876 src,
2877 g,
2878 b,
2879 mean,
2880 var,
2881 ..
2882 } => vec![*src, *g, *b, *mean, *var],
2883 Thunk::ResizeNearest2x { src, .. } => vec![*src],
2884 Thunk::Interpolate3d { src, .. } => vec![*src],
2885 Thunk::AxialRope2d { src, .. } => vec![*src],
2886 Thunk::FusedResidualLN {
2887 x, res, bias, g, b, ..
2888 } => vec![*x, *res, *bias, *g, *b],
2889 Thunk::FusedResidualRmsNorm {
2890 x, res, bias, g, b, ..
2891 } => vec![*x, *res, *bias, *g, *b],
2892 Thunk::RmsNorm { src, g, b, .. } => vec![*src, *g, *b],
2893 Thunk::AdaLayerNorm {
2894 src, scale, shift, ..
2895 } => vec![*src, *scale, *shift],
2896 Thunk::GatedResidual { x, y, gate, .. } => vec![*x, *y, *gate],
2897 Thunk::AdaLayerNormBackward {
2898 x,
2899 scale,
2900 shift,
2901 dy,
2902 ..
2903 } => vec![*x, *scale, *shift, *dy],
2904 Thunk::GatedResidualBackward { x, y, gate, dy, .. } => vec![*x, *y, *gate, *dy],
2905 Thunk::Softmax { data, .. } => vec![*data],
2906 Thunk::Cumsum { src, .. } => vec![*src],
2907 Thunk::CumScan { src, .. } => vec![*src],
2908 Thunk::Sample { logits, .. } => vec![*logits],
2909 Thunk::RngNormal { .. } | Thunk::RngUniform { .. } => vec![],
2910 Thunk::LoraMatMul { x, w, a, b, .. } => vec![*x, *w, *a, *b],
2911 Thunk::DequantMatMul {
2912 x, w_q, scale, zp, ..
2913 } => vec![*x, *w_q, *scale, *zp],
2914 Thunk::SynthMatMul {
2915 x,
2916 indices,
2917 codebook,
2918 ..
2919 } => vec![*x, *indices, *codebook],
2920 Thunk::SplineActivation { x, coeff, .. } => vec![*x, *coeff],
2921 Thunk::DequantMatMulGguf { x, w_q, .. } => vec![*x, *w_q],
2922 Thunk::DequantMatMulInt4 {
2923 x, w_q, scale, zp, ..
2924 } => vec![*x, *w_q, *scale, *zp],
2925 Thunk::DequantMatMulFp8 { x, w_q, scale, .. } => vec![*x, *w_q, *scale],
2926 Thunk::DequantMatMulNvfp4 {
2927 x,
2928 w_q,
2929 scale,
2930 global_scale,
2931 ..
2932 } => vec![*x, *w_q, *scale, *global_scale],
2933 Thunk::DequantMatMulMxFp4x2 { x, w_q, scale, .. } => vec![*x, *w_q, *scale],
2934 Thunk::DequantMatMulMlx {
2935 x, w_q, scale, zp, ..
2936 } => vec![*x, *w_q, *scale, *zp],
2937 Thunk::ScaledMatMul {
2938 lhs,
2939 rhs,
2940 lhs_scale,
2941 rhs_scale,
2942 bias,
2943 has_bias,
2944 ..
2945 } => {
2946 let mut v = vec![*lhs, *rhs, *lhs_scale, *rhs_scale];
2947 if *has_bias {
2948 v.push(*bias);
2949 }
2950 v
2951 }
2952 Thunk::ScaledQuantize { x, scale, .. } => vec![*x, *scale],
2953 Thunk::ScaledQuantScale { x, .. } => vec![*x],
2954 Thunk::ScaledDequantize { codes, scale, .. } => vec![*codes, *scale],
2955 Thunk::Conv2D1x1 { src, weight, .. } => vec![*src, *weight],
2956 Thunk::SelectiveScan {
2957 x, delta, a, b, c, ..
2958 } => vec![*x, *delta, *a, *b, *c],
2959 Thunk::GatedDeltaNet {
2960 q,
2961 k,
2962 v,
2963 g,
2964 beta,
2965 state,
2966 carry_state,
2967 ..
2968 } => {
2969 let mut v = vec![*q, *k, *v, *g, *beta];
2970 if *carry_state {
2971 v.push(*state);
2972 }
2973 v
2974 }
2975 Thunk::Attention { q, k, v, mask, .. } => vec![*q, *k, *v, *mask],
2976 Thunk::AttentionBackward {
2977 q, k, v, dy, mask, ..
2978 } => {
2979 let mut v = vec![*q, *k, *v, *dy];
2980 if *mask != 0 {
2981 v.push(*mask);
2982 }
2983 v
2984 }
2985 Thunk::Rope { src, cos, sin, .. } => vec![*src, *cos, *sin],
2986 Thunk::FusedAttnBlock {
2987 hidden,
2988 qkv_w,
2989 out_w,
2990 mask,
2991 qkv_b,
2992 out_b,
2993 cos,
2994 sin,
2995 ..
2996 } => vec![*hidden, *qkv_w, *out_w, *mask, *qkv_b, *out_b, *cos, *sin],
2997 Thunk::FusedSwiGLU { src, .. } => vec![*src],
2998 Thunk::Concat { inputs, .. } => inputs.iter().map(|(off, _, _)| *off).collect(),
2999 Thunk::ConcatF64 { inputs, .. } => inputs.iter().map(|(off, _, _)| *off).collect(),
3000 Thunk::Narrow { src, .. } => vec![*src],
3001 Thunk::Copy { src, .. } => vec![*src],
3002 Thunk::Gather { table, idx, .. } => vec![*table, *idx],
3003 Thunk::Histogram { src, .. } => vec![*src],
3004 // Anything not enumerated → return the dst as a "read" too,
3005 // forcing the fusion to bail (read_count >= 2 → skip). Keeps
3006 // this list safe to be incomplete.
3007 _ => vec![],
3008 }
3009}
3010
3011/// A scan body compiled into a self-contained CPU schedule plus the byte
3012/// offsets needed to drive it. Backends without a native `Op::Scan` kernel
3013/// (e.g. Metal via the unified-memory host fallback) build this once at compile
3014/// time and hand it to [`execute_scan_host`] at run time. Mirrors the inline
3015/// `Op::Scan` compilation used by the CPU executor.
3016pub struct ScanBodyPlan {
3017 pub body: ThunkSchedule,
3018 pub body_init: Vec<u8>,
3019 pub body_input_off: usize,
3020 pub body_output_off: usize,
3021 pub carry_bytes: usize,
3022 /// Body-arena offsets of the broadcast inputs, in declaration order.
3023 pub bcast_body_offs: Vec<usize>,
3024 /// Body-arena offsets of the per-step (`xs`) inputs, in declaration order.
3025 pub xs_body_offs: Vec<usize>,
3026}
3027
3028impl std::fmt::Debug for ScanBodyPlan {
3029 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3030 f.debug_struct("ScanBodyPlan")
3031 .field("carry_bytes", &self.carry_bytes)
3032 .field("body_input_off", &self.body_input_off)
3033 .field("body_output_off", &self.body_output_off)
3034 .field("num_bcast", &self.bcast_body_offs.len())
3035 .field("num_xs", &self.xs_body_offs.len())
3036 .finish()
3037 }
3038}