1use crate::pool::{matvec_rows, matvec_rows2, Pool};
15use cortiq_core::quant::{f16_to_f32, GROUP_SIZE, Q1_TILE, Q4_TILE};
16use cortiq_core::{CmfModel, TensorDtype};
17use std::sync::Arc;
18
19pub enum QTensor {
20 F32 {
21 data: Vec<f32>,
22 rows: usize,
23 cols: usize,
24 },
25 Mapped {
26 model: Arc<CmfModel>,
27 idx: usize,
29 dtype: TensorDtype,
30 rows: usize,
31 cols: usize,
32 row_scale: Vec<f32>,
34 col_field: Vec<f32>,
36 vbit_offsets: Vec<usize>,
40 repack: Vec<u8>,
48 },
49}
50
51fn repack_enabled() -> bool {
58 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
59 *ON.get_or_init(|| {
60 std::env::var("CMF_REPACK").map(|v| v == "1").unwrap_or(false)
61 })
62}
63
64fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
68 #[cfg(target_arch = "aarch64")]
69 let arch_ok = sdot_enabled();
70 #[cfg(not(target_arch = "aarch64"))]
71 let arch_ok = false;
72 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
73 return Vec::new();
74 }
75 q8_repack_layout(bytes, rows, cols)
76}
77
78fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
81 let groups = rows / 4;
82 let mut rep = vec![0u8; groups * 4 * cols];
83 for g in 0..groups {
84 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
85 for c in 0..cols / 16 {
86 for lane in 0..4 {
87 let src = (g * 4 + lane) * cols + c * 16;
88 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
89 .copy_from_slice(&bytes[src..src + 16]);
90 }
91 }
92 }
93 rep
94}
95
96fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
99 let ng = cols / GROUP_SIZE;
100 let bits = &bytes[..rows];
101 let mut offsets = Vec::with_capacity(rows + 1);
102 let mut off = rows + rows * ng * 2;
103 for r in 0..rows {
104 offsets.push(off);
105 off += (cols * bits[r] as usize + 7) / 8;
106 }
107 offsets.push(off);
108 offsets
109}
110
111impl QTensor {
112 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
113 debug_assert_eq!(data.len(), rows * cols);
114 Self::F32 { data, rows, cols }
115 }
116
117 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
120 let idx = model
123 .tensor_index(name)
124 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
125 let entry = &model.tensors[idx];
126 if entry.shape.len() != 2 {
127 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
128 }
129 let (rows, cols) = (entry.shape[0], entry.shape[1]);
130 let bytes = model.entry_bytes(entry);
131
132 match entry.dtype {
133 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
134 let n = rows * cols;
135 let scales_off = n;
136 let row_scale: Vec<f32> = (0..rows)
137 .map(|o| {
138 f16_to_f32(u16::from_le_bytes([
139 bytes[scales_off + o * 2],
140 bytes[scales_off + o * 2 + 1],
141 ]))
142 })
143 .collect();
144 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
145 let col_off = n + rows * 2;
146 (0..cols)
147 .map(|i| {
148 f16_to_f32(u16::from_le_bytes([
149 bytes[col_off + i * 2],
150 bytes[col_off + i * 2 + 1],
151 ]))
152 })
153 .collect()
154 } else {
155 Vec::new()
156 };
157 Ok(Self::Mapped {
158 model: model.clone(),
159 idx,
160 dtype: entry.dtype,
161 rows,
162 cols,
163 row_scale,
164 col_field,
165 vbit_offsets: Vec::new(),
166 repack: q8_repack(bytes, rows, cols),
167 })
168 }
169 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
171 model: model.clone(),
172 idx,
173 dtype: entry.dtype,
174 rows,
175 cols,
176 row_scale: Vec::new(),
177 col_field: Vec::new(),
178 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
179 repack: Vec::new(),
180 }),
181 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
185 let (_, off_off, packed_off) =
186 cortiq_core::quant::vbit_ro_sections(rows, cols);
187 let offsets: Vec<usize> = (0..=rows)
188 .map(|r| {
189 packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r)
190 })
191 .collect();
192 Ok(Self::Mapped {
193 model: model.clone(),
194 idx,
195 dtype: entry.dtype,
196 rows,
197 cols,
198 row_scale: Vec::new(),
199 col_field: Vec::new(),
200 vbit_offsets: offsets,
201 repack: Vec::new(),
202 })
203 }
204 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
210 model: model.clone(),
211 idx,
212 dtype: entry.dtype,
213 rows,
214 cols,
215 row_scale: Vec::new(),
216 col_field: Vec::new(),
217 vbit_offsets: Vec::new(),
218 repack: Vec::new(),
219 }),
220 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
221 model: model.clone(),
222 idx,
223 dtype: entry.dtype,
224 rows,
225 cols,
226 row_scale: Vec::new(),
227 col_field: Vec::new(),
228 vbit_offsets: Vec::new(),
229 repack: Vec::new(),
230 }),
231 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
233 model: model.clone(),
234 idx,
235 dtype: entry.dtype,
236 rows,
237 cols,
238 row_scale: Vec::new(),
239 col_field: Vec::new(),
240 vbit_offsets: Vec::new(),
241 repack: Vec::new(),
242 }),
243 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
247 model: model.clone(),
248 idx,
249 dtype: entry.dtype,
250 rows,
251 cols,
252 row_scale: Vec::new(),
253 col_field: Vec::new(),
254 vbit_offsets: Vec::new(),
255 repack: Vec::new(),
256 }),
257 _ => {
259 let mut data = vec![0.0f32; rows * cols];
260 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
261 Ok(Self::from_f32(data, rows, cols))
262 }
263 }
264 }
265
266 pub(crate) fn is_q1(&self) -> bool {
269 matches!(self, Self::Mapped { dtype: TensorDtype::Q1, .. })
270 }
271
272 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
275 match self {
276 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
277 _ => None,
278 }
279 }
280
281 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
288 match self {
289 Self::Mapped {
290 idx,
291 dtype: TensorDtype::Q1 | TensorDtype::Q1T | TensorDtype::Q4Block,
292 rows,
293 cols,
294 ..
295 } => Some((*idx, *rows, *cols)),
296 _ => None,
297 }
298 }
299
300 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
305 match self {
306 Self::Mapped {
307 idx,
308 dtype: TensorDtype::Q8Row,
309 rows,
310 cols,
311 row_scale,
312 col_field,
313 ..
314 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
315 _ => None,
316 }
317 }
318
319 pub fn rows(&self) -> usize {
320 match self {
321 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
322 }
323 }
324
325 pub fn cols(&self) -> usize {
326 match self {
327 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
328 }
329 }
330
331 pub fn as_f32(&self) -> Option<&[f32]> {
334 match self {
335 Self::F32 { data, .. } => Some(data),
336 Self::Mapped { .. } => None,
337 }
338 }
339
340 fn quant_bytes(&self) -> &[u8] {
341 match self {
342 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
343 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
344 }
345 }
346
347 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
349 let cols = self.cols();
350 debug_assert_eq!(dst.len(), cols);
351 match self {
352 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
353 Self::Mapped {
354 dtype,
355 row_scale,
356 col_field,
357 vbit_offsets,
358 ..
359 } => {
360 if *dtype == TensorDtype::Q4Tiled {
361 let bytes = self.quant_bytes();
362 let gpr = cols / GROUP_SIZE;
363 for gi in 0..gpr {
364 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
365 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
366 for (k, &b) in tile[2..].iter().enumerate() {
367 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
368 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
369 }
370 }
371 return;
372 }
373 if *dtype == TensorDtype::Q4Block {
374 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
375 let gpr = cols / GROUP_SIZE;
376 for gi in 0..gpr {
377 let g = r * gpr + gi;
378 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
379 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
380 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
381 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
382 }
383 }
384 return;
385 }
386 if *dtype == TensorDtype::Q1 {
387 let bytes = self.quant_bytes();
388 let gpr = cols / GROUP_SIZE;
389 for gi in 0..gpr {
390 let tile =
391 &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
392 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
393 for (j, &b) in tile[2..].iter().enumerate() {
394 for k in 0..8 {
395 dst[gi * GROUP_SIZE + j * 8 + k] =
396 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
397 }
398 }
399 }
400 return;
401 }
402 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
403 let bytes = self.quant_bytes();
404 let rows = self.rows();
405 let ng = cols / GROUP_SIZE;
406 let bits = &bytes[..rows];
407 let sc_off = rows;
408 let off = vbit_offsets[r];
411 let b = bits[r] as usize;
412 let l = ((1usize << (b - 1)) - 1) as f32;
413 let data = &bytes[off..];
414 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
415 for (i, d) in dst.iter_mut().enumerate() {
416 while nbits < b {
417 acc = (acc << 8) | data[idx] as u64;
418 idx += 1;
419 nbits += 8;
420 }
421 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
422 nbits -= b;
423 let so = (r * ng + i / GROUP_SIZE) * 2;
424 let sv = f16_to_f32(u16::from_le_bytes([
425 bytes[sc_off + so],
426 bytes[sc_off + so + 1],
427 ]));
428 *d = (u - l) * sv;
429 }
430 return;
431 }
432 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
433 let s = row_scale[r];
434 match dtype {
435 TensorDtype::Q8Row => {
436 for (d, &b) in dst.iter_mut().zip(q) {
437 *d = (b as i8) as f32 * s;
438 }
439 }
440 TensorDtype::Q8_2f => {
441 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
442 *d = (b as i8) as f32 * s * col_field[i];
443 }
444 }
445 _ => unreachable!(),
446 }
447 }
448 }
449 }
450
451 pub fn sparse_col_ok(&self) -> bool {
456 match self {
457 Self::F32 { .. } => true,
458 Self::Mapped { dtype, .. } => {
459 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
460 }
461 }
462 }
463
464 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
468 let inter = self.cols();
469 let hidden = self.rows();
470 debug_assert_eq!(out.len(), hidden);
471 match self {
472 Self::F32 { data, .. } => {
473 for (k, o) in out.iter_mut().enumerate() {
474 *o += w * data[k * inter + c];
475 }
476 }
477 Self::Mapped {
478 dtype, row_scale, col_field, ..
479 } => {
480 let q = self.quant_bytes();
481 let colf = if *dtype == TensorDtype::Q8_2f {
482 col_field[c]
483 } else {
484 1.0
485 };
486 let wc = w * colf;
487 for (k, o) in out.iter_mut().enumerate() {
488 let b = q[k * inter + c] as i8 as f32;
489 *o += wc * b * row_scale[k];
490 }
491 }
492 }
493 }
494
495 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
499 let cols = self.cols();
500 match self {
501 Self::F32 { data, .. } => {
502 let row = &data[r * cols..(r + 1) * cols];
503 row.iter().zip(x).map(|(w, v)| w * v).sum()
504 }
505 Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
506 TensorDtype::Q8Row => {
507 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
508 dot_i8_f32(q, x) * row_scale[r]
509 }
510 TensorDtype::Q8_2f => {
511 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
512 dot_i8_col_f32(q, x, col_field) * row_scale[r]
513 }
514 _ => {
515 self.row_f32(r, scratch);
516 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
517 }
518 },
519 }
520 }
521
522 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
525 match self {
526 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
527 Self::Mapped {
528 model,
529 idx,
530 dtype,
531 rows,
532 cols,
533 row_scale,
534 col_field,
535 vbit_offsets,
536 repack,
537 } => {
538 let _ = (model, idx);
539 if *dtype == TensorDtype::Q4Block {
540 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
544 let t0 = std::time::Instant::now();
545 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
546 crate::gpu::ProbeArm::Gpu => {
547 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
548 crate::gpu::probe_record(
549 crate::gpu::OpClass::Matvec,
550 true,
551 t0.elapsed(),
552 );
553 return;
554 }
555 }
556 crate::gpu::ProbeArm::CpuTimed => {
557 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
558 crate::gpu::probe_record(
559 crate::gpu::OpClass::Matvec,
560 false,
561 t0.elapsed(),
562 );
563 return;
564 }
565 crate::gpu::ProbeArm::Cpu => {}
566 }
567 }
568 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
569 return;
570 }
571 if *dtype == TensorDtype::Q4Tiled {
572 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
573 return;
574 }
575 if *dtype == TensorDtype::Q1 {
576 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
581 let t0 = std::time::Instant::now();
582 let arm = if crate::gpu::q1_force() {
583 crate::gpu::ProbeArm::Gpu
584 } else {
585 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
586 };
587 match arm {
588 crate::gpu::ProbeArm::Gpu => {
589 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
590 crate::gpu::probe_record(
591 crate::gpu::OpClass::Matvec,
592 true,
593 t0.elapsed(),
594 );
595 return;
596 }
597 }
598 crate::gpu::ProbeArm::CpuTimed => {
599 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
600 crate::gpu::probe_record(
601 crate::gpu::OpClass::Matvec,
602 false,
603 t0.elapsed(),
604 );
605 return;
606 }
607 crate::gpu::ProbeArm::Cpu => {}
608 }
609 }
610 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
611 return;
612 }
613 if *dtype == TensorDtype::Q1T {
614 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
618 let t0 = std::time::Instant::now();
619 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
620 crate::gpu::ProbeArm::Gpu => {
621 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
622 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
623 crate::gpu::probe_record(
624 crate::gpu::OpClass::Matvec,
625 true,
626 t0.elapsed(),
627 );
628 return;
629 }
630 }
631 crate::gpu::ProbeArm::CpuTimed => {
632 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
633 crate::gpu::probe_record(
634 crate::gpu::OpClass::Matvec,
635 false,
636 t0.elapsed(),
637 );
638 return;
639 }
640 crate::gpu::ProbeArm::Cpu => {}
641 }
642 }
643 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
644 return;
645 }
646 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
647 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
648 return;
649 }
650 let xs = prescale(x, col_field, *dtype);
651 if *rows >= crate::gpu::min_rows()
656 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
657 && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
658 && crate::gpu::enabled_here()
659 {
660 let t0 = std::time::Instant::now();
663 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
664 crate::gpu::ProbeArm::Gpu => {}
665 crate::gpu::ProbeArm::CpuTimed => {
666 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
667 crate::gpu::probe_record(
668 crate::gpu::OpClass::Matvec,
669 false,
670 t0.elapsed(),
671 );
672 return;
673 }
674 crate::gpu::ProbeArm::Cpu => {
675 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
676 return;
677 }
678 }
679 let frac = std::env::var("CMF_GPU_SPLIT")
680 .ok()
681 .and_then(|v| v.parse::<f32>().ok())
682 .unwrap_or(0.5)
683 .clamp(0.0, 1.0);
684 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
685 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
686 let bytes = self.quant_bytes();
687 let ok = std::thread::scope(|sc| {
688 let g = sc.spawn(|| {
689 crate::gpu::q8_matvec_range(
690 model,
691 *idx,
692 cpu_rows,
693 &row_scale[cpu_rows..],
694 &xs,
695 *rows - cpu_rows,
696 *cols,
697 out_gpu,
698 )
699 });
700 if cpu_rows > 0 {
701 let rep_cpu = if repack.is_empty() {
704 &[][..]
705 } else {
706 &repack[..(cpu_rows / 4) * 4 * *cols]
707 };
708 qmatvec(
709 &bytes[..cpu_rows * *cols],
710 rep_cpu,
711 &row_scale[..cpu_rows],
712 &xs,
713 cpu_rows,
714 *cols,
715 out_cpu,
716 pool,
717 );
718 }
719 g.join().unwrap_or(false)
720 });
721 if ok {
722 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
723 return;
724 }
725 qmatvec(
728 &bytes[cpu_rows * *cols..(*rows) * *cols],
729 &[],
730 &row_scale[cpu_rows..],
731 &xs,
732 *rows - cpu_rows,
733 *cols,
734 out_gpu,
735 pool,
736 );
737 return;
738 }
739 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
740 }
741 }
742 }
743
744 pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
746 match self {
747 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
748 Self::Mapped {
749 dtype,
750 rows,
751 cols,
752 row_scale,
753 col_field,
754 vbit_offsets,
755 ..
756 } => {
757 if *dtype == TensorDtype::Q4Block {
758 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
759 return;
760 }
761 if *dtype == TensorDtype::Q4Tiled {
762 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
763 return;
764 }
765 if *dtype == TensorDtype::Q1 {
766 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
767 return;
768 }
769 if *dtype == TensorDtype::Q1T {
770 q1t_matvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
775 q1t_matvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
776 return;
777 }
778 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
779 vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
780 return;
781 }
782 let x1s = prescale(x1, col_field, *dtype);
783 let x2s = prescale(x2, col_field, *dtype);
784 qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
785 }
786 }
787 }
788}
789
790impl QTensor {
791 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
796 let cols = self.cols();
797 let rows = self.rows();
798 debug_assert_eq!(xs_all.len(), b * cols);
799 debug_assert_eq!(out.len(), b * rows);
800 if crate::gptq_capture::capturing() {
804 if let Self::Mapped { model, idx, .. } = self {
805 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
806 }
807 }
808 match self {
809 Self::F32 { data, .. } => {
810 let out_addr = SendMut(out.as_mut_ptr());
811 let run = |start: usize, end: usize| {
812 for o in start..end {
813 let row = &data[o * cols..(o + 1) * cols];
814 for bi in 0..b {
815 let x = &xs_all[bi * cols..(bi + 1) * cols];
816 let mut acc = 0f32;
817 for j in 0..cols {
818 acc += row[j] * x[j];
819 }
820 unsafe { *out_addr.at(bi * rows + o) = acc };
821 }
822 }
823 };
824 dispatch_rows(pool, rows, &run);
825 }
826 Self::Mapped {
827 dtype,
828 row_scale,
829 col_field,
830 vbit_offsets,
831 ..
832 } => {
833 if *dtype == TensorDtype::Q4Block {
834 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
835 return;
836 }
837 if *dtype == TensorDtype::Q4Tiled {
838 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
839 return;
840 }
841 if *dtype == TensorDtype::Q1 {
842 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
843 return;
844 }
845 if *dtype == TensorDtype::Q1T {
846 if b >= 32
849 && b * rows * cols >= 128_000_000
850 && crate::gpu::enabled_here()
851 {
852 if let Self::Mapped { model, idx, .. } = self {
853 let t0 = std::time::Instant::now();
854 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
855 crate::gpu::ProbeArm::Gpu => {
856 if crate::gpu::q1t_matmat(model, *idx, xs_all, b, rows, cols, out) {
857 crate::gpu::probe_record(
858 crate::gpu::OpClass::Matmat,
859 true,
860 t0.elapsed(),
861 );
862 return;
863 }
864 }
865 crate::gpu::ProbeArm::CpuTimed => {
866 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
867 crate::gpu::probe_record(
868 crate::gpu::OpClass::Matmat,
869 false,
870 t0.elapsed(),
871 );
872 return;
873 }
874 crate::gpu::ProbeArm::Cpu => {}
875 }
876 }
877 }
878 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
879 return;
880 }
881 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
882 vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
883 return;
884 }
885 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
886 .map(|bi| {
887 prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
888 })
889 .collect();
890 if b >= 8
896 && b * rows * cols >= 128_000_000
897 && crate::gpu::enabled_here()
898 {
899 if let Self::Mapped { model, idx, .. } = self {
900 let t0 = std::time::Instant::now();
901 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
902 crate::gpu::ProbeArm::Gpu
903 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
904 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
905 {
906 let q = self.quant_bytes();
910 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
911 return;
912 }
913 crate::gpu::ProbeArm::Gpu => {
914 let flat: Vec<f32> =
915 pre.iter().flat_map(|v| v.iter().copied()).collect();
916 if crate::gpu::q8_matmat(
917 model, *idx, row_scale, &flat, b, rows, cols, out)
918 {
919 crate::gpu::probe_record(
920 crate::gpu::OpClass::Matmat,
921 true,
922 t0.elapsed(),
923 );
924 return;
925 }
926 }
927 crate::gpu::ProbeArm::CpuTimed => {
928 let q = self.quant_bytes();
929 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
930 crate::gpu::probe_record(
931 crate::gpu::OpClass::Matmat,
932 false,
933 t0.elapsed(),
934 );
935 return;
936 }
937 crate::gpu::ProbeArm::Cpu => {}
938 }
939 }
940 }
941 let q = self.quant_bytes();
942 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
943 }
944 }
945 }
946}
947
948impl QTensor {
949 pub fn matvec_many<const N: usize>(
956 ts: [&QTensor; N],
957 x: &[f32],
958 mut outs: [&mut [f32]; N],
959 pool: Option<&Pool>,
960 ) {
961 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
962 let uniform_q8 = ts.iter().all(|t| {
963 matches!(
964 t,
965 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
966 )
967 });
968 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
969 let uniform_q4 = ts
970 .iter()
971 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
972 let uniform_vbit = ts
973 .iter()
974 .all(|t| matches!(
975 t,
976 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
977 ));
978 let uniform_q1 = ts
979 .iter()
980 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q1, .. }));
981 let Some(pool) = pool else {
982 for (t, o) in ts.iter().zip(outs.iter_mut()) {
983 t.matvec(x, o, None);
984 }
985 return;
986 };
987 if total_rows < 256
988 || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit || uniform_q1)
989 {
990 for (t, o) in ts.iter().zip(outs.iter_mut()) {
991 t.matvec(x, o, Some(pool));
992 }
993 return;
994 }
995
996 if uniform_q1 {
997 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1000 if a8w8_enabled() {
1001 let act = split_act(x);
1002 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1003 let (act, gsum) = (&act, &gsum);
1004 let closures: [_; N] = std::array::from_fn(|i| {
1005 let (bytes, gpr, out) =
1006 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1007 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1008 });
1009 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1010 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1011 pool.run_many(&parts);
1012 } else {
1013 let closures: [_; N] = std::array::from_fn(|i| {
1014 let (bytes, gpr, out) =
1015 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1016 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1017 });
1018 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1019 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1020 pool.run_many(&parts);
1021 }
1022 return;
1023 }
1024
1025 if uniform_q4 || uniform_vbit {
1026 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1027 if a8w8_enabled() {
1029 let act = split_act(x);
1030 let act = &act;
1031 if uniform_q4 {
1032 let closures: [_; N] = std::array::from_fn(|i| {
1033 let (packed, scales) =
1034 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1035 let (gpr, cols, out) =
1036 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1037 move |s: usize, e: usize| {
1038 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1039 }
1040 });
1041 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1042 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1043 pool.run_many(&parts);
1044 } else {
1045 let closures: [_; N] = std::array::from_fn(|i| {
1046 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1047 let (bytes, rows, cols, out) =
1048 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
1049 move |s: usize, e: usize| {
1050 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1051 }
1052 });
1053 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1054 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1055 pool.run_many(&parts);
1056 }
1057 return;
1058 }
1059 if uniform_q4 {
1060 let closures: [_; N] = std::array::from_fn(|i| {
1061 let (packed, scales) =
1062 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1063 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1064 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1065 });
1066 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1067 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1068 pool.run_many(&parts);
1069 } else {
1070 let closures: [_; N] = std::array::from_fn(|i| {
1071 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1072 let (bytes, rows, cols, out) =
1073 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
1074 move |s: usize, e: usize| {
1075 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1076 }
1077 });
1078 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1079 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1080 pool.run_many(&parts);
1081 }
1082 return;
1083 }
1084
1085 if uniform_f32 {
1086 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1087 let closures: [_; N] = std::array::from_fn(|i| {
1088 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1089 let out = outs_addr[i];
1090 move |start: usize, end: usize| {
1091 for o in start..end {
1092 let row = &data[o * cols..(o + 1) * cols];
1093 let mut sum = 0.0f32;
1094 for j in 0..*cols {
1095 sum += row[j] * x[j];
1096 }
1097 unsafe { *out.at(o) = sum };
1099 }
1100 }
1101 });
1102 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1103 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1104 pool.run_many(&parts);
1105 return;
1106 }
1107
1108 struct Ctx<'a> {
1111 bytes: &'a [u8],
1112 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1113 rep: &'a [u8],
1114 row_scale: &'a [f32],
1115 cols: usize,
1116 xs: std::borrow::Cow<'a, [f32]>,
1117 }
1118 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1119 let Self::Mapped { dtype, cols, row_scale, col_field, repack, .. } = ts[i] else {
1120 unreachable!()
1121 };
1122 Ctx {
1123 bytes: ts[i].quant_bytes(),
1124 rep: repack,
1125 row_scale,
1126 cols: *cols,
1127 xs: prescale(x, col_field, *dtype),
1128 }
1129 });
1130 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1131 #[cfg(target_arch = "aarch64")]
1132 if sdot_enabled() {
1133 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1134 let closures: [_; N] = std::array::from_fn(|i| {
1135 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1136 move |start: usize, end: usize| {
1137 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1138 }
1139 });
1140 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1141 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1142 pool.run_many(&parts);
1143 return;
1144 }
1145 #[cfg(target_arch = "x86_64")]
1146 if avx2_a8w8_enabled() {
1147 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1148 let closures: [_; N] = std::array::from_fn(|i| {
1149 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1150 move |start: usize, end: usize| {
1151 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1152 }
1153 });
1154 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1155 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1156 pool.run_many(&parts);
1157 return;
1158 }
1159 let closures: [_; N] = std::array::from_fn(|i| {
1160 let (c, out) = (&ctxs[i], outs_addr[i]);
1161 move |start: usize, end: usize| {
1162 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1163 }
1164 });
1165 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1166 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1167 pool.run_many(&parts);
1168 }
1169}
1170
1171impl QTensor {
1172 #[allow(clippy::needless_range_loop)]
1177 pub fn matvec2_many<const N: usize>(
1178 ts: [&QTensor; N],
1179 x1: &[f32],
1180 x2: &[f32],
1181 mut o1s: [&mut [f32]; N],
1182 mut o2s: [&mut [f32]; N],
1183 pool: Option<&Pool>,
1184 ) {
1185 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1186 let uniform_q8 = ts.iter().all(|t| {
1187 matches!(
1188 t,
1189 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1190 )
1191 });
1192 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1193 let uniform_q4 = ts
1194 .iter()
1195 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1196 let uniform_vbit = ts
1197 .iter()
1198 .all(|t| matches!(
1199 t,
1200 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1201 ));
1202 let fusable = pool.is_some()
1203 && total_rows >= 256
1204 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1205 if !fusable {
1206 for i in 0..N {
1207 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1208 }
1209 return;
1210 }
1211 let pool = pool.unwrap();
1212
1213 if uniform_q4 || uniform_vbit {
1214 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1215 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1216 if a8w8_enabled() {
1218 let a1 = split_act(x1);
1219 let a2 = split_act(x2);
1220 let (a1, a2) = (&a1, &a2);
1221 if uniform_q4 {
1222 let closures: [_; N] = std::array::from_fn(|i| {
1223 let (packed, scales) =
1224 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1225 let (gpr, cols, o1, o2) =
1226 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1227 move |s: usize, e: usize| {
1228 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1229 }
1230 });
1231 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1232 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1233 pool.run_many(&parts);
1234 } else {
1235 let closures: [_; N] = std::array::from_fn(|i| {
1236 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1237 let (bytes, rows, cols, o1, o2) =
1238 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1239 move |s: usize, e: usize| {
1240 vbit_range2_a8w8(
1241 bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
1242 )
1243 }
1244 });
1245 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1246 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1247 pool.run_many(&parts);
1248 }
1249 return;
1250 }
1251 if uniform_q4 {
1252 let closures: [_; N] = std::array::from_fn(|i| {
1253 let (packed, scales) =
1254 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1255 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1256 move |s: usize, e: usize| {
1257 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1258 }
1259 });
1260 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1261 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1262 pool.run_many(&parts);
1263 } else {
1264 let closures: [_; N] = std::array::from_fn(|i| {
1265 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1266 let (bytes, rows, cols, o1, o2) =
1267 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1268 move |s: usize, e: usize| {
1269 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1270 }
1271 });
1272 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1273 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1274 pool.run_many(&parts);
1275 }
1276 return;
1277 }
1278
1279 if uniform_f32 {
1280 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1281 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1282 let closures: [_; N] = std::array::from_fn(|i| {
1283 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1284 let (o1, o2) = (p1[i], p2[i]);
1285 move |start: usize, end: usize| {
1286 for o in start..end {
1287 let row = &data[o * cols..(o + 1) * cols];
1288 let (mut s1, mut s2) = (0.0f32, 0.0f32);
1289 for j in 0..*cols {
1290 s1 += row[j] * x1[j];
1291 s2 += row[j] * x2[j];
1292 }
1293 unsafe {
1295 *o1.at(o) = s1;
1296 *o2.at(o) = s2;
1297 }
1298 }
1299 }
1300 });
1301 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1302 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1303 pool.run_many(&parts);
1304 return;
1305 }
1306
1307 struct Ctx<'a> {
1308 bytes: &'a [u8],
1309 row_scale: &'a [f32],
1310 cols: usize,
1311 xs1: std::borrow::Cow<'a, [f32]>,
1312 xs2: std::borrow::Cow<'a, [f32]>,
1313 }
1314 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1315 let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
1316 unreachable!()
1317 };
1318 Ctx {
1319 bytes: ts[i].quant_bytes(),
1320 row_scale,
1321 cols: *cols,
1322 xs1: prescale(x1, col_field, *dtype),
1323 xs2: prescale(x2, col_field, *dtype),
1324 }
1325 });
1326 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1327 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1328 #[cfg(target_arch = "aarch64")]
1329 if sdot_enabled() {
1330 let acts: [(SplitAct, SplitAct); N] =
1331 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1332 let closures: [_; N] = std::array::from_fn(|i| {
1333 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1334 move |start: usize, end: usize| {
1335 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1336 }
1337 });
1338 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1339 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1340 pool.run_many(&parts);
1341 return;
1342 }
1343 #[cfg(target_arch = "x86_64")]
1344 if avx2_a8w8_enabled() {
1345 let acts: [(SplitAct, SplitAct); N] =
1346 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1347 let closures: [_; N] = std::array::from_fn(|i| {
1348 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1349 move |start: usize, end: usize| {
1350 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1351 }
1352 });
1353 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1354 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1355 pool.run_many(&parts);
1356 return;
1357 }
1358 let closures: [_; N] = std::array::from_fn(|i| {
1359 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1360 move |start: usize, end: usize| {
1361 q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
1362 }
1363 });
1364 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1365 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1366 pool.run_many(&parts);
1367 }
1368}
1369
1370#[cfg(target_os = "macos")]
1375mod accel_blas {
1376 #[link(name = "Accelerate", kind = "framework")]
1377 unsafe extern "C" {
1378 pub fn cblas_sgemm(
1379 order: i32,
1380 trans_a: i32,
1381 trans_b: i32,
1382 m: i32,
1383 n: i32,
1384 k: i32,
1385 alpha: f32,
1386 a: *const f32,
1387 lda: i32,
1388 b: *const f32,
1389 ldb: i32,
1390 beta: f32,
1391 c: *mut f32,
1392 ldc: i32,
1393 );
1394 }
1395}
1396
1397#[cfg(target_os = "macos")]
1398pub(crate) fn accel_gemm_enabled() -> bool {
1399 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1400 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
1401}
1402
1403#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1406pub(crate) fn accel_gemm_enabled() -> bool {
1407 true
1408}
1409
1410#[cfg(target_arch = "aarch64")]
1417#[allow(clippy::too_many_arguments)]
1418pub(crate) fn neon_gemm_rm(
1419 m: usize,
1420 n: usize,
1421 k: usize,
1422 alpha: f32,
1423 a: &[f32],
1424 lda: usize,
1425 b_mat: &[f32],
1426 ldb: usize,
1427 b_rows_are_n: bool,
1428 c: &mut [f32],
1429 ldc: usize,
1430) {
1431 debug_assert!(a.len() >= (m - 1) * lda + k);
1432 debug_assert!(c.len() >= (m - 1) * ldc + n);
1433 unsafe {
1435 use core::arch::aarch64::*;
1436 let mut i = 0usize;
1437 while i < m {
1438 let mi = (m - i).min(4);
1439 let mut j = 0usize;
1440 while j < n {
1441 let nj = (n - j).min(8);
1442 if mi == 4 && nj == 8 {
1443 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1444 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1445 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1446 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1447 for p in 0..k {
1448 let (b0, b1) = if b_rows_are_n {
1449 let base = b_mat.as_ptr().add(j * ldb + p);
1452 let g = |o: usize| *base.add(o * ldb);
1453 (
1454 [g(0), g(1), g(2), g(3)],
1455 [g(4), g(5), g(6), g(7)],
1456 )
1457 } else {
1458 let base = b_mat.as_ptr().add(p * ldb + j);
1459 (
1460 [*base, *base.add(1), *base.add(2), *base.add(3)],
1461 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
1462 )
1463 };
1464 let bv0 = vld1q_f32(b0.as_ptr());
1465 let bv1 = vld1q_f32(b1.as_ptr());
1466 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
1467 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
1468 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
1469 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
1470 c0a = vfmaq_f32(c0a, a0, bv0);
1471 c0b = vfmaq_f32(c0b, a0, bv1);
1472 c1a = vfmaq_f32(c1a, a1, bv0);
1473 c1b = vfmaq_f32(c1b, a1, bv1);
1474 c2a = vfmaq_f32(c2a, a2, bv0);
1475 c2b = vfmaq_f32(c2b, a2, bv1);
1476 c3a = vfmaq_f32(c3a, a3, bv0);
1477 c3b = vfmaq_f32(c3b, a3, bv1);
1478 }
1479 let al = vdupq_n_f32(alpha);
1480 for (r, (ca, cb)) in
1481 [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)].iter().enumerate()
1482 {
1483 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
1484 vst1q_f32(dst, vmulq_f32(*ca, al));
1485 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
1486 }
1487 } else {
1488 for r in 0..mi {
1489 for q in 0..nj {
1490 let mut acc = 0f32;
1491 for p in 0..k {
1492 let bv = if b_rows_are_n {
1493 b_mat[(j + q) * ldb + p]
1494 } else {
1495 b_mat[p * ldb + j + q]
1496 };
1497 acc += a[(i + r) * lda + p] * bv;
1498 }
1499 c[(i + r) * ldc + j + q] = acc * alpha;
1500 }
1501 }
1502 }
1503 j += nj;
1504 }
1505 i += mi;
1506 }
1507 }
1508}
1509
1510#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1512#[allow(clippy::too_many_arguments)]
1513pub(crate) fn sgemm_rm(
1514 m: usize,
1515 n: usize,
1516 k: usize,
1517 alpha: f32,
1518 a: &[f32],
1519 lda: usize,
1520 b_mat: &[f32],
1521 ldb: usize,
1522 b_rows_are_n: bool,
1523 c: &mut [f32],
1524 ldc: usize,
1525) {
1526 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1527}
1528
1529#[cfg(target_os = "macos")]
1532#[allow(clippy::too_many_arguments)]
1533pub(crate) fn sgemm_rm(
1534 m: usize,
1535 n: usize,
1536 k: usize,
1537 alpha: f32,
1538 a: &[f32],
1539 lda: usize,
1540 b_mat: &[f32],
1541 ldb: usize,
1542 b_rows_are_n: bool,
1543 c: &mut [f32],
1544 ldc: usize,
1545) {
1546 debug_assert!(a.len() >= (m - 1) * lda + k);
1547 debug_assert!(c.len() >= (m - 1) * ldc + n);
1548 #[cfg(target_arch = "aarch64")]
1553 if std::env::var("CMF_FORCE_NEON_GEMM").map(|v| v == "1").unwrap_or(false) {
1554 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1555 }
1556 unsafe {
1557 accel_blas::cblas_sgemm(
1558 101, 111, if b_rows_are_n { 112 } else { 111 },
1561 m as i32,
1562 n as i32,
1563 k as i32,
1564 alpha,
1565 a.as_ptr(),
1566 lda as i32,
1567 b_mat.as_ptr(),
1568 ldb as i32,
1569 0.0,
1570 c.as_mut_ptr(),
1571 ldc as i32,
1572 );
1573 }
1574}
1575
1576#[cfg(target_os = "macos")]
1583fn qmatmat_accel(
1584 q: &[u8],
1585 row_scale: &[f32],
1586 pre: &[std::borrow::Cow<'_, [f32]>],
1587 rows: usize,
1588 cols: usize,
1589 out: &mut [f32],
1590 pool: Option<&Pool>,
1591) {
1592 const TR: usize = 2048;
1597 let b = pre.len();
1598 thread_local! {
1599 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1600 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1601 }
1602 XPANEL.with(|xp| {
1603 WTILE.with(|wt| {
1604 let mut xpanel = xp.borrow_mut();
1605 xpanel.clear();
1606 for x in pre {
1607 xpanel.extend_from_slice(x);
1608 }
1609 let mut wtile = wt.borrow_mut();
1610 wtile.resize(TR * cols, 0.0);
1611 let mut r0 = 0usize;
1612 while r0 < rows {
1613 let tr = TR.min(rows - r0);
1614 let wt_addr = SendMut(wtile.as_mut_ptr());
1616 let run = |start: usize, end: usize| {
1617 for r in start..end {
1618 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1619 let s = row_scale[r0 + r];
1620 let dst = unsafe {
1622 std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1623 };
1624 for (d, &v) in dst.iter_mut().zip(row) {
1625 *d = (v as i8) as f32 * s;
1626 }
1627 }
1628 };
1629 dispatch_rows(pool, tr, &run);
1630 unsafe {
1632 accel_blas::cblas_sgemm(
1633 101, 111, 112, b as i32,
1637 tr as i32,
1638 cols as i32,
1639 1.0,
1640 xpanel.as_ptr(),
1641 cols as i32,
1642 wtile.as_ptr(),
1643 cols as i32,
1644 0.0,
1645 out.as_mut_ptr().add(r0),
1646 rows as i32,
1647 );
1648 }
1649 r0 += tr;
1650 }
1651 })
1652 });
1653}
1654
1655fn qmatmat(
1656 q: &[u8],
1657 row_scale: &[f32],
1658 pre: &[std::borrow::Cow<'_, [f32]>],
1659 rows: usize,
1660 cols: usize,
1661 out: &mut [f32],
1662 pool: Option<&Pool>,
1663) {
1664 let b = pre.len();
1665 debug_assert_eq!(out.len(), b * rows);
1666 #[cfg(target_os = "macos")]
1671 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1672 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1673 return;
1674 }
1675 #[cfg(target_arch = "aarch64")]
1676 if sdot_enabled() {
1677 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1678 let out_addr = SendMut(out.as_mut_ptr());
1679 let blocked_ok =
1682 std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1683 let use_i8mm = i8mm_enabled();
1684 if blocked_ok {
1685 let run = |start: usize, end: usize| {
1686 let mut o = start;
1687 while o < end {
1688 if o + 2 <= end {
1689 let r0 = &q[o * cols..(o + 1) * cols];
1690 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1691 let mut bi = 0usize;
1692 while bi + 4 <= acts.len() {
1693 let xs = [
1694 acts[bi].xq.as_slice(),
1695 acts[bi + 1].xq.as_slice(),
1696 acts[bi + 2].xq.as_slice(),
1697 acts[bi + 3].xq.as_slice(),
1698 ];
1699 let d = if use_i8mm {
1700 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
1701 } else {
1702 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
1703 };
1704 for (r, row) in [r0, r1].into_iter().enumerate() {
1705 for k in 0..4 {
1706 let act = &acts[bi + k];
1707 let mut v = d[r][k] as f32 * act.sx;
1708 for &(j, xv) in &act.outliers {
1709 v += (row[j] as i8) as f32 * xv;
1710 }
1711 unsafe {
1712 *out_addr.at((bi + k) * rows + o + r) =
1713 v * row_scale[o + r]
1714 };
1715 }
1716 }
1717 bi += 4;
1718 }
1719 while bi < acts.len() {
1720 for (r, row) in [r0, r1].into_iter().enumerate() {
1721 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
1722 unsafe { *out_addr.at(bi * rows + o + r) = v };
1723 }
1724 bi += 1;
1725 }
1726 o += 2;
1727 } else {
1728 let row = &q[o * cols..(o + 1) * cols];
1729 for (bi, act) in acts.iter().enumerate() {
1730 let v = row_dot_sdot(row, act) * row_scale[o];
1731 unsafe { *out_addr.at(bi * rows + o) = v };
1732 }
1733 o += 1;
1734 }
1735 }
1736 };
1737 dispatch_rows(pool, rows, &run);
1738 return;
1739 }
1740 let run = |start: usize, end: usize| {
1741 for o in start..end {
1742 let row = &q[o * cols..(o + 1) * cols];
1743 for (bi, act) in acts.iter().enumerate() {
1744 let v = row_dot_sdot(row, act) * row_scale[o];
1745 unsafe { *out_addr.at(bi * rows + o) = v };
1746 }
1747 }
1748 };
1749 dispatch_rows(pool, rows, &run);
1750 return;
1751 }
1752 #[cfg(target_arch = "x86_64")]
1757 if avx2_a8w8_enabled() {
1758 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1759 let out_addr = SendMut(out.as_mut_ptr());
1760 let blocked_ok = std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1763 if !avx512vnni_enabled() && blocked_ok {
1764 let run = |start: usize, end: usize| {
1765 let mut o = start;
1766 while o < end {
1767 if o + 2 <= end {
1768 let r0 = &q[o * cols..(o + 1) * cols];
1769 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1770 let mut bi = 0usize;
1771 while bi + 4 <= acts.len() {
1772 let xs = [
1773 acts[bi].xq.as_slice(),
1774 acts[bi + 1].xq.as_slice(),
1775 acts[bi + 2].xq.as_slice(),
1776 acts[bi + 3].xq.as_slice(),
1777 ];
1778 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
1779 for (r, row) in [r0, r1].into_iter().enumerate() {
1780 for k in 0..4 {
1781 let act = &acts[bi + k];
1782 let mut v = d[r][k] as f32 * act.sx;
1783 for &(j, xv) in &act.outliers {
1784 v += (row[j] as i8) as f32 * xv;
1785 }
1786 unsafe {
1787 *out_addr.at((bi + k) * rows + o + r) =
1788 v * row_scale[o + r]
1789 };
1790 }
1791 }
1792 bi += 4;
1793 }
1794 while bi < acts.len() {
1795 for (r, row) in [r0, r1].into_iter().enumerate() {
1796 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
1797 unsafe { *out_addr.at(bi * rows + o + r) = v };
1798 }
1799 bi += 1;
1800 }
1801 o += 2;
1802 } else {
1803 let row = &q[o * cols..(o + 1) * cols];
1804 for (bi, act) in acts.iter().enumerate() {
1805 let v = row_dot_avx2(row, act) * row_scale[o];
1806 unsafe { *out_addr.at(bi * rows + o) = v };
1807 }
1808 o += 1;
1809 }
1810 }
1811 };
1812 dispatch_rows(pool, rows, &run);
1813 return;
1814 }
1815 let run = |start: usize, end: usize| {
1816 for o in start..end {
1817 let row = &q[o * cols..(o + 1) * cols];
1818 for (bi, act) in acts.iter().enumerate() {
1819 let v = row_dot_avx2(row, act) * row_scale[o];
1820 unsafe { *out_addr.at(bi * rows + o) = v };
1821 }
1822 }
1823 };
1824 dispatch_rows(pool, rows, &run);
1825 return;
1826 }
1827 let out_addr = SendMut(out.as_mut_ptr());
1828 let run = |start: usize, end: usize| {
1829 for o in start..end {
1830 let row = &q[o * cols..(o + 1) * cols];
1831 for (bi, x) in pre.iter().enumerate() {
1832 let mut acc = 0f32;
1833 for j in 0..cols {
1834 acc += (row[j] as i8) as f32 * x[j];
1835 }
1836 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1837 }
1838 }
1839 };
1840 dispatch_rows(pool, rows, &run);
1841}
1842
1843fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1846 match pool {
1847 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1848 _ => run(0, rows),
1849 }
1850}
1851
1852fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1854 let groups = rows * cols / GROUP_SIZE;
1855 bytes.split_at(groups * 16)
1856}
1857
1858#[inline]
1863fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1864 #[cfg(target_arch = "aarch64")]
1865 unsafe {
1866 return vbit_fill4_neon(data, buf);
1867 }
1868 #[cfg(target_arch = "x86_64")]
1869 if avx2_enabled() {
1870 return unsafe { vbit_fill4_avx2(data, buf) };
1871 }
1872 #[allow(unreachable_code)]
1873 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1874 let u = unpack8::<4>(&data[blk * 4..]);
1875 for k in 0..8 {
1876 chunk[k] = (u[k] - 7) as i8 as u8;
1877 }
1878 }
1879}
1880
1881#[cfg(target_arch = "aarch64")]
1882#[target_feature(enable = "neon")]
1883unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1884 unsafe {
1887 use core::arch::aarch64::*;
1888 let n = buf.len();
1889 let mask = vdupq_n_u8(0x0F);
1890 let seven = vdupq_n_s8(7);
1891 let mut g = 0usize;
1892 while g * 32 + 32 <= n {
1893 let b = vld1q_u8(data.as_ptr().add(g * 16));
1894 let hi = vshrq_n_u8::<4>(b);
1895 let lo = vandq_u8(b, mask);
1896 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1897 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1898 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1899 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1900 g += 1;
1901 }
1902 }
1903}
1904
1905#[cfg(target_arch = "x86_64")]
1906#[target_feature(enable = "avx2")]
1907unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1908 unsafe {
1910 use core::arch::x86_64::*;
1911 let n = buf.len();
1912 let mask = _mm_set1_epi8(0x0F);
1913 let seven = _mm256_set1_epi8(7);
1914 let mut g = 0usize;
1915 while g * 32 + 32 <= n {
1916 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1917 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1918 let lo = _mm_and_si128(b, mask);
1919 let z = _mm256_sub_epi8(
1920 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1921 seven,
1922 );
1923 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1924 g += 1;
1925 }
1926 }
1927}
1928
1929#[inline(always)]
1934fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1935 let mut acc = 0u64;
1936 for i in 0..B {
1937 acc = (acc << 8) | data[i] as u64;
1938 }
1939 let mask = (1u64 << B) - 1;
1940 let mut out = [0i32; 8];
1941 for (k, o) in out.iter_mut().enumerate() {
1942 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1943 }
1944 out
1945}
1946
1947#[allow(clippy::too_many_arguments)]
1953fn vbitmatvec(
1954 bytes: &[u8],
1955 offsets: &[usize],
1956 x: &[f32],
1957 rows: usize,
1958 cols: usize,
1959 out: &mut [f32],
1960 pool: Option<&Pool>,
1961) {
1962 debug_assert_eq!(out.len(), rows);
1963 debug_assert_eq!(offsets.len(), rows + 1);
1964
1965 if a8w8_enabled() {
1969 let act = split_act(x);
1970 let out_addr = SendMut(out.as_mut_ptr());
1971 let run = move |start: usize, end: usize| {
1972 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1973 };
1974 dispatch_rows(pool, rows, &run);
1975 return;
1976 }
1977
1978 let out_addr = SendMut(out.as_mut_ptr());
1979 let run = move |start: usize, end: usize| {
1980 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1981 };
1982 dispatch_rows(pool, rows, &run);
1983}
1984
1985#[allow(clippy::too_many_arguments)]
1989fn vbit_range_a8w8(
1990 bytes: &[u8],
1991 offsets: &[usize],
1992 x: &[f32],
1993 act: &SplitAct,
1994 rows: usize,
1995 cols: usize,
1996 out: SendMut,
1997 start: usize,
1998 end: usize,
1999) {
2000 let ng = cols / GROUP_SIZE;
2001 let bits = &bytes[..rows];
2002 let sc_off = rows;
2003 let row_dot = |r: usize| -> f32 {
2004 let b = bits[r] as usize;
2005 let l = ((1i32 << (b - 1)) - 1) as i32;
2006 let mask = (1u64 << b) - 1;
2007 let data = &bytes[offsets[r]..offsets[r + 1]];
2008 if b == 8 {
2009 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2011 let mut dot = 0f32;
2012 for g in 0..ng {
2013 let so = (r * ng + g) * 2;
2014 let sgf = f16_to_f32(u16::from_le_bytes([
2015 bytes[sc_off + so],
2016 bytes[sc_off + so + 1],
2017 ]));
2018 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2019 let mut gd = 0f32;
2020 for &xv in xg.iter() {
2021 if nbits < 8 {
2022 acc = (acc << 8) | data[idx] as u64;
2023 idx += 1;
2024 nbits += 8;
2025 }
2026 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2027 nbits -= 8;
2028 gd += (u - l) as f32 * xv;
2029 }
2030 dot += gd * sgf;
2031 }
2032 return dot;
2033 }
2034 thread_local! {
2038 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
2039 const { std::cell::RefCell::new(Vec::new()) };
2040 }
2041 #[inline(always)]
2042 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2043 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2044 let u = unpack8::<B>(&data[blk * B..]);
2045 for k in 0..8 {
2046 chunk[k] = (u[k] - l) as i8 as u8;
2047 }
2048 }
2049 }
2050 let _ = mask;
2051 VBIT_SCRATCH.with(|scratch| {
2052 let mut buf = scratch.borrow_mut();
2053 buf.resize(cols, 0);
2054 match b {
2055 3 => fill::<3>(data, l, &mut buf),
2056 4 => vbit_fill4(data, &mut buf),
2057 5 => fill::<5>(data, l, &mut buf),
2058 6 => fill::<6>(data, l, &mut buf),
2059 _ => unreachable!(),
2060 }
2061 let mut dot = 0f32;
2062 for g in 0..ng {
2063 let so = (r * ng + g) * 2;
2064 let s = f16_to_f32(u16::from_le_bytes([
2065 bytes[sc_off + so],
2066 bytes[sc_off + so + 1],
2067 ]));
2068 let d = dot_i8_i8(
2069 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2070 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2071 ) as f32
2072 * act.sx;
2073 dot += d * s;
2074 }
2075 for &(j, xv) in &act.outliers {
2076 let so = (r * ng + j / GROUP_SIZE) * 2;
2077 let s = f16_to_f32(u16::from_le_bytes([
2078 bytes[sc_off + so],
2079 bytes[sc_off + so + 1],
2080 ]));
2081 dot += (buf[j] as i8) as f32 * s * xv;
2083 }
2084 dot
2085 })
2086 };
2087 for r in start..end {
2088 unsafe { *out.at(r) = row_dot(r) };
2090 }
2091}
2092
2093#[allow(clippy::too_many_arguments)]
2095fn vbit_range_f32(
2096 bytes: &[u8],
2097 offsets: &[usize],
2098 x: &[f32],
2099 rows: usize,
2100 cols: usize,
2101 out: SendMut,
2102 start: usize,
2103 end: usize,
2104) {
2105 let ng = cols / GROUP_SIZE;
2106 let bits = &bytes[..rows];
2107 let sc_off = rows;
2108 #[inline(always)]
2112 fn dot_row<const B: usize>(
2113 data: &[u8],
2114 bytes: &[u8],
2115 sc_off: usize,
2116 r: usize,
2117 ng: usize,
2118 x: &[f32],
2119 ) -> f32 {
2120 let l = ((1i32 << (B - 1)) - 1) as f32;
2121 let gbytes = GROUP_SIZE * B / 8;
2122 let mut dot = 0f32;
2123 for g in 0..ng {
2124 let so = (r * ng + g) * 2;
2125 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2126 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2127 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2128 let mut gd = 0f32;
2129 for blk in 0..GROUP_SIZE / 8 {
2130 let u = unpack8::<B>(&gd0[blk * B..]);
2131 let xb = &xg[blk * 8..blk * 8 + 8];
2132 for k in 0..8 {
2133 gd += (u[k] as f32 - l) * xb[k];
2134 }
2135 }
2136 dot += gd * s;
2137 }
2138 dot
2139 }
2140 for r in start..end {
2141 let data = &bytes[offsets[r]..offsets[r + 1]];
2142 let v = match bits[r] {
2143 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2144 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2145 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2146 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2147 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2148 b => unreachable!("vbit bit-width {b} (validated at load)"),
2149 };
2150 unsafe { *out.at(r) = v };
2152 }
2153}
2154
2155#[allow(clippy::too_many_arguments)]
2160fn vbitmatvec2(
2161 bytes: &[u8],
2162 offsets: &[usize],
2163 x1: &[f32],
2164 x2: &[f32],
2165 rows: usize,
2166 cols: usize,
2167 o1: &mut [f32],
2168 o2: &mut [f32],
2169 pool: Option<&Pool>,
2170) {
2171 debug_assert_eq!(o1.len(), rows);
2172 debug_assert_eq!(o2.len(), rows);
2173
2174 if a8w8_enabled() {
2175 let a1 = split_act(x1);
2176 let a2 = split_act(x2);
2177 let p1 = SendMut(o1.as_mut_ptr());
2178 let p2 = SendMut(o2.as_mut_ptr());
2179 let run = move |start: usize, end: usize| {
2180 vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
2181 };
2182 dispatch_rows(pool, rows, &run);
2183 return;
2184 }
2185
2186 let p1 = SendMut(o1.as_mut_ptr());
2187 let p2 = SendMut(o2.as_mut_ptr());
2188 let run = move |start: usize, end: usize| {
2189 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
2190 };
2191 dispatch_rows(pool, rows, &run);
2192}
2193
2194#[allow(clippy::too_many_arguments)]
2198fn vbit_range2_a8w8(
2199 bytes: &[u8],
2200 offsets: &[usize],
2201 x1: &[f32],
2202 x2: &[f32],
2203 a1: &SplitAct,
2204 a2: &SplitAct,
2205 rows: usize,
2206 cols: usize,
2207 p1: SendMut,
2208 p2: SendMut,
2209 start: usize,
2210 end: usize,
2211) {
2212 let ng = cols / GROUP_SIZE;
2213 let bits = &bytes[..rows];
2214 let sc_off = rows;
2215 let row_dots = |r: usize| -> (f32, f32) {
2216 let b = bits[r] as usize;
2217 let l = (1i32 << (b - 1)) - 1;
2218 let data = &bytes[offsets[r]..offsets[r + 1]];
2219 if b == 8 {
2220 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2223 let (mut d1, mut d2) = (0f32, 0f32);
2224 for g in 0..ng {
2225 let so = (r * ng + g) * 2;
2226 let sgf = f16_to_f32(u16::from_le_bytes([
2227 bytes[sc_off + so],
2228 bytes[sc_off + so + 1],
2229 ]));
2230 let (mut g1, mut g2) = (0f32, 0f32);
2231 for k in 0..GROUP_SIZE {
2232 if nbits < 8 {
2233 acc = (acc << 8) | data[idx] as u64;
2234 idx += 1;
2235 nbits += 8;
2236 }
2237 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2238 nbits -= 8;
2239 let w = (u - l) as f32;
2240 g1 += w * x1[g * GROUP_SIZE + k];
2241 g2 += w * x2[g * GROUP_SIZE + k];
2242 }
2243 d1 += g1 * sgf;
2244 d2 += g2 * sgf;
2245 }
2246 return (d1, d2);
2247 }
2248 thread_local! {
2249 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
2250 const { std::cell::RefCell::new(Vec::new()) };
2251 }
2252 #[inline(always)]
2253 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2254 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2255 let u = unpack8::<B>(&data[blk * B..]);
2256 for k in 0..8 {
2257 chunk[k] = (u[k] - l) as i8 as u8;
2258 }
2259 }
2260 }
2261 VBIT_SCRATCH2.with(|scratch| {
2262 let mut buf = scratch.borrow_mut();
2263 buf.resize(cols, 0);
2264 match b {
2265 3 => fill::<3>(data, l, &mut buf),
2266 4 => vbit_fill4(data, &mut buf),
2267 5 => fill::<5>(data, l, &mut buf),
2268 6 => fill::<6>(data, l, &mut buf),
2269 _ => unreachable!(),
2270 }
2271 let (mut d1, mut d2) = (0f32, 0f32);
2272 for g in 0..ng {
2273 let so = (r * ng + g) * 2;
2274 let s = f16_to_f32(u16::from_le_bytes([
2275 bytes[sc_off + so],
2276 bytes[sc_off + so + 1],
2277 ]));
2278 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2279 let v1 =
2280 dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
2281 let v2 =
2282 dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
2283 d1 += v1 * s;
2284 d2 += v2 * s;
2285 }
2286 for &(j, xv) in &a1.outliers {
2287 let so = (r * ng + j / GROUP_SIZE) * 2;
2288 let s = f16_to_f32(u16::from_le_bytes([
2289 bytes[sc_off + so],
2290 bytes[sc_off + so + 1],
2291 ]));
2292 d1 += (buf[j] as i8) as f32 * s * xv;
2293 }
2294 for &(j, xv) in &a2.outliers {
2295 let so = (r * ng + j / GROUP_SIZE) * 2;
2296 let s = f16_to_f32(u16::from_le_bytes([
2297 bytes[sc_off + so],
2298 bytes[sc_off + so + 1],
2299 ]));
2300 d2 += (buf[j] as i8) as f32 * s * xv;
2301 }
2302 (d1, d2)
2303 })
2304 };
2305 for r in start..end {
2306 let (v1, v2) = row_dots(r);
2307 unsafe {
2309 *p1.at(r) = v1;
2310 *p2.at(r) = v2;
2311 }
2312 }
2313}
2314
2315#[allow(clippy::too_many_arguments)]
2319fn vbit_range2_f32(
2320 bytes: &[u8],
2321 offsets: &[usize],
2322 x1: &[f32],
2323 x2: &[f32],
2324 rows: usize,
2325 cols: usize,
2326 p1: SendMut,
2327 p2: SendMut,
2328 start: usize,
2329 end: usize,
2330) {
2331 let ng = cols / GROUP_SIZE;
2332 let bits = &bytes[..rows];
2333 let sc_off = rows;
2334 #[inline(always)]
2335 #[allow(clippy::too_many_arguments)]
2336 fn dot_row2<const B: usize>(
2337 data: &[u8],
2338 bytes: &[u8],
2339 sc_off: usize,
2340 r: usize,
2341 ng: usize,
2342 x1: &[f32],
2343 x2: &[f32],
2344 ) -> (f32, f32) {
2345 let l = ((1i32 << (B - 1)) - 1) as f32;
2346 let gbytes = GROUP_SIZE * B / 8;
2347 let (mut d1, mut d2) = (0f32, 0f32);
2348 for g in 0..ng {
2349 let so = (r * ng + g) * 2;
2350 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2351 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2352 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2353 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2354 let (mut g1, mut g2) = (0f32, 0f32);
2355 for blk in 0..GROUP_SIZE / 8 {
2356 let u = unpack8::<B>(&gd0[blk * B..]);
2357 for k in 0..8 {
2358 let w = u[k] as f32 - l;
2359 g1 += w * x1g[blk * 8 + k];
2360 g2 += w * x2g[blk * 8 + k];
2361 }
2362 }
2363 d1 += g1 * s;
2364 d2 += g2 * s;
2365 }
2366 (d1, d2)
2367 }
2368 for r in start..end {
2369 let data = &bytes[offsets[r]..offsets[r + 1]];
2370 let (v1, v2) = match bits[r] {
2371 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2372 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2373 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2374 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2375 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2376 b => unreachable!("vbit bit-width {b} (validated at load)"),
2377 };
2378 unsafe {
2380 *p1.at(r) = v1;
2381 *p2.at(r) = v2;
2382 }
2383 }
2384}
2385
2386#[inline]
2393#[allow(unreachable_code)]
2394fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2395 #[cfg(target_arch = "aarch64")]
2396 unsafe {
2397 return dot_q4t_row_sdot(bytes, r, gpr, xq);
2398 }
2399 #[cfg(target_arch = "x86_64")]
2400 unsafe {
2401 return dot_q4t_row_avx2(bytes, r, gpr, xq);
2402 }
2403 let mut acc = 0f32;
2404 for gi in 0..gpr {
2405 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2406 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2407 let mut d = 0i32;
2408 for (k, &b) in tile[2..].iter().enumerate() {
2409 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2410 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2411 }
2412 acc += d as f32 * s;
2413 }
2414 acc
2415}
2416
2417#[cfg(target_arch = "aarch64")]
2418#[target_feature(enable = "neon,dotprod")]
2419unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2420 unsafe {
2423 use core::arch::aarch64::*;
2424 use core::arch::asm;
2425 let lomask = vdupq_n_u8(0x0F);
2426 let eight = vdupq_n_s8(8);
2427 let mut acc = 0f32;
2428 for gi in 0..gpr {
2429 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2430 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2431 let b = vld1q_u8(t.add(2));
2432 let lo = vandq_u8(b, lomask);
2433 let hi = vshrq_n_u8::<4>(b);
2434 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2435 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2436 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2437 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2438 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2439 asm!(
2440 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2441 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2442 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2443 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2444 options(pure, nomem, nostack),
2445 );
2446 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2447 }
2448 acc
2449 }
2450}
2451
2452#[cfg(target_arch = "x86_64")]
2453#[target_feature(enable = "avx2")]
2454unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2455 unsafe {
2457 use core::arch::x86_64::*;
2458 let lomask = _mm_set1_epi8(0x0F);
2459 let eight = _mm256_set1_epi8(8);
2460 let ones = _mm256_set1_epi16(1);
2461 let mut acc = 0f32;
2462 for gi in 0..gpr {
2463 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2464 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2465 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2466 let lo = _mm_and_si128(b, lomask);
2467 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2468 let w = _mm256_sub_epi8(
2469 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2470 eight,
2471 );
2472 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2473 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2474 let d = _mm256_madd_epi16(p16, ones);
2475 let hi128 = _mm256_extracti128_si256::<1>(d);
2476 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2477 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2478 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2479 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2480 }
2481 acc
2482 }
2483}
2484
2485#[cfg(target_arch = "x86_64")]
2490#[target_feature(enable = "avx2")]
2491unsafe fn dot_q4t_row_1x4_avx2(
2492 bytes: &[u8],
2493 r: usize,
2494 gpr: usize,
2495 xs: [&[i8]; 4],
2496) -> [f32; 4] {
2497 unsafe {
2499 use core::arch::x86_64::*;
2500 let lomask = _mm_set1_epi8(0x0F);
2501 let eight = _mm256_set1_epi8(8);
2502 let ones = _mm256_set1_epi16(1);
2503 let mut acc = [0f32; 4];
2504 for gi in 0..gpr {
2505 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2506 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2507 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
2508 let lo = _mm_and_si128(bb, lomask);
2509 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
2510 let w = _mm256_sub_epi8(
2511 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2512 eight,
2513 );
2514 let aw = _mm256_abs_epi8(w);
2515 for (k, xq) in xs.iter().enumerate() {
2516 let x =
2517 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2518 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2519 let d = _mm256_madd_epi16(p16, ones);
2520 let hi128 = _mm256_extracti128_si256::<1>(d);
2521 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2522 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2523 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2524 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2525 }
2526 }
2527 acc
2528 }
2529}
2530
2531#[inline]
2533fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2534 let gi = j / GROUP_SIZE;
2535 let k = j % GROUP_SIZE;
2536 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2537 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2538 let byte = tile[2 + k / 2];
2539 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2540 ((nib as i32 - 8) as f32, s)
2541}
2542
2543#[inline]
2546fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2547 let mut acc = 0f32;
2548 for gi in 0..gpr {
2549 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2550 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2551 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2552 let mut ga = 0f32;
2553 for (k, &b) in tile[2..].iter().enumerate() {
2554 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2555 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2556 }
2557 acc += ga * s;
2558 }
2559 acc
2560}
2561
2562fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2564 debug_assert_eq!(out.len(), rows);
2565 let gpr = cols / GROUP_SIZE;
2566 let out_addr = SendMut(out.as_mut_ptr());
2567 if a8w8_enabled() {
2568 let act = split_act(x);
2569 let run = move |start: usize, end: usize| {
2570 for r in start..end {
2571 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2572 for &(j, xv) in &act.outliers {
2573 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2574 acc += w * s * xv;
2575 }
2576 unsafe { *out_addr.at(r) = acc };
2578 }
2579 };
2580 dispatch_rows(pool, rows, &run);
2581 return;
2582 }
2583 let run = move |start: usize, end: usize| {
2584 for r in start..end {
2585 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2587 }
2588 };
2589 dispatch_rows(pool, rows, &run);
2590}
2591
2592#[allow(clippy::too_many_arguments)]
2594fn q4t_matvec2(
2595 bytes: &[u8],
2596 x1: &[f32],
2597 x2: &[f32],
2598 rows: usize,
2599 cols: usize,
2600 o1: &mut [f32],
2601 o2: &mut [f32],
2602 pool: Option<&Pool>,
2603) {
2604 let gpr = cols / GROUP_SIZE;
2605 let p1 = SendMut(o1.as_mut_ptr());
2606 let p2 = SendMut(o2.as_mut_ptr());
2607 if a8w8_enabled() {
2608 let a1 = split_act(x1);
2609 let a2 = split_act(x2);
2610 let run = move |start: usize, end: usize| {
2611 for r in start..end {
2612 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2613 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2614 for &(j, xv) in &a1.outliers {
2615 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2616 v1 += w * s * xv;
2617 }
2618 for &(j, xv) in &a2.outliers {
2619 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2620 v2 += w * s * xv;
2621 }
2622 unsafe {
2624 *p1.at(r) = v1;
2625 *p2.at(r) = v2;
2626 }
2627 }
2628 };
2629 dispatch_rows(pool, rows, &run);
2630 return;
2631 }
2632 let run = move |start: usize, end: usize| {
2633 for r in start..end {
2634 unsafe {
2636 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2637 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2638 }
2639 }
2640 };
2641 dispatch_rows(pool, rows, &run);
2642}
2643
2644#[allow(clippy::too_many_arguments)]
2646fn q4t_matmat(
2647 bytes: &[u8],
2648 xs_all: &[f32],
2649 b: usize,
2650 rows: usize,
2651 cols: usize,
2652 out: &mut [f32],
2653 pool: Option<&Pool>,
2654) {
2655 debug_assert_eq!(out.len(), b * rows);
2656 let gpr = cols / GROUP_SIZE;
2657 let out_addr = SendMut(out.as_mut_ptr());
2658 if a8w8_enabled() {
2659 let acts: Vec<SplitAct> =
2660 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2661 let acts = &acts;
2662 #[cfg(target_arch = "x86_64")]
2663 let blocked_ok = avx2_enabled()
2664 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2665 #[cfg(not(target_arch = "x86_64"))]
2666 let blocked_ok = false;
2667 let run = move |start: usize, end: usize| {
2668 for r in start..end {
2669 let mut bi = 0usize;
2670 #[cfg(target_arch = "x86_64")]
2671 if blocked_ok {
2672 while bi + 4 <= acts.len() {
2673 let xs = [
2674 acts[bi].xq.as_slice(),
2675 acts[bi + 1].xq.as_slice(),
2676 acts[bi + 2].xq.as_slice(),
2677 acts[bi + 3].xq.as_slice(),
2678 ];
2679 let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
2680 for k in 0..4 {
2681 let act = &acts[bi + k];
2682 let mut acc = d[k] * act.sx;
2683 for &(j, xv) in &act.outliers {
2684 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
2685 acc += w * sc * xv;
2686 }
2687 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
2689 }
2690 bi += 4;
2691 }
2692 }
2693 let _ = blocked_ok;
2694 while bi < acts.len() {
2695 let act = &acts[bi];
2696 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2697 for &(j, xv) in &act.outliers {
2698 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2699 acc += w * s * xv;
2700 }
2701 unsafe { *out_addr.at(bi * rows + r) = acc };
2703 bi += 1;
2704 }
2705 }
2706 };
2707 dispatch_rows(pool, rows, &run);
2708 return;
2709 }
2710 let run = move |start: usize, end: usize| {
2711 for r in start..end {
2712 for bi in 0..b {
2713 let x = &xs_all[bi * cols..(bi + 1) * cols];
2714 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2716 }
2717 }
2718 };
2719 dispatch_rows(pool, rows, &run);
2720}
2721
2722fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2731 (0..gpr)
2732 .map(|gi| {
2733 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2734 .iter()
2735 .map(|&v| v as i32)
2736 .sum()
2737 })
2738 .collect()
2739}
2740
2741#[inline]
2745#[allow(unreachable_code)]
2746#[cfg(target_arch = "x86_64")]
2751#[target_feature(enable = "avx2")]
2752unsafe fn dot_q1_row_avx2(
2753 bytes: &[u8],
2754 r: usize,
2755 gpr: usize,
2756 xq: &[i8],
2757 gsum: &[i32],
2758) -> f32 {
2759 unsafe {
2761 use core::arch::x86_64::*;
2762 let expand = _mm256_setr_epi8(
2764 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2765 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2766 );
2767 let bitsel = _mm256_setr_epi8(
2768 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2769 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2770 );
2771 let ones8 = _mm256_set1_epi8(1);
2772 let ones16 = _mm256_set1_epi16(1);
2773 let mut acc = 0f32;
2774 for gi in 0..gpr {
2775 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2776 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2777 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2778 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2779 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2780 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2781 let sel = _mm256_and_si256(x, mask);
2782 let p16 = _mm256_maddubs_epi16(ones8, sel);
2784 let d32 = _mm256_madd_epi16(p16, ones16);
2785 let hi128 = _mm256_extracti128_si256::<1>(d32);
2786 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2787 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2788 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2789 let msum = _mm_cvtsi128_si32(s32);
2790 let d = 2 * msum - gsum[gi];
2793 acc += d as f32 * s;
2794 }
2795 acc
2796 }
2797}
2798
2799#[cfg(target_arch = "x86_64")]
2802#[target_feature(enable = "avx2")]
2803unsafe fn dot_q1_row_1x4_avx2(
2804 bytes: &[u8],
2805 r: usize,
2806 gpr: usize,
2807 xs: [&[i8]; 4],
2808 gsums: [&[i32]; 4],
2809) -> [f32; 4] {
2810 unsafe {
2812 use core::arch::x86_64::*;
2813 let expand = _mm256_setr_epi8(
2814 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2815 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2816 );
2817 let bitsel = _mm256_setr_epi8(
2818 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2819 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2820 );
2821 let ones8 = _mm256_set1_epi8(1);
2822 let ones16 = _mm256_set1_epi16(1);
2823 let mut acc = [0f32; 4];
2824 for gi in 0..gpr {
2825 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2826 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2827 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2828 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2829 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2830 for (k, xq) in xs.iter().enumerate() {
2831 let x =
2832 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2833 let sel = _mm256_and_si256(x, mask);
2834 let p16 = _mm256_maddubs_epi16(ones8, sel);
2835 let d32 = _mm256_madd_epi16(p16, ones16);
2836 let hi128 = _mm256_extracti128_si256::<1>(d32);
2837 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2838 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2839 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2840 let msum = _mm_cvtsi128_si32(s32);
2841 let d = 2 * msum - gsums[k][gi];
2842 acc[k] += d as f32 * s;
2843 }
2844 }
2845 acc
2846 }
2847}
2848
2849fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2850 #[cfg(target_arch = "aarch64")]
2851 unsafe {
2852 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2853 }
2854 #[cfg(target_arch = "x86_64")]
2855 if avx2_enabled() {
2856 unsafe {
2857 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
2858 }
2859 }
2860 let _ = gsum;
2861 let mut acc = 0f32;
2862 for gi in 0..gpr {
2863 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2864 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2865 let mut d = 0i32;
2866 for (j, &b) in tile[2..].iter().enumerate() {
2867 for k in 0..8 {
2868 let w = ((b >> k) & 1) as i32 * 2 - 1;
2869 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2870 }
2871 }
2872 acc += d as f32 * s;
2873 }
2874 acc
2875}
2876
2877#[cfg(target_arch = "aarch64")]
2886#[target_feature(enable = "neon,dotprod")]
2887unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2888 unsafe {
2891 use core::arch::aarch64::*;
2892 use core::arch::asm;
2893 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2894 let m = vld1q_u8(MASKS.as_ptr());
2895 macro_rules! tile_dot {
2897 ($t:expr, $x:expr) => {{
2898 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2899 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2900 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2901 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2902 let x0 = vld1q_s8($x);
2903 let x1 = vld1q_s8($x.add(16));
2904 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2905 asm!(
2906 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2907 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2908 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2909 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2910 options(pure, nomem, nostack),
2911 );
2912 vaddq_s32(a0, a1)
2913 }};
2914 }
2915 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
2924 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
2925 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
2926 const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
2927 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
2928 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
2929 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
2930 let isc = vld1_u8(ISC.as_ptr());
2931 macro_rules! tile_dot_tbl {
2933 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
2934 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
2935 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
2936 let x0 = vld1q_s8($x);
2937 let x1 = vld1q_s8($x.add(16));
2938 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2939 asm!(
2940 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2941 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2942 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2943 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2944 options(pure, nomem, nostack),
2945 );
2946 vaddq_s32(a0, a1)
2947 }};
2948 }
2949 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2950 let row_base = r * gpr * Q1_TILE;
2951 let abs_end = bytes.len();
2952 let xp = xq.as_ptr();
2953 let gp = gsum.as_ptr();
2954 let mut accv = vdupq_n_f32(0.0);
2955 let mut gi = 0;
2956 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
2959 let t0 = base.add(gi * Q1_TILE);
2960 let ld_a = vld1q_u8(t0);
2961 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
2962 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
2963 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
2964 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
2965 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
2966 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2968 let g = vld1q_s32(gp.add(gi));
2969 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2970 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
2971 let scf: float32x4_t;
2972 asm!(
2973 "fcvtl {o:v}.4s, {i:v}.4h",
2974 o = out(vreg) scf, i = in(vreg) sc16,
2975 options(pure, nomem, nostack),
2976 );
2977 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
2978 gi += 4;
2979 }
2980 let mut acc = vaddvq_f32(accv);
2981 while gi < gpr {
2982 let t = base.add(gi * Q1_TILE);
2983 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2984 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2985 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2986 gi += 1;
2987 }
2988 acc
2989 }
2990}
2991
2992#[cfg(target_arch = "aarch64")]
2997#[target_feature(enable = "neon,dotprod")]
2998unsafe fn dot_q1_row_1x4_sdot(
2999 bytes: &[u8],
3000 r: usize,
3001 gpr: usize,
3002 xs: [&[i8]; 4],
3003 gs: [&[i32]; 4],
3004) -> [f32; 4] {
3005 unsafe {
3007 use core::arch::aarch64::*;
3008 use core::arch::asm;
3009 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3010 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3011 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3012 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3013 const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
3014 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3015 let m = vld1q_u8(MASKS.as_ptr());
3016 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3017 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3018 let isc = vld1_u8(ISC.as_ptr());
3019 macro_rules! sdot2 {
3020 ($w0:expr, $w1:expr, $x:expr) => {{
3021 let x0 = vld1q_s8($x);
3022 let x1 = vld1q_s8($x.add(16));
3023 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3024 asm!(
3025 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3026 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3027 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3028 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
3029 options(pure, nomem, nostack),
3030 );
3031 vaddq_s32(a0, a1)
3032 }};
3033 }
3034 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3035 let row_base = r * gpr * Q1_TILE;
3036 let abs_end = bytes.len();
3037 let mut accv = [vdupq_n_f32(0.0); 4];
3038 let mut gi = 0;
3039 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3040 let t0 = base.add(gi * Q1_TILE);
3041 let ld_a = vld1q_u8(t0);
3042 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3043 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
3045 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
3046 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
3047 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
3048 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
3049 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
3050 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
3051 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
3052 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3053 let scf: float32x4_t;
3054 asm!(
3055 "fcvtl {o:v}.4s, {i:v}.4h",
3056 o = out(vreg) scf, i = in(vreg) sc16,
3057 options(pure, nomem, nostack),
3058 );
3059 for k in 0..4 {
3060 let xp = xs[k].as_ptr();
3061 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
3062 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
3063 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
3064 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
3065 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3066 let g = vld1q_s32(gs[k].as_ptr().add(gi));
3067 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3068 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
3069 }
3070 gi += 4;
3071 }
3072 let mut acc = [
3073 vaddvq_f32(accv[0]),
3074 vaddvq_f32(accv[1]),
3075 vaddvq_f32(accv[2]),
3076 vaddvq_f32(accv[3]),
3077 ];
3078 while gi < gpr {
3079 let t = base.add(gi * Q1_TILE);
3080 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3081 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
3082 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
3083 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
3084 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
3085 for k in 0..4 {
3086 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
3087 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
3088 }
3089 gi += 1;
3090 }
3091 acc
3092 }
3093}
3094
3095#[inline]
3097fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3098 let gi = j / GROUP_SIZE;
3099 let k = j % GROUP_SIZE;
3100 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3101 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3102 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
3103 ((bit as i32 * 2 - 1) as f32, s)
3104}
3105
3106#[inline]
3108fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3109 let mut acc = 0f32;
3110 for gi in 0..gpr {
3111 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3112 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3113 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3114 let mut ga = 0f32;
3115 for (j, &b) in tile[2..].iter().enumerate() {
3116 for k in 0..8 {
3117 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
3118 }
3119 }
3120 acc += ga * s;
3121 }
3122 acc
3123}
3124
3125#[allow(clippy::too_many_arguments)]
3128fn q1_range_a8w8(
3129 bytes: &[u8],
3130 gpr: usize,
3131 act: &SplitAct,
3132 gsum: &[i32],
3133 out: SendMut,
3134 start: usize,
3135 end: usize,
3136) {
3137 for r in start..end {
3138 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3139 for &(j, xv) in &act.outliers {
3140 let (w, s) = q1_outlier(bytes, r, gpr, j);
3141 acc += w * s * xv;
3142 }
3143 unsafe { *out.at(r) = acc };
3145 }
3146}
3147
3148fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
3150 for r in start..end {
3151 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
3153 }
3154}
3155
3156fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
3161 let entries = base_len + (rows + 1) * 4;
3162 (base_len, entries, entries <= bytes.len())
3163}
3164
3165#[inline]
3167fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
3168 let o = rp_off + r * 4;
3169 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
3170}
3171
3172const SIGN5: [[f32; 5]; 256] = {
3176 let mut lut = [[0.0f32; 5]; 256];
3177 let pow3 = [1u16, 3, 9, 27, 81];
3178 let mut byte = 0usize;
3179 while byte < 256 {
3180 let mut i = 0usize;
3181 while i < 5 {
3182 let code = (byte as u16 / pow3[i]) % 3;
3183 lut[byte][i] = if code == 1 {
3184 1.0
3185 } else if code == 2 {
3186 -1.0
3187 } else {
3188 0.0
3189 };
3190 i += 1;
3191 }
3192 byte += 1;
3193 }
3194 lut
3195};
3196
3197const SIGN5_I8: [[i8; 5]; 256] = {
3199 let mut lut = [[0i8; 5]; 256];
3200 let pow3 = [1u16, 3, 9, 27, 81];
3201 let mut byte = 0usize;
3202 while byte < 256 {
3203 let mut i = 0usize;
3204 while i < 5 {
3205 let code = (byte as u16 / pow3[i]) % 3;
3206 lut[byte][i] = if code == 1 {
3207 1
3208 } else if code == 2 {
3209 -1
3210 } else {
3211 0
3212 };
3213 i += 1;
3214 }
3215 byte += 1;
3216 }
3217 lut
3218};
3219
3220const SIGN5_U64: [u64; 256] = {
3226 let mut lut = [0u64; 256];
3227 let pow3 = [1u16, 3, 9, 27, 81];
3228 let mut byte = 0usize;
3229 while byte < 256 {
3230 let mut v = 0u64;
3231 let mut i = 0usize;
3232 while i < 5 {
3233 let code = (byte as u16 / pow3[i]) % 3;
3234 let s: u8 = if code == 1 {
3235 1
3236 } else if code == 2 {
3237 0xFF
3238 } else {
3239 0
3240 };
3241 v |= (s as u64) << (i * 8);
3242 i += 1;
3243 }
3244 lut[byte] = v;
3245 byte += 1;
3246 }
3247 lut
3248};
3249
3250#[inline]
3255fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
3256 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3257 let off = (r * gpr + j / GROUP_SIZE) * TILE;
3258 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3259 let within = j % GROUP_SIZE;
3260 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
3261}
3262
3263#[cfg(target_arch = "aarch64")]
3266#[target_feature(enable = "neon,dotprod")]
3267#[inline]
3268unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
3269 unsafe {
3271 use core::arch::aarch64::*;
3272 use core::arch::asm;
3273 let w0 = vld1q_s8(w);
3274 let w1 = vld1q_s8(w.add(16));
3275 let x0 = vld1q_s8(x);
3276 let x1 = vld1q_s8(x.add(16));
3277 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3278 asm!(
3279 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3280 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3281 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3282 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3283 options(pure, nomem, nostack),
3284 );
3285 vaddvq_s32(vaddq_s32(a0, a1))
3286 }
3287}
3288
3289#[cfg(target_arch = "x86_64")]
3292#[target_feature(enable = "avx2")]
3293#[inline]
3294unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
3295 unsafe {
3297 use core::arch::x86_64::*;
3298 let wv = _mm256_loadu_si256(w as *const __m256i);
3299 let xv = _mm256_loadu_si256(x as *const __m256i);
3300 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3301 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
3302 let hi128 = _mm256_extracti128_si256::<1>(d);
3303 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3304 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3305 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3306 _mm_cvtsi128_si32(s32)
3307 }
3308}
3309
3310#[inline]
3315fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
3316 debug_assert!(dst.len() >= 40);
3317 unsafe {
3320 let p = dst.as_mut_ptr();
3321 for bi in 0..7 {
3322 core::ptr::write_unaligned(
3323 p.add(bi * 5) as *mut u64,
3324 SIGN5_U64[*codes.add(bi) as usize],
3325 );
3326 }
3327 }
3328}
3329
3330#[inline]
3335fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
3336 #[cfg(target_arch = "aarch64")]
3337 unsafe {
3338 return sdot32_i8(w, x);
3339 }
3340 #[cfg(target_arch = "x86_64")]
3341 unsafe {
3342 return i8dot32_avx2(w, x);
3343 }
3344 #[allow(unreachable_code)]
3345 unsafe {
3346 let mut s = 0i32;
3347 for k in 0..GROUP_SIZE {
3348 s += *w.add(k) as i32 * *x.add(k) as i32;
3349 }
3350 s
3351 }
3352}
3353
3354#[cfg(target_arch = "aarch64")]
3358#[target_feature(enable = "neon,dotprod")]
3359unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3360 unsafe {
3362 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3363 let mut acc = 0f32;
3364 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3366 let off = (r * gpr + gi) * TILE;
3367 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3368 q1t_unpack_group_i8(bytes.as_ptr().add(off + 2), &mut sg);
3369 acc += sdot32_i8(sg.as_ptr(), xq.as_ptr().add(gi * GROUP_SIZE)) as f32 * s;
3370 }
3371 acc
3372 }
3373}
3374
3375#[cfg(target_arch = "x86_64")]
3377#[target_feature(enable = "avx2")]
3378unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3379 unsafe {
3381 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3382 let mut acc = 0f32;
3383 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3385 let off = (r * gpr + gi) * TILE;
3386 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3387 q1t_unpack_group_i8(bytes.as_ptr().add(off + 2), &mut sg);
3388 acc += i8dot32_avx2(sg.as_ptr(), xq.as_ptr().add(gi * GROUP_SIZE)) as f32 * s;
3389 }
3390 acc
3391 }
3392}
3393
3394#[inline]
3398fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3399 #[cfg(target_arch = "aarch64")]
3400 unsafe {
3401 return q1t_dot_row_sdot(bytes, r, gpr, xq);
3402 }
3403 #[cfg(target_arch = "x86_64")]
3404 unsafe {
3405 return q1t_dot_row_avx2(bytes, r, gpr, xq);
3406 }
3407 #[allow(unreachable_code)]
3408 {
3409 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3410 let mut acc = 0f32;
3411 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3413 let off = (r * gpr + gi) * TILE;
3414 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3415 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
3416 let mut d = 0i32;
3417 for k in 0..GROUP_SIZE {
3418 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
3419 }
3420 acc += d as f32 * s;
3421 }
3422 acc
3423 }
3424}
3425
3426fn q1t_row_outlier_correction(
3433 bytes: &[u8],
3434 r: usize,
3435 rp_off: usize,
3436 entries_off: usize,
3437 has_ov: bool,
3438 x: &[f32],
3439) -> f32 {
3440 if !has_ov {
3441 return 0.0;
3442 }
3443 let (c0, c1) = (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3444 let mut corr = 0f32;
3445 for p in c0..c1 {
3446 let e = entries_off + p * 4;
3447 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3448 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3449 corr += val * x[col];
3450 }
3451 corr
3452}
3453
3454fn q1t_dequant_row(
3458 bytes: &[u8],
3459 r: usize,
3460 gpr: usize,
3461 rp_off: usize,
3462 entries_off: usize,
3463 has_ov: bool,
3464 buf: &mut [f32],
3465) {
3466 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3467 for g in 0..gpr {
3468 let off = (r * gpr + g) * TILE;
3469 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3470 let codes = &bytes[off + 2..off + TILE];
3471 let bc = g * GROUP_SIZE;
3472 for bi in 0..6 {
3474 let lut = &SIGN5[codes[bi] as usize];
3475 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
3476 for i in 0..5 {
3477 d[i] = lut[i] * s;
3478 }
3479 }
3480 let lut = &SIGN5[codes[6] as usize];
3481 buf[bc + 30] = lut[0] * s;
3482 buf[bc + 31] = lut[1] * s;
3483 }
3484 if !has_ov {
3485 return;
3486 }
3487 let (c0, c1) = (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3488 for p in c0..c1 {
3489 let e = entries_off + p * 4;
3490 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3491 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3492 }
3493}
3494
3495fn q1t_add_overlay(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3499 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3500 let gpr = cols / GROUP_SIZE;
3501 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3502 if !has_ov {
3503 return;
3504 }
3505 let out_addr = SendMut(out.as_mut_ptr());
3506 let run = move |start: usize, end: usize| {
3507 for r in start..end {
3508 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3509 unsafe { *out_addr.at(r) += corr };
3511 }
3512 };
3513 dispatch_rows(pool, rows, &run);
3514}
3515
3516fn q1t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3520 debug_assert_eq!(out.len(), rows);
3521 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3522 let gpr = cols / GROUP_SIZE;
3523 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3524 let out_addr = SendMut(out.as_mut_ptr());
3525 if a8w8_enabled() {
3529 let act = split_act(x);
3530 let act = &act;
3531 let run = move |start: usize, end: usize| {
3532 for r in start..end {
3533 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3534 for &(j, xv) in &act.outliers {
3535 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
3536 }
3537 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3538 unsafe { *out_addr.at(r) = acc };
3540 }
3541 };
3542 dispatch_rows(pool, rows, &run);
3543 return;
3544 }
3545 let run = move |start: usize, end: usize| {
3546 let mut sg = [0f32; GROUP_SIZE];
3550 for r in start..end {
3551 let mut acc = 0f32;
3552 for g in 0..gpr {
3553 let off = (r * gpr + g) * TILE;
3554 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3555 let codes = &bytes[off + 2..off + TILE];
3556 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
3557 for bi in 0..6 {
3558 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
3559 }
3560 let lut = &SIGN5[codes[6] as usize];
3561 sg[30] = lut[0];
3562 sg[31] = lut[1];
3563 let mut gsum = 0f32;
3564 for k in 0..GROUP_SIZE {
3565 gsum += sg[k] * xg[k];
3566 }
3567 acc += s * gsum;
3568 }
3569 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3570 unsafe { *out_addr.at(r) = acc };
3571 }
3572 };
3573 dispatch_rows(pool, rows, &run);
3574}
3575
3576fn q1t_matmat(
3579 bytes: &[u8],
3580 xs: &[f32],
3581 b: usize,
3582 rows: usize,
3583 cols: usize,
3584 out: &mut [f32],
3585 pool: Option<&Pool>,
3586) {
3587 debug_assert_eq!(out.len(), b * rows);
3588 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3589 let gpr = cols / GROUP_SIZE;
3590 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3591 let out_addr = SendMut(out.as_mut_ptr());
3592 if a8w8_enabled() {
3596 let acts: Vec<SplitAct> = (0..b)
3597 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
3598 .collect();
3599 let acts = &acts;
3600 let run = move |start: usize, end: usize| {
3601 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
3605 for g in 0..gpr {
3606 let off = (r * gpr + g) * TILE;
3607 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3608 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg[g * GROUP_SIZE..]);
3609 }
3610 for bi in 0..b {
3611 let act = &acts[bi];
3612 let mut isum = 0f32;
3613 for g in 0..gpr {
3614 let d = q1t_i8dot32(
3615 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
3616 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
3617 );
3618 isum += d as f32 * sc[g];
3619 }
3620 let mut acc = isum * act.sx;
3621 for &(j, xv) in &act.outliers {
3622 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
3623 }
3624 accs[bi] = acc;
3625 }
3626 if has_ov {
3630 let (c0, c1) =
3631 (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3632 for p in c0..c1 {
3633 let e = ent_off + p * 4;
3634 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3635 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3636 for bi in 0..b {
3637 accs[bi] += val * xs[bi * cols + col];
3638 }
3639 }
3640 }
3641 for bi in 0..b {
3642 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
3643 }
3644 }
3645 };
3646 dispatch_rows(pool, rows, &run);
3647 return;
3648 }
3649 let run = move |start: usize, end: usize| {
3650 let mut buf = vec![0f32; cols];
3651 for r in start..end {
3652 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
3653 for bi in 0..b {
3654 let xr = &xs[bi * cols..(bi + 1) * cols];
3655 let mut acc = 0f32;
3656 for j in 0..cols {
3657 acc += buf[j] * xr[j];
3658 }
3659 unsafe { *out_addr.at(bi * rows + r) = acc };
3660 }
3661 }
3662 };
3663 dispatch_rows(pool, rows, &run);
3664}
3665
3666fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3667 debug_assert_eq!(out.len(), rows);
3668 let gpr = cols / GROUP_SIZE;
3669 let out_addr = SendMut(out.as_mut_ptr());
3670 if a8w8_enabled() {
3671 let act = split_act(x);
3672 let gsum = q1_group_sums(&act.xq, gpr);
3673 let (act, gsum) = (&act, &gsum);
3674 let run =
3675 move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
3676 dispatch_rows(pool, rows, &run);
3677 return;
3678 }
3679 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
3680 dispatch_rows(pool, rows, &run);
3681}
3682
3683#[allow(clippy::too_many_arguments)]
3685fn q1_matvec2(
3686 bytes: &[u8],
3687 x1: &[f32],
3688 x2: &[f32],
3689 rows: usize,
3690 cols: usize,
3691 o1: &mut [f32],
3692 o2: &mut [f32],
3693 pool: Option<&Pool>,
3694) {
3695 let gpr = cols / GROUP_SIZE;
3696 let p1 = SendMut(o1.as_mut_ptr());
3697 let p2 = SendMut(o2.as_mut_ptr());
3698 if a8w8_enabled() {
3699 let a1 = split_act(x1);
3700 let a2 = split_act(x2);
3701 let g1 = q1_group_sums(&a1.xq, gpr);
3702 let g2 = q1_group_sums(&a2.xq, gpr);
3703 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
3704 let run = move |start: usize, end: usize| {
3705 for r in start..end {
3706 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
3707 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
3708 for &(j, xv) in &a1.outliers {
3709 let (w, s) = q1_outlier(bytes, r, gpr, j);
3710 v1 += w * s * xv;
3711 }
3712 for &(j, xv) in &a2.outliers {
3713 let (w, s) = q1_outlier(bytes, r, gpr, j);
3714 v2 += w * s * xv;
3715 }
3716 unsafe {
3718 *p1.at(r) = v1;
3719 *p2.at(r) = v2;
3720 }
3721 }
3722 };
3723 dispatch_rows(pool, rows, &run);
3724 return;
3725 }
3726 let run = move |start: usize, end: usize| {
3727 for r in start..end {
3728 unsafe {
3730 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
3731 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
3732 }
3733 }
3734 };
3735 dispatch_rows(pool, rows, &run);
3736}
3737
3738#[allow(clippy::too_many_arguments)]
3740fn q1_matmat(
3741 bytes: &[u8],
3742 xs_all: &[f32],
3743 b: usize,
3744 rows: usize,
3745 cols: usize,
3746 out: &mut [f32],
3747 pool: Option<&Pool>,
3748) {
3749 debug_assert_eq!(out.len(), b * rows);
3750 let gpr = cols / GROUP_SIZE;
3751 let out_addr = SendMut(out.as_mut_ptr());
3752 if a8w8_enabled() {
3753 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
3754 .map(|bi| {
3755 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
3756 let gsum = q1_group_sums(&act.xq, gpr);
3757 (act, gsum)
3758 })
3759 .collect();
3760 let acts = &acts;
3761 #[cfg(target_arch = "x86_64")]
3762 let blocked_ok = avx2_enabled()
3763 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3764 #[cfg(target_arch = "aarch64")]
3765 let blocked_ok = sdot_enabled()
3766 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3767 let run = move |start: usize, end: usize| {
3768 for r in start..end {
3769 let mut bi = 0usize;
3770 #[cfg(target_arch = "aarch64")]
3773 if blocked_ok {
3774 while bi + 4 <= acts.len() {
3775 let xs = [
3776 acts[bi].0.xq.as_slice(),
3777 acts[bi + 1].0.xq.as_slice(),
3778 acts[bi + 2].0.xq.as_slice(),
3779 acts[bi + 3].0.xq.as_slice(),
3780 ];
3781 let gs = [
3782 acts[bi].1.as_slice(),
3783 acts[bi + 1].1.as_slice(),
3784 acts[bi + 2].1.as_slice(),
3785 acts[bi + 3].1.as_slice(),
3786 ];
3787 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
3788 for k in 0..4 {
3789 let (act, _) = &acts[bi + k];
3790 let mut acc = d[k] * act.sx;
3791 for &(j, xv) in &act.outliers {
3792 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3793 acc += w * sc * xv;
3794 }
3795 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3797 }
3798 bi += 4;
3799 }
3800 }
3801 #[cfg(target_arch = "x86_64")]
3802 if blocked_ok {
3803 while bi + 4 <= acts.len() {
3804 let xs = [
3805 acts[bi].0.xq.as_slice(),
3806 acts[bi + 1].0.xq.as_slice(),
3807 acts[bi + 2].0.xq.as_slice(),
3808 acts[bi + 3].0.xq.as_slice(),
3809 ];
3810 let gs = [
3811 acts[bi].1.as_slice(),
3812 acts[bi + 1].1.as_slice(),
3813 acts[bi + 2].1.as_slice(),
3814 acts[bi + 3].1.as_slice(),
3815 ];
3816 let d = unsafe { dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs) };
3817 for k in 0..4 {
3818 let (act, _) = &acts[bi + k];
3819 let mut acc = d[k] * act.sx;
3820 for &(j, xv) in &act.outliers {
3821 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3822 acc += w * sc * xv;
3823 }
3824 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3826 }
3827 bi += 4;
3828 }
3829 }
3830 while bi < acts.len() {
3831 let (act, gsum) = &acts[bi];
3832 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3833 for &(j, xv) in &act.outliers {
3834 let (w, s) = q1_outlier(bytes, r, gpr, j);
3835 acc += w * s * xv;
3836 }
3837 unsafe { *out_addr.at(bi * rows + r) = acc };
3839 bi += 1;
3840 }
3841 }
3842 };
3843 dispatch_rows(pool, rows, &run);
3844 return;
3845 }
3846 let run = move |start: usize, end: usize| {
3847 for r in start..end {
3848 for bi in 0..b {
3849 let x = &xs_all[bi * cols..(bi + 1) * cols];
3850 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
3852 }
3853 }
3854 };
3855 dispatch_rows(pool, rows, &run);
3856}
3857
3858fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3864 debug_assert_eq!(out.len(), rows);
3865 let (packed, scales) = q4_split(bytes, rows, cols);
3866 let gpr = cols / GROUP_SIZE;
3867 let out_addr = SendMut(out.as_mut_ptr());
3868
3869 if a8w8_enabled() {
3870 let act = split_act(x);
3871 let run = move |start: usize, end: usize| {
3872 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
3873 };
3874 dispatch_rows(pool, rows, &run);
3875 return;
3876 }
3877
3878 let run = move |start: usize, end: usize| {
3879 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
3880 };
3881 dispatch_rows(pool, rows, &run);
3882}
3883
3884#[inline]
3887#[allow(unreachable_code)]
3888#[cfg(target_arch = "x86_64")]
3893#[target_feature(enable = "avx2")]
3894unsafe fn dot_q4b_row_1x4_avx2(
3895 buf: &[u8],
3896 scales: &[u8],
3897 g0: usize,
3898 gpr: usize,
3899 xs: [&[i8]; 4],
3900) -> [f32; 4] {
3901 unsafe {
3903 use core::arch::x86_64::*;
3904 let ones = _mm256_set1_epi16(1);
3905 let mut acc = [0f32; 4];
3906 for gi in 0..gpr {
3907 let s = f16_to_f32(u16::from_le_bytes([
3908 scales[(g0 + gi) * 2],
3909 scales[(g0 + gi) * 2 + 1],
3910 ]));
3911 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3912 let aw = _mm256_abs_epi8(w);
3913 for (k, xq) in xs.iter().enumerate() {
3914 let x =
3915 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3916 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3917 let d = _mm256_madd_epi16(p16, ones);
3918 let hi128 = _mm256_extracti128_si256::<1>(d);
3919 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3920 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3921 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3922 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
3923 }
3924 }
3925 acc
3926 }
3927}
3928
3929#[cfg(target_arch = "x86_64")]
3935#[target_feature(enable = "avx2")]
3936unsafe fn dot_q4b_row_1x4_sx_avx2(
3937 buf: &[u8],
3938 scales: &[u8],
3939 g0: usize,
3940 gpr: usize,
3941 xs: [&[i8]; 4],
3942 sxs: [f32; 4],
3943) -> [f32; 4] {
3944 unsafe {
3946 use core::arch::x86_64::*;
3947 let ones = _mm256_set1_epi16(1);
3948 let mut acc = [0f32; 4];
3949 for gi in 0..gpr {
3950 let s = f16_to_f32(u16::from_le_bytes([
3951 scales[(g0 + gi) * 2],
3952 scales[(g0 + gi) * 2 + 1],
3953 ]));
3954 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3955 let aw = _mm256_abs_epi8(w);
3956 for (k, xq) in xs.iter().enumerate() {
3957 let x =
3958 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3959 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3960 let d = _mm256_madd_epi16(p16, ones);
3961 let hi128 = _mm256_extracti128_si256::<1>(d);
3962 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3963 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3964 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3965 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
3966 }
3967 }
3968 acc
3969 }
3970}
3971
3972fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3973 #[cfg(target_arch = "aarch64")]
3974 unsafe {
3975 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
3976 }
3977 #[cfg(target_arch = "x86_64")]
3978 unsafe {
3979 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
3980 }
3981 let mut acc = 0f32;
3982 for gi in 0..gpr {
3983 let g = g0 + gi;
3984 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3985 let mut d = 0i32;
3986 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3987 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3988 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3989 }
3990 acc += d as f32 * s;
3991 }
3992 acc
3993}
3994
3995#[inline]
3997#[allow(unreachable_code)]
3998fn dot_q4_row_i8_2(
3999 packed: &[u8],
4000 scales: &[u8],
4001 g0: usize,
4002 gpr: usize,
4003 xq1: &[i8],
4004 xq2: &[i8],
4005) -> (f32, f32) {
4006 #[cfg(target_arch = "aarch64")]
4007 unsafe {
4008 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
4009 }
4010 #[cfg(target_arch = "x86_64")]
4011 unsafe {
4012 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
4013 }
4014 (
4015 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
4016 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
4017 )
4018}
4019
4020#[allow(clippy::too_many_arguments)]
4023fn q4_range_a8w8(
4024 packed: &[u8],
4025 scales: &[u8],
4026 gpr: usize,
4027 cols: usize,
4028 act: &SplitAct,
4029 out: SendMut,
4030 start: usize,
4031 end: usize,
4032) {
4033 for r in start..end {
4034 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
4035 for &(j, xv) in &act.outliers {
4037 let flat = r * cols + j;
4038 let byte = packed[flat / 2];
4039 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4040 let s = f16_to_f32(u16::from_le_bytes([
4041 scales[(flat / GROUP_SIZE) * 2],
4042 scales[(flat / GROUP_SIZE) * 2 + 1],
4043 ]));
4044 acc += ((nib as i32 - 8) as f32) * s * xv;
4045 }
4046 unsafe { *out.at(r) = acc };
4048 }
4049}
4050
4051#[allow(clippy::too_many_arguments)]
4054fn q4_range2_a8w8(
4055 packed: &[u8],
4056 scales: &[u8],
4057 gpr: usize,
4058 cols: usize,
4059 a1: &SplitAct,
4060 a2: &SplitAct,
4061 p1: SendMut,
4062 p2: SendMut,
4063 start: usize,
4064 end: usize,
4065) {
4066 for r in start..end {
4067 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
4068 let mut acc1 = s1 * a1.sx;
4069 let mut acc2 = s2 * a2.sx;
4070 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
4072 for &(j, xv) in outliers {
4073 let flat = r * cols + j;
4074 let byte = packed[flat / 2];
4075 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4076 let s = f16_to_f32(u16::from_le_bytes([
4077 scales[(flat / GROUP_SIZE) * 2],
4078 scales[(flat / GROUP_SIZE) * 2 + 1],
4079 ]));
4080 *acc += ((nib as i32 - 8) as f32) * s * xv;
4081 }
4082 };
4083 fix(&a1.outliers, &mut acc1);
4084 fix(&a2.outliers, &mut acc2);
4085 unsafe {
4087 *p1.at(r) = acc1;
4088 *p2.at(r) = acc2;
4089 }
4090 }
4091}
4092
4093fn q4_range_f32(
4095 packed: &[u8],
4096 scales: &[u8],
4097 gpr: usize,
4098 x: &[f32],
4099 out: SendMut,
4100 start: usize,
4101 end: usize,
4102) {
4103 for r in start..end {
4104 let mut acc = 0f32;
4105 for gi in 0..gpr {
4106 let g = r * gpr + gi;
4107 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4108 let pk = &packed[g * 16..(g + 1) * 16];
4109 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4110 let mut ga = 0f32;
4111 for (k, &b) in pk.iter().enumerate() {
4112 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4113 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4114 }
4115 acc += ga * s;
4116 }
4117 unsafe { *out.at(r) = acc };
4119 }
4120}
4121
4122#[allow(clippy::too_many_arguments)]
4126fn q4matvec2(
4127 bytes: &[u8],
4128 x1: &[f32],
4129 x2: &[f32],
4130 rows: usize,
4131 cols: usize,
4132 o1: &mut [f32],
4133 o2: &mut [f32],
4134 pool: Option<&Pool>,
4135) {
4136 debug_assert_eq!(o1.len(), rows);
4137 debug_assert_eq!(o2.len(), rows);
4138 let (packed, scales) = q4_split(bytes, rows, cols);
4139 let gpr = cols / GROUP_SIZE;
4140
4141 if a8w8_enabled() {
4142 let a1 = split_act(x1);
4143 let a2 = split_act(x2);
4144 let p1 = SendMut(o1.as_mut_ptr());
4145 let p2 = SendMut(o2.as_mut_ptr());
4146 let run = move |start: usize, end: usize| {
4147 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
4148 };
4149 dispatch_rows(pool, rows, &run);
4150 return;
4151 }
4152
4153 let p1 = SendMut(o1.as_mut_ptr());
4154 let p2 = SendMut(o2.as_mut_ptr());
4155 let run = move |start: usize, end: usize| {
4156 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
4157 };
4158 dispatch_rows(pool, rows, &run);
4159}
4160
4161#[allow(clippy::too_many_arguments)]
4163fn q4_range2_f32(
4164 packed: &[u8],
4165 scales: &[u8],
4166 gpr: usize,
4167 x1: &[f32],
4168 x2: &[f32],
4169 p1: SendMut,
4170 p2: SendMut,
4171 start: usize,
4172 end: usize,
4173) {
4174 for r in start..end {
4175 let (mut acc1, mut acc2) = (0f32, 0f32);
4176 for gi in 0..gpr {
4177 let g = r * gpr + gi;
4178 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4179 let pk = &packed[g * 16..(g + 1) * 16];
4180 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4181 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4182 let (mut g1, mut g2) = (0f32, 0f32);
4183 for (k, &b) in pk.iter().enumerate() {
4184 let wl = (b & 0x0F) as f32 - 8.0;
4185 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
4186 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
4187 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
4188 }
4189 acc1 += g1 * s;
4190 acc2 += g2 * s;
4191 }
4192 unsafe {
4194 *p1.at(r) = acc1;
4195 *p2.at(r) = acc2;
4196 }
4197 }
4198}
4199
4200thread_local! {
4201 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
4204 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4205}
4206
4207#[allow(clippy::too_many_arguments)]
4213fn q4matmat(
4214 bytes: &[u8],
4215 xs_all: &[f32],
4216 b: usize,
4217 rows: usize,
4218 cols: usize,
4219 out: &mut [f32],
4220 pool: Option<&Pool>,
4221) {
4222 debug_assert_eq!(xs_all.len(), b * cols);
4223 debug_assert_eq!(out.len(), b * rows);
4224 let (packed, scales) = q4_split(bytes, rows, cols);
4225 let gpr = cols / GROUP_SIZE;
4226 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4227
4228 if a8w8_enabled() {
4229 let acts: Vec<SplitAct> =
4230 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
4231 let acts = &acts;
4232 let out_addr = SendMut(out.as_mut_ptr());
4233 let run = move |start: usize, end: usize| {
4234 ROW_I8.with(|rb| {
4235 let mut buf = rb.borrow_mut();
4236 buf.resize(cols, 0);
4237 for r in start..end {
4238 for gi in 0..gpr {
4242 let g = r * gpr + gi;
4243 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4244 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
4245 buf[gi * GROUP_SIZE + k * 2 + 1] =
4246 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
4247 }
4248 }
4249 let mut bi = 0usize;
4250 #[cfg(target_arch = "x86_64")]
4251 if avx2_enabled()
4252 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
4253 {
4254 while bi + 4 <= acts.len() {
4255 let xs = [
4256 acts[bi].xq.as_slice(),
4257 acts[bi + 1].xq.as_slice(),
4258 acts[bi + 2].xq.as_slice(),
4259 acts[bi + 3].xq.as_slice(),
4260 ];
4261 let d = unsafe {
4262 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
4263 };
4264 for k in 0..4 {
4265 let act = &acts[bi + k];
4266 let mut acc = d[k] * act.sx;
4267 for &(j, xv) in &act.outliers {
4268 acc += (buf[j] as i8) as f32
4269 * gscale((r * cols + j) / GROUP_SIZE)
4270 * xv;
4271 }
4272 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4274 }
4275 bi += 4;
4276 }
4277 }
4278 while bi < acts.len() {
4279 let act = &acts[bi];
4280 let mut acc = 0f32;
4281 for gi in 0..gpr {
4282 let d = dot_i8_i8(
4283 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
4284 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
4285 );
4286 acc += d as f32 * gscale(r * gpr + gi);
4287 }
4288 acc *= act.sx;
4289 for &(j, xv) in &act.outliers {
4291 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
4292 }
4293 unsafe { *out_addr.at(bi * rows + r) = acc };
4295 bi += 1;
4296 }
4297 }
4298 })
4299 };
4300 dispatch_rows(pool, rows, &run);
4301 return;
4302 }
4303
4304 let out_addr = SendMut(out.as_mut_ptr());
4305 let run = move |start: usize, end: usize| {
4306 ROW_F32.with(|rb| {
4307 let mut buf = rb.borrow_mut();
4308 buf.resize(cols, 0.0);
4309 for r in start..end {
4310 for gi in 0..gpr {
4313 let g = r * gpr + gi;
4314 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4315 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
4316 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
4317 }
4318 }
4319 for bi in 0..b {
4320 let x = &xs_all[bi * cols..(bi + 1) * cols];
4321 let mut acc = 0f32;
4322 for gi in 0..gpr {
4323 let mut ga = 0f32;
4324 for k in 0..GROUP_SIZE / 2 {
4329 let e = gi * GROUP_SIZE + k * 2;
4330 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
4331 }
4332 acc += ga * gscale(r * gpr + gi);
4333 }
4334 unsafe { *out_addr.at(bi * rows + r) = acc };
4336 }
4337 }
4338 })
4339 };
4340 dispatch_rows(pool, rows, &run);
4341}
4342
4343#[allow(clippy::too_many_arguments)]
4348fn vbitmatmat(
4349 bytes: &[u8],
4350 offsets: &[usize],
4351 xs_all: &[f32],
4352 b: usize,
4353 rows: usize,
4354 cols: usize,
4355 out: &mut [f32],
4356 pool: Option<&Pool>,
4357) {
4358 debug_assert_eq!(xs_all.len(), b * cols);
4359 debug_assert_eq!(out.len(), b * rows);
4360 debug_assert_eq!(offsets.len(), rows + 1);
4361 let ng = cols / GROUP_SIZE;
4362 let bits = &bytes[..rows];
4363 let sc_off = rows;
4364 let gscale = |r: usize, g: usize| {
4365 let so = (r * ng + g) * 2;
4366 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
4367 };
4368
4369 let decode_f32 = |r: usize, dst: &mut [f32]| {
4371 let bw = bits[r] as usize;
4372 let l = ((1i32 << (bw - 1)) - 1) as f32;
4373 let data = &bytes[offsets[r]..offsets[r + 1]];
4374 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4375 for d in dst.iter_mut() {
4376 while nbits < bw {
4377 acc = (acc << 8) | data[idx] as u64;
4378 idx += 1;
4379 nbits += 8;
4380 }
4381 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
4382 nbits -= bw;
4383 *d = u - l;
4384 }
4385 };
4386
4387 if a8w8_enabled() {
4388 let acts: Vec<SplitAct> =
4389 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
4390 let acts = &acts;
4391 let out_addr = SendMut(out.as_mut_ptr());
4392 let run = move |start: usize, end: usize| {
4393 for r in start..end {
4394 let bw = bits[r] as usize;
4395 if bw == 8 {
4396 ROW_F32.with(|rb| {
4399 let mut buf = rb.borrow_mut();
4400 buf.resize(cols, 0.0);
4401 decode_f32(r, &mut buf);
4402 for bi in 0..b {
4403 let x = &xs_all[bi * cols..(bi + 1) * cols];
4404 let mut dot = 0f32;
4405 for g in 0..ng {
4406 let mut gd = 0f32;
4407 for k in 0..GROUP_SIZE {
4408 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
4409 }
4410 dot += gd * gscale(r, g);
4411 }
4412 unsafe { *out_addr.at(bi * rows + r) = dot };
4414 }
4415 });
4416 continue;
4417 }
4418 let l = (1i32 << (bw - 1)) - 1;
4419 let data = &bytes[offsets[r]..offsets[r + 1]];
4420 ROW_I8.with(|rb| {
4421 let mut buf = rb.borrow_mut();
4422 buf.resize(cols, 0);
4423 #[inline(always)]
4424 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4425 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4426 let u = unpack8::<B>(&data[blk * B..]);
4427 for k in 0..8 {
4428 chunk[k] = (u[k] - l) as i8 as u8;
4429 }
4430 }
4431 }
4432 match bw {
4433 3 => fill::<3>(data, l, &mut buf),
4434 4 => vbit_fill4(data, &mut buf),
4435 5 => fill::<5>(data, l, &mut buf),
4436 6 => fill::<6>(data, l, &mut buf),
4437 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
4438 }
4439 let mut bi = 0usize;
4440 #[cfg(target_arch = "x86_64")]
4444 if avx2_enabled()
4445 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
4446 {
4447 while bi + 4 <= acts.len() {
4448 let xs = [
4449 acts[bi].xq.as_slice(),
4450 acts[bi + 1].xq.as_slice(),
4451 acts[bi + 2].xq.as_slice(),
4452 acts[bi + 3].xq.as_slice(),
4453 ];
4454 let sxs = [
4455 acts[bi].sx,
4456 acts[bi + 1].sx,
4457 acts[bi + 2].sx,
4458 acts[bi + 3].sx,
4459 ];
4460 let d = unsafe {
4461 dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
4462 };
4463 for k in 0..4 {
4464 let act = &acts[bi + k];
4465 let mut dot = d[k];
4466 for &(j, xv) in &act.outliers {
4467 dot +=
4468 (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
4469 }
4470 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
4472 }
4473 bi += 4;
4474 }
4475 }
4476 while bi < acts.len() {
4477 let act = &acts[bi];
4478 let mut dot = 0f32;
4479 for g in 0..ng {
4480 let d = dot_i8_i8(
4481 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4482 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4483 ) as f32
4484 * act.sx;
4485 dot += d * gscale(r, g);
4486 }
4487 for &(j, xv) in &act.outliers {
4488 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
4489 }
4490 unsafe { *out_addr.at(bi * rows + r) = dot };
4492 bi += 1;
4493 }
4494 });
4495 }
4496 };
4497 dispatch_rows(pool, rows, &run);
4498 return;
4499 }
4500
4501 let out_addr = SendMut(out.as_mut_ptr());
4502 let run = move |start: usize, end: usize| {
4503 ROW_F32.with(|rb| {
4504 let mut buf = rb.borrow_mut();
4505 buf.resize(cols, 0.0);
4506 for r in start..end {
4507 decode_f32(r, &mut buf);
4508 for bi in 0..b {
4509 let x = &xs_all[bi * cols..(bi + 1) * cols];
4510 let mut dot = 0f32;
4511 for g in 0..ng {
4512 let mut gd = 0f32;
4513 for k in 0..GROUP_SIZE {
4514 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
4515 }
4516 dot += gd * gscale(r, g);
4517 }
4518 unsafe { *out_addr.at(bi * rows + r) = dot };
4520 }
4521 }
4522 })
4523 };
4524 dispatch_rows(pool, rows, &run);
4525}
4526
4527pub(crate) fn gpu_batch_job<'a>(
4531 t: &'a QTensor,
4532 x: &[f32],
4533) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
4534 match t {
4535 QTensor::Mapped {
4536 model,
4537 idx,
4538 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
4539 rows,
4540 cols,
4541 row_scale,
4542 col_field,
4543 ..
4544 } => Some((
4545 model.clone(),
4546 crate::gpu::BatchJob {
4547 idx: *idx,
4548 rows: *rows,
4549 cols: *cols,
4550 row_scale,
4551 xs: prescale(x, col_field, *dt).into_owned(),
4552 q1: false,
4553 },
4554 )),
4555 QTensor::Mapped {
4557 model,
4558 idx,
4559 dtype: TensorDtype::Q1,
4560 rows,
4561 cols,
4562 ..
4563 } => Some((
4564 model.clone(),
4565 crate::gpu::BatchJob {
4566 idx: *idx,
4567 rows: *rows,
4568 cols: *cols,
4569 row_scale: &[],
4570 xs: x.to_vec(),
4571 q1: true,
4572 },
4573 )),
4574 _ => None,
4575 }
4576}
4577
4578pub(crate) fn prescale<'a>(
4582 x: &'a [f32],
4583 col_field: &[f32],
4584 dtype: TensorDtype,
4585) -> std::borrow::Cow<'a, [f32]> {
4586 if dtype == TensorDtype::Q8_2f {
4587 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
4588 } else {
4589 std::borrow::Cow::Borrowed(x)
4590 }
4591}
4592
4593#[cfg(target_arch = "x86_64")]
4598pub(crate) fn avx2_enabled() -> bool {
4599 use std::sync::OnceLock;
4600 static ON: OnceLock<bool> = OnceLock::new();
4601 *ON.get_or_init(|| {
4602 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
4603 && std::arch::is_x86_feature_detected!("avx2")
4604 && std::arch::is_x86_feature_detected!("fma")
4605 })
4606}
4607
4608#[cfg(target_arch = "x86_64")]
4613fn avx2_a8w8_enabled() -> bool {
4614 use std::sync::OnceLock;
4615 static ON: OnceLock<bool> = OnceLock::new();
4616 *ON.get_or_init(|| {
4617 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
4618 })
4619}
4620
4621#[inline]
4625pub(crate) fn a8w8_enabled() -> bool {
4626 #[cfg(target_arch = "aarch64")]
4627 {
4628 sdot_enabled()
4629 }
4630 #[cfg(target_arch = "x86_64")]
4631 {
4632 avx2_a8w8_enabled()
4633 }
4634 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
4635 {
4636 false
4637 }
4638}
4639
4640#[inline]
4643#[allow(unreachable_code)]
4644fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
4645 #[cfg(target_arch = "aarch64")]
4646 unsafe {
4647 return dot_i8_sdot(w, xq);
4648 }
4649 #[cfg(target_arch = "x86_64")]
4650 unsafe {
4651 if avx512vnni_enabled() {
4652 return dot_i8_i8_vnni(w, xq);
4653 }
4654 return dot_i8_i8_avx2(w, xq);
4655 }
4656 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
4657}
4658
4659#[cfg(target_arch = "x86_64")]
4663fn avx512vnni_enabled() -> bool {
4664 use std::sync::OnceLock;
4665 static ON: OnceLock<bool> = OnceLock::new();
4666 *ON.get_or_init(|| {
4667 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
4668 && std::arch::is_x86_feature_detected!("avx512f")
4669 && std::arch::is_x86_feature_detected!("avx512bw")
4670 && std::arch::is_x86_feature_detected!("avx512vl")
4671 && std::arch::is_x86_feature_detected!("avx512vnni")
4672 })
4673}
4674
4675#[cfg(target_arch = "x86_64")]
4680#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4681unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
4682 unsafe {
4684 use core::arch::x86_64::*;
4685 let n = w.len();
4686 let mut j = 0usize;
4687 let mut total: i32;
4688 {
4693 #[inline(always)]
4694 unsafe fn step(
4695 w: *const u8,
4696 x: *const i8,
4697 acc: core::arch::x86_64::__m512i,
4698 ) -> core::arch::x86_64::__m512i {
4699 unsafe {
4700 use core::arch::x86_64::*;
4701 let wv = _mm512_loadu_si512(w as *const _);
4702 let xv = _mm512_loadu_si512(x as *const _);
4703 let aw = _mm512_abs_epi8(wv);
4704 let neg = _mm512_movepi8_mask(wv);
4705 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
4706 _mm512_dpbusd_epi32(acc, aw, sx)
4707 }
4708 }
4709 let (mut a0, mut a1, mut a2, mut a3) = (
4710 _mm512_setzero_si512(),
4711 _mm512_setzero_si512(),
4712 _mm512_setzero_si512(),
4713 _mm512_setzero_si512(),
4714 );
4715 while j + 256 <= n {
4716 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4717 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
4718 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
4719 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
4720 j += 256;
4721 }
4722 while j + 64 <= n {
4723 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4724 j += 64;
4725 }
4726 let s01 = _mm512_add_epi32(a0, a1);
4727 let s23 = _mm512_add_epi32(a2, a3);
4728 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
4729 }
4730 if j + 32 <= n {
4732 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4733 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4734 let d = _mm256_dpbusd_epi32(
4735 _mm256_setzero_si256(),
4736 _mm256_abs_epi8(wv),
4737 _mm256_sign_epi8(xv, wv),
4738 );
4739 let hi128 = _mm256_extracti128_si256::<1>(d);
4740 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4741 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4742 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4743 total += _mm_cvtsi128_si32(s32);
4744 j += 32;
4745 }
4746 while j < n {
4747 total += (w[j] as i8) as i32 * xq[j] as i32;
4748 j += 1;
4749 }
4750 total
4751 }
4752}
4753
4754#[cfg(target_arch = "x86_64")]
4756#[target_feature(enable = "avx2,fma")]
4757unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
4758 unsafe {
4760 use core::arch::x86_64::*;
4761 let n = x.len();
4762 let wp = w.as_ptr();
4763 let xp = x.as_ptr();
4764 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
4765 let mut j = 0usize;
4766 while j + 16 <= n {
4767 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
4768 let lo = _mm256_cvtepi8_epi32(wb);
4769 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
4770 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
4771 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
4772 j += 16;
4773 }
4774 let acc = _mm256_add_ps(a0, a1);
4775 let hi128 = _mm256_extractf128_ps::<1>(acc);
4776 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
4777 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
4778 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
4779 let mut sum = _mm_cvtss_f32(s32);
4780 while j < n {
4781 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
4782 j += 1;
4783 }
4784 sum
4785 }
4786}
4787
4788#[cfg(target_arch = "x86_64")]
4793#[target_feature(enable = "avx2")]
4794unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
4795 unsafe {
4797 use core::arch::x86_64::*;
4798 let n = w.len();
4799 let ones = _mm256_set1_epi16(1);
4800 let mut acc = _mm256_setzero_si256();
4801 let mut j = 0usize;
4802 while j + 32 <= n {
4803 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4804 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4805 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4806 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
4807 j += 32;
4808 }
4809 let hi128 = _mm256_extracti128_si256::<1>(acc);
4810 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
4811 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4812 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4813 let mut s = _mm_cvtsi128_si32(s32);
4814 while j < n {
4815 s += (w[j] as i8) as i32 * xq[j] as i32;
4816 j += 1;
4817 }
4818 s
4819 }
4820}
4821
4822#[cfg(target_arch = "aarch64")]
4826#[target_feature(enable = "neon,i8mm")]
4827unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4828 unsafe {
4830 use core::arch::aarch64::*;
4831 use core::arch::asm;
4832 let n = w0.len();
4833 let w0p = w0.as_ptr() as *const i8;
4834 let w1p = w1.as_ptr() as *const i8;
4835 let mut acc01 = vdupq_n_s32(0);
4838 let mut acc23 = vdupq_n_s32(0);
4839 let mut i = 0usize;
4840 while i + 8 <= n {
4841 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
4842 let xb01 =
4843 vcombine_s8(vld1_s8(xs[0].as_ptr().add(i)), vld1_s8(xs[1].as_ptr().add(i)));
4844 let xb23 =
4845 vcombine_s8(vld1_s8(xs[2].as_ptr().add(i)), vld1_s8(xs[3].as_ptr().add(i)));
4846 asm!(
4847 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
4848 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
4849 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
4850 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
4851 options(pure, nomem, nostack),
4852 );
4853 i += 8;
4854 }
4855 let mut out = [[0i32; 4]; 2];
4856 let a01: [i32; 4] = core::mem::transmute(acc01);
4857 let a23: [i32; 4] = core::mem::transmute(acc23);
4858 out[0][0] = a01[0];
4859 out[0][1] = a01[1];
4860 out[1][0] = a01[2];
4861 out[1][1] = a01[3];
4862 out[0][2] = a23[0];
4863 out[0][3] = a23[1];
4864 out[1][2] = a23[2];
4865 out[1][3] = a23[3];
4866 if i < n {
4867 for (k, x) in xs.iter().enumerate() {
4868 for j in i..n {
4869 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4870 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4871 }
4872 }
4873 }
4874 out
4875 }
4876}
4877
4878#[cfg(target_arch = "aarch64")]
4882#[target_feature(enable = "neon,dotprod")]
4883unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4884 unsafe {
4886 use core::arch::aarch64::*;
4887 use core::arch::asm;
4888 let n = w0.len();
4889 let w0p = w0.as_ptr() as *const i8;
4890 let w1p = w1.as_ptr() as *const i8;
4891 let mut acc = [[vdupq_n_s32(0); 4]; 2];
4892 let mut i = 0usize;
4893 while i + 16 <= n {
4894 let wv0 = vld1q_s8(w0p.add(i));
4895 let wv1 = vld1q_s8(w1p.add(i));
4896 for (k, x) in xs.iter().enumerate() {
4897 let xv = vld1q_s8(x.as_ptr().add(i));
4898 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
4899 asm!(
4900 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
4901 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
4902 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4903 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
4904 options(pure, nomem, nostack),
4905 );
4906 acc[0][k] = a0;
4907 acc[1][k] = a1;
4908 }
4909 i += 16;
4910 }
4911 let mut out = [[0i32; 4]; 2];
4912 for r in 0..2 {
4913 for k in 0..4 {
4914 out[r][k] = vaddvq_s32(acc[r][k]);
4915 }
4916 }
4917 if i < n {
4918 for (k, x) in xs.iter().enumerate() {
4919 for j in i..n {
4920 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4921 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4922 }
4923 }
4924 }
4925 out
4926 }
4927}
4928
4929#[cfg(target_arch = "x86_64")]
4935#[target_feature(enable = "avx2")]
4936unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4937 unsafe {
4939 use core::arch::x86_64::*;
4940 let n = w0.len();
4941 let ones = _mm256_set1_epi16(1);
4942 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
4943 let mut j = 0usize;
4944 while j + 32 <= n {
4945 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
4946 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
4947 let aw0 = _mm256_abs_epi8(wv0);
4948 let aw1 = _mm256_abs_epi8(wv1);
4949 for (k, x) in xs.iter().enumerate() {
4950 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
4951 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
4952 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
4953 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
4954 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
4955 }
4956 j += 32;
4957 }
4958 let mut out = [[0i32; 4]; 2];
4959 for r in 0..2 {
4960 for k in 0..4 {
4961 let a = acc[r][k];
4962 let hi128 = _mm256_extracti128_si256::<1>(a);
4963 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
4964 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4965 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4966 out[r][k] = _mm_cvtsi128_si32(s32);
4967 }
4968 }
4969 if j < n {
4970 for (k, x) in xs.iter().enumerate() {
4971 for i in j..n {
4972 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
4973 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
4974 }
4975 }
4976 }
4977 out
4978 }
4979}
4980
4981#[cfg(target_arch = "x86_64")]
4986#[inline]
4987fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
4988 let dot = if avx512vnni_enabled() && row.len() >= 64 {
4989 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
4990 } else {
4991 unsafe { dot_i8_i8_avx2(row, &act.xq) }
4992 };
4993 let mut acc = dot as f32 * act.sx;
4994 for &(j, xv) in &act.outliers {
4995 acc += (row[j] as i8) as f32 * xv;
4996 }
4997 acc
4998}
4999
5000#[cfg(target_arch = "x86_64")]
5004#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5005unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
5006 unsafe {
5008 use core::arch::x86_64::*;
5009 let n = w.len();
5010 let flip = _mm512_set1_epi8(-128); #[inline(always)]
5012 unsafe fn step(
5013 w: *const u8,
5014 x: *const i8,
5015 flip: core::arch::x86_64::__m512i,
5016 acc: core::arch::x86_64::__m512i,
5017 ) -> core::arch::x86_64::__m512i {
5018 unsafe {
5019 use core::arch::x86_64::*;
5020 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
5021 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
5022 }
5023 }
5024 let (mut a0, mut a1, mut a2, mut a3) = (
5025 _mm512_setzero_si512(),
5026 _mm512_setzero_si512(),
5027 _mm512_setzero_si512(),
5028 _mm512_setzero_si512(),
5029 );
5030 let mut j = 0usize;
5031 while j + 256 <= n {
5032 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5033 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
5034 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
5035 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
5036 j += 256;
5037 }
5038 while j + 64 <= n {
5039 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5040 j += 64;
5041 }
5042 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
5043 _mm512_add_epi32(a0, a1),
5044 _mm512_add_epi32(a2, a3),
5045 ));
5046 while j < n {
5048 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
5049 j += 1;
5050 }
5051 total
5052 }
5053}
5054
5055#[cfg(target_arch = "x86_64")]
5061#[target_feature(enable = "avx2")]
5062unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5063 unsafe {
5066 use core::arch::x86_64::*;
5067 let lomask = _mm_set1_epi8(0x0F);
5068 let eight = _mm256_set1_epi8(8);
5069 let ones = _mm256_set1_epi16(1);
5070 let mut acc = 0f32;
5071 for gi in 0..gpr {
5072 let g = g0 + gi;
5073 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5074 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5075 let lo = _mm_and_si128(b, lomask);
5076 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5077 let w = _mm256_sub_epi8(
5078 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5079 eight,
5080 );
5081 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5082 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5083 let d = _mm256_madd_epi16(p16, ones);
5084 let hi128 = _mm256_extracti128_si256::<1>(d);
5085 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5086 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5087 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5088 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5089 }
5090 acc
5091 }
5092}
5093
5094#[cfg(target_arch = "x86_64")]
5097#[target_feature(enable = "avx2")]
5098unsafe fn dot_q4_row_avx2_2(
5099 packed: &[u8],
5100 scales: &[u8],
5101 g0: usize,
5102 gpr: usize,
5103 xq1: &[i8],
5104 xq2: &[i8],
5105) -> (f32, f32) {
5106 unsafe {
5108 use core::arch::x86_64::*;
5109 let lomask = _mm_set1_epi8(0x0F);
5110 let eight = _mm256_set1_epi8(8);
5111 let ones = _mm256_set1_epi16(1);
5112 let (mut acc1, mut acc2) = (0f32, 0f32);
5113 #[inline(always)]
5114 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
5115 unsafe {
5116 use core::arch::x86_64::*;
5117 let hi128 = _mm256_extracti128_si256::<1>(d);
5118 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5119 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5120 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5121 _mm_cvtsi128_si32(s32)
5122 }
5123 }
5124 for gi in 0..gpr {
5125 let g = g0 + gi;
5126 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5127 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5128 let lo = _mm_and_si128(b, lomask);
5129 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5130 let w = _mm256_sub_epi8(
5131 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5132 eight,
5133 );
5134 let aw = _mm256_abs_epi8(w);
5135 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5136 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5137 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
5138 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
5139 acc1 += hsum(d1) as f32 * s;
5140 acc2 += hsum(d2) as f32 * s;
5141 }
5142 (acc1, acc2)
5143 }
5144}
5145
5146#[cfg(target_arch = "x86_64")]
5148fn q8_range_avx2(
5149 q: &[u8],
5150 row_scale: &[f32],
5151 act: &SplitAct,
5152 cols: usize,
5153 out_addr: SendMut,
5154 start: usize,
5155 end: usize,
5156) {
5157 for o in start..end {
5158 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5159 unsafe { *out_addr.at(o) = v };
5161 }
5162}
5163
5164#[cfg(target_arch = "x86_64")]
5166#[allow(clippy::too_many_arguments)]
5167fn q8_range2_avx2(
5168 q: &[u8],
5169 row_scale: &[f32],
5170 a1: &SplitAct,
5171 a2: &SplitAct,
5172 cols: usize,
5173 p1: SendMut,
5174 p2: SendMut,
5175 start: usize,
5176 end: usize,
5177) {
5178 for o in start..end {
5179 let row = &q[o * cols..(o + 1) * cols];
5180 unsafe {
5182 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
5183 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
5184 }
5185 }
5186}
5187
5188#[cfg(target_arch = "aarch64")]
5199fn i8mm_enabled() -> bool {
5200 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5201 *ON.get_or_init(|| {
5202 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
5203 && std::arch::is_aarch64_feature_detected!("i8mm")
5204 })
5205}
5206
5207#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
5211fn sdot_enabled() -> bool {
5212 use std::sync::OnceLock;
5213 static ON: OnceLock<bool> = OnceLock::new();
5214 *ON.get_or_init(|| {
5215 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
5216 #[cfg(target_arch = "aarch64")]
5217 {
5218 want && std::arch::is_aarch64_feature_detected!("dotprod")
5219 }
5220 #[cfg(not(target_arch = "aarch64"))]
5221 {
5222 let _ = want;
5223 false
5224 }
5225 })
5226}
5227
5228struct SplitAct {
5233 xq: Vec<i8>,
5234 sx: f32,
5235 outliers: Vec<(usize, f32)>,
5236 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
5239 xsum: i32,
5240}
5241
5242thread_local! {
5243 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
5246 const { std::cell::RefCell::new(Vec::new()) };
5247}
5248
5249impl Drop for SplitAct {
5250 fn drop(&mut self) {
5251 let buf = std::mem::take(&mut self.xq);
5252 if buf.capacity() > 0 {
5253 XQ_FREE.with(|f| {
5254 let mut f = f.borrow_mut();
5255 if f.len() < 16 {
5256 f.push(buf);
5257 }
5258 });
5259 }
5260 }
5261}
5262
5263fn split_act(x: &[f32]) -> SplitAct {
5264 let n = x.len();
5265 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
5266 let thr = 8.0 * rms;
5267 let mut outliers: Vec<(usize, f32)> = Vec::new();
5271 let mut amax = 0f32;
5272 for (j, &v) in x.iter().enumerate() {
5273 let a = v.abs();
5274 if a > thr {
5275 outliers.push((j, v));
5276 } else if a > amax {
5277 amax = a;
5278 }
5279 }
5280 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
5281 let inv = 1.0 / sx;
5282 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
5283 xq.clear();
5284 xq.reserve(n);
5285 if outliers.is_empty() {
5286 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
5287 } else {
5288 xq.extend(x.iter().map(|&v| {
5290 if v.abs() > thr {
5291 0
5292 } else {
5293 (v * inv).round().clamp(-127.0, 127.0) as i8
5294 }
5295 }));
5296 }
5297 let xsum = xq.iter().map(|&v| v as i32).sum();
5298 SplitAct { xq, sx, outliers, xsum }
5299}
5300
5301#[cfg(target_arch = "aarch64")]
5304#[target_feature(enable = "neon,dotprod")]
5305unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
5306 unsafe {
5308 use core::arch::aarch64::*;
5309 use core::arch::asm;
5310 let wp = w.as_ptr() as *const i8;
5311 let n = w.len();
5312 let (mut a0, mut a1, mut a2, mut a3) =
5313 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5314 let mut i = 0;
5315 while i + 64 <= n {
5316 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
5317 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
5318 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
5319 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
5320 asm!(
5321 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5322 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5323 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
5324 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
5325 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5326 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5327 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
5328 options(pure, nomem, nostack),
5329 );
5330 i += 64;
5331 }
5332 while i + 16 <= n {
5333 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
5334 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
5335 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
5336 i += 16;
5337 }
5338 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
5339 while i < n {
5340 s += (*wp.add(i)) as i32 * xq[i] as i32;
5341 i += 1;
5342 }
5343 s
5344}
5345}
5346
5347#[cfg(target_arch = "aarch64")]
5351#[target_feature(enable = "neon,dotprod")]
5352unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
5353 unsafe {
5355 use core::arch::aarch64::*;
5356 use core::arch::asm;
5357 let n = xq.len();
5358 let px = xq.as_ptr();
5359 let (p0, p1, p2, p3) = (
5360 w0.as_ptr() as *const i8,
5361 w1.as_ptr() as *const i8,
5362 w2.as_ptr() as *const i8,
5363 w3.as_ptr() as *const i8,
5364 );
5365 let (mut a0, mut a1, mut a2, mut a3) =
5366 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5367 let mut i = 0;
5368 while i + 16 <= n {
5369 let x = vld1q_s8(px.add(i));
5370 let v0 = vld1q_s8(p0.add(i));
5371 let v1 = vld1q_s8(p1.add(i));
5372 let v2 = vld1q_s8(p2.add(i));
5373 let v3 = vld1q_s8(p3.add(i));
5374 asm!(
5375 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
5376 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
5377 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
5378 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
5379 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5380 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
5381 options(pure, nomem, nostack),
5382 );
5383 i += 16;
5384 }
5385 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
5386 while i < n {
5387 let xi = *px.add(i) as i32;
5388 r[0] += (*p0.add(i)) as i32 * xi;
5389 r[1] += (*p1.add(i)) as i32 * xi;
5390 r[2] += (*p2.add(i)) as i32 * xi;
5391 r[3] += (*p3.add(i)) as i32 * xi;
5392 i += 1;
5393 }
5394 r
5395}
5396}
5397
5398#[cfg(target_arch = "aarch64")]
5405#[target_feature(enable = "neon,dotprod")]
5406unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
5407 unsafe {
5410 use core::arch::aarch64::*;
5411 use core::arch::asm;
5412 let n = xq.len();
5413 let px = xq.as_ptr();
5414 let pg = g.as_ptr() as *const i8;
5415 let (mut a0, mut a1, mut a2, mut a3) =
5416 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5417 let mut i = 0;
5418 while i + 16 <= n {
5419 let x = vld1q_s8(px.add(i));
5420 let base = pg.add(4 * i);
5421 let v0 = vld1q_s8(base);
5422 let v1 = vld1q_s8(base.add(16));
5423 let v2 = vld1q_s8(base.add(32));
5424 let v3 = vld1q_s8(base.add(48));
5425 asm!(
5426 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
5427 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
5428 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
5429 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
5430 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5431 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
5432 options(pure, nomem, nostack),
5433 );
5434 i += 16;
5435 }
5436 [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
5437 }
5438}
5439
5440#[cfg(target_arch = "aarch64")]
5446fn q8_range_sdot(
5447 q: &[u8],
5448 rep: &[u8],
5449 row_scale: &[f32],
5450 act: &SplitAct,
5451 cols: usize,
5452 out_addr: SendMut,
5453 start: usize,
5454 end: usize,
5455) {
5456 let mut o = start;
5457 if !rep.is_empty() {
5460 while o < end && o % 4 != 0 {
5461 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5462 unsafe { *out_addr.at(o) = v };
5463 o += 1;
5464 }
5465 }
5466 while o + 4 <= end {
5467 let r = if rep.is_empty() {
5468 unsafe {
5469 dot_i8_sdot_4rows(
5470 &q[o * cols..(o + 1) * cols],
5471 &q[(o + 1) * cols..(o + 2) * cols],
5472 &q[(o + 2) * cols..(o + 3) * cols],
5473 &q[(o + 3) * cols..(o + 4) * cols],
5474 &act.xq,
5475 )
5476 }
5477 } else {
5478 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
5479 };
5480 for k in 0..4 {
5481 let mut acc = r[k] as f32 * act.sx;
5482 for &(j, xv) in &act.outliers {
5483 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
5484 }
5485 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
5487 }
5488 o += 4;
5489 }
5490 while o < end {
5491 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5492 unsafe { *out_addr.at(o) = v };
5493 o += 1;
5494 }
5495}
5496
5497#[cfg(target_arch = "aarch64")]
5500#[allow(clippy::too_many_arguments)]
5501fn q8_range2_sdot(
5502 q: &[u8],
5503 row_scale: &[f32],
5504 a1: &SplitAct,
5505 a2: &SplitAct,
5506 cols: usize,
5507 p1: SendMut,
5508 p2: SendMut,
5509 start: usize,
5510 end: usize,
5511) {
5512 for o in start..end {
5513 let row = &q[o * cols..(o + 1) * cols];
5514 unsafe {
5516 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
5517 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
5518 }
5519 }
5520}
5521
5522#[allow(clippy::too_many_arguments)]
5524fn q8_range2_f32(
5525 q: &[u8],
5526 row_scale: &[f32],
5527 x1: &[f32],
5528 x2: &[f32],
5529 cols: usize,
5530 p1: SendMut,
5531 p2: SendMut,
5532 start: usize,
5533 end: usize,
5534) {
5535 for o in start..end {
5536 let row = &q[o * cols..(o + 1) * cols];
5537 unsafe {
5539 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
5540 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
5541 }
5542 }
5543}
5544
5545fn q8_range_f32(
5547 q: &[u8],
5548 row_scale: &[f32],
5549 xs: &[f32],
5550 cols: usize,
5551 out_addr: SendMut,
5552 start: usize,
5553 end: usize,
5554) {
5555 for o in start..end {
5556 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
5557 unsafe { *out_addr.at(o) = v };
5559 }
5560}
5561
5562#[cfg(target_arch = "aarch64")]
5565#[inline]
5566fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
5567 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
5568 for &(j, xv) in &act.outliers {
5569 acc += (row[j] as i8) as f32 * xv;
5570 }
5571 acc
5572}
5573
5574#[cfg(target_arch = "aarch64")]
5582#[target_feature(enable = "neon,dotprod")]
5583unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5584 unsafe {
5587 use core::arch::aarch64::*;
5588 use core::arch::asm;
5589 let lomask = vdupq_n_u8(0x0F);
5590 let eight = vdupq_n_s8(8);
5591 let mut acc = 0f32;
5592 for gi in 0..gpr {
5593 let g = g0 + gi;
5594 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5595 let b = vld1q_u8(packed.as_ptr().add(g * 16));
5596 let lo = vandq_u8(b, lomask);
5597 let hi = vshrq_n_u8::<4>(b);
5598 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5599 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5600 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5601 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5602 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5603 asm!(
5604 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5605 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5606 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5607 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5608 options(pure, nomem, nostack),
5609 );
5610 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5611 }
5612 acc
5613 }
5614}
5615
5616#[cfg(target_arch = "aarch64")]
5621#[target_feature(enable = "neon,dotprod")]
5622unsafe fn dot_q4_row_sdot2(
5623 packed: &[u8],
5624 scales: &[u8],
5625 g0: usize,
5626 gpr: usize,
5627 xq1: &[i8],
5628 xq2: &[i8],
5629) -> (f32, f32) {
5630 unsafe {
5633 use core::arch::aarch64::*;
5634 use core::arch::asm;
5635 let lomask = vdupq_n_u8(0x0F);
5636 let eight = vdupq_n_s8(8);
5637 let (mut acc1, mut acc2) = (0f32, 0f32);
5638 for gi in 0..gpr {
5639 let g = g0 + gi;
5640 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5641 let b = vld1q_u8(packed.as_ptr().add(g * 16));
5642 let lo = vandq_u8(b, lomask);
5643 let hi = vshrq_n_u8::<4>(b);
5644 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5645 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5646 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
5647 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
5648 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
5649 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
5650 let (mut a0, mut a1, mut b0, mut b1) =
5651 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5652 asm!(
5653 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
5654 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
5655 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
5656 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
5657 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5658 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
5659 e0 = in(vreg) e0, e1 = in(vreg) e1,
5660 x10 = in(vreg) x10, x11 = in(vreg) x11,
5661 x20 = in(vreg) x20, x21 = in(vreg) x21,
5662 options(pure, nomem, nostack),
5663 );
5664 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5665 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
5666 }
5667 (acc1, acc2)
5668 }
5669}
5670
5671#[inline]
5676pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
5677 #[cfg(target_arch = "aarch64")]
5678 unsafe {
5679 return axpy_i8_f32_neon(acc, row, w);
5680 }
5681 #[cfg(target_arch = "x86_64")]
5682 if avx2_enabled() {
5683 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
5684 }
5685 #[allow(unreachable_code)]
5686 {
5687 for (a, &b) in acc.iter_mut().zip(row) {
5688 *a += w * b as f32;
5689 }
5690 }
5691}
5692
5693#[cfg(target_arch = "x86_64")]
5695#[target_feature(enable = "avx2,fma")]
5696unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
5697 unsafe {
5699 use core::arch::x86_64::*;
5700 let n = acc.len().min(row.len());
5701 let ap = acc.as_mut_ptr();
5702 let rp = row.as_ptr();
5703 let wv = _mm256_set1_ps(w);
5704 let mut j = 0usize;
5705 while j + 16 <= n {
5706 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
5707 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
5708 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
5709 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
5710 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
5711 _mm256_storeu_ps(ap.add(j), v0);
5712 _mm256_storeu_ps(ap.add(j + 8), v1);
5713 j += 16;
5714 }
5715 while j < n {
5716 *ap.add(j) += w * (*rp.add(j)) as f32;
5717 j += 1;
5718 }
5719 }
5720}
5721
5722#[cfg(target_arch = "aarch64")]
5723#[target_feature(enable = "neon")]
5724unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
5725 unsafe {
5727 use core::arch::aarch64::*;
5728 let n = acc.len().min(row.len());
5729 let ap = acc.as_mut_ptr();
5730 let rp = row.as_ptr();
5731 let wv = vdupq_n_f32(w);
5732 let mut j = 0usize;
5733 while j + 16 <= n {
5734 let rb = vld1q_s8(rp.add(j));
5735 let lo = vmovl_s8(vget_low_s8(rb));
5736 let hi = vmovl_s8(vget_high_s8(rb));
5737 for (off, half) in [(0, lo), (8, hi)] {
5738 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
5739 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
5740 let o = j + off;
5741 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
5742 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
5743 }
5744 j += 16;
5745 }
5746 while j < n {
5747 *ap.add(j) += w * (*rp.add(j)) as f32;
5748 j += 1;
5749 }
5750}
5751}
5752
5753#[inline]
5756pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
5757 #[cfg(target_arch = "aarch64")]
5758 unsafe {
5759 return dot_i8_f32_neon(w, x);
5760 }
5761 #[cfg(target_arch = "x86_64")]
5762 if avx2_enabled() {
5763 return unsafe { dot_i8_f32_avx2(w, x) };
5764 }
5765 #[allow(unreachable_code)]
5766 {
5767 let mut sum = 0.0f32;
5768 for (j, &b) in w.iter().enumerate() {
5769 sum += (b as i8) as f32 * x[j];
5770 }
5771 sum
5772 }
5773}
5774
5775#[inline]
5779fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5780 #[cfg(target_arch = "aarch64")]
5781 unsafe {
5782 return dot_i8_col_f32_neon(w, x, col);
5783 }
5784 #[allow(unreachable_code)]
5785 {
5786 let mut sum = 0.0f32;
5787 for (j, &b) in w.iter().enumerate() {
5788 sum += (b as i8) as f32 * x[j] * col[j];
5789 }
5790 sum
5791 }
5792}
5793
5794#[cfg(target_arch = "aarch64")]
5795#[target_feature(enable = "neon")]
5796unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5797 unsafe {
5799 use core::arch::aarch64::*;
5800 let n = x.len();
5801 let wp = w.as_ptr() as *const i8;
5802 let xp = x.as_ptr();
5803 let cp = col.as_ptr();
5804 let (mut a0, mut a1, mut a2, mut a3) =
5805 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5806 let mut j = 0usize;
5807 while j + 16 <= n {
5808 let wb = vld1q_s8(wp.add(j));
5809 let lo = vmovl_s8(vget_low_s8(wb));
5810 let hi = vmovl_s8(vget_high_s8(wb));
5811 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5812 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5813 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5814 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5815 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
5816 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
5817 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
5818 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
5819 j += 16;
5820 }
5821 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5822 while j < n {
5823 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
5824 j += 1;
5825 }
5826 sum
5827}
5828}
5829
5830#[cfg(target_arch = "aarch64")]
5831#[target_feature(enable = "neon")]
5832unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
5833 unsafe {
5835 use core::arch::aarch64::*;
5836 let n = x.len();
5837 let wp = w.as_ptr() as *const i8;
5838 let xp = x.as_ptr();
5839 let (mut a0, mut a1, mut a2, mut a3) =
5840 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5841 let mut j = 0usize;
5842 while j + 16 <= n {
5843 let wb = vld1q_s8(wp.add(j));
5844 let lo = vmovl_s8(vget_low_s8(wb));
5845 let hi = vmovl_s8(vget_high_s8(wb));
5846 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5847 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5848 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5849 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5850 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
5851 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
5852 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
5853 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
5854 j += 16;
5855 }
5856 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5857 while j < n {
5858 sum += (*wp.add(j)) as f32 * *xp.add(j);
5859 j += 1;
5860 }
5861 sum
5862}
5863}
5864
5865#[allow(clippy::too_many_arguments)]
5866fn qmatvec(
5867 q: &[u8],
5868 rep: &[u8],
5869 row_scale: &[f32],
5870 xs: &[f32],
5871 rows: usize,
5872 cols: usize,
5873 out: &mut [f32],
5874 pool: Option<&Pool>,
5875) {
5876 debug_assert_eq!(out.len(), rows);
5877 #[cfg(not(target_arch = "aarch64"))]
5878 let _ = rep;
5879
5880 #[cfg(target_arch = "aarch64")]
5881 if sdot_enabled() {
5882 let act = split_act(xs);
5883 let out_addr = SendMut(out.as_mut_ptr());
5884 let run_range = |start: usize, end: usize| {
5885 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
5886 };
5887 match pool {
5888 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5889 _ => run_range(0, rows),
5890 }
5891 return;
5892 }
5893 #[cfg(target_arch = "x86_64")]
5896 if avx2_a8w8_enabled() {
5897 let act = split_act(xs);
5898 let out_addr = SendMut(out.as_mut_ptr());
5899 let run_range =
5900 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
5901 match pool {
5902 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5903 _ => run_range(0, rows),
5904 }
5905 return;
5906 }
5907
5908 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
5909 match pool {
5910 Some(pool) if rows >= 256 => {
5911 let out_addr = SendMut(out.as_mut_ptr());
5912 pool.run(&move |widx, n| {
5913 let chunk = rows.div_ceil(n);
5914 let start = widx * chunk;
5915 let end = (start + chunk).min(rows);
5916 for o in start..end {
5917 unsafe { *out_addr.at(o) = row_dot(o) };
5919 }
5920 });
5921 }
5922 _ => {
5923 for (o, dst) in out.iter_mut().enumerate() {
5924 *dst = row_dot(o);
5925 }
5926 }
5927 }
5928}
5929
5930#[allow(clippy::too_many_arguments)]
5931fn qmatvec2(
5932 q: &[u8],
5933 row_scale: &[f32],
5934 x1: &[f32],
5935 x2: &[f32],
5936 rows: usize,
5937 cols: usize,
5938 o1: &mut [f32],
5939 o2: &mut [f32],
5940 pool: Option<&Pool>,
5941) {
5942 #[cfg(target_arch = "aarch64")]
5943 if sdot_enabled() {
5944 let a1s = split_act(x1);
5945 let a2s = split_act(x2);
5946 let p1 = SendMut(o1.as_mut_ptr());
5947 let p2 = SendMut(o2.as_mut_ptr());
5948 let run_range = |start: usize, end: usize| {
5949 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
5950 };
5951 match pool {
5952 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5953 _ => run_range(0, rows),
5954 }
5955 return;
5956 }
5957 #[cfg(target_arch = "x86_64")]
5958 if avx2_a8w8_enabled() {
5959 let a1s = split_act(x1);
5960 let a2s = split_act(x2);
5961 let p1 = SendMut(o1.as_mut_ptr());
5962 let p2 = SendMut(o2.as_mut_ptr());
5963 let run_range = |start: usize, end: usize| {
5964 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
5965 };
5966 match pool {
5967 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5968 _ => run_range(0, rows),
5969 }
5970 return;
5971 }
5972
5973 let row_dots = |o: usize| -> (f32, f32) {
5974 let row = &q[o * cols..(o + 1) * cols];
5975 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
5976 };
5977 match pool {
5978 Some(pool) if rows >= 256 => {
5979 let p1 = SendMut(o1.as_mut_ptr());
5980 let p2 = SendMut(o2.as_mut_ptr());
5981 pool.run(&move |widx, n| {
5982 let chunk = rows.div_ceil(n);
5983 let start = widx * chunk;
5984 let end = (start + chunk).min(rows);
5985 for o in start..end {
5986 let (s1, s2) = row_dots(o);
5987 unsafe {
5989 *p1.at(o) = s1;
5990 *p2.at(o) = s2;
5991 }
5992 }
5993 });
5994 }
5995 _ => {
5996 for o in 0..rows {
5997 let (s1, s2) = row_dots(o);
5998 o1[o] = s1;
5999 o2[o] = s2;
6000 }
6001 }
6002 }
6003}
6004
6005#[derive(Clone, Copy)]
6006struct SendMut(*mut f32);
6007unsafe impl Send for SendMut {}
6008unsafe impl Sync for SendMut {}
6009
6010impl SendMut {
6011 #[inline]
6012 fn at(self, i: usize) -> *mut f32 {
6013 unsafe { self.0.add(i) }
6014 }
6015}
6016
6017#[cfg(test)]
6018mod tests {
6019 use super::*;
6020
6021 #[test]
6022 fn f32_matvec_matches_matvec_rows_bitexact() {
6023 let (rows, cols) = (300, 40);
6024 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
6025 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
6026 let qt = QTensor::from_f32(w.clone(), rows, cols);
6027
6028 let mut a = vec![0.0f32; rows];
6029 matvec_rows(None, &w, &x, &mut a);
6030 let mut b = vec![0.0f32; rows];
6031 qt.matvec(&x, &mut b, None);
6032 assert_eq!(a, b);
6033 }
6034
6035 #[test]
6036 fn sdot_kernel_exact_on_grid() {
6037 eprintln!("sdot_enabled = {}", sdot_enabled());
6042 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
6044 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6045 .collect();
6046 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
6047 let x: Vec<f32> = (0..cols)
6048 .map(|i| match i % 3 {
6049 0 => 1.0,
6050 1 => -1.0,
6051 _ => 0.0,
6052 })
6053 .collect();
6054 let mut a = vec![0.0f32; rows];
6055 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
6056 for o in 0..rows {
6057 let mut acc = 0.0f32;
6058 for j in 0..cols {
6059 acc += (w[o * cols + j] as i8) as f32 * x[j];
6060 }
6061 let expect = acc * scales[o];
6062 assert!(
6063 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6064 "row {o}: {} vs {expect}",
6065 a[o]
6066 );
6067 }
6068 }
6069
6070 #[test]
6071 fn q1_tbl_fast_path_matches_reference() {
6072 let (rows, cols) = (5, 256);
6077 let gpr = cols / GROUP_SIZE;
6078 let mut bytes = Vec::new();
6079 for t in 0..rows * gpr {
6080 let s = 0.007 + (t % 11) as f32 * 0.004;
6081 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6082 for j in 0..4 {
6083 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
6084 }
6085 }
6086 let x: Vec<f32> = (0..cols)
6087 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
6088 .collect();
6089 let mut w = vec![0.0f32; rows * cols];
6090 cortiq_core::quant::dequant_q1(&bytes, &mut w);
6091 let mut got = vec![0.0f32; rows];
6092 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
6093 for o in 0..rows {
6094 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
6095 assert!(
6096 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6097 "row {o}: {} vs {expect}",
6098 got[o]
6099 );
6100 }
6101 let b = 5usize;
6104 let mut xs_all = Vec::new();
6105 for bi in 0..b {
6106 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
6107 }
6108 let mut mm = vec![0.0f32; b * rows];
6109 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
6110 for bi in 0..b {
6111 let mut single = vec![0.0f32; rows];
6112 q1_matvec(&bytes, &xs_all[bi * cols..(bi + 1) * cols], rows, cols, &mut single, None);
6113 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
6114 }
6115 }
6116
6117 #[test]
6118 fn q1_kernels_match_exact_reference() {
6119 let (rows, cols) = (7, 96);
6121 let gpr = cols / GROUP_SIZE;
6122 let mut bytes = Vec::new();
6123 for t in 0..rows * gpr {
6124 let s = 0.01 + (t % 13) as f32 * 0.003;
6125 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6126 for j in 0..4 {
6127 bytes.push(((t * 31 + j * 97) % 251) as u8);
6128 }
6129 }
6130 let x: Vec<f32> = (0..cols)
6132 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
6133 .collect();
6134 let mut w = vec![0.0f32; rows * cols];
6136 cortiq_core::quant::dequant_q1(&bytes, &mut w);
6137 let mut expect = vec![0.0f32; rows];
6138 for o in 0..rows {
6139 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
6140 }
6141 let mut got = vec![0.0f32; rows];
6142 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
6143 for o in 0..rows {
6144 assert!(
6145 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
6146 "row {o}: {} vs {}",
6147 got[o],
6148 expect[o]
6149 );
6150 }
6151 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
6153 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
6154 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
6155 assert_eq!(a1, got);
6156 let mut xs = x.clone();
6157 xs.extend_from_slice(&x2);
6158 let mut mm = vec![0.0f32; 2 * rows];
6159 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
6160 assert_eq!(&mm[..rows], got.as_slice());
6161 assert_eq!(&mm[rows..], a2.as_slice());
6162 }
6163
6164 #[test]
6165 fn repack_is_bit_identical() {
6166 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
6172 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
6173 .collect();
6174 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
6175 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
6176 let rep = q8_repack_layout(&w, rows, cols);
6177 for g in 0..rows / 4 {
6179 for c in 0..cols / 16 {
6180 for lane in 0..4 {
6181 assert_eq!(
6182 &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
6183 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
6184 );
6185 }
6186 }
6187 }
6188 let mut a = vec![0.0f32; rows];
6189 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
6190 let mut b = vec![0.0f32; rows];
6191 qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
6192 assert_eq!(a, b, "full-range repack output diverged");
6193
6194 #[cfg(target_arch = "aarch64")]
6195 if sdot_enabled() {
6196 let act = split_act(&x);
6198 let mut c1 = vec![0.0f32; rows];
6199 let mut c2 = vec![0.0f32; rows];
6200 q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
6201 q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
6202 assert_eq!(c1, c2, "unaligned-range repack output diverged");
6203 }
6204 }
6205
6206 #[test]
6207 fn sdot_a8w8_noise_is_bounded() {
6208 let (rows, cols) = (16, 512);
6212 let w: Vec<u8> = (0..rows * cols)
6213 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6214 .collect();
6215 let scales = vec![0.01f32; rows];
6216 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
6217 let mut a = vec![0.0f32; rows];
6218 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
6219 let (mut num, mut den) = (0f64, 0f64);
6220 for o in 0..rows {
6221 let mut acc = 0.0f32;
6222 for j in 0..cols {
6223 acc += (w[o * cols + j] as i8) as f32 * x[j];
6224 }
6225 let expect = acc * scales[o];
6226 num += ((a[o] - expect) as f64).powi(2);
6227 den += (expect as f64).powi(2);
6228 }
6229 let rel = (num / den.max(1e-12)).sqrt();
6230 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
6231 }
6232
6233 #[test]
6234 fn i8_dot_neon_matches_scalar() {
6235 let n = 100;
6236 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
6237 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
6238 let mut scalar = 0.0f32;
6239 for j in 0..n {
6240 scalar += (w[j] as i8) as f32 * x[j];
6241 }
6242 let fast = dot_i8_f32(&w, &x);
6243 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
6244 }
6245
6246 #[test]
6248 fn vbitmatvec_matches_full_dequant() {
6249 let (rows, cols) = (6, 64);
6250 let ng = cols / GROUP_SIZE;
6251 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
6253 let mut bytes = bits.clone();
6254 for g in 0..rows * ng {
6255 let s = 0.02 + 0.001 * g as f32;
6256 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6257 }
6258 for r in 0..rows {
6259 let b = bits[r] as usize;
6260 let (mut acc, mut nb) = (0u64, 0usize);
6261 let mut rowbytes = Vec::new();
6262 for i in 0..cols {
6263 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
6264 acc = (acc << b) | v;
6265 nb += b;
6266 while nb >= 8 {
6267 nb -= 8;
6268 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6269 }
6270 }
6271 if nb > 0 {
6272 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6273 }
6274 bytes.extend_from_slice(&rowbytes);
6275 }
6276 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
6277
6278 let mut reference = vec![0f32; rows * cols];
6279 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
6280 let mut expect = vec![0f32; rows];
6281 for r in 0..rows {
6282 expect[r] = reference[r * cols..(r + 1) * cols]
6283 .iter()
6284 .zip(&x)
6285 .map(|(w, xv)| w * xv)
6286 .sum();
6287 }
6288 let mut got = vec![0f32; rows];
6289 let offsets = vbit_row_offsets(&bytes, rows, cols);
6290 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
6291 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
6295 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
6296 for r in 0..rows {
6297 assert!(
6298 (got[r] - expect[r]).abs() < tol * scale,
6299 "row {r}: {} vs {}",
6300 got[r],
6301 expect[r]
6302 );
6303 }
6304 }
6305
6306 #[test]
6311 #[cfg(target_arch = "x86_64")]
6312 fn vbit_matmat_blocked_matches_per_row() {
6313 let (rows, cols, b) = (64usize, 128usize, 9usize);
6314 let ng = cols / GROUP_SIZE;
6315 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
6316 let mut bytes = bits.clone();
6317 for g in 0..rows * ng {
6318 let sc = 0.02 + 0.0005 * g as f32;
6319 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
6320 }
6321 for r in 0..rows {
6322 let bw = bits[r] as usize;
6323 let (mut acc, mut nb) = (0u64, 0usize);
6324 let mut rowbytes = Vec::new();
6325 for i in 0..cols {
6326 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
6327 acc = (acc << bw) | v;
6328 nb += bw;
6329 while nb >= 8 {
6330 nb -= 8;
6331 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6332 }
6333 }
6334 if nb > 0 {
6335 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6336 }
6337 bytes.extend_from_slice(&rowbytes);
6338 }
6339 let x: Vec<f32> =
6340 (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
6341 let offsets = vbit_row_offsets(&bytes, rows, cols);
6342 let mut y_a = vec![0f32; b * rows];
6343 let mut y_b = vec![0f32; b * rows];
6344 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
6345 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
6346 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
6347 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
6348 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
6349 let max_d =
6350 y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
6351 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
6352 }
6353
6354 #[test]
6355 fn q4matvec_matches_full_dequant() {
6356 let (rows, cols) = (8, 64);
6357 let groups = rows * cols / GROUP_SIZE;
6358 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6360 for i in 0..groups * 16 {
6361 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6362 }
6363 for g in 0..groups {
6364 let s = 0.01 + 0.003 * g as f32;
6365 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6366 }
6367 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6368
6369 let mut reference = vec![0.0f32; rows * cols];
6370 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
6371 let mut expect = vec![0.0f32; rows];
6372 for r in 0..rows {
6373 expect[r] = reference[r * cols..(r + 1) * cols]
6374 .iter()
6375 .zip(&x)
6376 .map(|(w, xv)| w * xv)
6377 .sum();
6378 }
6379
6380 let mut got = vec![0.0f32; rows];
6381 q4matvec(&bytes, &x, rows, cols, &mut got, None);
6382 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
6386 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6387 for r in 0..rows {
6388 assert!(
6389 (got[r] - expect[r]).abs() < tol * scale,
6390 "row {r}: {} vs {}",
6391 got[r],
6392 expect[r]
6393 );
6394 }
6395 }
6396
6397 #[test]
6400 fn vbitmatvec2_equals_two_singles() {
6401 let (rows, cols) = (6, 64);
6402 let ng = cols / GROUP_SIZE;
6403 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
6404 let mut bytes = bits.clone();
6405 for g in 0..rows * ng {
6406 let s = 0.02 + 0.001 * g as f32;
6407 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6408 }
6409 for r in 0..rows {
6410 let b = bits[r] as usize;
6411 let (mut acc, mut nb) = (0u64, 0usize);
6412 let mut rowbytes = Vec::new();
6413 for i in 0..cols {
6414 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
6415 acc = (acc << b) | v;
6416 nb += b;
6417 while nb >= 8 {
6418 nb -= 8;
6419 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6420 }
6421 }
6422 if nb > 0 {
6423 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6424 }
6425 bytes.extend_from_slice(&rowbytes);
6426 }
6427 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
6428 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
6429 let offsets = vbit_row_offsets(&bytes, rows, cols);
6430
6431 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6432 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
6433 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
6434 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
6435 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
6436 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
6437 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
6438 }
6439
6440 #[test]
6442 fn q4matvec2_equals_two_singles() {
6443 let (rows, cols) = (8, 128);
6444 let groups = rows * cols / GROUP_SIZE;
6445 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6446 for i in 0..groups * 16 {
6447 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6448 }
6449 for g in 0..groups {
6450 let s = 0.01 + 0.003 * g as f32;
6451 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6452 }
6453 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6456 x1[9] = 250.0;
6457 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
6458
6459 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6460 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
6461 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
6462 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
6463 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
6464 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
6465 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
6466 }
6467
6468 #[test]
6471 fn matvec_many_equals_separate_matvecs() {
6472 use crate::pool::Pool;
6473 let (r1, r2, cols) = (300, 200, 64);
6474 let mk = |salt: usize, rows: usize| {
6475 QTensor::from_f32(
6476 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
6477 rows,
6478 cols,
6479 )
6480 };
6481 let (a, b) = (mk(1, r1), mk(5, r2));
6482 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
6483 let pool = Pool::new(3);
6484
6485 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
6486 a.matvec(&x, &mut ea, Some(&pool));
6487 b.matvec(&x, &mut eb, Some(&pool));
6488 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
6489 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
6490 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
6491 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
6492 }
6493
6494 #[test]
6497 fn batched_matmat_equals_per_position_matvec() {
6498 let (rows, cols, b) = (8, 64, 5);
6499 let groups = rows * cols / GROUP_SIZE;
6501 let mut q4 = Vec::new();
6502 for i in 0..groups * 16 {
6503 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6504 }
6505 for g in 0..groups {
6506 q4.extend_from_slice(
6507 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
6508 );
6509 }
6510 let ng = cols / GROUP_SIZE;
6512 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
6513 let mut vb = bits.clone();
6514 for g in 0..rows * ng {
6515 vb.extend_from_slice(
6516 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
6517 );
6518 }
6519 for r in 0..rows {
6520 let bw = bits[r] as usize;
6521 let (mut acc, mut nb) = (0u64, 0usize);
6522 let mut rowbytes = Vec::new();
6523 for i in 0..cols {
6524 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
6525 acc = (acc << bw) | v;
6526 nb += bw;
6527 while nb >= 8 {
6528 nb -= 8;
6529 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6530 }
6531 }
6532 if nb > 0 {
6533 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6534 }
6535 vb.extend_from_slice(&rowbytes);
6536 }
6537 let offsets = vbit_row_offsets(&vb, rows, cols);
6538
6539 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
6540
6541 let mut got = vec![0f32; b * rows];
6543 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
6544 for bi in 0..b {
6545 let mut expect = vec![0f32; rows];
6546 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
6547 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
6548 }
6549
6550 let mut got = vec![0f32; b * rows];
6552 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
6553 for bi in 0..b {
6554 let mut expect = vec![0f32; rows];
6555 vbitmatvec(
6556 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
6557 );
6558 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
6559 }
6560 }
6561
6562 #[test]
6566 fn q4_tiled_matches_q4_block_bitexact() {
6567 let (rows, cols, b) = (8usize, 128usize, 3usize);
6568 let groups = rows * cols / GROUP_SIZE;
6569 let mut split = Vec::with_capacity(groups * 18);
6570 for i in 0..groups * 16 {
6571 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6572 }
6573 for g in 0..groups {
6574 split.extend_from_slice(
6575 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
6576 );
6577 }
6578 let (packed, scales) = split.split_at(groups * 16);
6580 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
6581 for g in 0..groups {
6582 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
6583 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
6584 }
6585
6586 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6587 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
6589
6590 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
6591 q4matvec(&split, &x1, rows, cols, &mut a, None);
6592 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
6593 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
6594
6595 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6596 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
6597 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
6598 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
6599 assert_eq!(a1, t1);
6600 assert_eq!(a2, t2);
6601
6602 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
6603 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
6604 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
6605 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
6606 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
6607 }
6608
6609 #[test]
6616 fn q4matvec_sdot_outlier_exact() {
6617 let (rows, cols) = (4, 128);
6618 let groups = rows * cols / GROUP_SIZE;
6619 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6620 for i in 0..groups * 16 {
6621 bytes.push(((i * 11 + 5) % 256) as u8);
6622 }
6623 for g in 0..groups {
6624 let s = 0.02 + 0.002 * g as f32;
6625 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6626 }
6627 let mut x: Vec<f32> = (0..cols)
6628 .map(|i| match i % 3 {
6629 0 => 1.0,
6630 1 => -1.0,
6631 _ => 0.0,
6632 })
6633 .collect();
6634 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
6637 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
6638 let mut expect = vec![0.0f32; rows];
6639 for r in 0..rows {
6640 expect[r] = reference[r * cols..(r + 1) * cols]
6641 .iter()
6642 .zip(&x)
6643 .map(|(w, xv)| w * xv)
6644 .sum();
6645 }
6646 let mut got = vec![0.0f32; rows];
6647 q4matvec(&bytes, &x, rows, cols, &mut got, None);
6648 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6649 for r in 0..rows {
6650 assert!(
6651 (got[r] - expect[r]).abs() < 2e-3 * scale,
6652 "row {r}: {} vs {} (outlier term must be exact)",
6653 got[r],
6654 expect[r]
6655 );
6656 }
6657 }
6658
6659 #[test]
6663 fn q1t_matvec_matches_reference() {
6664 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
6665 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
6667 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
6668 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
6670 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
6671 let mut bytes = Vec::new();
6672 for r in 0..rows {
6673 for g in 0..gpr {
6674 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
6675 let mut c = [0u8; 7];
6676 for k in 0..GROUP_SIZE {
6677 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
6679 0
6680 } else {
6681 ((k + r * 3 + g) % 3) as u8 };
6683 cortiq_core::quant::q1t_pack(&mut c, k, code);
6684 }
6685 bytes.extend_from_slice(&c);
6686 }
6687 }
6688 let mut row_ptr = vec![0u32; rows + 1];
6691 for &(idx, _) in &outliers {
6692 row_ptr[idx as usize / cols + 1] += 1;
6693 }
6694 for r in 0..rows {
6695 row_ptr[r + 1] += row_ptr[r];
6696 }
6697 for &p in &row_ptr {
6698 bytes.extend_from_slice(&p.to_le_bytes());
6699 }
6700 for &(idx, v) in &outliers {
6701 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
6702 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
6703 }
6704
6705 let mut refw = vec![0f32; rows * cols];
6706 dequant_q1t(&bytes, rows, cols, &mut refw);
6707 let x: Vec<f32> = (0..cols).map(|j| if j % 3 == 0 { 1.0 } else { -1.0 }).collect();
6710 let mut expect = vec![0f32; rows];
6711 for r in 0..rows {
6712 let mut a = 0.0f32;
6713 for j in 0..cols {
6714 a += refw[r * cols + j] * x[j];
6715 }
6716 expect[r] = a;
6717 }
6718 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
6719 let mut got = vec![0f32; rows];
6720 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
6721 for r in 0..rows {
6722 assert!((got[r] - expect[r]).abs() < tol(expect[r]), "row {r}: {} vs {}", got[r], expect[r]);
6723 }
6724 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
6726 let mut gm = vec![0f32; 2 * rows];
6727 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
6728 for r in 0..rows {
6729 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
6730 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
6731 }
6732 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
6737 q1t_matvec(&bytes, &x, rows, cols, &mut p1, None);
6738 q1t_matvec(&bytes, &x, rows, cols, &mut p2, None);
6739 for r in 0..rows {
6740 assert!((p1[r] - expect[r]).abs() < tol(expect[r]) && (p2[r] - expect[r]).abs() < tol(expect[r]));
6741 }
6742 }
6743
6744 #[test]
6748 #[ignore]
6749 fn q1t_matvec_speed() {
6750 use cortiq_core::quant::{f32_to_f16, q1t_code, q1t_pack, Q1T_TILE};
6751 use std::time::Instant;
6752 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
6754 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
6755 for r in 0..rows {
6756 for g in 0..gpr {
6757 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
6758 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
6759 let mut c = [0u8; 7];
6760 for k in 0..GROUP_SIZE {
6761 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
6762 }
6763 bytes.extend_from_slice(&c);
6764 }
6765 }
6766 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
6768 let mut idx = 0usize;
6769 while idx < n {
6770 row_ptr[idx / cols + 1] += 1;
6771 idx += stride;
6772 }
6773 for r in 0..rows {
6774 row_ptr[r + 1] += row_ptr[r];
6775 }
6776 for &p in &row_ptr {
6777 bytes.extend_from_slice(&p.to_le_bytes());
6778 }
6779 let mut idx = 0usize;
6780 while idx < n {
6781 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
6782 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
6783 idx += stride;
6784 }
6785 let x: Vec<f32> = (0..cols).map(|j| if j % 3 == 0 { 1.0 } else { -1.0 }).collect();
6788 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
6789
6790 let slow = |out: &mut [f32]| {
6792 let mut buf = vec![0f32; cols];
6793 for r in 0..rows {
6794 for g in 0..gpr {
6795 let off = (r * gpr + g) * Q1T_TILE;
6796 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6797 let codes = &bytes[off + 2..off + Q1T_TILE];
6798 for k in 0..GROUP_SIZE {
6799 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
6800 1 => s,
6801 2 => -s,
6802 _ => 0.0,
6803 };
6804 }
6805 }
6806 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
6807 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
6808 }
6809 };
6810 let iters = 5;
6811 let mut a = vec![0f32; rows];
6812 slow(&mut a); let t = Instant::now();
6814 for _ in 0..iters {
6815 slow(&mut a);
6816 }
6817 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
6818
6819 let mut b = vec![0f32; rows];
6820 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
6822 for _ in 0..iters {
6823 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
6824 }
6825 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
6826
6827 for r in 0..rows {
6828 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
6829 }
6830 println!(
6831 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
6832 slow_ms / fast_ms
6833 );
6834 }
6835}