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