rlx_cpu/thunk/exec_dispatch.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/// Execute a thunk schedule on a raw arena buffer.
8/// Fastest executor: call pre-compiled closures sequentially.
9/// Zero match dispatch — each closure is a direct kernel call.
10pub fn execute_compiled(schedule: &ThunkSchedule, arena_buf: &mut [u8]) {
11 let base = arena_buf.as_mut_ptr();
12 // `RLX_CPU_WATCH=<byte_offset>:<n_floats>` reports every thunk that changes
13 // the watched region. A node read by one thunk and written by a LATER one
14 // is a scheduling hazard the post-run arena dump cannot see: it shows the
15 // final value, so the reader looks like it was fed data it never got.
16 if let Some(spec) = rlx_ir::env::var("RLX_CPU_WATCH") {
17 let mut it = spec
18 .split(':')
19 .filter_map(|t| t.trim().parse::<usize>().ok());
20 if let (Some(off), Some(len)) = (it.next(), it.next()) {
21 let mut prev: Option<Vec<f32>> = None;
22 for (i, f) in schedule.compiled_fns.iter().enumerate() {
23 f(base);
24 let cur: Vec<f32> =
25 unsafe { std::slice::from_raw_parts(base.add(off) as *const f32, len) }
26 .to_vec();
27 if prev.as_deref() != Some(cur.as_slice()) {
28 let tabs: f64 = cur.iter().map(|&v| v.abs() as f64).sum();
29 eprintln!(
30 "[rlx-cpu-watch] thunk {i} {:?} changed watch: tabs={tabs:.6}",
31 schedule.thunks.get(i).map(std::mem::discriminant)
32 );
33 prev = Some(cur);
34 }
35 }
36 return;
37 }
38 }
39 for f in &schedule.compiled_fns {
40 f(base);
41 }
42}
43
44/// Active-extent execution stub. The runtime calls this when it has an
45/// active-extent hint set. CPU doesn't implement per-thunk active-extent
46/// scaling yet — return false so the caller falls back to the full
47/// `execute_thunks` path.
48pub fn execute_thunks_active(
49 schedule: &ThunkSchedule,
50 _arena_buf: &mut [u8],
51 _actual: usize,
52 _upper: usize,
53) -> bool {
54 let _ = schedule;
55 false
56}
57
58/// Match-based executor (fallback, used by tests).
59pub(crate) struct MoeResidencyGuard;
60impl Drop for MoeResidencyGuard {
61 fn drop(&mut self) {
62 if let Some(stats) = crate::moe_residency::take_stats() {
63 crate::moe_residency::stash_last_forward_stats(stats);
64 } else {
65 crate::moe_residency::clear_mask();
66 }
67 }
68}
69
70/// Contiguous AVX2 Add/Mul/Sub. Large buffers are split across Rayon so we
71/// keep SIMD *and* multi-core (serial AVX2 alone lost to 16-wide scalar rayon
72/// on F5-TTS DiT residuals).
73#[cfg(target_arch = "x86_64")]
74#[inline]
75fn binary_contig_f32(l: &[f32], r: &[f32], o: &mut [f32], op: BinaryOp) -> bool {
76 if !std::arch::is_x86_feature_detected!("avx2") {
77 return false;
78 }
79 if !matches!(op, BinaryOp::Add | BinaryOp::Mul | BinaryOp::Sub) {
80 return false;
81 }
82 let len = o.len();
83 let l_ptr = l.as_ptr() as usize;
84 let r_ptr = r.as_ptr() as usize;
85 let o_ptr = o.as_mut_ptr() as usize;
86 let run = |i0: usize, i1: usize| unsafe {
87 use std::arch::x86_64::*;
88 let l = l_ptr as *const f32;
89 let r = r_ptr as *const f32;
90 let o = o_ptr as *mut f32;
91 let mut i = i0;
92 // Align to 8-wide when possible within the chunk.
93 while i + 8 <= i1 {
94 let a = _mm256_loadu_ps(l.add(i));
95 let b = _mm256_loadu_ps(r.add(i));
96 let res = match op {
97 BinaryOp::Add => _mm256_add_ps(a, b),
98 BinaryOp::Sub => _mm256_sub_ps(a, b),
99 BinaryOp::Mul => _mm256_mul_ps(a, b),
100 _ => unreachable!(),
101 };
102 _mm256_storeu_ps(o.add(i), res);
103 i += 8;
104 }
105 while i < i1 {
106 let a = *l.add(i);
107 let b = *r.add(i);
108 *o.add(i) = match op {
109 BinaryOp::Add => a + b,
110 BinaryOp::Sub => a - b,
111 BinaryOp::Mul => a * b,
112 _ => unreachable!(),
113 };
114 i += 1;
115 }
116 };
117 if len >= 8192 && crate::pool::num_threads() > 1 {
118 crate::pool::par_for(len, crate::pool::chunk_floor(len), &|off, cnt| {
119 run(off, off + cnt);
120 });
121 } else {
122 run(0, len);
123 }
124 true
125}
126
127#[cfg(not(target_arch = "x86_64"))]
128#[inline]
129#[allow(dead_code)]
130fn binary_contig_f32(_l: &[f32], _r: &[f32], _o: &mut [f32], _op: BinaryOp) -> bool {
131 false
132}
133
134/// Row-tiled rhs broadcast: `o[i] = op(l[i], r[i % rl])` with optional AVX2.
135#[inline]
136fn binary_row_bcast_f32(l: &[f32], r: &[f32], o: &mut [f32], op: BinaryOp, rl: usize) -> bool {
137 let len = o.len();
138 if rl == 0 || rl >= len || !len.is_multiple_of(rl) || l.len() < len || r.len() < rl {
139 return false;
140 }
141 let rows = len / rl;
142 let l_ptr = l.as_ptr() as usize;
143 let r_ptr = r.as_ptr() as usize;
144 let o_ptr = o.as_mut_ptr() as usize;
145
146 #[cfg(target_arch = "x86_64")]
147 let use_avx2 = rl >= 8
148 && rl.is_multiple_of(8)
149 && matches!(op, BinaryOp::Add | BinaryOp::Mul | BinaryOp::Sub)
150 && std::arch::is_x86_feature_detected!("avx2");
151 #[cfg(not(target_arch = "x86_64"))]
152 let _use_avx2 = false;
153
154 let run_rows = |row0: usize, row1: usize| unsafe {
155 let l = l_ptr as *const f32;
156 let r = r_ptr as *const f32;
157 let o = o_ptr as *mut f32;
158 #[cfg(target_arch = "x86_64")]
159 if use_avx2 {
160 use std::arch::x86_64::*;
161 let chunks = rl / 8;
162 for row in row0..row1 {
163 let base = row * rl;
164 for c in 0..chunks {
165 let off = base + c * 8;
166 let roff = c * 8;
167 let a = _mm256_loadu_ps(l.add(off));
168 let b = _mm256_loadu_ps(r.add(roff));
169 let res = match op {
170 BinaryOp::Add => _mm256_add_ps(a, b),
171 BinaryOp::Sub => _mm256_sub_ps(a, b),
172 BinaryOp::Mul => _mm256_mul_ps(a, b),
173 _ => unreachable!(),
174 };
175 _mm256_storeu_ps(o.add(off), res);
176 }
177 }
178 return;
179 }
180 for row in row0..row1 {
181 let base = row * rl;
182 for j in 0..rl {
183 let i = base + j;
184 let a = *l.add(i);
185 let b = *r.add(j);
186 *o.add(i) = match op {
187 BinaryOp::Add => a + b,
188 BinaryOp::Sub => a - b,
189 BinaryOp::Mul => a * b,
190 BinaryOp::Div => a / b,
191 BinaryOp::Max => a.max(b),
192 BinaryOp::Min => a.min(b),
193 BinaryOp::Pow => a.powf(b),
194 BinaryOp::Mod => a % b,
195 BinaryOp::Atan2 => a.atan2(b),
196 BinaryOp::BitAnd => ((a as i64) & (b as i64)) as f32,
197 BinaryOp::BitOr => ((a as i64) | (b as i64)) as f32,
198 BinaryOp::BitXor => ((a as i64) ^ (b as i64)) as f32,
199 BinaryOp::Shl => (a as i64).wrapping_shl(b as u32) as f32,
200 BinaryOp::Shr => (a as i64).wrapping_shr(b as u32) as f32,
201 };
202 }
203 }
204 };
205 if rows >= 4 && crate::pool::num_threads() > 1 && len >= 8192 {
206 crate::pool::par_for(rows, 1, &|off, cnt| run_rows(off, off + cnt));
207 } else {
208 run_rows(0, rows);
209 }
210 true
211}
212
213pub(crate) fn thunk_kind_name(t: &Thunk) -> &'static str {
214 // Zero-alloc variant name from the derived `IntoStaticStr` on `Thunk`
215 // (see its definition). Covers EVERY variant automatically, so a newly
216 // added thunk is labeled by name instead of silently bucketing into
217 // "Other" — which had hidden the pathologically slow AttentionBackward.
218 <&'static str>::from(t)
219}
220
221/// Per-thunk-kind wall-time accumulator, populated only when the env var
222/// `RLX_PROFILE_THUNKS` is set. Used to see which ops dominate a step so the
223/// optimizer/kernels can target the real hotspots rather than guesses.
224pub(crate) static THUNK_PROFILE: std::sync::Mutex<
225 Option<std::collections::BTreeMap<&'static str, (u128, u64)>>,
226> = std::sync::Mutex::new(None);
227
228#[inline]
229pub(crate) fn profile_record(name: &'static str, d: std::time::Duration) {
230 let mut g = THUNK_PROFILE.lock().unwrap();
231 let map = g.get_or_insert_with(std::collections::BTreeMap::new);
232 let e = map.entry(name).or_insert((0, 0));
233 e.0 += d.as_nanos();
234 e.1 += 1;
235}
236
237/// Print and clear the per-thunk-kind time profile gathered under
238/// `RLX_PROFILE_THUNKS`. Call after a run to see where the time went.
239pub fn dump_thunk_profile() {
240 let mut g = THUNK_PROFILE.lock().unwrap();
241 if let Some(map) = g.take() {
242 let mut v: Vec<_> = map.into_iter().collect();
243 v.sort_by_key(|b| std::cmp::Reverse(b.1.0));
244 let total: u128 = v.iter().map(|(_, (ns, _))| *ns).sum();
245 eprintln!(
246 "[thunk-profile] total {:.1}ms across kinds:",
247 total as f64 / 1e6
248 );
249 for (name, (ns, c)) in v.iter().take(25) {
250 eprintln!(" {name:<28} {:>8.1}ms ({c} calls)", *ns as f64 / 1e6);
251 }
252 }
253}
254
255/// Do the arena byte ranges `[a, a + a_elems*4)` and `[c, c + c_elems*4)`
256/// overlap at all?
257///
258/// The BLAS aliasing guard has to compare *ranges*, not pointers. When a
259/// fusion folds a `Transpose` into a matmul, the surviving operand is the
260/// transpose's *source* buffer, which the planner scored as dead at the
261/// transpose — so its slot is free to be reused for the matmul's own output.
262/// That leaves an operand sitting strictly inside the destination (WavTokenizer
263/// attention: `probs @ V` reads a 256 B `probs` from the middle of its own
264/// 24 KB output slot). `ptr::eq` sees two different addresses and waves it
265/// through, after which cblas streams `probs` while overwriting it. Accelerate
266/// happens to pack its operands up front and survives; OpenBLAS does not, which
267/// is why this only ever corrupted x86_64.
268#[inline]
269fn arena_ranges_overlap(a: usize, a_elems: usize, c: usize, c_elems: usize) -> bool {
270 a < c + c_elems * 4 && c < a + a_elems * 4
271}
272
273pub fn execute_thunks(schedule: &ThunkSchedule, arena_buf: &mut [u8]) {
274 crate::moe_residency::reset_gmm_counters();
275 if let Some(layers) = schedule.moe_resident_layers.clone() {
276 crate::moe_residency::set_per_layer_masks(Some(layers));
277 } else {
278 crate::moe_residency::set_mask(schedule.moe_resident.clone());
279 }
280 if let Some(cap) = schedule.moe_topk_capture.as_ref() {
281 cap.clear();
282 }
283 let _moe_guard = MoeResidencyGuard;
284 let base = arena_buf.as_mut_ptr();
285 let mask_thr = schedule.mask_threshold;
286 let mask_neg = schedule.mask_neg_inf;
287 let score_thr = schedule.score_skip;
288 let thunks = &schedule.thunks;
289 let len = thunks.len();
290
291 // Pre-allocate ALL reusable buffers once (zero per-call allocation)
292 let max_h = thunks
293 .iter()
294 .filter_map(|t| match t {
295 Thunk::FusedResidualLN { h, .. }
296 | Thunk::FusedResidualRmsNorm { h, .. }
297 | Thunk::LayerNorm { h, .. } => Some(*h as usize),
298 _ => None,
299 })
300 .max()
301 .unwrap_or(0);
302 let zero_bias = vec![0f32; max_h];
303
304 // Pre-allocate per-(batch,head) score buffers for parallel SDPA.
305 // Q/K/V/out are accessed via strided BLAS — no deinterleave copy needed.
306 let max_sdpa = thunks
307 .iter()
308 .filter_map(|t| match t {
309 Thunk::Attention {
310 batch,
311 seq,
312 kv_seq,
313 heads,
314 head_dim,
315 ..
316 } => Some((
317 *batch as usize,
318 (*seq as usize).max(*kv_seq as usize),
319 *heads as usize,
320 *head_dim as usize,
321 )),
322 _ => None,
323 })
324 .fold((0, 0, 0, 0), |(mb, ms, mh, md), (b, s, h, d)| {
325 (mb.max(b), ms.max(s), mh.max(h), md.max(d))
326 });
327 let (max_batch, max_seq, max_heads, _max_dh) = max_sdpa;
328 let max_units = max_batch * max_heads;
329 let mut sdpa_scores = vec![0f32; max_units * max_seq * max_seq];
330
331 // Pre-allocate fused layer buffers (reused across all 12+ layers — zero malloc per layer)
332 let fl = thunks
333 .iter()
334 .filter_map(|t| match t {
335 Thunk::FusedBertLayer {
336 batch,
337 seq,
338 hs,
339 int_dim,
340 ..
341 } => {
342 let m = (*batch as usize) * (*seq as usize);
343 let h = *hs as usize;
344 let id = *int_dim as usize;
345 Some((m, h, id, m * (*seq as usize)))
346 }
347 Thunk::FusedNomicLayer {
348 batch,
349 seq,
350 hs,
351 int_dim,
352 ..
353 } => {
354 let m = (*batch as usize) * (*seq as usize);
355 let h = *hs as usize;
356 let id = *int_dim as usize;
357 Some((m, h, id, m * (*seq as usize)))
358 }
359 _ => None,
360 })
361 .fold((0, 0, 0, 0), |(mm, mh, mi, ms), (m, h, id, ss)| {
362 (mm.max(m), mh.max(h), mi.max(id), ms.max(ss))
363 });
364 let (fl_m, fl_h, fl_int, fl_ss) = fl;
365 let mut fl_qkv = vec![0f32; fl_m * 3 * fl_h];
366 let mut fl_attn = vec![0f32; fl_m * fl_h];
367 let mut fl_res = vec![0f32; fl_m * fl_h];
368 let mut fl_normed = vec![0f32; fl_m * fl_h];
369 let mut fl_ffn = vec![0f32; fl_m * fl_int.max(2 * fl_int)]; // Nomic needs 2×int for fused fc11+fc12
370 let mut fl_sc = vec![0f32; fl_ss.max(1)];
371
372 let trace_thunks = rlx_ir::env::var_os("RLX_TRACE_THUNK").is_some();
373 if trace_thunks {
374 eprintln!(
375 "[thunk] prealloc max_h={max_h} sdpa={} fl_m={fl_m} fl_h={fl_h} fl_int={fl_int}",
376 max_units * max_seq * max_seq
377 );
378 }
379 let profile = rlx_ir::env::var_os("RLX_PROFILE_THUNKS").is_some();
380 // Time the previous thunk at the top of each iteration (avoids touching the
381 // giant match's many arms). The last thunk's tail is folded into the next
382 // step's first sample — negligible over a training run.
383 let mut prof_prev: Option<(&'static str, std::time::Instant)> = None;
384 // Byte-range write tracer (`RLX_ARENA_WRITE_TRACE=1`). Snapshots the arena
385 // around every thunk and reports the exact range each one dirtied, so a
386 // kernel that writes past its planned slot — the classic SIMD-tail overrun
387 // that corrupts whichever buffer the planner reused next — is named
388 // directly instead of inferred. O(arena) per thunk, so diagnostics only.
389 let write_trace = rlx_ir::env::var_os("RLX_ARENA_WRITE_TRACE").is_some();
390 let mut wt_before: Vec<u8> = if write_trace {
391 arena_buf.to_vec()
392 } else {
393 Vec::new()
394 };
395 // `RLX_CPU_WATCH=<byte_offset>:<n_floats>` reports the watched region's
396 // contents whenever they change, naming the thunk about to run — so the
397 // writer is the thunk *before* the one printed. A node written after its
398 // reader is a scheduling hazard invisible to the post-run arena dump,
399 // which only ever shows the final value.
400 let watch = rlx_ir::env::var("RLX_CPU_WATCH").and_then(|spec| {
401 let mut it = spec
402 .split(':')
403 .filter_map(|t| t.trim().parse::<usize>().ok());
404 Some((it.next()?, it.next()?))
405 });
406 let mut watch_prev: Option<Vec<f32>> = None;
407 // `RLX_CPU_DUMP_THUNKS=<lo>:<hi>` prints the compiled thunks in that index
408 // range, with their arena offsets and dimensions — the arguments a kernel
409 // actually receives, as opposed to the ones the graph implies.
410 if let Some(spec) = rlx_ir::env::var("RLX_CPU_DUMP_THUNKS") {
411 let mut it = spec
412 .split(':')
413 .filter_map(|t| t.trim().parse::<usize>().ok());
414 if let (Some(lo), Some(hi)) = (it.next(), it.next()) {
415 for (i, t) in thunks.iter().enumerate().take(hi.min(len)).skip(lo) {
416 let detail = match t {
417 Thunk::FusedMmBiasAct {
418 a,
419 w,
420 bias,
421 c,
422 m,
423 k,
424 n,
425 act,
426 } => format!(" a={a} w={w} bias={bias} c={c} m={m} k={k} n={n} act={act:?}"),
427 Thunk::Sgemm { a, b, c, m, k, n } => {
428 format!(" a={a} b={b} c={c} m={m} k={k} n={n}")
429 }
430 _ => String::new(),
431 };
432 eprintln!("[rlx-cpu-thunk] {i} {}{detail}", thunk_kind_name(t));
433 }
434 }
435 }
436 for i in 0..len {
437 if profile {
438 if let Some((pn, pt)) = prof_prev.take() {
439 profile_record(pn, pt.elapsed());
440 }
441 }
442 let thunk = unsafe { thunks.get_unchecked(i) };
443 if let Some((woff, wlen)) = watch {
444 let cur: Vec<f32> =
445 unsafe { std::slice::from_raw_parts(base.add(woff) as *const f32, wlen) }.to_vec();
446 if watch_prev.as_deref() != Some(cur.as_slice()) {
447 let tabs: f64 = cur.iter().map(|&v| v.abs() as f64).sum();
448 eprintln!(
449 "[rlx-cpu-watch] before thunk {i} ({}) tabs={tabs:.6}",
450 thunk_kind_name(thunk)
451 );
452 watch_prev = Some(cur);
453 }
454 }
455 if trace_thunks && (i < 120 || i % 200 == 0 || i + 1 == len) {
456 eprintln!("[thunk {i}/{len}] {}", thunk_kind_name(thunk));
457 }
458 let trace_done = trace_thunks && i < 120;
459 if profile {
460 prof_prev = Some((thunk_kind_name(thunk), std::time::Instant::now()));
461 }
462 match thunk {
463 Thunk::Nop => exec_nop(thunk),
464 Thunk::ElementwiseRegion { .. } => exec_elementwise_region(thunk, base),
465 Thunk::GaussianSplatRender { .. } => exec_gaussian_splat_render(thunk, base),
466 Thunk::GaussianSplatRenderBackward { .. } => {
467 exec_gaussian_splat_render_backward(thunk, base)
468 }
469 Thunk::GaussianSplatPrepare { .. } => exec_gaussian_splat_prepare(thunk, base),
470 Thunk::GaussianSplatRasterize { .. } => exec_gaussian_splat_rasterize(thunk, base),
471 Thunk::Fft1d { .. } => exec_fft1d(thunk, base),
472 Thunk::Fft1dQ { .. } => exec_fft1d_q(thunk, base),
473 Thunk::FftButterflyStage { .. } => exec_fft_butterfly_stage(thunk, base),
474 Thunk::LogMel { .. } => exec_log_mel(thunk, base),
475 Thunk::LogMelBackward { .. } => exec_log_mel_backward(thunk, base),
476 Thunk::WelchPeaks { .. } => exec_welch_peaks(thunk, base),
477 Thunk::CustomFn { .. } => exec_custom_fn(thunk, base),
478 Thunk::Sgemm { a, b, c, m, k, n } => {
479 let (m, k, n) = (*m as usize, *k as usize, *n as usize);
480 if trace_thunks {
481 eprintln!("[sgemm] m={m} k={k} n={n} a={} b={} c={}", *a, *b, *c);
482 }
483 let c_len = m.saturating_mul(n);
484 let a_len = m.saturating_mul(k);
485 let b_len = k.saturating_mul(n);
486 let arena_len = arena_buf.len();
487 let max_a = (arena_len.saturating_sub(*a)) / 4;
488 let max_b = (arena_len.saturating_sub(*b)) / 4;
489 let max_c = (arena_len.saturating_sub(*c)) / 4;
490 let a_len = a_len.min(max_a);
491 let b_len = b_len.min(max_b);
492 let c_len = c_len.min(max_c);
493 unsafe {
494 let a_sl = sl(*a, base, a_len);
495 let b_sl = sl(*b, base, b_len);
496 let c_sl = sl_mut(*c, base, c_len);
497 if arena_ranges_overlap(*a, a_len, *c, c_len)
498 || arena_ranges_overlap(*b, b_len, *c, c_len)
499 {
500 let mut tmp = vec![0.0f32; c_len];
501 crate::blas::sgemm_auto(a_sl, b_sl, &mut tmp, m, k, n);
502 c_sl.copy_from_slice(&tmp);
503 } else {
504 crate::blas::sgemm_auto(a_sl, b_sl, c_sl, m, k, n);
505 }
506 }
507 }
508
509 Thunk::SgemmBf16 { a, b, c, m, k, n } => {
510 let (m, k, n) = (*m as usize, *k as usize, *n as usize);
511 let c_len = m.saturating_mul(n);
512 let a_len = m.saturating_mul(k);
513 let bn = k.saturating_mul(n);
514 unsafe {
515 let a_sl = sl(*a, base, a_len);
516 let b16 = sl_typed::<u16>(*b, base, bn);
517 let c_sl = sl_mut(*c, base, c_len);
518 crate::blas::sgemm_bf16_rhs(a_sl, b16, c_sl, m, k, n);
519 }
520 }
521
522 Thunk::SgemmF16 { a, b, c, m, k, n } => {
523 let (m, k, n) = (*m as usize, *k as usize, *n as usize);
524 let c_len = m.saturating_mul(n);
525 let a_len = m.saturating_mul(k);
526 let bn = k.saturating_mul(n);
527 unsafe {
528 let a_sl = sl(*a, base, a_len);
529 let b16 = sl_typed::<u16>(*b, base, bn);
530 let c_sl = sl_mut(*c, base, c_len);
531 crate::blas::sgemm_f16_rhs(a_sl, b16, c_sl, m, k, n);
532 }
533 }
534
535 Thunk::SgemmT {
536 a,
537 b,
538 c,
539 m,
540 k,
541 n,
542 ta,
543 tb,
544 } => {
545 // C[m,n] = op(A) @ op(B). RowMajor cblas: lda/ldb = stored
546 // row-length of each operand → m if A is transposed else k;
547 // k if B is transposed else n. Element counts are m*k / k*n
548 // regardless of layout.
549 let (m, k, n) = (*m as usize, *k as usize, *n as usize);
550 let lda = if *ta { m } else { k };
551 let ldb = if *tb { k } else { n };
552 if trace_thunks {
553 eprintln!(
554 "[sgemmT] m={m} k={k} n={n} ta={ta} tb={tb} a={a}..{} b={b}..{} c={c}..{}",
555 a + m * k * 4,
556 b + k * n * 4,
557 c + m * n * 4
558 );
559 }
560 let arena_len = arena_buf.len();
561 let a_len = (m * k).min((arena_len.saturating_sub(*a)) / 4);
562 let b_len = (k * n).min((arena_len.saturating_sub(*b)) / 4);
563 let c_len = (m * n).min((arena_len.saturating_sub(*c)) / 4);
564 unsafe {
565 let a_sl = sl(*a, base, a_len);
566 let b_sl = sl(*b, base, b_len);
567 let c_sl = sl_mut(*c, base, c_len);
568 let (ap, bp) = (a_sl.as_ptr(), b_sl.as_ptr());
569 if arena_ranges_overlap(*a, a_len, *c, c_len)
570 || arena_ranges_overlap(*b, b_len, *c, c_len)
571 {
572 let mut tmp = vec![0.0f32; c_len];
573 crate::blas::sgemm_general(
574 ap,
575 bp,
576 tmp.as_mut_ptr(),
577 m,
578 n,
579 k,
580 1.0,
581 0.0,
582 lda,
583 ldb,
584 n,
585 *ta,
586 *tb,
587 );
588 c_sl.copy_from_slice(&tmp);
589 } else {
590 crate::blas::sgemm_general(
591 ap,
592 bp,
593 c_sl.as_mut_ptr(),
594 m,
595 n,
596 k,
597 1.0,
598 0.0,
599 lda,
600 ldb,
601 n,
602 *ta,
603 *tb,
604 );
605 }
606 }
607 }
608
609 Thunk::SgdMomentum { .. } => exec_sgd_momentum(thunk, base),
610 Thunk::CgemmC64 { .. } => exec_cgemm_c64(thunk, base),
611 Thunk::DenseSolveF64 { .. } => exec_dense_solve_f64(thunk, base),
612 Thunk::DenseSolveF32 { .. } => exec_dense_solve_f32(thunk, base),
613 Thunk::CholeskyF32 { .. } => exec_cholesky(thunk, base),
614 Thunk::TriangularSolveF32 { .. } => exec_triangular_solve(thunk, base),
615 Thunk::DetF32 { .. } => exec_det(thunk, base),
616 Thunk::SortF32 { .. } => exec_sort(thunk, base),
617 Thunk::SvdF32 { .. } => exec_svd(thunk, base),
618 Thunk::QrF32 { .. } => exec_qr(thunk, base),
619 Thunk::BatchedDenseSolveF64 { .. } => exec_batched_dense_solve_f64(thunk, base),
620 Thunk::BatchedDenseSolveF32 { .. } => exec_batched_dense_solve_f32(thunk, base),
621 Thunk::BatchedDgemmF64 { .. } => exec_batched_dgemm_f64(thunk, base),
622 Thunk::BatchedSgemm {
623 a,
624 b,
625 c,
626 batch,
627 m,
628 k,
629 n,
630 a_bcast,
631 b_bcast,
632 } => {
633 let (b_, m_, k_, n_) = (*batch as usize, *m as usize, *k as usize, *n as usize);
634 if trace_thunks {
635 eprintln!(
636 "[batched-sgemm] batch={b_} m={m_} k={k_} n={n_} a_bcast={a_bcast} b_bcast={b_bcast} a={} b={} c={}",
637 *a, *b, *c
638 );
639 }
640 let a_mat = m_.saturating_mul(k_); // per-matrix element count
641 let b_mat = k_.saturating_mul(n_);
642 let c_mat = m_.saturating_mul(n_);
643 // A broadcast operand (batch dim 1) has batch stride 0 → reuse
644 // matrix 0 for every output batch; else stride by its matrix size.
645 let a_bstride = if *a_bcast { 0 } else { a_mat };
646 let b_bstride = if *b_bcast { 0 } else { b_mat };
647 let arena_len = arena_buf.len();
648 let a_cap = (arena_len.saturating_sub(*a)) / 4;
649 let b_cap = (arena_len.saturating_sub(*b)) / 4;
650 let c_cap = (arena_len.saturating_sub(*c)) / 4;
651 let a_count = if *a_bcast { 1 } else { b_ };
652 let b_count = if *b_bcast { 1 } else { b_ };
653 let a_elems = (a_count * a_mat).min(a_cap);
654 let b_elems = (b_count * b_mat).min(b_cap);
655 let c_elems = (b_ * c_mat).min(c_cap);
656 unsafe {
657 let a_full = sl(*a, base, a_elems);
658 let b_full = sl(*b, base, b_elems);
659 let c_full = sl_mut(*c, base, c_elems);
660 // Parallelize independent batch GEMMs with Rayon (BLAS
661 // stays 1-thread per worker). Serial for tiny batches.
662 if b_ >= 2 && crate::pool::num_threads() > 1 {
663 let a_ptr = a_full.as_ptr() as usize;
664 let b_ptr = b_full.as_ptr() as usize;
665 let c_ptr = c_full.as_mut_ptr() as usize;
666 let a_len = a_full.len();
667 let b_len = b_full.len();
668 let c_len = c_full.len();
669 crate::pool::par_for(b_, 1, &|off, cnt| {
670 for bi in off..off + cnt {
671 let a0 = bi * a_bstride;
672 let b0 = bi * b_bstride;
673 let c0 = bi * c_mat;
674 if a0 + a_mat > a_len || b0 + b_mat > b_len || c0 + c_mat > c_len {
675 break;
676 }
677 // SAFETY: pointers from disjoint batch slices; c
678 // ranges don't overlap across workers.
679 let a_slice = std::slice::from_raw_parts(
680 (a_ptr as *const f32).add(a0),
681 a_mat,
682 );
683 let b_slice = std::slice::from_raw_parts(
684 (b_ptr as *const f32).add(b0),
685 b_mat,
686 );
687 let c_slice = std::slice::from_raw_parts_mut(
688 (c_ptr as *mut f32).add(c0),
689 c_mat,
690 );
691 if std::ptr::eq(a_slice.as_ptr(), c_slice.as_mut_ptr())
692 || std::ptr::eq(b_slice.as_ptr(), c_slice.as_mut_ptr())
693 {
694 let mut tmp = vec![0.0f32; c_mat];
695 crate::blas::sgemm(a_slice, b_slice, &mut tmp, m_, k_, n_);
696 c_slice.copy_from_slice(&tmp);
697 } else {
698 crate::blas::sgemm(a_slice, b_slice, c_slice, m_, k_, n_);
699 }
700 }
701 });
702 } else {
703 for bi in 0..b_ {
704 let a0 = bi * a_bstride;
705 let b0 = bi * b_bstride;
706 let c0 = bi * c_mat;
707 if a0 + a_mat > a_full.len()
708 || b0 + b_mat > b_full.len()
709 || c0 + c_mat > c_full.len()
710 {
711 break;
712 }
713 let a_slice = &a_full[a0..a0 + a_mat];
714 let b_slice = &b_full[b0..b0 + b_mat];
715 let c_slice = &mut c_full[c0..c0 + c_mat];
716 if std::ptr::eq(a_slice.as_ptr(), c_slice.as_mut_ptr())
717 || std::ptr::eq(b_slice.as_ptr(), c_slice.as_mut_ptr())
718 {
719 let mut tmp = vec![0.0f32; c_mat];
720 crate::blas::sgemm_auto(a_slice, b_slice, &mut tmp, m_, k_, n_);
721 c_slice.copy_from_slice(&tmp);
722 } else {
723 crate::blas::sgemm_auto(a_slice, b_slice, c_slice, m_, k_, n_);
724 }
725 }
726 }
727 }
728 }
729
730 Thunk::Dgemm { .. } => exec_dgemm(thunk, base),
731 Thunk::TransposeF64 { .. } => exec_transpose_f64(thunk, base),
732 Thunk::ActivationF64 { .. } => exec_activation_f64(thunk, base),
733 Thunk::ReduceSumF64 { .. } => exec_reduce_sum_f64(thunk, base),
734 Thunk::CopyF64 { src, dst, len } => {
735 let mut len = *len as usize;
736 if *src == *dst || len == 0 {
737 continue;
738 }
739 let arena_len = arena_buf.len();
740 let max_from_src = (arena_len.saturating_sub(*src)) / 8;
741 let max_from_dst = (arena_len.saturating_sub(*dst)) / 8;
742 len = len.min(max_from_src).min(max_from_dst);
743 if len == 0 {
744 continue;
745 }
746 let byte_len = len.saturating_mul(8);
747 unsafe {
748 std::ptr::copy(base.add(*src), base.add(*dst), byte_len);
749 }
750 }
751
752 Thunk::CopyI64 { src, dst, len } => {
753 let mut len = *len as usize;
754 if *src == *dst || len == 0 {
755 continue;
756 }
757 let arena_len = arena_buf.len();
758 let max_from_src = (arena_len.saturating_sub(*src)) / 8;
759 let max_from_dst = (arena_len.saturating_sub(*dst)) / 8;
760 len = len.min(max_from_src).min(max_from_dst);
761 if len == 0 {
762 continue;
763 }
764 let byte_len = len.saturating_mul(8);
765 unsafe {
766 std::ptr::copy(base.add(*src), base.add(*dst), byte_len);
767 }
768 }
769
770 Thunk::CastF32ToI64 { src, dst, len } => {
771 let len = *len as usize;
772 if len == 0 {
773 continue;
774 }
775 unsafe {
776 let inp = sl(*src, base, len);
777 let out = sl_mut_i64(*dst, base, len);
778 // ONNX Cast float→int truncates toward zero (not round-half).
779 for i in 0..len {
780 out[i] = inp[i] as i64;
781 }
782 }
783 }
784
785 Thunk::CastF32ToF64 { src, dst, len } => {
786 let len = *len as usize;
787 if len == 0 {
788 continue;
789 }
790 unsafe {
791 let inp = sl(*src, base, len);
792 let out = sl_mut_f64(*dst, base, len);
793 for i in 0..len {
794 out[i] = inp[i] as f64;
795 }
796 }
797 }
798
799 Thunk::CastF64ToF32 { src, dst, len } => {
800 let len = *len as usize;
801 if len == 0 {
802 continue;
803 }
804 unsafe {
805 let inp = sl_f64(*src, base, len);
806 let out = sl_mut(*dst, base, len);
807 // f64 → f32 narrowing (round-to-nearest), not a bitcast.
808 for i in 0..len {
809 out[i] = inp[i] as f32;
810 }
811 }
812 }
813
814 Thunk::CastF64ToI64 { src, dst, len } => {
815 let len = *len as usize;
816 if len == 0 {
817 continue;
818 }
819 unsafe {
820 let inp = sl_f64(*src, base, len);
821 let out = sl_mut_i64(*dst, base, len);
822 // ONNX Cast float→int truncates toward zero.
823 for i in 0..len {
824 out[i] = inp[i] as i64;
825 }
826 }
827 }
828
829 Thunk::CastF64ToI32 { src, dst, len } => {
830 let len = *len as usize;
831 if len == 0 {
832 continue;
833 }
834 unsafe {
835 let inp = sl_f64(*src, base, len);
836 let out = sl_mut_i32(*dst, base, len);
837 // ONNX Cast float→int truncates toward zero.
838 for i in 0..len {
839 out[i] = inp[i] as i32;
840 }
841 }
842 }
843
844 Thunk::CastF64ToBool { src, dst, len } => {
845 let len = *len as usize;
846 if len == 0 {
847 continue;
848 }
849 // f64 = 8 bytes, bool = 1 byte — copy the input out first to
850 // avoid aliasing the arena while writing the narrower output.
851 let inp = unsafe { sl_f64(*src, base, len).to_vec() };
852 for i in 0..len {
853 arena_buf[*dst + i] = u8::from(inp[i] != 0.0);
854 }
855 }
856
857 Thunk::CastI64ToI32 { src, dst, len } => {
858 let len = *len as usize;
859 if len == 0 {
860 continue;
861 }
862 unsafe {
863 let inp = sl_i64(*src, base, len).to_vec();
864 let out = sl_mut_i32(*dst, base, len);
865 // Wrapping narrowing (matches Rust `as`, ONNX int→int).
866 for i in 0..len {
867 out[i] = inp[i] as i32;
868 }
869 }
870 }
871
872 Thunk::CastI64ToF64 { src, dst, len } => {
873 let len = *len as usize;
874 if len == 0 {
875 continue;
876 }
877 unsafe {
878 let inp = sl_i64(*src, base, len).to_vec();
879 let out = sl_mut_f64(*dst, base, len);
880 for i in 0..len {
881 out[i] = inp[i] as f64;
882 }
883 }
884 }
885
886 Thunk::CastI32ToF64 { src, dst, len } => {
887 let len = *len as usize;
888 if len == 0 {
889 continue;
890 }
891 unsafe {
892 let inp = sl_i32(*src, base, len).to_vec();
893 let out = sl_mut_f64(*dst, base, len);
894 for i in 0..len {
895 out[i] = inp[i] as f64;
896 }
897 }
898 }
899
900 Thunk::CastF32ToF16 { src, dst, len } => {
901 let len = *len as usize;
902 if len == 0 {
903 continue;
904 }
905 // f32 (4 bytes) → f16 (2 bytes). No sl_f16 accessor; copy the
906 // input out then write the LE f16 bytes into the arena.
907 let inp = unsafe { sl(*src, base, len).to_vec() };
908 for i in 0..len {
909 let h = half::f16::from_f32(inp[i]);
910 arena_buf[*dst + i * 2..*dst + i * 2 + 2].copy_from_slice(&h.to_le_bytes());
911 }
912 }
913
914 Thunk::CastF16ToF32 { src, dst, len } => {
915 let len = *len as usize;
916 if len == 0 {
917 continue;
918 }
919 // f16 (2 bytes) → f32 (4 bytes). Read the LE f16 bytes first,
920 // then write the widened f32 output.
921 let bytes: Vec<u8> = arena_buf[*src..*src + len * 2].to_vec();
922 unsafe {
923 let out = sl_mut(*dst, base, len);
924 for i in 0..len {
925 out[i] =
926 half::f16::from_le_bytes([bytes[i * 2], bytes[i * 2 + 1]]).to_f32();
927 }
928 }
929 }
930
931 Thunk::CastF32ToBf16 { src, dst, len } => {
932 let len = *len as usize;
933 if len == 0 {
934 continue;
935 }
936 let inp = unsafe { sl(*src, base, len).to_vec() };
937 for i in 0..len {
938 let h = half::bf16::from_f32(inp[i]);
939 arena_buf[*dst + i * 2..*dst + i * 2 + 2].copy_from_slice(&h.to_le_bytes());
940 }
941 }
942
943 Thunk::CastBf16ToF32 { src, dst, len } => {
944 let len = *len as usize;
945 if len == 0 {
946 continue;
947 }
948 let bytes: Vec<u8> = arena_buf[*src..*src + len * 2].to_vec();
949 unsafe {
950 let out = sl_mut(*dst, base, len);
951 for i in 0..len {
952 out[i] =
953 half::bf16::from_le_bytes([bytes[i * 2], bytes[i * 2 + 1]]).to_f32();
954 }
955 }
956 }
957
958 Thunk::CastF32ToI32 { src, dst, len } => {
959 let len = *len as usize;
960 if len == 0 {
961 continue;
962 }
963 unsafe {
964 let inp = sl(*src, base, len);
965 let out = sl_mut_i32(*dst, base, len);
966 // ONNX Cast float→int truncates toward zero (not round-half).
967 for i in 0..len {
968 out[i] = inp[i] as i32;
969 }
970 }
971 }
972
973 Thunk::CastI64ToF32 { src, dst, len } => {
974 let len = *len as usize;
975 if len == 0 {
976 continue;
977 }
978 unsafe {
979 let inp = sl_i64(*src, base, len);
980 let out = sl_mut(*dst, base, len);
981 for i in 0..len {
982 out[i] = inp[i] as f32;
983 }
984 }
985 }
986
987 Thunk::CastBoolToI32 { src, dst, len } => {
988 let len = *len as usize;
989 if len == 0 {
990 continue;
991 }
992 unsafe {
993 let inp = &arena_buf[*src..*src + len];
994 let out = sl_mut_i32(*dst, base, len);
995 for i in 0..len {
996 out[i] = i32::from(inp[i] != 0);
997 }
998 }
999 }
1000
1001 Thunk::CastI32ToF32 { src, dst, len } => {
1002 let len = *len as usize;
1003 if len == 0 {
1004 continue;
1005 }
1006 unsafe {
1007 let inp = sl_i32(*src, base, len);
1008 let out = sl_mut(*dst, base, len);
1009 for i in 0..len {
1010 out[i] = inp[i] as f32;
1011 }
1012 }
1013 }
1014
1015 Thunk::CastI32ToI64 { src, dst, len } => {
1016 let len = *len as usize;
1017 if len == 0 {
1018 continue;
1019 }
1020 unsafe {
1021 let inp = sl_i32(*src, base, len);
1022 let out = sl_mut_i64(*dst, base, len);
1023 for i in 0..len {
1024 out[i] = inp[i] as i64;
1025 }
1026 }
1027 }
1028
1029 Thunk::CastI32ToBool { src, dst, len } => {
1030 let len = *len as usize;
1031 if len == 0 {
1032 continue;
1033 }
1034 // src/dst are byte offsets; i32 = 4 bytes, bool = 1 byte. Copy the
1035 // input first to avoid aliasing the arena while writing the output.
1036 let bytes: Vec<u8> = arena_buf[*src..*src + len * 4].to_vec();
1037 for i in 0..len {
1038 let v = i32::from_le_bytes([
1039 bytes[i * 4],
1040 bytes[i * 4 + 1],
1041 bytes[i * 4 + 2],
1042 bytes[i * 4 + 3],
1043 ]);
1044 arena_buf[*dst + i] = u8::from(v != 0);
1045 }
1046 }
1047
1048 Thunk::CastI64ToBool { src, dst, len } => {
1049 let len = *len as usize;
1050 if len == 0 {
1051 continue;
1052 }
1053 // i64 = 8 bytes, bool = 1 byte — copy first to avoid aliasing.
1054 let bytes: Vec<u8> = arena_buf[*src..*src + len * 8].to_vec();
1055 for i in 0..len {
1056 let v = i64::from_le_bytes([
1057 bytes[i * 8],
1058 bytes[i * 8 + 1],
1059 bytes[i * 8 + 2],
1060 bytes[i * 8 + 3],
1061 bytes[i * 8 + 4],
1062 bytes[i * 8 + 5],
1063 bytes[i * 8 + 6],
1064 bytes[i * 8 + 7],
1065 ]);
1066 arena_buf[*dst + i] = u8::from(v != 0);
1067 }
1068 }
1069
1070 Thunk::CastBoolToI64 { src, dst, len } => {
1071 let len = *len as usize;
1072 if len == 0 {
1073 continue;
1074 }
1075 let bools: Vec<u8> = arena_buf[*src..*src + len].to_vec();
1076 for i in 0..len {
1077 let v = (bools[i] != 0) as i64;
1078 arena_buf[*dst + i * 8..*dst + i * 8 + 8].copy_from_slice(&v.to_le_bytes());
1079 }
1080 }
1081
1082 Thunk::CastBoolToF32 { src, dst, len } => {
1083 let len = *len as usize;
1084 if len == 0 {
1085 continue;
1086 }
1087 unsafe {
1088 let inp = &arena_buf[*src..*src + len];
1089 let out = sl_mut(*dst, base, len);
1090 for i in 0..len {
1091 out[i] = if inp[i] != 0 { 1.0 } else { 0.0 };
1092 }
1093 }
1094 }
1095
1096 Thunk::CastGeneric {
1097 src,
1098 dst,
1099 len,
1100 src_dtype,
1101 dst_dtype,
1102 } => {
1103 let len = *len as usize;
1104 if len == 0 {
1105 continue;
1106 }
1107 exec_cast_generic(*src, *dst, len, *src_dtype, *dst_dtype, base);
1108 }
1109
1110 Thunk::CastF32ToBool { src, dst, len } => {
1111 let len = *len as usize;
1112 if len == 0 {
1113 continue;
1114 }
1115 unsafe {
1116 let inp = sl(*src, base, len).to_vec();
1117 let out = &mut arena_buf[*dst..*dst + len];
1118 for i in 0..len {
1119 out[i] = u8::from(inp[i] != 0.0);
1120 }
1121 }
1122 }
1123
1124 Thunk::BinaryFullF64 { .. } => exec_binary_full_f64(thunk, base),
1125 Thunk::BinaryFullC64 { .. } => exec_binary_full_c64(thunk, base),
1126 Thunk::ComplexNormSqF32 { .. } => exec_complex_norm_sq_f32(thunk, base),
1127 Thunk::ComplexNormSqBackwardF32 { .. } => {
1128 exec_complex_norm_sq_backward_f32(thunk, base)
1129 }
1130 Thunk::ConjugateC64 { .. } => exec_conjugate_c64(thunk, base),
1131 Thunk::ActivationC64 { .. } => exec_activation_c64(thunk, base),
1132 Thunk::Scan { .. } => exec_scan(thunk, base),
1133 Thunk::ScanBackward {
1134 body_vjp,
1135 body_init,
1136 body_carry_in_off,
1137 body_x_offs,
1138 body_d_output_off,
1139 body_dcarry_out_off,
1140 outer_init_off,
1141 outer_traj_off,
1142 outer_upstream_off,
1143 outer_xs_offs,
1144 outer_dinit_off,
1145 length,
1146 carry_bytes,
1147 save_trajectory,
1148 num_checkpoints,
1149 forward_body,
1150 forward_body_init,
1151 forward_body_carry_in_off,
1152 forward_body_output_off,
1153 forward_body_x_offs,
1154 carry_elem_size,
1155 } => {
1156 // Two backward paths share the same per-iteration body
1157 // (body_vjp run + dcarry threading). The "All" path
1158 // reads the carry directly from the saved trajectory
1159 // each step. The "Recursive checkpointing" path stores
1160 // only K saved checkpoints and reconstructs intermediate
1161 // carries via Griewank-style recursive subdivision —
1162 // see [`griewank_process_segment`]. Auxiliary memory
1163 // is `O(log(segment_size) · carry_bytes)` for the
1164 // recursion stack, vs the old segment-cache scheme's
1165 // `O(segment_size · carry_bytes)`. Total recompute work
1166 // grows from `O(length)` to `O(length · log)`, which
1167 // is the canonical Griewank trade.
1168 let cb = *carry_bytes as usize;
1169 let n_steps = *length as usize;
1170 let k_total = *num_checkpoints as usize;
1171 let is_recursive = k_total != 0 && k_total != n_steps;
1172 let checkpoint_t_for_k = |k: usize| -> usize {
1173 ((k + 1) * n_steps)
1174 .div_ceil(k_total)
1175 .saturating_sub(1)
1176 .min(n_steps - 1)
1177 };
1178
1179 let mut fwd_buf: Vec<u8> = if is_recursive {
1180 (**forward_body_init.as_ref().unwrap()).clone()
1181 } else {
1182 Vec::new()
1183 };
1184
1185 let mut dcarry: Vec<u8> = vec![0u8; cb];
1186 if !*save_trajectory {
1187 unsafe {
1188 std::ptr::copy_nonoverlapping(
1189 base.add(*outer_upstream_off),
1190 dcarry.as_mut_ptr(),
1191 cb,
1192 );
1193 }
1194 }
1195
1196 let mut body_buf: Vec<u8> = (**body_init).clone();
1197
1198 // Per-iteration backward action — shared between the
1199 // direct-trajectory (All) and Griewank (Recursive) paths.
1200 // Both feed the same body_vjp run with carry-at-t,
1201 // x_t_i, and d_output, then thread dcarry backward.
1202 let process_iter =
1203 |t: usize, carry_in: &[u8], dcarry: &mut Vec<u8>, body_buf: &mut Vec<u8>| {
1204 if *save_trajectory {
1205 unsafe {
1206 let up_off = *outer_upstream_off + t * cb;
1207 match *carry_elem_size {
1208 4 => {
1209 let up_ptr = base.add(up_off) as *const f32;
1210 let dc_ptr = dcarry.as_mut_ptr() as *mut f32;
1211 let n_elems = cb / 4;
1212 for i in 0..n_elems {
1213 *dc_ptr.add(i) += *up_ptr.add(i);
1214 }
1215 }
1216 8 => {
1217 let up_ptr = base.add(up_off) as *const f64;
1218 let dc_ptr = dcarry.as_mut_ptr() as *mut f64;
1219 let n_elems = cb / 8;
1220 for i in 0..n_elems {
1221 *dc_ptr.add(i) += *up_ptr.add(i);
1222 }
1223 }
1224 other => panic!(
1225 "ScanBackward: unsupported carry elem size {other} \
1226 (only f32/f64 carries are supported today)"
1227 ),
1228 }
1229 }
1230 }
1231 body_buf[*body_carry_in_off..*body_carry_in_off + cb]
1232 .copy_from_slice(carry_in);
1233 unsafe {
1234 for (i, body_x_off) in body_x_offs.iter().enumerate() {
1235 let (outer_xs_off, per_step_bytes) = outer_xs_offs[i];
1236 let psb = per_step_bytes as usize;
1237 std::ptr::copy_nonoverlapping(
1238 base.add(outer_xs_off + t * psb),
1239 body_buf.as_mut_ptr().add(*body_x_off),
1240 psb,
1241 );
1242 }
1243 std::ptr::copy_nonoverlapping(
1244 dcarry.as_ptr(),
1245 body_buf.as_mut_ptr().add(*body_d_output_off),
1246 cb,
1247 );
1248 }
1249 execute_thunks(body_vjp, body_buf);
1250 unsafe {
1251 std::ptr::copy_nonoverlapping(
1252 body_buf.as_ptr().add(*body_dcarry_out_off),
1253 dcarry.as_mut_ptr(),
1254 cb,
1255 );
1256 }
1257 };
1258
1259 if is_recursive {
1260 // Griewank treeverse path. Process saved-checkpoint
1261 // segments from highest-t to lowest-t; within each,
1262 // recursive binary subdivision via
1263 // `griewank_process_segment`. Auxiliary memory:
1264 // O(log(seg_size) · cb) for the recursion stack
1265 // (vs O(seg_size · cb) for the older segment-cache
1266 // scheme); recompute work: O(seg_size · log).
1267 let leaf_threshold = 4usize;
1268 let fb_sched = forward_body.as_ref().unwrap();
1269 let fb_init = forward_body_init.as_ref().unwrap().as_slice();
1270 let mut segment_end = n_steps - 1;
1271 for seg_k in (0..k_total).rev() {
1272 let segment_start = if seg_k == 0 {
1273 0
1274 } else {
1275 checkpoint_t_for_k(seg_k - 1) + 1
1276 };
1277 let mut anchor: Vec<u8> = vec![0u8; cb];
1278 unsafe {
1279 let src = if seg_k == 0 {
1280 base.add(*outer_init_off)
1281 } else {
1282 base.add(*outer_traj_off + (seg_k - 1) * cb)
1283 };
1284 std::ptr::copy_nonoverlapping(src, anchor.as_mut_ptr(), cb);
1285 }
1286 // Closure adapter for the helper's signature
1287 // (mutably re-borrows dcarry / body_buf each call).
1288 let mut leaf_action = |t: usize, carry_in: &[u8]| {
1289 process_iter(t, carry_in, &mut dcarry, &mut body_buf);
1290 };
1291 unsafe {
1292 griewank_process_segment(
1293 segment_start,
1294 segment_end,
1295 &anchor,
1296 cb,
1297 fb_sched,
1298 fb_init,
1299 *forward_body_carry_in_off,
1300 *forward_body_output_off,
1301 forward_body_x_offs,
1302 base,
1303 outer_xs_offs,
1304 &mut fwd_buf,
1305 leaf_threshold,
1306 &mut leaf_action,
1307 );
1308 }
1309 if seg_k == 0 {
1310 break;
1311 }
1312 segment_end = segment_start - 1;
1313 }
1314 } else {
1315 // All-trajectory path: read each carry directly
1316 // from the saved trajectory buffer.
1317 let mut carry_buf: Vec<u8> = vec![0u8; cb];
1318 for t in (0..n_steps).rev() {
1319 unsafe {
1320 let src = if t == 0 {
1321 base.add(*outer_init_off)
1322 } else {
1323 base.add(*outer_traj_off + (t - 1) * cb)
1324 };
1325 std::ptr::copy_nonoverlapping(src, carry_buf.as_mut_ptr(), cb);
1326 }
1327 process_iter(t, &carry_buf, &mut dcarry, &mut body_buf);
1328 }
1329 }
1330
1331 unsafe {
1332 std::ptr::copy_nonoverlapping(dcarry.as_ptr(), base.add(*outer_dinit_off), cb);
1333 }
1334 }
1335
1336 Thunk::ScanBackwardXs { .. } => exec_scan_backward_xs(thunk, base),
1337 Thunk::FusedMmBiasAct { .. } => exec_fused_mm_bias_act(thunk, base),
1338 Thunk::FusedResidualLN {
1339 x,
1340 res,
1341 bias,
1342 g,
1343 b,
1344 out,
1345 rows,
1346 h,
1347 eps,
1348 has_bias,
1349 } => {
1350 let (rows, h) = (*rows as usize, *h as usize);
1351 unsafe {
1352 let zero = &zero_bias[..h];
1353 let bi = if *has_bias { sl(*bias, base, h) } else { zero };
1354 let x_ptr = sl(*x, base, rows * h).as_ptr() as usize;
1355 let r_ptr = sl(*res, base, rows * h).as_ptr() as usize;
1356 let o_ptr = sl_mut(*out, base, rows * h).as_mut_ptr() as usize;
1357 let bi_ptr = bi.as_ptr() as usize;
1358 let g_ptr = sl(*g, base, h).as_ptr() as usize;
1359 let b_ptr = sl(*b, base, h).as_ptr() as usize;
1360 let e = *eps;
1361 crate::pool::par_for(rows, 4, &|off, cnt| {
1362 let xs =
1363 std::slice::from_raw_parts((x_ptr as *const f32).add(off * h), cnt * h);
1364 let rs =
1365 std::slice::from_raw_parts((r_ptr as *const f32).add(off * h), cnt * h);
1366 let os = std::slice::from_raw_parts_mut(
1367 (o_ptr as *mut f32).add(off * h),
1368 cnt * h,
1369 );
1370 let bi = std::slice::from_raw_parts(bi_ptr as *const f32, h);
1371 let g = std::slice::from_raw_parts(g_ptr as *const f32, h);
1372 let b = std::slice::from_raw_parts(b_ptr as *const f32, h);
1373 crate::kernels::residual_bias_layer_norm(xs, rs, bi, g, b, os, cnt, h, e);
1374 });
1375 }
1376 }
1377
1378 Thunk::FusedResidualRmsNorm {
1379 x,
1380 res,
1381 bias,
1382 g,
1383 b,
1384 out,
1385 rows,
1386 h,
1387 eps,
1388 has_bias,
1389 } => {
1390 let (rows, h) = (*rows as usize, *h as usize);
1391 unsafe {
1392 let zero = &zero_bias[..h];
1393 let bi = if *has_bias { sl(*bias, base, h) } else { zero };
1394 let x_ptr = sl(*x, base, rows * h).as_ptr() as usize;
1395 let r_ptr = sl(*res, base, rows * h).as_ptr() as usize;
1396 let o_ptr = sl_mut(*out, base, rows * h).as_mut_ptr() as usize;
1397 let bi_ptr = bi.as_ptr() as usize;
1398 let g_ptr = sl(*g, base, h).as_ptr() as usize;
1399 let b_ptr = sl(*b, base, h).as_ptr() as usize;
1400 let e = *eps;
1401 crate::pool::par_for(rows, 4, &|off, cnt| {
1402 let xs =
1403 std::slice::from_raw_parts((x_ptr as *const f32).add(off * h), cnt * h);
1404 let rs =
1405 std::slice::from_raw_parts((r_ptr as *const f32).add(off * h), cnt * h);
1406 let os = std::slice::from_raw_parts_mut(
1407 (o_ptr as *mut f32).add(off * h),
1408 cnt * h,
1409 );
1410 let bi = std::slice::from_raw_parts(bi_ptr as *const f32, h);
1411 let g = std::slice::from_raw_parts(g_ptr as *const f32, h);
1412 let b = std::slice::from_raw_parts(b_ptr as *const f32, h);
1413 crate::kernels::residual_bias_rms_norm(xs, rs, bi, g, b, os, cnt, h, e);
1414 });
1415 }
1416 }
1417
1418 Thunk::BiasAdd { .. } => exec_bias_add(thunk, base),
1419 Thunk::BinaryFull {
1420 lhs,
1421 rhs,
1422 dst,
1423 len,
1424 lhs_len,
1425 rhs_len,
1426 op,
1427 out_dims_bcast,
1428 bcast_lhs_strides,
1429 bcast_rhs_strides,
1430 elem_bytes,
1431 } => {
1432 let len = *len as usize;
1433 let ll = (*lhs_len as usize).max(1);
1434 let rl = (*rhs_len as usize).max(1);
1435 let eb = (*elem_bytes).max(1) as usize;
1436 let arena_len = arena_buf.len();
1437 let ll = ll.min((arena_len.saturating_sub(*lhs)) / eb);
1438 let rl = rl.min((arena_len.saturating_sub(*rhs)) / eb);
1439 let len = len.min((arena_len.saturating_sub(*dst)) / eb);
1440 unsafe {
1441 if eb == 8 {
1442 let l = sl_i64(*lhs, base, ll);
1443 let r = sl_i64(*rhs, base, rl);
1444 let o = sl_mut_i64(*dst, base, len);
1445 // Same fused-index + hoisted-op + parallel treatment as
1446 // the f32 path below (bit-exact; elements independent).
1447 let rank = out_dims_bcast.len();
1448 let odb = &out_dims_bcast[..];
1449 let lstr = &bcast_lhs_strides[..];
1450 let rstr = &bcast_rhs_strides[..];
1451 let idx = |i: usize| -> (usize, usize) {
1452 if rank == 0 {
1453 let li = if ll == 1 { 0 } else { i % ll };
1454 let ri = if rl == 1 { 0 } else { i % rl };
1455 (li, ri)
1456 } else {
1457 let mut rem = i;
1458 let (mut li, mut ri) = (0usize, 0usize);
1459 for ax in (0..rank).rev() {
1460 let sz = odb[ax] as usize;
1461 let c = rem % sz;
1462 rem /= sz;
1463 li += c * lstr[ax] as usize;
1464 ri += c * rstr[ax] as usize;
1465 }
1466 (li, ri)
1467 }
1468 };
1469 macro_rules! bini64 {
1470 ($f:expr) => {{
1471 let f = $f;
1472 if len >= 8192 {
1473 use rayon::prelude::*;
1474 o.par_iter_mut().enumerate().for_each(|(i, out)| {
1475 let (li, ri) = idx(i);
1476 *out = f(l[li], r[ri]);
1477 });
1478 } else {
1479 for i in 0..len {
1480 let (li, ri) = idx(i);
1481 o[i] = f(l[li], r[ri]);
1482 }
1483 }
1484 }};
1485 }
1486 match op {
1487 BinaryOp::Add => bini64!(|a: i64, b: i64| a.wrapping_add(b)),
1488 BinaryOp::Sub => bini64!(|a: i64, b: i64| a.wrapping_sub(b)),
1489 BinaryOp::Mul => bini64!(|a: i64, b: i64| a.wrapping_mul(b)),
1490 BinaryOp::Div => {
1491 bini64!(|a: i64, b: i64| if b == 0 { 0 } else { a / b })
1492 }
1493 BinaryOp::Max => bini64!(|a: i64, b: i64| a.max(b)),
1494 BinaryOp::Min => bini64!(|a: i64, b: i64| a.min(b)),
1495 BinaryOp::Pow => bini64!(|a: i64, b: i64| a.pow(b.max(0) as u32)),
1496 BinaryOp::Mod => bini64!(|a: i64, b: i64| a % b),
1497 // atan2 is float-only; unreachable on integer tensors.
1498 BinaryOp::Atan2 => {
1499 bini64!(|a: i64, b: i64| (a as f64).atan2(b as f64) as i64)
1500 }
1501 BinaryOp::BitAnd => bini64!(|a: i64, b: i64| a & b),
1502 BinaryOp::BitOr => bini64!(|a: i64, b: i64| a | b),
1503 BinaryOp::BitXor => bini64!(|a: i64, b: i64| a ^ b),
1504 BinaryOp::Shl => bini64!(|a: i64, b: i64| a.wrapping_shl(b as u32)),
1505 BinaryOp::Shr => bini64!(|a: i64, b: i64| a.wrapping_shr(b as u32)),
1506 }
1507 } else {
1508 let l = sl(*lhs, base, ll);
1509 let r = sl(*rhs, base, rl);
1510 let o = sl_mut(*dst, base, len);
1511 if ll == len && rl == len {
1512 #[cfg(target_arch = "aarch64")]
1513 if matches!(op, BinaryOp::Add | BinaryOp::Mul) {
1514 use std::arch::aarch64::*;
1515 let chunks = len / 4;
1516 for c in 0..chunks {
1517 let off = c * 4;
1518 let vl = vld1q_f32(l.as_ptr().add(off));
1519 let vr = vld1q_f32(r.as_ptr().add(off));
1520 let res = match op {
1521 BinaryOp::Add => vaddq_f32(vl, vr),
1522 BinaryOp::Mul => vmulq_f32(vl, vr),
1523 _ => unreachable!(),
1524 };
1525 vst1q_f32(o.as_mut_ptr().add(off), res);
1526 }
1527 for i in (chunks * 4)..len {
1528 o[i] = match op {
1529 BinaryOp::Add => l[i] + r[i],
1530 BinaryOp::Mul => l[i] * r[i],
1531 _ => unreachable!(),
1532 };
1533 }
1534 continue;
1535 }
1536 // x86: contiguous same-shape Add/Mul/Sub — AVX2 when
1537 // available, else a plain parallel loop (no N-D index walk).
1538 #[cfg(target_arch = "x86_64")]
1539 if matches!(op, BinaryOp::Add | BinaryOp::Mul | BinaryOp::Sub) {
1540 let used = binary_contig_f32(l, r, o, *op);
1541 if used {
1542 continue;
1543 }
1544 }
1545 // Contiguous fallback for other ops / arches: skip
1546 // broadcast index math entirely.
1547 macro_rules! bin_contig {
1548 ($f:expr) => {{
1549 let f = $f;
1550 if len >= 8192 {
1551 use rayon::prelude::*;
1552 o.par_iter_mut()
1553 .zip(l.par_iter())
1554 .zip(r.par_iter())
1555 .for_each(|((out, a), b)| *out = f(*a, *b));
1556 } else {
1557 for i in 0..len {
1558 o[i] = f(l[i], r[i]);
1559 }
1560 }
1561 }};
1562 }
1563 match op {
1564 BinaryOp::Add => bin_contig!(|a: f32, b: f32| a + b),
1565 BinaryOp::Sub => bin_contig!(|a: f32, b: f32| a - b),
1566 BinaryOp::Mul => bin_contig!(|a: f32, b: f32| a * b),
1567 BinaryOp::Div => bin_contig!(|a: f32, b: f32| a / b),
1568 BinaryOp::Max => bin_contig!(|a: f32, b: f32| a.max(b)),
1569 BinaryOp::Min => bin_contig!(|a: f32, b: f32| a.min(b)),
1570 BinaryOp::Pow => bin_contig!(|a: f32, b: f32| a.powf(b)),
1571 BinaryOp::Mod => bin_contig!(|a: f32, b: f32| a % b),
1572 BinaryOp::Atan2 => bin_contig!(|a: f32, b: f32| a.atan2(b)),
1573 BinaryOp::BitAnd => {
1574 bin_contig!(|a: f32, b: f32| ((a as i64) & (b as i64)) as f32)
1575 }
1576 BinaryOp::BitOr => {
1577 bin_contig!(|a: f32, b: f32| ((a as i64) | (b as i64)) as f32)
1578 }
1579 BinaryOp::BitXor => {
1580 bin_contig!(|a: f32, b: f32| ((a as i64) ^ (b as i64)) as f32)
1581 }
1582 BinaryOp::Shl => bin_contig!(|a: f32, b: f32| (a as i64)
1583 .wrapping_shl(b as u32)
1584 as f32),
1585 BinaryOp::Shr => bin_contig!(|a: f32, b: f32| (a as i64)
1586 .wrapping_shr(b as u32)
1587 as f32),
1588 }
1589 continue;
1590 }
1591 // Trailing-row broadcast: rhs tiles every `rl` elements
1592 // (bias/scale along last dim). Covers empty dims and the
1593 // common `[…, D]` rhs with leading broadcast strides 0.
1594 if ll == len && rl > 0 && rl < len && len.is_multiple_of(rl) {
1595 let rhs_tile = out_dims_bcast.is_empty()
1596 || (bcast_rhs_strides.len() == out_dims_bcast.len()
1597 && bcast_rhs_strides.last() == Some(&1)
1598 && bcast_rhs_strides
1599 [..bcast_rhs_strides.len().saturating_sub(1)]
1600 .iter()
1601 .all(|&s| s == 0));
1602 if rhs_tile {
1603 let used = binary_row_bcast_f32(l, r, o, *op, rl);
1604 if used {
1605 continue;
1606 }
1607 }
1608 }
1609 // Broadcast / small-operand path. This scalar loop was
1610 // ~88% of a supertonic subgraph's CPU time. Three fixes,
1611 // all bit-exact (each output element is independent):
1612 // 1) fuse the coord decomposition with the stride
1613 // dot-product — one pass, no per-call `coords` Vec,
1614 // no second loop;
1615 // 2) hoist the op-match out of the per-element loop
1616 // (one branch per call, not per element);
1617 // 3) parallelize across cores for large outputs.
1618 let rank = out_dims_bcast.len();
1619 let odb = &out_dims_bcast[..];
1620 let lstr = &bcast_lhs_strides[..];
1621 let rstr = &bcast_rhs_strides[..];
1622 let idx = |i: usize| -> (usize, usize) {
1623 if rank == 0 {
1624 let li = if ll == 1 { 0 } else { i % ll };
1625 let ri = if rl == 1 { 0 } else { i % rl };
1626 (li, ri)
1627 } else {
1628 let mut rem = i;
1629 let (mut li, mut ri) = (0usize, 0usize);
1630 for ax in (0..rank).rev() {
1631 let sz = odb[ax] as usize;
1632 let c = rem % sz;
1633 rem /= sz;
1634 li += c * lstr[ax] as usize;
1635 ri += c * rstr[ax] as usize;
1636 }
1637 (li, ri)
1638 }
1639 };
1640 macro_rules! binf32 {
1641 ($f:expr) => {{
1642 let f = $f;
1643 if len >= 8192 {
1644 use rayon::prelude::*;
1645 o.par_iter_mut().enumerate().for_each(|(i, out)| {
1646 let (li, ri) = idx(i);
1647 *out = f(l[li], r[ri]);
1648 });
1649 } else {
1650 for i in 0..len {
1651 let (li, ri) = idx(i);
1652 o[i] = f(l[li], r[ri]);
1653 }
1654 }
1655 }};
1656 }
1657 match op {
1658 BinaryOp::Add => binf32!(|a: f32, b: f32| a + b),
1659 BinaryOp::Sub => binf32!(|a: f32, b: f32| a - b),
1660 BinaryOp::Mul => binf32!(|a: f32, b: f32| a * b),
1661 BinaryOp::Div => binf32!(|a: f32, b: f32| a / b),
1662 BinaryOp::Max => binf32!(|a: f32, b: f32| a.max(b)),
1663 BinaryOp::Min => binf32!(|a: f32, b: f32| a.min(b)),
1664 BinaryOp::Pow => binf32!(|a: f32, b: f32| a.powf(b)),
1665 BinaryOp::Mod => binf32!(|a: f32, b: f32| a % b),
1666 BinaryOp::Atan2 => binf32!(|a: f32, b: f32| a.atan2(b)),
1667 BinaryOp::BitAnd => {
1668 binf32!(|a: f32, b: f32| ((a as i64) & (b as i64)) as f32)
1669 }
1670 BinaryOp::BitOr => {
1671 binf32!(|a: f32, b: f32| ((a as i64) | (b as i64)) as f32)
1672 }
1673 BinaryOp::BitXor => {
1674 binf32!(|a: f32, b: f32| ((a as i64) ^ (b as i64)) as f32)
1675 }
1676 BinaryOp::Shl => {
1677 binf32!(|a: f32, b: f32| (a as i64).wrapping_shl(b as u32) as f32)
1678 }
1679 BinaryOp::Shr => {
1680 binf32!(|a: f32, b: f32| (a as i64).wrapping_shr(b as u32) as f32)
1681 }
1682 }
1683 }
1684 }
1685 }
1686
1687 Thunk::Gather { .. } => exec_gather(thunk, base),
1688 Thunk::Narrow {
1689 src,
1690 dst,
1691 outer,
1692 src_stride,
1693 dst_stride,
1694 inner,
1695 elem_bytes,
1696 } => {
1697 let (outer, ss, ds, inner, eb) = (
1698 *outer as usize,
1699 *src_stride as usize,
1700 *dst_stride as usize,
1701 *inner as usize,
1702 *elem_bytes as usize,
1703 );
1704 let row_bytes = inner.saturating_mul(eb);
1705 let src_row_stride = ss.saturating_mul(eb);
1706 let dst_row_stride = ds.saturating_mul(eb);
1707 if trace_thunks {
1708 eprintln!(
1709 "[narrow] src={} dst={} outer={outer} ss={ss} ds={ds} inner={inner} eb={eb} row={row_bytes} arena={}",
1710 *src,
1711 *dst,
1712 arena_buf.len()
1713 );
1714 }
1715 if row_bytes > 0 && *src != *dst {
1716 let arena_len = arena_buf.len();
1717 // Parallelize independent row copies for large narrows
1718 // (attention head splits, DiT reshape slices).
1719 if outer >= 4
1720 && row_bytes >= 64
1721 && crate::pool::num_threads() > 1
1722 && crate::pool::should_parallelize(outer.saturating_mul(row_bytes / 4))
1723 {
1724 let base_addr = base as usize;
1725 let src0 = *src;
1726 let dst0 = *dst;
1727 crate::pool::par_for(outer, 1, &|off, cnt| {
1728 for o in off..off + cnt {
1729 let s_off = src0 + o * src_row_stride;
1730 let d_off = dst0 + o * dst_row_stride;
1731 if s_off == d_off {
1732 continue;
1733 }
1734 if s_off.saturating_add(row_bytes) > arena_len
1735 || d_off.saturating_add(row_bytes) > arena_len
1736 {
1737 break;
1738 }
1739 unsafe {
1740 std::ptr::copy_nonoverlapping(
1741 (base_addr as *const u8).add(s_off),
1742 (base_addr as *mut u8).add(d_off),
1743 row_bytes,
1744 );
1745 }
1746 }
1747 });
1748 } else {
1749 for o in 0..outer {
1750 let s_off = *src + o * src_row_stride;
1751 let d_off = *dst + o * dst_row_stride;
1752 if s_off == d_off {
1753 continue;
1754 }
1755 if s_off.saturating_add(row_bytes) > arena_len
1756 || d_off.saturating_add(row_bytes) > arena_len
1757 {
1758 break;
1759 }
1760 unsafe {
1761 std::ptr::copy_nonoverlapping(
1762 base.add(s_off),
1763 base.add(d_off),
1764 row_bytes,
1765 );
1766 }
1767 }
1768 }
1769 }
1770 }
1771
1772 Thunk::Copy { src, dst, len } => {
1773 let mut len = *len as usize;
1774 if *src == *dst || len == 0 {
1775 continue;
1776 }
1777 let arena_len = arena_buf.len();
1778 let max_from_src = (arena_len.saturating_sub(*src)) / 4;
1779 let max_from_dst = (arena_len.saturating_sub(*dst)) / 4;
1780 len = len.min(max_from_src).min(max_from_dst);
1781 if len == 0 {
1782 continue;
1783 }
1784 let byte_len = len.saturating_mul(4);
1785 // Parallel memcpy for huge arena moves (DiT residual buffers).
1786 if len >= 262_144 && crate::pool::num_threads() > 1 {
1787 let base_addr = base as usize;
1788 let src0 = *src;
1789 let dst0 = *dst;
1790 crate::pool::par_for(len, crate::pool::chunk_floor(len), &|off, cnt| {
1791 let n = cnt.saturating_mul(4);
1792 unsafe {
1793 std::ptr::copy(
1794 (base_addr as *const u8).add(src0 + off * 4),
1795 (base_addr as *mut u8).add(dst0 + off * 4),
1796 n,
1797 );
1798 }
1799 });
1800 } else {
1801 unsafe {
1802 std::ptr::copy(base.add(*src), base.add(*dst), byte_len);
1803 }
1804 }
1805 }
1806
1807 Thunk::LayerNorm { .. } => exec_layer_norm(thunk, base),
1808 Thunk::GroupNorm { .. } => exec_group_norm(thunk, base),
1809 Thunk::BatchNormInference { .. } => exec_batch_norm_inference(thunk, base),
1810 Thunk::LayerNorm2d { .. } => exec_layer_norm2d(thunk, base),
1811 Thunk::ConvTranspose2d { .. } => exec_conv_transpose2d(thunk, base),
1812 Thunk::ResizeNearest2x { .. } => exec_resize_nearest2x(thunk, base),
1813 Thunk::Interpolate3d { .. } => exec_interpolate3d(thunk, base),
1814 Thunk::AxialRope2d { .. } => exec_axial_rope2d(thunk, base),
1815 Thunk::RmsNorm { .. } => exec_rms_norm(thunk, base),
1816 Thunk::AdaLayerNorm { .. } => exec_ada_layer_norm(thunk, base),
1817 Thunk::GatedResidual { .. } => exec_gated_residual(thunk, base),
1818 Thunk::AdaLayerNormBackward { .. } => exec_ada_layer_norm_backward(thunk, base),
1819 Thunk::GatedResidualBackward { .. } => exec_gated_residual_backward(thunk, base),
1820 Thunk::Softmax { .. } => exec_softmax(thunk, base),
1821 Thunk::Cumsum { .. } => exec_cumsum(thunk, base),
1822 Thunk::CumScan { .. } => exec_cum_scan(thunk, base),
1823 Thunk::Sample { .. } => exec_sample(thunk, base),
1824 Thunk::RngNormal {
1825 dst,
1826 len,
1827 mean,
1828 scale,
1829 key,
1830 op_seed,
1831 } => {
1832 let n = *len as usize;
1833 unsafe {
1834 let out = sl_mut(*dst, base, n);
1835 let opts = *schedule.rng.read().unwrap();
1836 rlx_ir::fill_normal_like(out, *mean, *scale, opts, *key, *op_seed);
1837 }
1838 }
1839
1840 Thunk::RngUniform {
1841 dst,
1842 len,
1843 low,
1844 high,
1845 key,
1846 op_seed,
1847 } => {
1848 let n = *len as usize;
1849 unsafe {
1850 let out = sl_mut(*dst, base, n);
1851 let opts = *schedule.rng.read().unwrap();
1852 rlx_ir::fill_uniform_like(out, *low, *high, opts, *key, *op_seed);
1853 }
1854 }
1855
1856 Thunk::GatedDeltaNet { .. } => exec_gated_delta_net(thunk, base),
1857 Thunk::GatedDeltaNetBackward { .. } => exec_gated_delta_net_backward(thunk, base),
1858 Thunk::Lstm { .. } => exec_lstm(thunk, base),
1859 Thunk::Gru { .. } => exec_gru(thunk, base),
1860 Thunk::Rnn { .. } => exec_rnn(thunk, base),
1861 Thunk::Mamba2 { .. } => exec_mamba2(thunk, base),
1862 Thunk::SelectiveScan { .. } => exec_selective_scan(thunk, base),
1863 Thunk::DequantMatMul { .. } => exec_dequant_mat_mul(thunk, base),
1864 Thunk::SynthMatMul { .. } => exec_synth_mat_mul(thunk, base),
1865 Thunk::SplineActivation { .. } => exec_spline_activation(thunk, base),
1866 Thunk::DequantMatMulGguf { .. } => exec_dequant_mat_mul_gguf(thunk, base),
1867 Thunk::DequantMatMulInt4 { .. } => exec_dequant_mat_mul_int4(thunk, base),
1868 Thunk::DequantMatMulFp8 { .. } => exec_dequant_mat_mul_fp8(thunk, base),
1869 Thunk::DequantMatMulNvfp4 { .. } => exec_dequant_mat_mul_nvfp4(thunk, base),
1870 Thunk::DequantMatMulMxFp4x2 { .. } => exec_dequant_mat_mul_mxfp4x2(thunk, base),
1871 Thunk::DequantMatMulMlx { .. } => exec_dequant_mat_mul_mlx(thunk, base),
1872 Thunk::ScaledMatMul { .. } => exec_scaled_mat_mul(thunk, base),
1873 Thunk::ScaledGroupedMatMul { .. } => exec_scaled_grouped_mat_mul(thunk, base),
1874 Thunk::ScaledQuantize { .. } => exec_scaled_quantize(thunk, base),
1875 Thunk::ScaledQuantScale { .. } => exec_scaled_quant_scale(thunk, base),
1876 Thunk::ScaledDequantize { .. } => exec_scaled_dequantize(thunk, base),
1877 Thunk::LoraMatMul { .. } => exec_lora_mat_mul(thunk, base),
1878 Thunk::Attention {
1879 q,
1880 k,
1881 v,
1882 mask,
1883 out,
1884 batch,
1885 seq,
1886 kv_seq,
1887 heads,
1888 head_dim,
1889 v_head_dim,
1890 mask_kind,
1891 scale,
1892 softcap,
1893 q_row_stride,
1894 k_row_stride,
1895 v_row_stride,
1896 bhsd,
1897 kv_heads,
1898 } => {
1899 let (b, q_s, k_s, nh, dh) = (
1900 *batch as usize,
1901 *seq as usize,
1902 *kv_seq as usize,
1903 *heads as usize,
1904 *head_dim as usize,
1905 );
1906 // V/output per-head width (== dh for symmetric SDPA; MLA differs).
1907 let dh_v = *v_head_dim as usize;
1908 let nkv = (*kv_heads as usize).max(1);
1909 let group = (nh / nkv).max(1); // query heads per KV head (GQA/MQA)
1910 let hs_v = nh * dh_v; // output row stride (V-width)
1911 // For [B, H, S, D] layout each (b, h) tile is dense
1912 // contiguous; the qrs/krs/vrs strides are not used.
1913 let (qrs, krs, vrs) = if *bhsd {
1914 (dh, dh, dh_v)
1915 } else {
1916 (
1917 *q_row_stride as usize,
1918 *k_row_stride as usize,
1919 *v_row_stride as usize,
1920 )
1921 };
1922 let bhsd = *bhsd;
1923 let _ = (q_row_stride, k_row_stride, v_row_stride);
1924 let scale = *scale;
1925 let ss = q_s * k_s;
1926 let cfg = crate::config::RuntimeConfig::global();
1927 unsafe {
1928 // Slice lengths cover the strided span. When Q/K/V
1929 // alias the parent QKV (post-#46-fusion), the same
1930 // bytes back all three slices — compiler bounds
1931 // checks see the right size. For [B, H, S, D] the
1932 // buffer is densely B*H*S*D elements; the row
1933 // strides aren't used.
1934 let q_len = if bhsd {
1935 b * nh * q_s * dh
1936 } else {
1937 b * q_s * qrs
1938 };
1939 let k_len = if bhsd {
1940 b * nkv * k_s * dh
1941 } else {
1942 b * k_s * krs
1943 };
1944 let v_len = if bhsd {
1945 b * nkv * k_s * dh_v
1946 } else {
1947 b * k_s * vrs
1948 };
1949 let q_data = sl(*q, base, q_len);
1950 let k_data = sl(*k, base, k_len);
1951 let v_data = sl(*v, base, v_len);
1952 let mask_data: &[f32] = match mask_kind {
1953 rlx_ir::op::MaskKind::Custom => sl(*mask, base, b * k_s),
1954 rlx_ir::op::MaskKind::Bias => sl(*mask, base, b * nh * q_s * k_s),
1955 _ => &[],
1956 };
1957 let out_len = if bhsd {
1958 b * nh * q_s * dh_v
1959 } else {
1960 b * q_s * hs_v
1961 };
1962 let out_data = sl_mut(*out, base, out_len);
1963
1964 // ── [B, H, S, D] fallback ──────────────────────
1965 // The NEON / strided-BLAS specializations below
1966 // are written for the [B, S, H, D] layout. When
1967 // the input is head-major ([B, H, S, D] —
1968 // matching rlx-cuda / rlx-rocm / rlx-tpu), bypass
1969 // them and run a simple (correct but slower)
1970 // scalar implementation. Production-CPU inference
1971 // graphs use [B, S, H, D] so they still hit the
1972 // hot path; cross-backend parity tests use
1973 // [B, H, S, D] and land here.
1974 if bhsd {
1975 let scores = &mut sdpa_scores[..ss];
1976 for bi in 0..b {
1977 for hi in 0..nh {
1978 let kv_hi = hi / group; // GQA/MQA: shared KV head
1979 let q_head_base = bi * nh * q_s * dh + hi * q_s * dh;
1980 let k_head_base = bi * nkv * k_s * dh + kv_hi * k_s * dh;
1981 // Q@K^T
1982 for qi in 0..q_s {
1983 let q_base = q_head_base + qi * dh;
1984 for ki in 0..k_s {
1985 let k_base = k_head_base + ki * dh;
1986 let mut dot = 0f32;
1987 for d in 0..dh {
1988 dot += q_data[q_base + d] * k_data[k_base + d];
1989 }
1990 scores[qi * k_s + ki] = dot * scale;
1991 if matches!(mask_kind, rlx_ir::op::MaskKind::Custom)
1992 && !mask_data.is_empty()
1993 && mask_data[bi * k_s + ki] < mask_thr
1994 {
1995 scores[qi * k_s + ki] = mask_neg;
1996 }
1997 }
1998 }
1999 if matches!(mask_kind, rlx_ir::op::MaskKind::Bias) {
2000 let off = (bi * nh + hi) * q_s * k_s;
2001 for i in 0..q_s * k_s {
2002 scores[i] += mask_data[off + i];
2003 }
2004 }
2005 apply_synthetic_mask(scores, q_s, k_s, *mask_kind);
2006 // Gemma 2 attention logit soft-cap (post-mask, pre-softmax).
2007 if *softcap > 0.0 {
2008 for s in scores.iter_mut() {
2009 *s = *softcap * (*s / *softcap).tanh();
2010 }
2011 }
2012 crate::kernels::neon_softmax(scores, q_s, k_s);
2013 // score @ V — V and output are dh_v-wide per head.
2014 let o_head_base = bi * nh * q_s * dh_v + hi * q_s * dh_v;
2015 let v_head_base = bi * nkv * k_s * dh_v + kv_hi * k_s * dh_v;
2016 for qi in 0..q_s {
2017 let o_base = o_head_base + qi * dh_v;
2018 for d in 0..dh_v {
2019 out_data[o_base + d] = 0.0;
2020 }
2021 for ki in 0..k_s {
2022 let sc = scores[qi * k_s + ki];
2023 if sc > score_thr {
2024 let v_base = v_head_base + ki * dh_v;
2025 for d in 0..dh_v {
2026 out_data[o_base + d] += sc * v_data[v_base + d];
2027 }
2028 }
2029 }
2030 }
2031 }
2032 }
2033 continue;
2034 }
2035
2036 // ── Auto-select kernel: NEON dots vs strided BLAS ───
2037 // For tiny inputs (batch=1, short seq), per-head BLAS call
2038 // overhead (~0.5µs × 2 calls × num_heads × num_layers)
2039 // exceeds the NEON compute cost. Use direct strided NEON
2040 // with zero dispatch overhead.
2041 // For batch≥2: always BLAS + par_for (parallelism wins).
2042 if b == 1 && q_s.max(k_s) <= cfg.sdpa_seq_threshold {
2043 // ── Sequential NEON path (zero overhead) ──
2044 let scores = &mut sdpa_scores[..ss];
2045 #[cfg(target_arch = "aarch64")]
2046 let neon_chunks = dh / 4;
2047 #[cfg(target_arch = "aarch64")]
2048 let neon_chunks_v = dh_v / 4;
2049
2050 for bi in 0..b {
2051 for hi in 0..nh {
2052 let kv_hi = hi / group; // GQA/MQA: shared KV head
2053 // Q@K^T via strided NEON dot products
2054 for qi in 0..q_s {
2055 let q_off = bi * q_s * qrs + qi * qrs + hi * dh;
2056 for ki in 0..k_s {
2057 let k_off = bi * k_s * krs + ki * krs + kv_hi * dh;
2058 #[cfg(target_arch = "aarch64")]
2059 let mut dot;
2060 #[cfg(not(target_arch = "aarch64"))]
2061 let mut dot = 0f32;
2062 #[cfg(target_arch = "aarch64")]
2063 {
2064 use std::arch::aarch64::*;
2065 let mut acc = vdupq_n_f32(0.0);
2066 for c in 0..neon_chunks {
2067 let vq =
2068 vld1q_f32(q_data.as_ptr().add(q_off + c * 4));
2069 let vk =
2070 vld1q_f32(k_data.as_ptr().add(k_off + c * 4));
2071 acc = vfmaq_f32(acc, vq, vk);
2072 }
2073 dot = vaddvq_f32(acc);
2074 for d in (neon_chunks * 4)..dh {
2075 dot += q_data[q_off + d] * k_data[k_off + d];
2076 }
2077 }
2078 #[cfg(not(target_arch = "aarch64"))]
2079 for d in 0..dh {
2080 dot += q_data[q_off + d] * k_data[k_off + d];
2081 }
2082 scores[qi * k_s + ki] = dot * scale;
2083 // Inner-loop Custom mask check —
2084 // Causal / SlidingWindow / None
2085 // apply outside the loop below.
2086 // Skip for Bias — that mask is a
2087 // per-head additive tensor, not a
2088 // 0/1 key-padding mask.
2089 if matches!(mask_kind, rlx_ir::op::MaskKind::Custom)
2090 && !mask_data.is_empty()
2091 && mask_data[bi * k_s + ki] < mask_thr
2092 {
2093 scores[qi * k_s + ki] = mask_neg;
2094 }
2095 }
2096 }
2097
2098 if matches!(mask_kind, rlx_ir::op::MaskKind::Bias) {
2099 let off = (bi * nh + hi) * q_s * k_s;
2100 for i in 0..q_s * k_s {
2101 scores[i] += mask_data[off + i];
2102 }
2103 }
2104 apply_synthetic_mask(scores, q_s, k_s, *mask_kind);
2105 crate::kernels::neon_softmax(scores, q_s, k_s);
2106
2107 // Score@V via strided NEON accumulation (zero-copy).
2108 // V/output are dh_v-wide per head (== dh unless MLA).
2109 for qi in 0..q_s {
2110 let o_off = bi * q_s * hs_v + qi * hs_v + hi * dh_v;
2111 // Zero output for this head position
2112 for d in 0..dh_v {
2113 out_data[o_off + d] = 0.0;
2114 }
2115 for ki in 0..k_s {
2116 let sc = scores[qi * k_s + ki];
2117 if sc > score_thr {
2118 let v_off = bi * k_s * vrs + ki * vrs + kv_hi * dh_v;
2119 #[cfg(target_arch = "aarch64")]
2120 {
2121 use std::arch::aarch64::*;
2122 let vsc = vdupq_n_f32(sc);
2123 for c in 0..neon_chunks_v {
2124 let off = c * 4;
2125 let vo = vld1q_f32(
2126 out_data.as_ptr().add(o_off + off),
2127 );
2128 let vv =
2129 vld1q_f32(v_data.as_ptr().add(v_off + off));
2130 vst1q_f32(
2131 out_data.as_mut_ptr().add(o_off + off),
2132 vfmaq_f32(vo, vsc, vv),
2133 );
2134 }
2135 // scalar tail if dh_v not a multiple of 4
2136 for d in (neon_chunks_v * 4)..dh_v {
2137 out_data[o_off + d] += sc * v_data[v_off + d];
2138 }
2139 }
2140 #[cfg(not(target_arch = "aarch64"))]
2141 for d in 0..dh_v {
2142 out_data[o_off + d] += sc * v_data[v_off + d];
2143 }
2144 }
2145 }
2146 }
2147 }
2148 }
2149 } else {
2150 // ── Parallel strided BLAS path (high throughput) ──
2151 let total_work = b * nh;
2152 let q_addr = q_data.as_ptr() as usize;
2153 let k_addr = k_data.as_ptr() as usize;
2154 let v_addr = v_data.as_ptr() as usize;
2155 let m_addr = mask_data.as_ptr() as usize;
2156 let o_addr = out_data.as_mut_ptr() as usize;
2157 let sc_addr = sdpa_scores.as_mut_ptr() as usize;
2158
2159 crate::pool::par_for(total_work, 1, &|off, cnt| {
2160 for idx in off..off + cnt {
2161 let bi = idx / nh;
2162 let hi = idx % nh;
2163 let kv_hi = hi / group; // GQA/MQA: shared KV head
2164
2165 let q_start = (q_addr as *const f32).add(bi * q_s * qrs + hi * dh);
2166 let k_start =
2167 (k_addr as *const f32).add(bi * k_s * krs + kv_hi * dh);
2168 // V/output are dh_v-wide per head (== dh unless MLA).
2169 let v_start =
2170 (v_addr as *const f32).add(bi * k_s * vrs + kv_hi * dh_v);
2171 let o_start = (o_addr as *mut f32).add(bi * q_s * hs_v + hi * dh_v);
2172 let sc = std::slice::from_raw_parts_mut(
2173 (sc_addr as *mut f32).add(idx * ss),
2174 ss,
2175 );
2176
2177 // LDA = qrs, LDB = krs (parent row strides
2178 // when fused; hs otherwise).
2179 crate::blas::sgemm_general(
2180 q_start,
2181 k_start,
2182 sc.as_mut_ptr(),
2183 q_s,
2184 k_s,
2185 dh,
2186 scale,
2187 0.0,
2188 qrs,
2189 krs,
2190 k_s,
2191 false,
2192 true,
2193 );
2194
2195 match mask_kind {
2196 rlx_ir::op::MaskKind::Custom => {
2197 let mask_bi = std::slice::from_raw_parts(
2198 (m_addr as *const f32).add(bi * k_s),
2199 k_s,
2200 );
2201 for ki in 0..k_s {
2202 if mask_bi[ki] < mask_thr {
2203 for qi in 0..q_s {
2204 sc[qi * k_s + ki] = mask_neg;
2205 }
2206 }
2207 }
2208 }
2209 rlx_ir::op::MaskKind::Bias => {
2210 // Per-head additive bias slice.
2211 let bias = std::slice::from_raw_parts(
2212 (m_addr as *const f32).add((bi * nh + hi) * q_s * k_s),
2213 q_s * k_s,
2214 );
2215 for i in 0..q_s * k_s {
2216 sc[i] += bias[i];
2217 }
2218 }
2219 _ => apply_synthetic_mask(sc, q_s, k_s, *mask_kind),
2220 }
2221
2222 crate::kernels::neon_softmax(sc, q_s, k_s);
2223
2224 // score@V: N = dh_v (output/V width), LDB = vrs
2225 // (V row stride), LDC = hs_v (output row stride).
2226 crate::blas::sgemm_general(
2227 sc.as_ptr(),
2228 v_start,
2229 o_start,
2230 q_s,
2231 dh_v,
2232 k_s,
2233 1.0,
2234 0.0,
2235 k_s,
2236 vrs,
2237 hs_v,
2238 false,
2239 false,
2240 );
2241 }
2242 });
2243 }
2244 }
2245 }
2246
2247 Thunk::AttentionBackward { .. } => exec_attention_backward(thunk, base),
2248 Thunk::AttentionBackwardAll { .. } => exec_attention_backward_all(thunk, base),
2249 Thunk::ActivationInPlace { .. } => exec_activation_in_place(thunk, base),
2250 Thunk::FusedAttnBlock {
2251 hidden,
2252 qkv_w,
2253 out_w,
2254 mask,
2255 mask_kind,
2256 out,
2257 qkv_b,
2258 out_b,
2259 cos,
2260 sin,
2261 cos_len,
2262 batch,
2263 seq,
2264 hs,
2265 nh,
2266 dh,
2267 in_dim,
2268 out_dim,
2269 has_bias,
2270 has_rope,
2271 interleaved,
2272 } => {
2273 let (b, s) = (*batch as usize, *seq as usize);
2274 let (h, n_h, d_h) = (*hs as usize, *nh as usize, *dh as usize);
2275 // QKV-proj input width and out-proj output width — differ from
2276 // `h`(=hs) for MoonViT-style attention (hidden ≠ nh·dh).
2277 let (in_dim, out_dim) = (*in_dim as usize, *out_dim as usize);
2278 let interleaved = *interleaved;
2279 let m = b * s;
2280 let scale = (d_h as f32).powf(-0.5);
2281 let half = d_h / 2;
2282 // Only `Custom` consumes the per-key padding buffer; `Causal` /
2283 // `SlidingWindow` are synthesized from (qi, ki) below, and have
2284 // no mask buffer (so reading one would touch unrelated arena
2285 // bytes). q_seq == kv_seq here (guaranteed at fusion time), so
2286 // the absolute query position is just `qi`.
2287 let use_custom_mask = matches!(mask_kind, rlx_ir::op::MaskKind::Custom);
2288 unsafe {
2289 let inp = sl(*hidden, base, m * in_dim);
2290 let wq = sl(*qkv_w, base, in_dim * 3 * h);
2291 let wo = sl(*out_w, base, h * out_dim);
2292 let mk = if use_custom_mask {
2293 sl(*mask, base, b * s)
2294 } else {
2295 &[]
2296 };
2297 let dst = sl_mut(*out, base, m * out_dim);
2298
2299 // Stack-allocated intermediates — all fit in L1 cache for small batch
2300 let mut qkv = vec![0f32; m * 3 * h];
2301 let mut attn_out = vec![0f32; m * h];
2302
2303 // 1. QKV projection: [m, in_dim] @ [in_dim, 3h] → [m, 3h]
2304 crate::blas::par_sgemm(inp, wq, &mut qkv, m, in_dim, 3 * h);
2305 if *has_bias {
2306 let bias = sl(*qkv_b, base, 3 * h);
2307 crate::blas::bias_add(&mut qkv, bias, m, 3 * h);
2308 }
2309
2310 // 2. Multi-head SDPA — PARALLEL across (batch, head). Each
2311 // task owns its scores buffer and writes a DISJOINT
2312 // head-slice of attn_out (`o_base` is unique per (bi,hi)),
2313 // so the raw-pointer writes never alias; Q/K/V/cos/sin/mask
2314 // are read-only. Serial per-head SDPA lost ~3× to the
2315 // unfused Attention thunk (which is per-head parallel) on
2316 // wide heads — e.g. MoonViT 12×128 @ seq 64.
2317 #[cfg(target_arch = "aarch64")]
2318 let neon_chunks = d_h / 4;
2319 let qkv_ptr = qkv.as_ptr() as usize;
2320 let out_ptr = attn_out.as_mut_ptr() as usize;
2321 let mk_ptr = mk.as_ptr() as usize;
2322 let mk_len = mk.len();
2323 let base_addr = base as usize;
2324 let (cos_off, sin_off, cl) = (*cos, *sin, *cos_len as usize);
2325 let has_rope_c = *has_rope;
2326 let mask_kind_c = *mask_kind;
2327 let n_tasks = b * n_h;
2328 crate::pool::par_for(n_tasks, 1, &|off, cnt| {
2329 // Read-only views (multiple immutable aliases are fine).
2330 let qkv = std::slice::from_raw_parts(qkv_ptr as *const f32, m * 3 * h);
2331 let out = out_ptr as *mut f32;
2332 let base = base_addr as *mut u8;
2333 let (cos_tab, sin_tab): (&[f32], &[f32]) = if has_rope_c {
2334 (sl(cos_off, base, cl), sl(sin_off, base, cl))
2335 } else {
2336 (&[], &[])
2337 };
2338 let mk = std::slice::from_raw_parts(mk_ptr as *const f32, mk_len);
2339 let mut scores = vec![0f32; s * s];
2340 for t in off..off + cnt {
2341 let bi = t / n_h;
2342 let hi = t % n_h;
2343 for qi in 0..s {
2344 let q_base = bi * s * 3 * h + qi * 3 * h + hi * d_h;
2345 for ki in 0..s {
2346 let k_base = bi * s * 3 * h + ki * 3 * h + h + hi * d_h;
2347 let mut dot = 0f32;
2348 if has_rope_c {
2349 // Inline RoPE per pair, then dot. Pair offsets:
2350 // NeoX (i, i+half) GPT-J (2i, 2i+1).
2351 let q_cos = qi * half;
2352 let k_cos = ki * half;
2353 for i in 0..half {
2354 let (qo1, qo2, ko1, ko2) = if interleaved {
2355 (2 * i, 2 * i + 1, 2 * i, 2 * i + 1)
2356 } else {
2357 (i, half + i, i, half + i)
2358 };
2359 let q1 = qkv[q_base + qo1];
2360 let q2 = qkv[q_base + qo2];
2361 let k1 = qkv[k_base + ko1];
2362 let k2 = qkv[k_base + ko2];
2363 let c_q = cos_tab[q_cos + i];
2364 let s_q = sin_tab[q_cos + i];
2365 let c_k = cos_tab[k_cos + i];
2366 let s_k = sin_tab[k_cos + i];
2367 let qr1 = q1 * c_q - q2 * s_q;
2368 let kr1 = k1 * c_k - k2 * s_k;
2369 let qr2 = q2 * c_q + q1 * s_q;
2370 let kr2 = k2 * c_k + k1 * s_k;
2371 dot += qr1 * kr1 + qr2 * kr2;
2372 }
2373 } else {
2374 #[cfg(target_arch = "aarch64")]
2375 {
2376 use std::arch::aarch64::*;
2377 let mut acc = vdupq_n_f32(0.0);
2378 for c in 0..neon_chunks {
2379 let vq =
2380 vld1q_f32(qkv.as_ptr().add(q_base + c * 4));
2381 let vk =
2382 vld1q_f32(qkv.as_ptr().add(k_base + c * 4));
2383 acc = vfmaq_f32(acc, vq, vk);
2384 }
2385 dot = vaddvq_f32(acc);
2386 for d in (neon_chunks * 4)..d_h {
2387 dot += qkv[q_base + d] * qkv[k_base + d];
2388 }
2389 }
2390 #[cfg(not(target_arch = "aarch64"))]
2391 for d in 0..d_h {
2392 dot += qkv[q_base + d] * qkv[k_base + d];
2393 }
2394 }
2395 let mut sv = dot * scale;
2396 let pos_masked = match mask_kind_c {
2397 rlx_ir::op::MaskKind::Causal => ki > qi,
2398 rlx_ir::op::MaskKind::SlidingWindow(w) => {
2399 ki > qi || ki + w < qi
2400 }
2401 _ => false,
2402 };
2403 if pos_masked || (use_custom_mask && mk[bi * s + ki] < mask_thr)
2404 {
2405 sv = mask_neg;
2406 }
2407 scores[qi * s + ki] = sv;
2408 }
2409 }
2410 crate::kernels::neon_softmax(&mut scores[..s * s], s, s);
2411 // Score @ V accumulation (V at offset 2h in QKV).
2412 for qi in 0..s {
2413 let o_base = bi * s * h + qi * h + hi * d_h;
2414 for d in 0..d_h {
2415 *out.add(o_base + d) = 0.0;
2416 }
2417 for ki in 0..s {
2418 let sc = scores[qi * s + ki];
2419 if sc > score_thr {
2420 let v_base = bi * s * 3 * h + ki * 3 * h + 2 * h + hi * d_h;
2421 #[cfg(target_arch = "aarch64")]
2422 {
2423 use std::arch::aarch64::*;
2424 let vsc = vdupq_n_f32(sc);
2425 for c in 0..neon_chunks {
2426 let o = c * 4;
2427 let vo = vld1q_f32(out.add(o_base + o));
2428 let vv = vld1q_f32(qkv.as_ptr().add(v_base + o));
2429 vst1q_f32(
2430 out.add(o_base + o),
2431 vfmaq_f32(vo, vsc, vv),
2432 );
2433 }
2434 }
2435 #[cfg(not(target_arch = "aarch64"))]
2436 for d in 0..d_h {
2437 *out.add(o_base + d) += sc * qkv[v_base + d];
2438 }
2439 }
2440 }
2441 }
2442 }
2443 });
2444
2445 // 3. Output projection: [m, h] @ [h, out_dim] → dst
2446 crate::blas::par_sgemm(&attn_out, wo, dst, m, h, out_dim);
2447 if *has_bias {
2448 let bias = sl(*out_b, base, out_dim);
2449 crate::blas::bias_add(dst, bias, m, out_dim);
2450 }
2451 }
2452 }
2453
2454 Thunk::Rope { .. } => exec_rope(thunk, base),
2455 Thunk::FusedBertLayer {
2456 hidden,
2457 qkv_w,
2458 qkv_b,
2459 out_w,
2460 out_b,
2461 mask,
2462 ln1_g,
2463 ln1_b,
2464 eps1,
2465 fc1_w,
2466 fc1_b,
2467 fc2_w,
2468 fc2_b,
2469 ln2_g,
2470 ln2_b,
2471 eps2,
2472 out,
2473 batch,
2474 seq,
2475 hs,
2476 nh,
2477 dh,
2478 int_dim,
2479 } => {
2480 let (b, s, h, n_h, d_h) = (
2481 *batch as usize,
2482 *seq as usize,
2483 *hs as usize,
2484 *nh as usize,
2485 *dh as usize,
2486 );
2487 let m = b * s;
2488 let id = *int_dim as usize;
2489 let scale = (d_h as f32).powf(-0.5);
2490 let _half = d_h / 2;
2491 #[cfg(target_arch = "aarch64")]
2492 let neon_chunks = d_h / 4;
2493 unsafe {
2494 let inp = sl(*hidden, base, m * h);
2495 let dst = sl_mut(*out, base, m * h);
2496 let mk = sl(*mask, base, b * s);
2497
2498 // Pre-allocated buffers (zero malloc per layer — allocated once before thunk loop)
2499 let qkv = std::slice::from_raw_parts_mut(fl_qkv.as_mut_ptr(), m * 3 * h);
2500 let attn = std::slice::from_raw_parts_mut(fl_attn.as_mut_ptr(), m * h);
2501 let res = std::slice::from_raw_parts_mut(fl_res.as_mut_ptr(), m * h);
2502 let normed = std::slice::from_raw_parts_mut(fl_normed.as_mut_ptr(), m * h);
2503 let ffn = std::slice::from_raw_parts_mut(fl_ffn.as_mut_ptr(), m * id);
2504 let sc = std::slice::from_raw_parts_mut(fl_sc.as_mut_ptr(), s * s);
2505
2506 // QKV (parallelized across cores — multiple AMX coprocessors)
2507 crate::blas::par_sgemm_bias(
2508 inp,
2509 sl(*qkv_w, base, h * 3 * h),
2510 sl(*qkv_b, base, 3 * h),
2511 qkv,
2512 m,
2513 h,
2514 3 * h,
2515 );
2516
2517 // SDPA per head (sequential NEON, inline — zero overhead)
2518 for bi in 0..b {
2519 for hi in 0..n_h {
2520 for qi in 0..s {
2521 for ki in 0..s {
2522 let q_base = bi * s * 3 * h + qi * 3 * h + hi * d_h;
2523 let k_base = bi * s * 3 * h + ki * 3 * h + h + hi * d_h;
2524 #[cfg(target_arch = "aarch64")]
2525 let dot;
2526 #[cfg(not(target_arch = "aarch64"))]
2527 let mut dot = 0f32;
2528 #[cfg(target_arch = "aarch64")]
2529 {
2530 use std::arch::aarch64::*;
2531 let mut acc = vdupq_n_f32(0.0);
2532 for c in 0..neon_chunks {
2533 acc = vfmaq_f32(
2534 acc,
2535 vld1q_f32(qkv.as_ptr().add(q_base + c * 4)),
2536 vld1q_f32(qkv.as_ptr().add(k_base + c * 4)),
2537 );
2538 }
2539 dot = vaddvq_f32(acc);
2540 }
2541 #[cfg(not(target_arch = "aarch64"))]
2542 for d in 0..d_h {
2543 dot += qkv[q_base + d] * qkv[k_base + d];
2544 }
2545 sc[qi * s + ki] = dot * scale;
2546 if mk[bi * s + ki] < mask_thr {
2547 sc[qi * s + ki] = mask_neg;
2548 }
2549 }
2550 }
2551 crate::kernels::neon_softmax(&mut sc[..s * s], s, s);
2552 for qi in 0..s {
2553 let o = bi * s * h + qi * h + hi * d_h;
2554 for d in 0..d_h {
2555 attn[o + d] = 0.0;
2556 }
2557 for ki in 0..s {
2558 let w = sc[qi * s + ki];
2559 if w > score_thr {
2560 let v = bi * s * 3 * h + ki * 3 * h + 2 * h + hi * d_h;
2561 #[cfg(target_arch = "aarch64")]
2562 {
2563 use std::arch::aarch64::*;
2564 let vw = vdupq_n_f32(w);
2565 for c in 0..neon_chunks {
2566 let off = c * 4;
2567 vst1q_f32(
2568 attn.as_mut_ptr().add(o + off),
2569 vfmaq_f32(
2570 vld1q_f32(attn.as_ptr().add(o + off)),
2571 vw,
2572 vld1q_f32(qkv.as_ptr().add(v + off)),
2573 ),
2574 );
2575 }
2576 }
2577 #[cfg(not(target_arch = "aarch64"))]
2578 for d in 0..d_h {
2579 attn[o + d] += w * qkv[v + d];
2580 }
2581 }
2582 }
2583 }
2584 }
2585 }
2586
2587 // Out proj (sgemm + bias fused) + residual add with NEON
2588 crate::blas::sgemm_bias(
2589 attn,
2590 sl(*out_w, base, h * h),
2591 sl(*out_b, base, h),
2592 res,
2593 m,
2594 h,
2595 h,
2596 );
2597 #[cfg(target_arch = "aarch64")]
2598 {
2599 use std::arch::aarch64::*;
2600 let chunks_h = (m * h) / 4;
2601 for c in 0..chunks_h {
2602 let off = c * 4;
2603 vst1q_f32(
2604 res.as_mut_ptr().add(off),
2605 vaddq_f32(
2606 vld1q_f32(res.as_ptr().add(off)),
2607 vld1q_f32(inp.as_ptr().add(off)),
2608 ),
2609 );
2610 }
2611 for i in (chunks_h * 4)..(m * h) {
2612 res[i] += inp[i];
2613 }
2614 }
2615 #[cfg(not(target_arch = "aarch64"))]
2616 for i in 0..m * h {
2617 res[i] += inp[i];
2618 }
2619
2620 // LN1 (fused residual already done above — just normalize)
2621 let g1 = sl(*ln1_g, base, h);
2622 let b1 = sl(*ln1_b, base, h);
2623 for r in 0..m {
2624 crate::kernels::layer_norm_row(
2625 &res[r * h..(r + 1) * h],
2626 g1,
2627 b1,
2628 &mut normed[r * h..(r + 1) * h],
2629 h,
2630 *eps1,
2631 );
2632 }
2633
2634 // FFN: fc1 (parallel across cores) + GELU
2635 crate::blas::par_sgemm_bias(
2636 normed,
2637 sl(*fc1_w, base, h * id),
2638 sl(*fc1_b, base, id),
2639 ffn,
2640 m,
2641 h,
2642 id,
2643 );
2644 crate::kernels::par_gelu_inplace(ffn);
2645
2646 // fc2 + bias (parallel across cores) + residual with NEON
2647 crate::blas::par_sgemm_bias(
2648 ffn,
2649 sl(*fc2_w, base, id * h),
2650 sl(*fc2_b, base, h),
2651 res,
2652 m,
2653 id,
2654 h,
2655 );
2656 #[cfg(target_arch = "aarch64")]
2657 {
2658 use std::arch::aarch64::*;
2659 let chunks_h = (m * h) / 4;
2660 for c in 0..chunks_h {
2661 let off = c * 4;
2662 vst1q_f32(
2663 res.as_mut_ptr().add(off),
2664 vaddq_f32(
2665 vld1q_f32(res.as_ptr().add(off)),
2666 vld1q_f32(normed.as_ptr().add(off)),
2667 ),
2668 );
2669 }
2670 for i in (chunks_h * 4)..(m * h) {
2671 res[i] += normed[i];
2672 }
2673 }
2674 #[cfg(not(target_arch = "aarch64"))]
2675 for i in 0..m * h {
2676 res[i] += normed[i];
2677 }
2678
2679 // LN2 → output
2680 let g2 = sl(*ln2_g, base, h);
2681 let b2 = sl(*ln2_b, base, h);
2682 for r in 0..m {
2683 crate::kernels::layer_norm_row(
2684 &res[r * h..(r + 1) * h],
2685 g2,
2686 b2,
2687 &mut dst[r * h..(r + 1) * h],
2688 h,
2689 *eps2,
2690 );
2691 }
2692 }
2693 }
2694
2695 Thunk::FusedNomicLayer {
2696 hidden,
2697 qkv_w,
2698 out_w,
2699 mask,
2700 cos,
2701 sin,
2702 cos_len,
2703 ln1_g,
2704 ln1_b,
2705 eps1,
2706 fc11_w,
2707 fc12_w: _,
2708 fc2_w,
2709 ln2_g,
2710 ln2_b,
2711 eps2,
2712 out,
2713 batch,
2714 seq,
2715 hs,
2716 nh,
2717 dh,
2718 int_dim,
2719 interleaved,
2720 } => {
2721 let interleaved = *interleaved;
2722 let (b, s, h, n_h, d_h) = (
2723 *batch as usize,
2724 *seq as usize,
2725 *hs as usize,
2726 *nh as usize,
2727 *dh as usize,
2728 );
2729 let m = b * s;
2730 let id = *int_dim as usize;
2731 let scale = (d_h as f32).powf(-0.5);
2732 let half_dh = d_h / 2;
2733 #[cfg(target_arch = "aarch64")]
2734 let neon_chunks = d_h / 4;
2735 unsafe {
2736 let inp = sl(*hidden, base, m * h);
2737 let dst = sl_mut(*out, base, m * h);
2738 let mk = sl(*mask, base, b * s);
2739 let cos_tab = sl(*cos, base, *cos_len as usize);
2740 let sin_tab = sl(*sin, base, *cos_len as usize);
2741 // fc11_w is the fused [h, 2*int_dim] weight (fc11 || fc12 concatenated)
2742 let fused_fc_w = sl(*fc11_w, base, h * 2 * id);
2743
2744 let mut qkv = vec![0f32; m * 3 * h];
2745 let mut attn = vec![0f32; m * h];
2746 let mut res = vec![0f32; m * h];
2747 let mut normed = vec![0f32; m * h];
2748 let mut ffn_concat = vec![0f32; m * 2 * id]; // fc11||fc12 output
2749 let mut sc = vec![0f32; s * s];
2750
2751 // QKV (no bias)
2752 crate::blas::sgemm(inp, sl(*qkv_w, base, h * 3 * h), &mut qkv, m, h, 3 * h);
2753
2754 // SDPA with inline RoPE
2755 for bi in 0..b {
2756 for hi in 0..n_h {
2757 for qi in 0..s {
2758 for ki in 0..s {
2759 let q_base = bi * s * 3 * h + qi * 3 * h + hi * d_h;
2760 let k_base = bi * s * 3 * h + ki * 3 * h + h + hi * d_h;
2761 let mut dot = 0f32;
2762 for i in 0..half_dh {
2763 // NeoX pairs (i, i+half); GPT-J pairs (2i, 2i+1).
2764 let (o1, o2) = if interleaved {
2765 (2 * i, 2 * i + 1)
2766 } else {
2767 (i, half_dh + i)
2768 };
2769 let q1 = qkv[q_base + o1];
2770 let q2 = qkv[q_base + o2];
2771 let k1 = qkv[k_base + o1];
2772 let k2 = qkv[k_base + o2];
2773 let cq = cos_tab[qi * half_dh + i];
2774 let sq = sin_tab[qi * half_dh + i];
2775 let ck = cos_tab[ki * half_dh + i];
2776 let sk = sin_tab[ki * half_dh + i];
2777 dot += (q1 * cq - q2 * sq) * (k1 * ck - k2 * sk)
2778 + (q2 * cq + q1 * sq) * (k2 * ck + k1 * sk);
2779 }
2780 sc[qi * s + ki] = dot * scale;
2781 if mk[bi * s + ki] < mask_thr {
2782 sc[qi * s + ki] = mask_neg;
2783 }
2784 }
2785 }
2786 crate::kernels::neon_softmax(&mut sc[..s * s], s, s);
2787 for qi in 0..s {
2788 let o = bi * s * h + qi * h + hi * d_h;
2789 for d in 0..d_h {
2790 attn[o + d] = 0.0;
2791 }
2792 for ki in 0..s {
2793 let w = sc[qi * s + ki];
2794 if w > score_thr {
2795 let v = bi * s * 3 * h + ki * 3 * h + 2 * h + hi * d_h;
2796 #[cfg(target_arch = "aarch64")]
2797 {
2798 use std::arch::aarch64::*;
2799 let vw = vdupq_n_f32(w);
2800 for c in 0..neon_chunks {
2801 let off = c * 4;
2802 vst1q_f32(
2803 attn.as_mut_ptr().add(o + off),
2804 vfmaq_f32(
2805 vld1q_f32(attn.as_ptr().add(o + off)),
2806 vw,
2807 vld1q_f32(qkv.as_ptr().add(v + off)),
2808 ),
2809 );
2810 }
2811 }
2812 #[cfg(not(target_arch = "aarch64"))]
2813 for d in 0..d_h {
2814 attn[o + d] += w * qkv[v + d];
2815 }
2816 }
2817 }
2818 }
2819 }
2820 }
2821
2822 // Out proj (no bias) + residual
2823 crate::blas::sgemm(&attn, sl(*out_w, base, h * h), &mut res, m, h, h);
2824 for i in 0..m * h {
2825 res[i] += inp[i];
2826 }
2827
2828 // LN1
2829 let g1 = sl(*ln1_g, base, h);
2830 let b1 = sl(*ln1_b, base, h);
2831 for r in 0..m {
2832 crate::kernels::layer_norm_row(
2833 &res[r * h..(r + 1) * h],
2834 g1,
2835 b1,
2836 &mut normed[r * h..(r + 1) * h],
2837 h,
2838 *eps1,
2839 );
2840 }
2841
2842 // SwiGLU: fused fc11+fc12 sgemm, then split, silu, mul
2843 crate::blas::sgemm(&normed, fused_fc_w, &mut ffn_concat, m, h, 2 * id);
2844 // Split: first id cols = fc11 (up), second id cols = fc12 (gate)
2845 // SiLU on gate, then multiply up * gate → store in up region
2846 for row in 0..m {
2847 let bo = row * 2 * id;
2848 // SiLU in-place on gate portion
2849 for j in 0..id {
2850 let x = ffn_concat[bo + id + j];
2851 ffn_concat[bo + id + j] = x / (1.0 + (-x).exp());
2852 }
2853 // Multiply: up[j] *= gate[j]
2854 for j in 0..id {
2855 ffn_concat[bo + j] *= ffn_concat[bo + id + j];
2856 }
2857 }
2858
2859 // fc2 (no bias) + residual. The up*silu(gate) product lives in
2860 // the FIRST `id` cols of each 2*id-wide ffn_concat row; gather
2861 // it contiguous and run the SAME `sgemm` dispatch the unfused
2862 // path uses (a strided `sgemm_general` here would force the
2863 // BLAS/scalar path and diverge from the unfused NEON sgemm).
2864 let mut swiglu_contig = vec![0f32; m * id];
2865 for row in 0..m {
2866 let bo = row * 2 * id;
2867 swiglu_contig[row * id..(row + 1) * id]
2868 .copy_from_slice(&ffn_concat[bo..bo + id]);
2869 }
2870 crate::blas::sgemm(
2871 &swiglu_contig,
2872 sl(*fc2_w, base, id * h),
2873 &mut res,
2874 m,
2875 id,
2876 h,
2877 );
2878 for i in 0..m * h {
2879 res[i] += normed[i];
2880 }
2881
2882 // LN2 → output
2883 let g2 = sl(*ln2_g, base, h);
2884 let b2 = sl(*ln2_b, base, h);
2885 for r in 0..m {
2886 crate::kernels::layer_norm_row(
2887 &res[r * h..(r + 1) * h],
2888 g2,
2889 b2,
2890 &mut dst[r * h..(r + 1) * h],
2891 h,
2892 *eps2,
2893 );
2894 }
2895 }
2896 }
2897
2898 Thunk::FusedSwiGLU { .. } => exec_fused_swi_g_l_u(thunk, base),
2899 Thunk::Concat { .. } => exec_concat(thunk, base),
2900 Thunk::KvAppend { .. } => exec_kv_append(thunk, base),
2901 Thunk::ConcatF64 { .. } => exec_concat_f64(thunk, base),
2902 Thunk::Compare {
2903 lhs,
2904 rhs,
2905 dst,
2906 len,
2907 op,
2908 inputs_i64,
2909 inputs_elem_bytes,
2910 dst_elem_bytes,
2911 lhs_scalar,
2912 rhs_scalar,
2913 lhs_len,
2914 rhs_len,
2915 out_dims_bcast,
2916 bcast_lhs_strides,
2917 bcast_rhs_strides,
2918 } => {
2919 let len = *len as usize;
2920 let arena_len = arena_buf.len();
2921 let elem = (*inputs_elem_bytes).max(1) as usize;
2922 let dst_eb = (*dst_elem_bytes).max(1) as usize;
2923 let max_l = (arena_len.saturating_sub(*lhs)) / elem;
2924 let max_r = (arena_len.saturating_sub(*rhs)) / elem;
2925 let max_d = (arena_len.saturating_sub(*dst)) / dst_eb;
2926 if !out_dims_bcast.is_empty() {
2927 // Shape-aware broadcast (partial size-1 axes, e.g. `[.,L,1]`
2928 // vs `[.,L,L]`). Read each operand at its TRUE element count
2929 // and index via row-major coords · strides — the flat path
2930 // below would read past a partial operand into adjacent
2931 // arena memory. Mirrors `BinaryFull`'s broadcast loop.
2932 let len = len.min(max_d);
2933 let l_read = (*lhs_len as usize).min(max_l).max(1);
2934 let r_read = (*rhs_len as usize).min(max_r).max(1);
2935 let rank = out_dims_bcast.len();
2936 let bidx = |i: usize| -> (usize, usize) {
2937 let mut rem = i;
2938 let (mut li, mut ri) = (0usize, 0usize);
2939 for ax in (0..rank).rev() {
2940 let sz = out_dims_bcast[ax] as usize;
2941 let c = rem % sz;
2942 rem /= sz;
2943 li += c * bcast_lhs_strides[ax] as usize;
2944 ri += c * bcast_rhs_strides[ax] as usize;
2945 }
2946 (li, ri)
2947 };
2948 macro_rules! do_cmp {
2949 ($a:expr, $b:expr) => {
2950 match op {
2951 CmpOp::Eq => $a == $b,
2952 CmpOp::Ne => $a != $b,
2953 CmpOp::Lt => $a < $b,
2954 CmpOp::Le => $a <= $b,
2955 CmpOp::Gt => $a > $b,
2956 CmpOp::Ge => $a >= $b,
2957 }
2958 };
2959 }
2960 macro_rules! put {
2961 ($i:expr, $v:expr) => {
2962 if *dst_elem_bytes == 1 {
2963 arena_buf[*dst + $i] = u8::from($v);
2964 } else {
2965 let o = sl_mut(*dst, base, len);
2966 o[$i] = if $v { 1.0 } else { 0.0 };
2967 }
2968 };
2969 }
2970 if elem == 1 {
2971 let l = arena_buf[*lhs..*lhs + l_read].to_vec();
2972 let r = arena_buf[*rhs..*rhs + r_read].to_vec();
2973 for i in 0..len {
2974 let (li, ri) = bidx(i);
2975 let v = do_cmp!(l[li], r[ri]);
2976 unsafe { put!(i, v) }
2977 }
2978 } else if *inputs_i64 != 0 {
2979 unsafe {
2980 let l = sl_i64(*lhs, base, l_read);
2981 let r = sl_i64(*rhs, base, r_read);
2982 for i in 0..len {
2983 let (li, ri) = bidx(i);
2984 let v = do_cmp!(l[li], r[ri]);
2985 put!(i, v)
2986 }
2987 }
2988 } else {
2989 unsafe {
2990 let l = sl(*lhs, base, l_read);
2991 let r = sl(*rhs, base, r_read);
2992 for i in 0..len {
2993 let (li, ri) = bidx(i);
2994 let v = do_cmp!(l[li], r[ri]);
2995 put!(i, v)
2996 }
2997 }
2998 }
2999 } else {
3000 let l_n = if *lhs_scalar { 1 } else { len };
3001 let r_n = if *rhs_scalar { 1 } else { len };
3002 // Keep full `len` when broadcasting scalars — only the
3003 // non-scalar operands (and dst) may shrink the loop.
3004 let mut len = len.min(max_d);
3005 if *lhs_scalar {
3006 if max_l < 1 {
3007 len = 0;
3008 }
3009 } else {
3010 len = len.min(max_l);
3011 }
3012 if *rhs_scalar {
3013 if max_r < 1 {
3014 len = 0;
3015 }
3016 } else {
3017 len = len.min(max_r);
3018 }
3019 if trace_thunks && len > 0 {
3020 eprintln!(
3021 "[compare] len={len} lhs={} rhs={} dst={} ls={} rs={}",
3022 *lhs, *rhs, *dst, *lhs_scalar, *rhs_scalar
3023 );
3024 }
3025 if elem == 1 {
3026 let l = arena_buf[*lhs..*lhs + l_n.min(max_l).max(1)].to_vec();
3027 let r = arena_buf[*rhs..*rhs + r_n.min(max_r).max(1)].to_vec();
3028 for i in 0..len {
3029 let li = if *lhs_scalar { 0 } else { i };
3030 let ri = if *rhs_scalar { 0 } else { i };
3031 let v = match op {
3032 CmpOp::Eq => l[li] == r[ri],
3033 CmpOp::Ne => l[li] != r[ri],
3034 CmpOp::Lt => l[li] < r[ri],
3035 CmpOp::Le => l[li] <= r[ri],
3036 CmpOp::Gt => l[li] > r[ri],
3037 CmpOp::Ge => l[li] >= r[ri],
3038 };
3039 if *dst_elem_bytes == 1 {
3040 arena_buf[*dst + i] = u8::from(v);
3041 } else {
3042 unsafe {
3043 let o = sl_mut(*dst, base, len);
3044 o[i] = if v { 1.0 } else { 0.0 };
3045 }
3046 }
3047 }
3048 } else if *inputs_i64 != 0 {
3049 unsafe {
3050 let l = sl_i64(*lhs, base, l_n.min(max_l).max(1));
3051 let r = sl_i64(*rhs, base, r_n.min(max_r).max(1));
3052 for i in 0..len {
3053 let li = if *lhs_scalar { 0 } else { i };
3054 let ri = if *rhs_scalar { 0 } else { i };
3055 let v = match op {
3056 CmpOp::Eq => l[li] == r[ri],
3057 CmpOp::Ne => l[li] != r[ri],
3058 CmpOp::Lt => l[li] < r[ri],
3059 CmpOp::Le => l[li] <= r[ri],
3060 CmpOp::Gt => l[li] > r[ri],
3061 CmpOp::Ge => l[li] >= r[ri],
3062 };
3063 if *dst_elem_bytes == 1 {
3064 arena_buf[*dst + i] = u8::from(v);
3065 } else {
3066 let o = sl_mut(*dst, base, len);
3067 o[i] = if v { 1.0 } else { 0.0 };
3068 }
3069 }
3070 }
3071 } else {
3072 unsafe {
3073 let l = sl(*lhs, base, l_n.min(max_l).max(1));
3074 let r = sl(*rhs, base, r_n.min(max_r).max(1));
3075 for i in 0..len {
3076 let li = if *lhs_scalar { 0 } else { i };
3077 let ri = if *rhs_scalar { 0 } else { i };
3078 let v = match op {
3079 CmpOp::Eq => l[li] == r[ri],
3080 CmpOp::Ne => l[li] != r[ri],
3081 CmpOp::Lt => l[li] < r[ri],
3082 CmpOp::Le => l[li] <= r[ri],
3083 CmpOp::Gt => l[li] > r[ri],
3084 CmpOp::Ge => l[li] >= r[ri],
3085 };
3086 if *dst_elem_bytes == 1 {
3087 arena_buf[*dst + i] = u8::from(v);
3088 } else {
3089 let o = sl_mut(*dst, base, len);
3090 o[i] = if v { 1.0 } else { 0.0 };
3091 }
3092 }
3093 }
3094 }
3095 }
3096 }
3097
3098 Thunk::Where {
3099 cond,
3100 on_true,
3101 on_false,
3102 dst,
3103 len,
3104 elem_bytes,
3105 cond_elem_bytes,
3106 cond_scalar,
3107 true_scalar,
3108 false_scalar,
3109 } => {
3110 let len = *len as usize;
3111 let eb = *elem_bytes as usize;
3112 let cond_eb = (*cond_elem_bytes).max(1) as usize;
3113 let arena_len = arena_buf.len();
3114 let c_n = if *cond_scalar { 1 } else { len };
3115 let t_n = if *true_scalar { 1 } else { len };
3116 let f_n = if *false_scalar { 1 } else { len };
3117 let max_c = (arena_len.saturating_sub(*cond)) / cond_eb;
3118 let max_t = (arena_len.saturating_sub(*on_true)) / eb;
3119 let max_f = (arena_len.saturating_sub(*on_false)) / eb;
3120 let max_d = (arena_len.saturating_sub(*dst)) / eb;
3121 let mut len = len.min(max_d);
3122 if *cond_scalar {
3123 if max_c < 1 {
3124 len = 0;
3125 }
3126 } else {
3127 len = len.min(max_c);
3128 }
3129 if *true_scalar {
3130 if max_t < 1 {
3131 len = 0;
3132 }
3133 } else {
3134 len = len.min(max_t);
3135 }
3136 if *false_scalar {
3137 if max_f < 1 {
3138 len = 0;
3139 }
3140 } else {
3141 len = len.min(max_f);
3142 }
3143 unsafe {
3144 if *elem_bytes == 8 {
3145 let t = sl_i64(*on_true, base, t_n.min(max_t).max(1));
3146 let e = sl_i64(*on_false, base, f_n.min(max_f).max(1));
3147 let o = sl_mut_i64(*dst, base, len);
3148 if *cond_elem_bytes == 1 {
3149 let c = &arena_buf[*cond..*cond + c_n.min(max_c).max(1)];
3150 for i in 0..len {
3151 let ci = if *cond_scalar { 0 } else { i };
3152 let ti = if *true_scalar { 0 } else { i };
3153 let ei = if *false_scalar { 0 } else { i };
3154 o[i] = if c[ci] != 0 { t[ti] } else { e[ei] };
3155 }
3156 } else if *cond_elem_bytes == 4 {
3157 // Bool-as-f32 masks (f32-uniform arena).
3158 let c = sl(*cond, base, c_n.min(max_c).max(1));
3159 for i in 0..len {
3160 let ci = if *cond_scalar { 0 } else { i };
3161 let ti = if *true_scalar { 0 } else { i };
3162 let ei = if *false_scalar { 0 } else { i };
3163 o[i] = if c[ci] != 0.0 { t[ti] } else { e[ei] };
3164 }
3165 } else {
3166 let c = sl_i64(*cond, base, c_n.min(max_c).max(1));
3167 for i in 0..len {
3168 let ci = if *cond_scalar { 0 } else { i };
3169 let ti = if *true_scalar { 0 } else { i };
3170 let ei = if *false_scalar { 0 } else { i };
3171 o[i] = if c[ci] != 0 { t[ti] } else { e[ei] };
3172 }
3173 }
3174 } else if *cond_elem_bytes == 1 {
3175 let c = &arena_buf[*cond..*cond + c_n.min(max_c).max(1)];
3176 let t = sl(*on_true, base, t_n.min(max_t).max(1));
3177 let e = sl(*on_false, base, f_n.min(max_f).max(1));
3178 let o = sl_mut(*dst, base, len);
3179 for i in 0..len {
3180 let ci = if *cond_scalar { 0 } else { i };
3181 let ti = if *true_scalar { 0 } else { i };
3182 let ei = if *false_scalar { 0 } else { i };
3183 o[i] = if c[ci] != 0 { t[ti] } else { e[ei] };
3184 }
3185 } else {
3186 let c = sl(*cond, base, c_n.min(max_c).max(1));
3187 let t = sl(*on_true, base, t_n.min(max_t).max(1));
3188 let e = sl(*on_false, base, f_n.min(max_f).max(1));
3189 let o = sl_mut(*dst, base, len);
3190 for i in 0..len {
3191 let ci = if *cond_scalar { 0 } else { i };
3192 let ti = if *true_scalar { 0 } else { i };
3193 let ei = if *false_scalar { 0 } else { i };
3194 o[i] = if c[ci] != 0.0 { t[ti] } else { e[ei] };
3195 }
3196 }
3197 }
3198 }
3199
3200 Thunk::Fma {
3201 a,
3202 b,
3203 c,
3204 dst,
3205 len,
3206 elem_bytes,
3207 } => {
3208 let len = *len as usize;
3209 let eb = (*elem_bytes).max(1) as usize;
3210 let arena_len = arena_buf.len();
3211 let len = len
3212 .min(arena_len.saturating_sub(*a) / eb)
3213 .min(arena_len.saturating_sub(*b) / eb)
3214 .min(arena_len.saturating_sub(*c) / eb)
3215 .min(arena_len.saturating_sub(*dst) / eb);
3216 unsafe {
3217 if *elem_bytes == 8 {
3218 let av = sl_f64(*a, base, len);
3219 let bv = sl_f64(*b, base, len);
3220 let cv = sl_f64(*c, base, len);
3221 let o = sl_mut_f64(*dst, base, len);
3222 for i in 0..len {
3223 o[i] = av[i].mul_add(bv[i], cv[i]);
3224 }
3225 } else {
3226 let av = sl(*a, base, len);
3227 let bv = sl(*b, base, len);
3228 let cv = sl(*c, base, len);
3229 let o = sl_mut(*dst, base, len);
3230 for i in 0..len {
3231 o[i] = av[i].mul_add(bv[i], cv[i]);
3232 }
3233 }
3234 }
3235 }
3236
3237 Thunk::ScatterAdd { .. } => exec_scatter_add(thunk, base),
3238 Thunk::ScatterNd { .. } => exec_scatter_nd(thunk, base),
3239 Thunk::ScatterElements { .. } => exec_scatter_elements(thunk, base),
3240 Thunk::GatherNd { .. } => exec_gather_nd(thunk, base),
3241 Thunk::GatherElements { .. } => exec_gather_elements(thunk, base),
3242 Thunk::GroupedMatMul {
3243 input,
3244 weight,
3245 expert_idx,
3246 dst,
3247 m,
3248 k_dim,
3249 n,
3250 num_experts,
3251 w_transposed,
3252 } => {
3253 let m = *m as usize;
3254 let k_dim = *k_dim as usize;
3255 let n = *n as usize;
3256 let num_experts = *num_experts as usize;
3257 let w_transposed = *w_transposed;
3258 unsafe {
3259 let inp = sl(*input, base, m * k_dim);
3260 let wt = sl(*weight, base, num_experts * k_dim * n);
3261 let ids = sl(*expert_idx, base, m);
3262 let out = sl_mut(*dst, base, m * n);
3263
3264 // Counting-sort tokens by their assigned expert.
3265 // counts[e] = how many tokens routed to expert e.
3266 let mut counts = vec![0usize; num_experts];
3267 for i in 0..m {
3268 let e = ids[i] as usize;
3269 debug_assert!(
3270 e < num_experts,
3271 "expert_idx out of range: {e} >= {num_experts}"
3272 );
3273 counts[e] += 1;
3274 }
3275 // Cumulative offsets into the packed buffer.
3276 let mut offsets = vec![0usize; num_experts + 1];
3277 for e in 0..num_experts {
3278 offsets[e + 1] = offsets[e] + counts[e];
3279 }
3280 // Pack: each expert's rows land contiguously in `packed_in`.
3281 // `original_pos[packed_idx] = original_token_idx` for the
3282 // unpermute step at the end.
3283 let mut packed_in = vec![0f32; m * k_dim];
3284 let mut original_pos = vec![0usize; m];
3285 let mut write_idx = vec![0usize; num_experts];
3286 for i in 0..m {
3287 let e = ids[i] as usize;
3288 let dst_row = offsets[e] + write_idx[e];
3289 packed_in[dst_row * k_dim..(dst_row + 1) * k_dim]
3290 .copy_from_slice(&inp[i * k_dim..(i + 1) * k_dim]);
3291 original_pos[dst_row] = i;
3292 write_idx[e] += 1;
3293 }
3294
3295 // One BLAS sgemm per expert. Skip experts with no
3296 // tokens — common at the tail when M is much smaller
3297 // than num_experts × k.
3298 let mut packed_out = vec![0f32; m * n];
3299 let expert_stride = k_dim * n;
3300 let gmm_ord = crate::moe_residency::next_gmm_ord();
3301 let moe_layer = gmm_ord / 3;
3302 for e in 0..num_experts {
3303 let count = counts[e];
3304 if count == 0 {
3305 continue;
3306 }
3307 crate::moe_residency::record_expert_tokens(moe_layer, e, count);
3308 let in_start = offsets[e];
3309 let in_slice = &packed_in[in_start * k_dim..(in_start + count) * k_dim];
3310 let w_slab: &[f32] =
3311 if !crate::moe_residency::expert_on_device_for_layer(moe_layer, e) {
3312 if let Some(ptr) =
3313 crate::moe_residency::host_expert_weight_ptr(gmm_ord, e)
3314 {
3315 // Residency hands over `[K, N]` slabs by
3316 // construction. A folded transpose means the
3317 // bank is `[N, K]`, so the two do not compose
3318 // — and silently using one as the other reads
3319 // a correctly-sized block of the wrong values,
3320 // which no shape check would catch. Refuse.
3321 assert!(
3322 !w_transposed,
3323 "rlx-cpu: MoE host residency supplies [K, N] expert \
3324 slabs, but this GroupedMatMul folded a bank transpose \
3325 and reads [N, K]. The two are not composed; unbind \
3326 residency or disable the fold."
3327 );
3328 std::slice::from_raw_parts(ptr, expert_stride)
3329 } else {
3330 &wt[e * expert_stride..(e + 1) * expert_stride]
3331 }
3332 } else {
3333 &wt[e * expert_stride..(e + 1) * expert_stride]
3334 };
3335 let out_slice = &mut packed_out[in_start * n..(in_start + count) * n];
3336 // Measured: routing this through `par_sgemm` REGRESSES decode
3337 // (0.069 -> 0.085 s/token). At m == 1 the per-call work is a
3338 // 1.5 MFLOP matvec, and decode issues 368 of these per token,
3339 // so ~368 Rayon hand-offs cost more than the extra memory
3340 // parallelism buys. Keep the serial call.
3341 if w_transposed {
3342 // Bank is `[E, N, K]`: B-transposed GEMM, no copy.
3343 crate::blas::sgemm_bt(
3344 in_slice, w_slab, out_slice, count, k_dim, n, 1.0,
3345 );
3346 } else {
3347 crate::blas::sgemm(in_slice, w_slab, out_slice, count, k_dim, n);
3348 }
3349 }
3350
3351 // Unpermute back to original token order.
3352 for packed_idx in 0..m {
3353 let i = original_pos[packed_idx];
3354 out[i * n..(i + 1) * n]
3355 .copy_from_slice(&packed_out[packed_idx * n..(packed_idx + 1) * n]);
3356 }
3357 }
3358 }
3359
3360 Thunk::DequantGroupedMatMulGguf { .. } => {
3361 exec_dequant_grouped_mat_mul_gguf(thunk, base)
3362 }
3363 Thunk::DequantGroupedMatMulMlx { .. } => exec_dequant_grouped_mat_mul_mlx(thunk, base),
3364 Thunk::DequantMoEWeightsGguf { .. } => exec_dequant_mo_e_weights_gguf(thunk, base),
3365 Thunk::TopK {
3366 src,
3367 dst,
3368 outer,
3369 axis_dim,
3370 k,
3371 indices_i64,
3372 } => {
3373 let outer = *outer as usize;
3374 let axis_dim = *axis_dim as usize;
3375 let k = *k as usize;
3376 unsafe {
3377 let inp = sl(*src, base, outer * axis_dim);
3378 // Repeated argmax with masking. O(k * axis_dim) per row;
3379 // good enough for small k (MoE typical k=2–8). For larger
3380 // k a partial heap would win.
3381 let mut row_buf: Vec<f32> = vec![0.0; axis_dim];
3382 if *indices_i64 != 0 {
3383 let out = sl_mut_i64(*dst, base, outer * k);
3384 for o in 0..outer {
3385 row_buf.copy_from_slice(&inp[o * axis_dim..(o + 1) * axis_dim]);
3386 for ki in 0..k {
3387 let mut best_i = 0usize;
3388 let mut best_v = row_buf[0];
3389 for i in 1..axis_dim {
3390 let v = row_buf[i];
3391 if v > best_v {
3392 best_v = v;
3393 best_i = i;
3394 }
3395 }
3396 out[o * k + ki] = best_i as i64;
3397 row_buf[best_i] = f32::NEG_INFINITY;
3398 }
3399 }
3400 } else {
3401 let out = sl_mut(*dst, base, outer * k);
3402 for o in 0..outer {
3403 row_buf.copy_from_slice(&inp[o * axis_dim..(o + 1) * axis_dim]);
3404 for ki in 0..k {
3405 let mut best_i = 0usize;
3406 let mut best_v = row_buf[0];
3407 for i in 1..axis_dim {
3408 let v = row_buf[i];
3409 if v > best_v {
3410 best_v = v;
3411 best_i = i;
3412 }
3413 }
3414 out[o * k + ki] = best_i as f32;
3415 row_buf[best_i] = f32::NEG_INFINITY;
3416 }
3417 }
3418 if let Some(cap) = schedule.moe_topk_capture.as_ref() {
3419 cap.push_topk_f32(&out[..outer * k], axis_dim);
3420 }
3421 }
3422 }
3423 }
3424
3425 Thunk::Reduce { .. } => exec_reduce(thunk, base),
3426 Thunk::Histogram { .. } => exec_histogram(thunk, base),
3427 Thunk::ArgReduce { .. } => exec_arg_reduce(thunk, base),
3428 Thunk::Conv2D1x1 { .. } => exec_conv2_d1x1(thunk, base),
3429 Thunk::Conv2D { .. } => exec_conv2_d(thunk, base),
3430 Thunk::Conv3d { .. } => exec_conv3d(thunk, base),
3431 Thunk::ConvTranspose3d { .. } => exec_conv_transpose3d(thunk, base),
3432 Thunk::Pool2D {
3433 src,
3434 dst,
3435 n,
3436 c,
3437 h,
3438 w,
3439 h_out,
3440 w_out,
3441 kh,
3442 kw,
3443 sh,
3444 sw,
3445 ph,
3446 pw,
3447 kind,
3448 } => {
3449 let n = *n as usize;
3450 let c = *c as usize;
3451 let h = *h as usize;
3452 let w = *w as usize;
3453 let h_out = *h_out as usize;
3454 let w_out = *w_out as usize;
3455 let kh = *kh as usize;
3456 let kw = *kw as usize;
3457 let sh = *sh as usize;
3458 let sw = *sw as usize;
3459 let ph = *ph as usize;
3460 let pw = *pw as usize;
3461 let kernel_area = (kh * kw) as f32;
3462 unsafe {
3463 let inp = sl(*src, base, n * c * h * w);
3464 let out = sl_mut(*dst, base, n * c * h_out * w_out);
3465 // Each (n, c) plane is independent and writes a disjoint
3466 // output region, so pooling fans out over the channel-batch
3467 // when RLX_FAST_CONV is set.
3468 let out_addr = out.as_mut_ptr() as usize;
3469 let is_max = matches!(kind, ReduceOp::Max);
3470 let is_mean = matches!(kind, ReduceOp::Mean);
3471 // No-padding windows are *usually* fully in-bounds (the
3472 // conv-net case), so the hot path drops the per-element
3473 // bounds branches and hoists the reduce-op choice out of the
3474 // loop. "Usually" is not "always": under ONNX `ceil_mode=1`
3475 // the final window may start inside the input and overhang
3476 // its end (259 frames, kernel/stride 100 -> 3 windows, the
3477 // last spanning 200..300), and the fast path then indexes
3478 // past the buffer. Require every window to fit before taking
3479 // it; the general path below already treats the overhang as
3480 // padding, which is what ONNX specifies.
3481 let windows_fit = h_out.saturating_sub(1) * sh + kh <= h
3482 && w_out.saturating_sub(1) * sw + kw <= w;
3483 let nopad = ph == 0 && pw == 0 && windows_fit;
3484 let pool_plane = |nc: usize| {
3485 let ni = nc / c;
3486 let ci = nc % c;
3487 let in_chan = ni * c * h * w + ci * h * w;
3488 let out_chan = ni * c * h_out * w_out + ci * h_out * w_out;
3489 let op = out_addr as *mut f32;
3490 for ho in 0..h_out {
3491 for wo in 0..w_out {
3492 let mut valid = 0.0f32;
3493 let acc = if nopad {
3494 let row0 = in_chan + (ho * sh) * w + wo * sw;
3495 let mut a = if is_max { f32::NEG_INFINITY } else { 0.0 };
3496 for ki in 0..kh {
3497 let row = row0 + ki * w;
3498 if is_max {
3499 for kj in 0..kw {
3500 a = a.max(inp[row + kj]);
3501 }
3502 } else {
3503 for kj in 0..kw {
3504 a += inp[row + kj];
3505 }
3506 }
3507 }
3508 a
3509 } else {
3510 let mut a = if is_max { f32::NEG_INFINITY } else { 0.0 };
3511 for ki in 0..kh {
3512 for kj in 0..kw {
3513 let hi = ho * sh + ki;
3514 let wi = wo * sw + kj;
3515 if hi < ph || wi < pw {
3516 continue;
3517 }
3518 let hi = hi - ph;
3519 let wi = wi - pw;
3520 if hi >= h || wi >= w {
3521 continue;
3522 }
3523 let v = inp[in_chan + hi * w + wi];
3524 if is_max {
3525 a = a.max(v);
3526 } else {
3527 a += v;
3528 valid += 1.0;
3529 }
3530 }
3531 }
3532 a
3533 };
3534 // Average over the elements that actually exist,
3535 // not over the nominal window. ONNX's default is
3536 // `count_include_pad = 0`, and a `ceil_mode`
3537 // window that overhangs the input is not padding
3538 // at all — those positions simply are not there.
3539 // ChatterBox's speaker encoder pools 259 frames
3540 // with kernel/stride 100: the third window holds
3541 // 59 real frames, and dividing it by 100 scaled
3542 // it by 0.59 against the reference.
3543 let denom = if nopad { kernel_area } else { valid.max(1.0) };
3544 let acc = if is_mean { acc / denom } else { acc };
3545 *op.add(out_chan + ho * w_out + wo) = acc;
3546 }
3547 }
3548 };
3549 if fast_conv_enabled() && crate::pool::should_parallelize(n * c * h_out * w_out)
3550 {
3551 crate::pool::par_for(
3552 n * c,
3553 crate::pool::outer_chunk(n * c),
3554 &|off, cnt| {
3555 for nc in off..off + cnt {
3556 pool_plane(nc);
3557 }
3558 },
3559 );
3560 } else {
3561 for nc in 0..n * c {
3562 pool_plane(nc);
3563 }
3564 }
3565 }
3566 }
3567
3568 Thunk::Pool3D {
3569 src,
3570 dst,
3571 n,
3572 c,
3573 d,
3574 h,
3575 w,
3576 d_out,
3577 h_out,
3578 w_out,
3579 kd,
3580 kh,
3581 kw,
3582 sd,
3583 sh,
3584 sw,
3585 pd,
3586 ph,
3587 pw,
3588 kind,
3589 } => {
3590 let n = *n as usize;
3591 let c = *c as usize;
3592 let d = *d as usize;
3593 let h = *h as usize;
3594 let w = *w as usize;
3595 let d_out = *d_out as usize;
3596 let h_out = *h_out as usize;
3597 let w_out = *w_out as usize;
3598 let kd = *kd as usize;
3599 let kh = *kh as usize;
3600 let kw = *kw as usize;
3601 let sd = *sd as usize;
3602 let sh = *sh as usize;
3603 let sw = *sw as usize;
3604 let pd = *pd as usize;
3605 let ph = *ph as usize;
3606 let pw = *pw as usize;
3607 let kernel_vol = (kd * kh * kw) as f32;
3608 let is_max = matches!(kind, ReduceOp::Max);
3609 let is_mean = matches!(kind, ReduceOp::Mean);
3610 unsafe {
3611 let inp = sl(*src, base, n * c * d * h * w);
3612 let out = sl_mut(*dst, base, n * c * d_out * h_out * w_out);
3613 for ni in 0..n {
3614 for ci in 0..c {
3615 let in_base = (ni * c + ci) * d * h * w;
3616 let out_base = (ni * c + ci) * d_out * h_out * w_out;
3617 for od in 0..d_out {
3618 for ho in 0..h_out {
3619 for wo in 0..w_out {
3620 let mut acc = if is_max { f32::NEG_INFINITY } else { 0.0 };
3621 for kz in 0..kd {
3622 for ki in 0..kh {
3623 for kj in 0..kw {
3624 let di = od * sd + kz;
3625 let hi = ho * sh + ki;
3626 let wi = wo * sw + kj;
3627 if di < pd || hi < ph || wi < pw {
3628 continue;
3629 }
3630 let di = di - pd;
3631 let hi = hi - ph;
3632 let wi = wi - pw;
3633 if di >= d || hi >= h || wi >= w {
3634 continue;
3635 }
3636 let v = inp[in_base + (di * h + hi) * w + wi];
3637 if is_max {
3638 acc = acc.max(v);
3639 } else {
3640 acc += v;
3641 }
3642 }
3643 }
3644 }
3645 if is_mean {
3646 acc /= kernel_vol;
3647 }
3648 out[out_base + (od * h_out + ho) * w_out + wo] = acc;
3649 }
3650 }
3651 }
3652 }
3653 }
3654 }
3655 }
3656
3657 Thunk::ReluBackward { .. } => exec_relu_backward(thunk, base),
3658 Thunk::ReluBackwardF64 { .. } => exec_relu_backward_f64(thunk, base),
3659 Thunk::QMatMul { .. } => exec_q_mat_mul(thunk, base),
3660 Thunk::QConv2d {
3661 x,
3662 w,
3663 bias,
3664 out,
3665 n,
3666 c_in,
3667 h,
3668 w_in,
3669 c_out,
3670 h_out,
3671 w_out,
3672 kh,
3673 kw,
3674 sh,
3675 sw,
3676 ph,
3677 pw,
3678 dh,
3679 dw,
3680 groups,
3681 x_zp,
3682 w_zp,
3683 out_zp,
3684 mult,
3685 } => {
3686 let n = *n as usize;
3687 let c_in = *c_in as usize;
3688 let h = *h as usize;
3689 let w_in = *w_in as usize;
3690 let c_out = *c_out as usize;
3691 let h_out = *h_out as usize;
3692 let w_out = *w_out as usize;
3693 let kh = *kh as usize;
3694 let kw = *kw as usize;
3695 let sh = *sh as usize;
3696 let sw = *sw as usize;
3697 let ph = *ph as usize;
3698 let pw = *pw as usize;
3699 let dh = *dh as usize;
3700 let dw = *dw as usize;
3701 let groups = *groups as usize;
3702 let c_in_per_g = c_in / groups;
3703 let c_out_per_g = c_out / groups;
3704 unsafe {
3705 let x_ptr = base.add(*x) as *const i8;
3706 let w_ptr = base.add(*w) as *const i8;
3707 let bias_ptr = base.add(*bias) as *const i32;
3708 let out_ptr = base.add(*out) as *mut i8;
3709 for ni in 0..n {
3710 for co in 0..c_out {
3711 let g = co / c_out_per_g;
3712 let ci_start = g * c_in_per_g;
3713 for ho in 0..h_out {
3714 for wo in 0..w_out {
3715 let mut acc: i32 = *bias_ptr.add(co);
3716 for ci_off in 0..c_in_per_g {
3717 let ci = ci_start + ci_off;
3718 let in_chan = ((ni * c_in) + ci) * h * w_in;
3719 let wt_chan = ((co * c_in_per_g) + ci_off) * kh * kw;
3720 for ki in 0..kh {
3721 for kj in 0..kw {
3722 let hi = ho * sh + ki * dh;
3723 let wi = wo * sw + kj * dw;
3724 if hi < ph || wi < pw {
3725 continue;
3726 }
3727 let hi = hi - ph;
3728 let wi = wi - pw;
3729 if hi >= h || wi >= w_in {
3730 continue;
3731 }
3732 let xv = *x_ptr.add(in_chan + hi * w_in + wi)
3733 as i32
3734 - *x_zp;
3735 let wv = *w_ptr.add(wt_chan + ki * kw + kj) as i32
3736 - *w_zp;
3737 acc += xv * wv;
3738 }
3739 }
3740 }
3741 let r = (acc as f32 * *mult).round() as i32 + *out_zp;
3742 let r = r.clamp(-128, 127) as i8;
3743 let dst = ((ni * c_out) + co) * h_out * w_out + ho * w_out + wo;
3744 *out_ptr.add(dst) = r;
3745 }
3746 }
3747 }
3748 }
3749 }
3750 }
3751
3752 Thunk::Quantize { .. } => exec_quantize(thunk, base),
3753 Thunk::Dequantize { .. } => exec_dequantize(thunk, base),
3754 Thunk::FakeQuantize { .. } => exec_fake_quantize(thunk, base),
3755 Thunk::ActivationBackward { .. } => exec_activation_backward(thunk, base),
3756 Thunk::ActivationBackwardF64 { .. } => exec_activation_backward_f64(thunk, base),
3757 Thunk::FakeQuantizeLSQ { .. } => exec_fake_quantize_l_s_q(thunk, base),
3758 Thunk::FakeQuantizeLSQBackwardX { .. } => {
3759 exec_fake_quantize_l_s_q_backward_x(thunk, base)
3760 }
3761 Thunk::FakeQuantizeLSQBackwardScale { .. } => {
3762 exec_fake_quantize_l_s_q_backward_scale(thunk, base)
3763 }
3764 Thunk::FakeQuantizeBackward { .. } => exec_fake_quantize_backward(thunk, base),
3765 Thunk::LayerNormBackwardInput { .. } => exec_layer_norm_backward_input(thunk, base),
3766 Thunk::BatchNormInferenceBackwardInput { .. } => {
3767 exec_batch_norm_inference_backward_input(thunk, base)
3768 }
3769 Thunk::BatchNormInferenceBackwardGamma { .. } => {
3770 exec_batch_norm_inference_backward_gamma(thunk, base)
3771 }
3772 Thunk::BatchNormInferenceBackwardBeta { .. } => {
3773 exec_batch_norm_inference_backward_beta(thunk, base)
3774 }
3775 Thunk::LayerNormBackwardGamma { .. } => exec_layer_norm_backward_gamma(thunk, base),
3776 Thunk::RmsNormBackwardInput { .. } => exec_rms_norm_backward_input(thunk, base),
3777 Thunk::RmsNormBackwardGamma { .. } => exec_rms_norm_backward_gamma(thunk, base),
3778 Thunk::RmsNormBackwardBeta { .. } => exec_rms_norm_backward_beta(thunk, base),
3779 Thunk::RopeBackward { .. } => exec_rope_backward(thunk, base),
3780 Thunk::CumsumBackward { .. } => exec_cumsum_backward(thunk, base),
3781 Thunk::GroupNormBackwardInput { .. } => exec_group_norm_backward_input(thunk, base),
3782 Thunk::GroupNormBackwardGamma { .. } => exec_group_norm_backward_gamma(thunk, base),
3783 Thunk::GroupNormBackwardBeta { .. } => exec_group_norm_backward_beta(thunk, base),
3784 Thunk::GatherBackward { .. } => exec_gather_backward(thunk, base),
3785 Thunk::MaxPool2dBackward { .. } => exec_max_pool2d_backward(thunk, base),
3786 Thunk::Conv2dBackwardInput { .. } => exec_conv2d_backward_input(thunk, base),
3787 Thunk::Conv2dBackwardWeight { .. } => exec_conv2d_backward_weight(thunk, base),
3788 Thunk::MaxPool3dBackward { .. } => exec_max_pool3d_backward(thunk, base),
3789 Thunk::Conv3dBackwardInput { .. } => exec_conv3d_backward_input(thunk, base),
3790 Thunk::Conv3dBackwardWeight { .. } => exec_conv3d_backward_weight(thunk, base),
3791 Thunk::Im2Col { .. } => exec_im2_col(thunk, base),
3792 Thunk::SoftmaxCrossEntropyDense { .. } => exec_softmax_cross_entropy_dense(thunk, base),
3793 Thunk::SoftmaxCrossEntropy { .. } => exec_softmax_cross_entropy(thunk, base),
3794 Thunk::SoftmaxCrossEntropyBackward { .. } => {
3795 exec_softmax_cross_entropy_backward(thunk, base)
3796 }
3797 Thunk::GatherAxis { .. } => exec_gather_axis(thunk, base),
3798 Thunk::Transpose {
3799 src,
3800 dst,
3801 in_total,
3802 out_dims,
3803 in_strides,
3804 elem_bytes,
3805 } => {
3806 // N-D index walk: for each output flat index, decompose into
3807 // multi-dim coords using out_dims, then dot with in_strides
3808 // to find the source flat index. Stride 0 = broadcast (read
3809 // the same input element repeatedly along that dim).
3810 let rank = out_dims.len();
3811 let total: usize = out_dims.iter().map(|&d| d as usize).product();
3812 // Empty output (e.g. Zipformer downsample pad `Expand` → `[0,1,C]`
3813 // when T is already a multiple of the stride) — nothing to write.
3814 if total == 0 {
3815 // fall through
3816 } else {
3817 let in_total = *in_total as usize;
3818 unsafe {
3819 if *elem_bytes == 1 {
3820 // 1-byte dtypes (Bool / I8 / U8). Without this branch the
3821 // `else` path below reads/writes 4 bytes per element via the
3822 // f32 slice, corrupting e.g. a broadcast of the VITS attention
3823 // mask (Bool, expanded over heads) — masking wrong positions.
3824 let inp = arena_buf[*src..*src + in_total].to_vec();
3825 let out = &mut arena_buf[*dst..*dst + total];
3826 let mut idx = vec![0usize; rank];
3827 for o in 0..total {
3828 let mut src_idx = 0usize;
3829 for d in 0..rank {
3830 src_idx += idx[d] * in_strides[d] as usize;
3831 }
3832 out[o] = inp[broadcast_src_index(src_idx, in_total)];
3833 for d in (0..rank).rev() {
3834 idx[d] += 1;
3835 if idx[d] < out_dims[d] as usize {
3836 break;
3837 }
3838 idx[d] = 0;
3839 }
3840 }
3841 } else if *elem_bytes == 8 {
3842 let inp = sl_i64(*src, base, in_total);
3843 let out = sl_mut_i64(*dst, base, total);
3844 let mut idx = vec![0usize; rank];
3845 for o in 0..total {
3846 let mut src_idx = 0usize;
3847 for d in 0..rank {
3848 src_idx += idx[d] * in_strides[d] as usize;
3849 }
3850 out[o] = inp[broadcast_src_index(src_idx, in_total)];
3851 for d in (0..rank).rev() {
3852 idx[d] += 1;
3853 if idx[d] < out_dims[d] as usize {
3854 break;
3855 }
3856 idx[d] = 0;
3857 }
3858 }
3859 } else if *elem_bytes == 16 {
3860 // 16-byte dtypes (C128 = 2×f64 native / 4×f32 lanes).
3861 // Copy each element as two i64 words; without this
3862 // branch the `else` f32 path below reads 4 bytes per
3863 // element and shatters the complex element (keeps only
3864 // the first f32, drops the other 12 bytes).
3865 let inp = sl_i64(*src, base, in_total * 2);
3866 let out = sl_mut_i64(*dst, base, total * 2);
3867 let mut idx = vec![0usize; rank];
3868 for o in 0..total {
3869 let mut src_idx = 0usize;
3870 for d in 0..rank {
3871 src_idx += idx[d] * in_strides[d] as usize;
3872 }
3873 let s = broadcast_src_index(src_idx, in_total);
3874 out[o * 2] = inp[s * 2];
3875 out[o * 2 + 1] = inp[s * 2 + 1];
3876 for d in (0..rank).rev() {
3877 idx[d] += 1;
3878 if idx[d] < out_dims[d] as usize {
3879 break;
3880 }
3881 idx[d] = 0;
3882 }
3883 }
3884 } else {
3885 let inp = sl(*src, base, in_total);
3886 let out = sl_mut(*dst, base, total);
3887 if rank == 4
3888 && in_strides[0] == 0
3889 && in_strides[2] == 0
3890 && in_strides[3] == 0
3891 && in_strides[1] != 0
3892 {
3893 // Per-channel broadcast out[n,c,h,w] = in[c*sc] — how
3894 // a conv bias `[C]` reaches `[N,C,H,W]` (and its
3895 // recompute in the backward graph). Fill each (n,c)
3896 // plane with its scalar instead of an N-D index walk
3897 // over every element; parallel over the channel-batch.
3898 let d1 = out_dims[1] as usize;
3899 let sc = in_strides[1] as usize;
3900 let plane = (out_dims[2] as usize) * (out_dims[3] as usize);
3901 let nc_total = (out_dims[0] as usize) * d1;
3902 let out_addr = out.as_mut_ptr() as usize;
3903 let fill = |nc0: usize, nc1: usize| {
3904 let op = out_addr as *mut f32;
3905 for nc in nc0..nc1 {
3906 let v = inp[(nc % d1) * sc];
3907 let base_off = nc * plane;
3908 for k in 0..plane {
3909 *op.add(base_off + k) = v;
3910 }
3911 }
3912 };
3913 if fast_conv_enabled() && crate::pool::should_parallelize(total) {
3914 crate::pool::par_for(
3915 nc_total,
3916 crate::pool::outer_chunk(nc_total),
3917 &|off, cnt| fill(off, off + cnt),
3918 );
3919 } else {
3920 fill(0, nc_total);
3921 }
3922 } else if rank == 2 && in_strides[0] != 0 && in_strides[1] != 0 {
3923 // Fast 2D transpose (the common matmul-backward case:
3924 // xᵀ, wᵀ). out[i,j] = in[i*s0 + j*s1]; tiled over
3925 // columns for write-locality, parallel over rows. Far
3926 // cheaper than the general per-element index walk.
3927 let d0 = out_dims[0] as usize;
3928 let d1 = out_dims[1] as usize;
3929 let s0 = in_strides[0] as usize;
3930 let s1 = in_strides[1] as usize;
3931 let out_addr = out.as_mut_ptr() as usize;
3932 let tile = |i0: usize, i1: usize| {
3933 let op = out_addr as *mut f32;
3934 const T: usize = 32;
3935 let mut j0 = 0;
3936 while j0 < d1 {
3937 let j1 = (j0 + T).min(d1);
3938 for i in i0..i1 {
3939 let inb = i * s0;
3940 let outb = i * d1;
3941 for j in j0..j1 {
3942 *op.add(outb + j) = inp[inb + j * s1];
3943 }
3944 }
3945 j0 = j1;
3946 }
3947 };
3948 if fast_conv_enabled() && crate::pool::should_parallelize(total) {
3949 crate::pool::par_for(
3950 d0,
3951 crate::pool::outer_chunk(d0),
3952 &|off, cnt| tile(off, off + cnt),
3953 );
3954 } else {
3955 tile(0, d0);
3956 }
3957 } else if rank >= 3
3958 && in_strides[rank - 2] == 1
3959 && in_strides[rank - 1] as usize == out_dims[rank - 2] as usize
3960 && {
3961 // Leading axes untouched and contiguous, so
3962 // each batch plane is one dense 2-D block.
3963 let plane = (out_dims[rank - 2] as usize)
3964 * (out_dims[rank - 1] as usize);
3965 let mut expect = plane;
3966 let mut ok = true;
3967 for d in (0..rank - 2).rev() {
3968 if in_strides[d] as usize != expect {
3969 ok = false;
3970 break;
3971 }
3972 expect *= out_dims[d] as usize;
3973 }
3974 ok
3975 }
3976 {
3977 // BATCHED 2-D transpose: the permutation swaps
3978 // only the last two axes.
3979 //
3980 // Neither branch above catches this. The
3981 // innermost stride is not 1 (it is the input's
3982 // inner extent), and the rank is not 2, so it
3983 // used to fall through to the general
3984 // per-element index walk — measured at
3985 // **103 MB/s** on an M4 Pro, against ~20 GB/s of
3986 // memory bandwidth. A `[8, 2048, 1024]` f32
3987 // bank took 653 ms.
3988 //
3989 // It is not a rare shape: `GroupedMatMul` wants
3990 // `[E, in, out]` while GGUF stores expert banks
3991 // as `[E, out, in]`, so every dense MoE layer
3992 // transposes all three of its banks, and
3993 // attention permutes `[B, H, S, D]` the same
3994 // way. Reuse the tiling the rank-2 branch uses,
3995 // once per plane.
3996 let d_i = out_dims[rank - 2] as usize;
3997 let d_o = out_dims[rank - 1] as usize;
3998 let plane = d_i * d_o;
3999 let batches = total / plane.max(1);
4000 let out_addr = out.as_mut_ptr() as usize;
4001 let in_addr = inp.as_ptr() as usize;
4002 let plane_tp = |b0: usize, b1: usize| {
4003 const T: usize = 32;
4004 for b in b0..b1 {
4005 let ip = (in_addr as *const f32).add(b * plane);
4006 let op = (out_addr as *mut f32).add(b * plane);
4007 let mut i0 = 0;
4008 while i0 < d_i {
4009 let i1 = (i0 + T).min(d_i);
4010 let mut o0 = 0;
4011 while o0 < d_o {
4012 let o1 = (o0 + T).min(d_o);
4013 for i in i0..i1 {
4014 for o in o0..o1 {
4015 *op.add(i * d_o + o) = *ip.add(o * d_i + i);
4016 }
4017 }
4018 o0 = o1;
4019 }
4020 i0 = i1;
4021 }
4022 }
4023 };
4024 if crate::pool::should_parallelize(total) && batches >= 2 {
4025 crate::pool::par_for(
4026 batches,
4027 crate::pool::outer_chunk(batches),
4028 &|off, cnt| plane_tp(off, off + cnt),
4029 );
4030 } else {
4031 plane_tp(0, batches);
4032 }
4033 } else if rank >= 3
4034 && *in_strides.last().unwrap_or(&0) == 1
4035 && out_dims[rank - 1] as usize >= 8
4036 {
4037 // Innermost dim contiguous in both layouts: copy one
4038 // row at a time (memcpy) instead of per-element index
4039 // walk. Covers attention BSHD↔BHSD and most Expand/
4040 // permute patterns in DiT.
4041 let row = out_dims[rank - 1] as usize;
4042 let planes = total / row;
4043 let s_last = in_strides[rank - 1] as usize; // 1
4044 let _ = s_last;
4045 let out_addr = out.as_mut_ptr() as usize;
4046 let in_addr = inp.as_ptr() as usize;
4047 let dims = out_dims.to_vec();
4048 let strides = in_strides.to_vec();
4049 let copy_planes = |p0: usize, p1: usize| {
4050 let mut idx = vec![0usize; rank];
4051 let mut rem = p0;
4052 // Decode plane index into leading coords (exclude last dim).
4053 for d in (0..rank - 1).rev() {
4054 let dim = dims[d] as usize;
4055 idx[d] = rem % dim;
4056 rem /= dim;
4057 }
4058 for p in p0..p1 {
4059 let mut src = 0usize;
4060 for d in 0..rank - 1 {
4061 src += idx[d] * strides[d] as usize;
4062 }
4063 std::ptr::copy_nonoverlapping(
4064 (in_addr as *const f32).add(src),
4065 (out_addr as *mut f32).add(p * row),
4066 row,
4067 );
4068 for d in (0..rank - 1).rev() {
4069 idx[d] += 1;
4070 if idx[d] < dims[d] as usize {
4071 break;
4072 }
4073 idx[d] = 0;
4074 }
4075 }
4076 };
4077 if crate::pool::should_parallelize(total) && planes >= 4 {
4078 crate::pool::par_for(
4079 planes,
4080 crate::pool::outer_chunk(planes),
4081 &|off, cnt| copy_planes(off, off + cnt),
4082 );
4083 } else {
4084 copy_planes(0, planes);
4085 }
4086 } else if fast_conv_enabled() && crate::pool::should_parallelize(total)
4087 {
4088 // Parallel: each chunk seeds its starting multi-index
4089 // from `off`, then walks incrementally. Output writes
4090 // are disjoint per `o`.
4091 let out_addr = out.as_mut_ptr() as usize;
4092 crate::pool::par_for(
4093 total,
4094 crate::pool::chunk_floor(total),
4095 &|off, cnt| {
4096 let mut idx = vec![0usize; rank];
4097 let mut rem = off;
4098 for d in (0..rank).rev() {
4099 let dim = out_dims[d] as usize;
4100 idx[d] = rem % dim;
4101 rem /= dim;
4102 }
4103 for o in off..off + cnt {
4104 let mut src_idx = 0usize;
4105 for d in 0..rank {
4106 src_idx += idx[d] * in_strides[d] as usize;
4107 }
4108 let v = inp[broadcast_src_index(src_idx, in_total)];
4109 *((out_addr as *mut f32).add(o)) = v;
4110 for d in (0..rank).rev() {
4111 idx[d] += 1;
4112 if idx[d] < out_dims[d] as usize {
4113 break;
4114 }
4115 idx[d] = 0;
4116 }
4117 }
4118 },
4119 );
4120 } else {
4121 let mut idx = vec![0usize; rank];
4122 for o in 0..total {
4123 let mut src_idx = 0usize;
4124 for d in 0..rank {
4125 src_idx += idx[d] * in_strides[d] as usize;
4126 }
4127 out[o] = inp[broadcast_src_index(src_idx, in_total)];
4128 for d in (0..rank).rev() {
4129 idx[d] += 1;
4130 if idx[d] < out_dims[d] as usize {
4131 break;
4132 }
4133 idx[d] = 0;
4134 }
4135 }
4136 }
4137 }
4138 }
4139 } // total != 0
4140 }
4141
4142 Thunk::CustomOp { .. } => exec_custom_op(thunk, base, arena_buf.len()),
4143 Thunk::Reverse { .. } => exec_reverse(thunk, base),
4144 }
4145 if trace_done {
4146 eprintln!("[thunk {i} done]");
4147 }
4148 if write_trace {
4149 let lo = wt_before
4150 .iter()
4151 .zip(arena_buf.iter())
4152 .position(|(a, b)| a != b);
4153 match lo {
4154 None => eprintln!("[wt {i}] {} wrote nothing", thunk_kind_name(thunk)),
4155 Some(lo) => {
4156 let hi = wt_before.len()
4157 - wt_before
4158 .iter()
4159 .zip(arena_buf.iter())
4160 .rev()
4161 .position(|(a, b)| a != b)
4162 .unwrap();
4163 // Hash the CONTENT, not the offsets: under
4164 // `RLX_ARENA_NO_REUSE` every slot moves, so only a
4165 // content checksum lines up thunk-for-thunk between the
4166 // two layouts and pins down the first op that diverges.
4167 let h = arena_buf[lo..hi]
4168 .iter()
4169 .fold(0xcbf2_9ce4_8422_2325u64, |a, b| {
4170 (a ^ *b as u64).wrapping_mul(0x1000_0000_01b3)
4171 });
4172 eprintln!(
4173 "[wt {i}] {} wrote [{lo},{hi}) len={} hash={h:016x}",
4174 thunk_kind_name(thunk),
4175 hi - lo
4176 );
4177 wt_before.copy_from_slice(arena_buf);
4178 }
4179 }
4180 }
4181 }
4182 if profile {
4183 if let Some((pn, pt)) = prof_prev.take() {
4184 profile_record(pn, pt.elapsed());
4185 }
4186 // Auto-dump under RLX_PROFILE_THUNKS so any binary shows where the run's
4187 // time went (per execute_thunks call — the hot subgraph stands out).
4188 dump_thunk_profile();
4189 }
4190}
4191
4192#[inline(always)]
4193pub(crate) fn exec_nop(t: &Thunk) {
4194 let Thunk::Nop = t else { unreachable!() };
4195 {}
4196}