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