1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16 GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17 q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23 F32 {
24 data: Vec<f32>,
25 rows: usize,
26 cols: usize,
27 },
28 Mapped {
29 model: Arc<CmfModel>,
30 idx: usize,
32 dtype: TensorDtype,
33 rows: usize,
34 cols: usize,
35 row_scale: Vec<f32>,
37 col_field: Vec<f32>,
39 vbit_offsets: Vec<usize>,
43 repack: Vec<u8>,
51 },
52}
53
54fn repack_enabled() -> bool {
61 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62 *ON.get_or_init(|| {
63 std::env::var("CMF_REPACK")
64 .map(|v| v == "1")
65 .unwrap_or(cfg!(target_os = "android"))
66 })
67}
68
69fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73 #[cfg(target_arch = "aarch64")]
74 let arch_ok = sdot_enabled();
75 #[cfg(not(target_arch = "aarch64"))]
76 let arch_ok = false;
77 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78 return Vec::new();
79 }
80 q8_repack_layout(bytes, rows, cols)
81}
82
83fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86 let groups = rows / 4;
87 let mut rep = vec![0u8; groups * 4 * cols];
88 for g in 0..groups {
89 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90 for c in 0..cols / 16 {
91 for lane in 0..4 {
92 let src = (g * 4 + lane) * cols + c * 16;
93 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94 .copy_from_slice(&bytes[src..src + 16]);
95 }
96 }
97 }
98 rep
99}
100
101fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104 let ng = cols / GROUP_SIZE;
105 let bits = &bytes[..rows];
106 let mut offsets = Vec::with_capacity(rows + 1);
107 let mut off = rows + rows * ng * 2;
108 for r in 0..rows {
109 offsets.push(off);
110 off += (cols * bits[r] as usize).div_ceil(8);
111 }
112 offsets.push(off);
113 offsets
114}
115
116fn blocked_enabled() -> bool {
122 use std::sync::atomic::Ordering::Relaxed;
123 match BLOCKED_OVERRIDE.load(Relaxed) {
124 1 => false,
125 2 => true,
126 _ => {
127 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128 *ON.get_or_init(|| {
129 std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
130 })
131 }
132 }
133}
134
135static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
136
137pub fn set_blocked_override(on: Option<bool>) {
144 let v = match on {
145 None => 0,
146 Some(false) => 1,
147 Some(true) => 2,
148 };
149 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
150}
151
152fn gpu_lmhead_enabled() -> bool {
153 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
154 *ON.get_or_init(|| std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true))
155}
156
157fn gpu_split_frac() -> f32 {
158 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
159 *F.get_or_init(|| {
160 std::env::var("CMF_GPU_SPLIT")
161 .ok()
162 .and_then(|v| v.parse::<f32>().ok())
163 .unwrap_or(0.5)
164 .clamp(0.0, 1.0)
165 })
166}
167
168impl QTensor {
169 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
170 debug_assert_eq!(data.len(), rows * cols);
171 Self::F32 { data, rows, cols }
172 }
173
174 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
177 let idx = model
180 .tensor_index(name)
181 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
182 let entry = &model.tensors[idx];
183 if entry.shape.len() != 2 {
184 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
185 }
186 let (rows, cols) = (entry.shape[0], entry.shape[1]);
187 let bytes = model.entry_bytes(entry);
188
189 match entry.dtype {
190 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
191 let n = rows * cols;
192 let scales_off = n;
193 let row_scale: Vec<f32> = (0..rows)
194 .map(|o| {
195 f16_to_f32(u16::from_le_bytes([
196 bytes[scales_off + o * 2],
197 bytes[scales_off + o * 2 + 1],
198 ]))
199 })
200 .collect();
201 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
202 let col_off = n + rows * 2;
203 (0..cols)
204 .map(|i| {
205 f16_to_f32(u16::from_le_bytes([
206 bytes[col_off + i * 2],
207 bytes[col_off + i * 2 + 1],
208 ]))
209 })
210 .collect()
211 } else {
212 Vec::new()
213 };
214 Ok(Self::Mapped {
215 model: model.clone(),
216 idx,
217 dtype: entry.dtype,
218 rows,
219 cols,
220 row_scale,
221 col_field,
222 vbit_offsets: Vec::new(),
223 repack: q8_repack(bytes, rows, cols),
224 })
225 }
226 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
228 model: model.clone(),
229 idx,
230 dtype: entry.dtype,
231 rows,
232 cols,
233 row_scale: Vec::new(),
234 col_field: Vec::new(),
235 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
236 repack: Vec::new(),
237 }),
238 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
242 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
243 let offsets: Vec<usize> = (0..=rows)
244 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
245 .collect();
246 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: offsets,
255 repack: Vec::new(),
256 })
257 }
258 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
264 model: model.clone(),
265 idx,
266 dtype: entry.dtype,
267 rows,
268 cols,
269 row_scale: Vec::new(),
270 col_field: Vec::new(),
271 vbit_offsets: Vec::new(),
272 repack: Vec::new(),
273 }),
274 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
277 model: model.clone(),
278 idx,
279 dtype: entry.dtype,
280 rows,
281 cols,
282 row_scale: Vec::new(),
283 col_field: Vec::new(),
284 vbit_offsets: Vec::new(),
285 repack: Vec::new(),
286 }),
287 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
289 model: model.clone(),
290 idx,
291 dtype: entry.dtype,
292 rows,
293 cols,
294 row_scale: Vec::new(),
295 col_field: Vec::new(),
296 vbit_offsets: Vec::new(),
297 repack: Vec::new(),
298 }),
299 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
300 model: model.clone(),
301 idx,
302 dtype: entry.dtype,
303 rows,
304 cols,
305 row_scale: Vec::new(),
306 col_field: Vec::new(),
307 vbit_offsets: Vec::new(),
308 repack: Vec::new(),
309 }),
310 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
312 model: model.clone(),
313 idx,
314 dtype: entry.dtype,
315 rows,
316 cols,
317 row_scale: Vec::new(),
318 col_field: Vec::new(),
319 vbit_offsets: Vec::new(),
320 repack: Vec::new(),
321 }),
322 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
326 model: model.clone(),
327 idx,
328 dtype: entry.dtype,
329 rows,
330 cols,
331 row_scale: Vec::new(),
332 col_field: Vec::new(),
333 vbit_offsets: Vec::new(),
334 repack: Vec::new(),
335 }),
336 _ => {
338 let mut data = vec![0.0f32; rows * cols];
339 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
340 Ok(Self::from_f32(data, rows, cols))
341 }
342 }
343 }
344
345 pub(crate) fn is_q1(&self) -> bool {
348 matches!(
349 self,
350 Self::Mapped {
351 dtype: TensorDtype::Q1,
352 ..
353 }
354 )
355 }
356
357 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
360 match self {
361 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
362 _ => None,
363 }
364 }
365
366 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
373 match self {
374 #[cfg(target_os = "macos")]
375 Self::Mapped {
376 dtype: TensorDtype::Q1T,
377 ..
378 } if !crate::gpu::metal_q1t_enabled() => None,
379 Self::Mapped {
380 idx,
381 dtype:
382 TensorDtype::Q1
383 | TensorDtype::Q1T
384 | TensorDtype::Q4Block
385 | TensorDtype::Q4Tiled
386 | TensorDtype::Q4TiledP
390 | TensorDtype::Q8Row
391 | TensorDtype::Q8_2f,
392 rows,
393 cols,
394 ..
395 } => Some((*idx, *rows, *cols)),
396 _ => None,
397 }
398 }
399
400 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
406 match self {
407 Self::Mapped {
408 idx,
409 dtype: TensorDtype::Q4Tiled,
410 rows,
411 cols,
412 ..
413 } => Some((*idx, *rows, *cols)),
414 _ => None,
415 }
416 }
417
418 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
422 match self {
423 Self::Mapped {
424 idx,
425 dtype: TensorDtype::Q4TiledP,
426 rows,
427 cols,
428 ..
429 } => Some((*idx, *rows, *cols)),
430 _ => None,
431 }
432 }
433
434 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
439 match self {
440 Self::Mapped {
441 idx,
442 dtype: TensorDtype::Q8Row,
443 rows,
444 cols,
445 row_scale,
446 col_field,
447 ..
448 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
449 _ => None,
450 }
451 }
452
453 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
457 match self {
458 Self::Mapped { dtype, .. } => Some(*dtype),
459 _ => None,
460 }
461 }
462
463 pub fn model_idx(&self) -> Option<usize> {
468 match self {
469 Self::Mapped { idx, .. } => Some(*idx),
470 _ => None,
471 }
472 }
473
474 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
479 match self {
480 Self::Mapped { model, .. } => Some(model.clone()),
481 _ => None,
482 }
483 }
484
485 pub fn rows(&self) -> usize {
486 match self {
487 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
488 }
489 }
490
491 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
494 match self {
495 Self::Mapped {
496 model,
497 idx,
498 dtype: TensorDtype::Q4Tiled,
499 ..
500 } => Some((model, *idx)),
501 _ => None,
502 }
503 }
504
505 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
508 match self {
509 Self::Mapped {
510 model,
511 idx,
512 dtype: TensorDtype::Q4TiledP,
513 ..
514 } => Some((model, *idx)),
515 _ => None,
516 }
517 }
518
519 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
522 match self {
523 Self::Mapped {
524 model,
525 idx,
526 dtype: TensorDtype::Q2TiledP,
527 ..
528 } => Some((model, *idx)),
529 _ => None,
530 }
531 }
532
533 pub fn cols(&self) -> usize {
534 match self {
535 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
536 }
537 }
538
539 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
542 match self {
543 Self::Mapped {
544 model,
545 idx,
546 dtype: TensorDtype::Q1,
547 ..
548 } => Some((model, *idx)),
549 _ => None,
550 }
551 }
552
553 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
563 match self {
564 Self::Mapped {
565 model,
566 idx,
567 dtype: TensorDtype::Q8Row,
568 row_scale,
569 ..
570 } => Some((model, *idx, 0, row_scale.as_slice())),
571 Self::Mapped {
572 model,
573 idx,
574 dtype: TensorDtype::Q1,
575 ..
576 } => Some((model, *idx, 1, &[])),
577 Self::Mapped {
582 model,
583 idx,
584 dtype: TensorDtype::Q4Tiled,
585 ..
586 } => Some((model, *idx, 5, &[])),
587 Self::Mapped {
591 model,
592 idx,
593 dtype: TensorDtype::Q4TiledP,
594 ..
595 } => Some((model, *idx, 6, &[])),
596 Self::Mapped {
597 model,
598 idx,
599 dtype: TensorDtype::Q4Block,
600 ..
601 } => Some((model, *idx, 2, &[])),
602 Self::Mapped {
607 model,
608 idx,
609 dtype: TensorDtype::Q8_2f,
610 ..
611 } => Some((model, *idx, 7, &[])),
612 Self::Mapped {
613 model,
614 idx,
615 dtype: TensorDtype::Q1T,
616 ..
617 } => Some((model, *idx, 3, &[])),
618 _ => None,
619 }
620 }
621
622 pub fn as_f32(&self) -> Option<&[f32]> {
625 match self {
626 Self::F32 { data, .. } => Some(data),
627 Self::Mapped { .. } => None,
628 }
629 }
630
631 fn quant_bytes(&self) -> &[u8] {
632 match self {
633 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
634 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
635 }
636 }
637
638 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
640 let cols = self.cols();
641 debug_assert_eq!(dst.len(), cols);
642 match self {
643 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
644 Self::Mapped {
645 dtype,
646 row_scale,
647 col_field,
648 vbit_offsets,
649 ..
650 } => {
651 if *dtype == TensorDtype::Q4Tiled {
652 let bytes = self.quant_bytes();
653 let gpr = cols / GROUP_SIZE;
654 for gi in 0..gpr {
655 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
656 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
657 for (k, &b) in tile[2..].iter().enumerate() {
658 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
659 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
660 }
661 }
662 return;
663 }
664 if *dtype == TensorDtype::Q4TiledP {
665 let bytes = self.quant_bytes();
666 let gpr = cols / GROUP_SIZE;
667 let v = Q4tpView::new(bytes, self.rows(), cols);
668 let mut sc = vec![0f32; gpr];
669 v.scales_into(r, gpr, &mut sc);
670 for gi in 0..gpr {
671 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
672 let s = sc[gi];
673 for (k, &b) in tile.iter().enumerate() {
674 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
675 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
676 }
677 }
678 return;
679 }
680 if *dtype == TensorDtype::Q2TiledP {
681 let bytes = self.quant_bytes();
682 let gpr = cols / GROUP_SIZE;
683 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
684 let mut sc = vec![0f32; gpr];
685 v.scales_into(r, gpr, &mut sc);
686 for gi in 0..gpr {
687 let ch =
688 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
689 let s = sc[gi];
690 for (k, &b) in ch.iter().enumerate() {
691 for j in 0..4 {
692 dst[gi * GROUP_SIZE + k * 4 + j] =
693 (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
694 }
695 }
696 }
697 return;
698 }
699 if *dtype == TensorDtype::Q4Block {
700 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
701 let gpr = cols / GROUP_SIZE;
702 for gi in 0..gpr {
703 let g = r * gpr + gi;
704 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
705 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
706 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
707 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
708 }
709 }
710 return;
711 }
712 if *dtype == TensorDtype::Q1 {
713 let bytes = self.quant_bytes();
714 let gpr = cols / GROUP_SIZE;
715 for gi in 0..gpr {
716 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
717 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
718 for (j, &b) in tile[2..].iter().enumerate() {
719 for k in 0..8 {
720 dst[gi * GROUP_SIZE + j * 8 + k] =
721 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
722 }
723 }
724 }
725 return;
726 }
727 if *dtype == TensorDtype::Q1T {
728 let bytes = self.quant_bytes();
729 let gpr = cols / GROUP_SIZE;
730 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
731 for gi in 0..gpr {
732 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
733 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
734 bytes[off],
735 bytes[off + 1],
736 ]));
737 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
738 for k in 0..GROUP_SIZE {
739 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
740 {
741 1 => s,
742 2 => -s,
743 _ => 0.0,
744 };
745 }
746 }
747 let rows = self.rows();
749 let entries = base_len + (rows + 1) * 4;
750 if entries <= bytes.len() {
751 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
752 let r0 = u32::from_le_bytes([
753 ptrs[r * 4],
754 ptrs[r * 4 + 1],
755 ptrs[r * 4 + 2],
756 ptrs[r * 4 + 3],
757 ]) as usize;
758 let r1 = u32::from_le_bytes([
759 ptrs[(r + 1) * 4],
760 ptrs[(r + 1) * 4 + 1],
761 ptrs[(r + 1) * 4 + 2],
762 ptrs[(r + 1) * 4 + 3],
763 ]) as usize;
764 let off = entries + r0 * 4;
765 for i in 0..r1 - r0 {
766 let item = &bytes[off + i * 4..off + i * 4 + 4];
767 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
768 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
769 item[2], item[3],
770 ]));
771 if c < cols {
772 dst[c] = v;
773 }
774 }
775 }
776 return;
777 }
778 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
779 let bytes = self.quant_bytes();
780 let rows = self.rows();
781 let ng = cols / GROUP_SIZE;
782 let bits = &bytes[..rows];
783 let sc_off = rows;
784 let off = vbit_offsets[r];
787 let b = bits[r] as usize;
788 let l = ((1usize << (b - 1)) - 1) as f32;
789 let data = &bytes[off..];
790 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
791 for (i, d) in dst.iter_mut().enumerate() {
792 while nbits < b {
793 acc = (acc << 8) | data[idx] as u64;
794 idx += 1;
795 nbits += 8;
796 }
797 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
798 nbits -= b;
799 let so = (r * ng + i / GROUP_SIZE) * 2;
800 let sv = f16_to_f32(u16::from_le_bytes([
801 bytes[sc_off + so],
802 bytes[sc_off + so + 1],
803 ]));
804 *d = (u - l) * sv;
805 }
806 return;
807 }
808 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
809 let s = row_scale[r];
810 match dtype {
811 TensorDtype::Q8Row => {
812 for (d, &b) in dst.iter_mut().zip(q) {
813 *d = (b as i8) as f32 * s;
814 }
815 }
816 TensorDtype::Q8_2f => {
817 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
818 *d = (b as i8) as f32 * s * col_field[i];
819 }
820 }
821 _ => unreachable!(),
822 }
823 }
824 }
825 }
826
827 pub fn sparse_col_ok(&self) -> bool {
832 match self {
833 Self::F32 { .. } => true,
834 Self::Mapped { dtype, .. } => {
835 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
836 }
837 }
838 }
839
840 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
844 let inter = self.cols();
845 let hidden = self.rows();
846 debug_assert_eq!(out.len(), hidden);
847 match self {
848 Self::F32 { data, .. } => {
849 for (k, o) in out.iter_mut().enumerate() {
850 *o += w * data[k * inter + c];
851 }
852 }
853 Self::Mapped {
854 dtype,
855 row_scale,
856 col_field,
857 ..
858 } => {
859 let q = self.quant_bytes();
860 let colf = if *dtype == TensorDtype::Q8_2f {
861 col_field[c]
862 } else {
863 1.0
864 };
865 let wc = w * colf;
866 for (k, o) in out.iter_mut().enumerate() {
867 let b = q[k * inter + c] as i8 as f32;
868 *o += wc * b * row_scale[k];
869 }
870 }
871 }
872 }
873
874 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
878 let cols = self.cols();
879 match self {
880 Self::F32 { data, .. } => {
881 let row = &data[r * cols..(r + 1) * cols];
882 row.iter().zip(x).map(|(w, v)| w * v).sum()
883 }
884 Self::Mapped {
885 dtype,
886 row_scale,
887 col_field,
888 ..
889 } => match dtype {
890 TensorDtype::Q8Row => {
891 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
892 dot_i8_f32(q, x) * row_scale[r]
893 }
894 TensorDtype::Q8_2f => {
895 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
896 dot_i8_col_f32(q, x, col_field) * row_scale[r]
897 }
898 _ => {
899 self.row_f32(r, scratch);
900 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
901 }
902 },
903 }
904 }
905
906 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
909 match self {
910 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
914 Self::Mapped {
915 model,
916 idx,
917 dtype,
918 rows,
919 cols,
920 row_scale,
921 col_field,
922 vbit_offsets,
923 repack,
924 } => {
925 let _ = (model, idx);
926 assert!(
936 out.len() >= *rows && x.len() >= *cols,
937 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
938 out.len(),
939 x.len(),
940 );
941 if *dtype == TensorDtype::Q4Block {
942 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
946 let t0 = std::time::Instant::now();
947 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
948 crate::gpu::ProbeArm::Gpu => {
949 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
950 crate::gpu::probe_record(
951 crate::gpu::OpClass::Matvec,
952 true,
953 t0.elapsed(),
954 );
955 return;
956 }
957 }
958 crate::gpu::ProbeArm::CpuTimed => {
959 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
960 crate::gpu::probe_record(
961 crate::gpu::OpClass::Matvec,
962 false,
963 t0.elapsed(),
964 );
965 return;
966 }
967 crate::gpu::ProbeArm::Cpu => {}
968 }
969 }
970 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
971 return;
972 }
973 if *dtype == TensorDtype::Q4Tiled {
974 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
979 let t0 = std::time::Instant::now();
980 let cls = crate::gpu::matvec_class(*rows, *cols);
981 match crate::gpu::probe_arm(cls) {
982 crate::gpu::ProbeArm::Gpu => {
983 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
984 crate::gpu::probe_record(cls, true, t0.elapsed());
985 return;
986 }
987 }
988 crate::gpu::ProbeArm::CpuTimed => {
989 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
990 crate::gpu::probe_record(cls, false, t0.elapsed());
991 return;
992 }
993 crate::gpu::ProbeArm::Cpu => {}
994 }
995 }
996 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
997 return;
998 }
999 if *dtype == TensorDtype::Q4TiledP {
1000 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1006 let t0 = std::time::Instant::now();
1007 let cls = crate::gpu::matvec_class(*rows, *cols);
1008 match crate::gpu::probe_arm(cls) {
1009 crate::gpu::ProbeArm::Gpu => {
1010 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1011 crate::gpu::probe_record(cls, true, t0.elapsed());
1012 return;
1013 }
1014 }
1015 crate::gpu::ProbeArm::CpuTimed => {
1016 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1017 crate::gpu::probe_record(cls, false, t0.elapsed());
1018 return;
1019 }
1020 crate::gpu::ProbeArm::Cpu => {}
1021 }
1022 }
1023 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1024 return;
1025 }
1026 if *dtype == TensorDtype::Q2TiledP {
1027 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1028 return;
1029 }
1030 if *dtype == TensorDtype::Q1 {
1031 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1036 let t0 = std::time::Instant::now();
1037 let arm = if crate::gpu::q1_force() {
1038 crate::gpu::ProbeArm::Gpu
1039 } else {
1040 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1041 };
1042 match arm {
1043 crate::gpu::ProbeArm::Gpu => {
1044 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1045 crate::gpu::probe_record(
1046 crate::gpu::OpClass::Matvec,
1047 true,
1048 t0.elapsed(),
1049 );
1050 return;
1051 }
1052 }
1053 crate::gpu::ProbeArm::CpuTimed => {
1054 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1055 crate::gpu::probe_record(
1056 crate::gpu::OpClass::Matvec,
1057 false,
1058 t0.elapsed(),
1059 );
1060 return;
1061 }
1062 crate::gpu::ProbeArm::Cpu => {}
1063 }
1064 }
1065 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1066 return;
1067 }
1068 if *dtype == TensorDtype::Q1T {
1069 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1073 let t0 = std::time::Instant::now();
1074 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1075 crate::gpu::ProbeArm::Gpu => {
1076 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1077 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1078 crate::gpu::probe_record(
1079 crate::gpu::OpClass::Matvec,
1080 true,
1081 t0.elapsed(),
1082 );
1083 return;
1084 }
1085 }
1086 crate::gpu::ProbeArm::CpuTimed => {
1087 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1088 crate::gpu::probe_record(
1089 crate::gpu::OpClass::Matvec,
1090 false,
1091 t0.elapsed(),
1092 );
1093 return;
1094 }
1095 crate::gpu::ProbeArm::Cpu => {}
1096 }
1097 }
1098 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1099 return;
1100 }
1101 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1102 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1103 return;
1104 }
1105 let xs = prescale(x, col_field, *dtype);
1106 if *rows >= crate::gpu::min_rows()
1111 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1112 && gpu_lmhead_enabled()
1113 && crate::gpu::enabled_here()
1114 {
1115 let t0 = std::time::Instant::now();
1118 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1119 crate::gpu::ProbeArm::Gpu => {}
1120 crate::gpu::ProbeArm::CpuTimed => {
1121 qmatvec(
1122 self.quant_bytes(),
1123 repack,
1124 row_scale,
1125 x,
1126 col_field,
1127 *dtype,
1128 *rows,
1129 *cols,
1130 out,
1131 pool,
1132 );
1133 crate::gpu::probe_record(
1134 crate::gpu::OpClass::Matvec,
1135 false,
1136 t0.elapsed(),
1137 );
1138 return;
1139 }
1140 crate::gpu::ProbeArm::Cpu => {
1141 qmatvec(
1142 self.quant_bytes(),
1143 repack,
1144 row_scale,
1145 x,
1146 col_field,
1147 *dtype,
1148 *rows,
1149 *cols,
1150 out,
1151 pool,
1152 );
1153 return;
1154 }
1155 }
1156 let frac = gpu_split_frac();
1157 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1158 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1159 let bytes = self.quant_bytes();
1160 let ok = std::thread::scope(|sc| {
1161 let g = sc.spawn(|| {
1162 crate::gpu::q8_matvec_range(
1163 model,
1164 *idx,
1165 cpu_rows,
1166 &row_scale[cpu_rows..],
1167 &xs,
1168 *rows - cpu_rows,
1169 *cols,
1170 out_gpu,
1171 )
1172 });
1173 if cpu_rows > 0 {
1174 let rep_cpu = if repack.is_empty() {
1177 &[][..]
1178 } else {
1179 &repack[..(cpu_rows / 4) * 4 * *cols]
1180 };
1181 qmatvec(
1182 &bytes[..cpu_rows * *cols],
1183 rep_cpu,
1184 &row_scale[..cpu_rows],
1185 x,
1186 col_field,
1187 *dtype,
1188 cpu_rows,
1189 *cols,
1190 out_cpu,
1191 pool,
1192 );
1193 }
1194 g.join().unwrap_or(false)
1195 });
1196 if ok {
1197 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1198 return;
1199 }
1200 qmatvec(
1203 &bytes[cpu_rows * *cols..(*rows) * *cols],
1204 &[],
1205 &row_scale[cpu_rows..],
1206 x,
1207 col_field,
1208 *dtype,
1209 *rows - cpu_rows,
1210 *cols,
1211 out_gpu,
1212 pool,
1213 );
1214 return;
1215 }
1216 qmatvec(
1217 self.quant_bytes(),
1218 repack,
1219 row_scale,
1220 x,
1221 col_field,
1222 *dtype,
1223 *rows,
1224 *cols,
1225 out,
1226 pool,
1227 );
1228 }
1229 }
1230 }
1231
1232 pub fn matvec2(
1234 &self,
1235 x1: &[f32],
1236 x2: &[f32],
1237 o1: &mut [f32],
1238 o2: &mut [f32],
1239 pool: Option<&Pool>,
1240 ) {
1241 match self {
1242 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1243 Self::Mapped {
1244 dtype,
1245 rows,
1246 cols,
1247 row_scale,
1248 col_field,
1249 vbit_offsets,
1250 ..
1251 } => {
1252 if *dtype == TensorDtype::Q4Block {
1253 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1254 return;
1255 }
1256 if *dtype == TensorDtype::Q4Tiled {
1257 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1258 return;
1259 }
1260 if *dtype == TensorDtype::Q4TiledP {
1261 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1262 return;
1263 }
1264 if *dtype == TensorDtype::Q2TiledP {
1265 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1266 return;
1267 }
1268 if *dtype == TensorDtype::Q1 {
1269 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1270 return;
1271 }
1272 if *dtype == TensorDtype::Q1T {
1273 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1279 return;
1280 }
1281 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1282 vbitmatvec2(
1283 self.quant_bytes(),
1284 vbit_offsets,
1285 x1,
1286 x2,
1287 *rows,
1288 *cols,
1289 o1,
1290 o2,
1291 pool,
1292 );
1293 return;
1294 }
1295 qmatvec2(
1296 self.quant_bytes(),
1297 row_scale,
1298 x1,
1299 x2,
1300 col_field,
1301 *dtype,
1302 *rows,
1303 *cols,
1304 o1,
1305 o2,
1306 pool,
1307 );
1308 }
1309 }
1310 }
1311}
1312
1313impl QTensor {
1314 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1322 match self {
1323 Self::Mapped { model, idx, dtype, .. } if *dtype == TensorDtype::Q4TiledP => {
1324 Some((model, *idx))
1325 }
1326 _ => None,
1327 }
1328 }
1329
1330 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1331 let cols = self.cols();
1332 let rows = self.rows();
1333 debug_assert_eq!(xs_all.len(), b * cols);
1334 debug_assert_eq!(out.len(), b * rows);
1335 if crate::gptq_capture::capturing() {
1339 if let Self::Mapped { model, idx, .. } = self {
1340 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1341 }
1342 }
1343 match self {
1344 Self::F32 { data, .. } => {
1345 let out_addr = SendMut(out.as_mut_ptr());
1346 let run = |start: usize, end: usize| {
1347 for o in start..end {
1348 let row = &data[o * cols..(o + 1) * cols];
1349 for bi in 0..b {
1350 let x = &xs_all[bi * cols..(bi + 1) * cols];
1351 let mut acc = 0f32;
1352 for j in 0..cols {
1353 acc += row[j] * x[j];
1354 }
1355 unsafe { *out_addr.at(bi * rows + o) = acc };
1356 }
1357 }
1358 };
1359 dispatch_rows(pool, rows, &run);
1360 }
1361 Self::Mapped {
1362 dtype,
1363 row_scale,
1364 col_field,
1365 vbit_offsets,
1366 ..
1367 } => {
1368 if *dtype == TensorDtype::Q4Block {
1369 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1370 return;
1371 }
1372 if *dtype == TensorDtype::Q4TiledP {
1373 if b >= 32
1385 && b * rows * cols >= 128_000_000
1386 && cols % 32 == 0
1387 && !crate::gpu::mm_killed()
1388 && crate::gpu::enabled_here()
1389 {
1390 let class = if b >= 128 {
1391 crate::gpu::OpClass::MatmatWide
1392 } else {
1393 crate::gpu::OpClass::Matmat
1394 };
1395 if let Self::Mapped { model, idx, .. } = self {
1396 if crate::mm_ab::on() {
1408 let mut g = vec![0f32; b * rows];
1409 let t = std::time::Instant::now();
1410 let took = crate::gpu::q4tp_matmat(
1411 model, *idx, xs_all, b, rows, cols, &mut g,
1412 );
1413 let dg = t.elapsed();
1414 let t = std::time::Instant::now();
1415 q4tp_matmat(
1416 self.quant_bytes(), xs_all, b, rows, cols, out, pool,
1417 );
1418 let dc = t.elapsed();
1419 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1420 return;
1421 }
1422 let t0 = std::time::Instant::now();
1423 let resident = crate::gpu::weight_is_resident(model, *idx);
1427 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1428 crate::gpu::ProbeArm::Gpu => {
1429 if crate::gpu::q4tp_matmat(
1430 model, *idx, xs_all, b, rows, cols, out,
1431 ) {
1432 let el = t0.elapsed();
1433 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1443 let budget = std::time::Duration::from_secs_f64(
1444 flops / 1.5e12 * 8.0 + 0.020,
1445 );
1446 if el > budget && !crate::gpu::probe_was_cold() {
1447 tracing::warn!(
1448 "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1449 device contended, CPU for the rest of the process"
1450 );
1451 crate::gpu::mm_kill();
1452 }
1453 crate::gpu::probe_record(class, true, el);
1454 return;
1455 }
1456 }
1457 crate::gpu::ProbeArm::CpuTimed => {
1458 q4tp_matmat(
1459 self.quant_bytes(),
1460 xs_all,
1461 b,
1462 rows,
1463 cols,
1464 out,
1465 pool,
1466 );
1467 crate::gpu::probe_record(class, false, t0.elapsed());
1468 return;
1469 }
1470 crate::gpu::ProbeArm::Cpu => {}
1471 }
1472 }
1473 }
1474 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1475 return;
1476 }
1477 if *dtype == TensorDtype::Q2TiledP {
1478 if b >= 32
1484 && b * rows * cols >= 128_000_000
1485 && cols % 32 == 0
1486 && !crate::gpu::mm_killed()
1487 && crate::gpu::enabled_here()
1488 {
1489 let class = if b >= 128 {
1490 crate::gpu::OpClass::MatmatWide
1491 } else {
1492 crate::gpu::OpClass::Matmat
1493 };
1494 if let Self::Mapped { model, idx, .. } = self {
1495 let t0 = std::time::Instant::now();
1496 match crate::gpu::probe_arm(class) {
1497 crate::gpu::ProbeArm::Gpu => {
1498 if crate::gpu::q2tp_matmat(
1499 model, *idx, xs_all, b, rows, cols, out,
1500 ) {
1501 crate::gpu::probe_record(class, true, t0.elapsed());
1502 return;
1503 }
1504 }
1505 crate::gpu::ProbeArm::CpuTimed => {
1506 q2tp_matmat(
1507 self.quant_bytes(),
1508 xs_all,
1509 b,
1510 rows,
1511 cols,
1512 out,
1513 pool,
1514 );
1515 crate::gpu::probe_record(class, false, t0.elapsed());
1516 return;
1517 }
1518 crate::gpu::ProbeArm::Cpu => {}
1519 }
1520 }
1521 }
1522 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1527 return;
1528 }
1529 if *dtype == TensorDtype::Q4Tiled {
1530 if b >= 32
1542 && b * rows * cols >= 128_000_000
1543 && cols % 32 == 0
1544 && !crate::gpu::mm_killed()
1545 && crate::gpu::enabled_here()
1546 {
1547 let class = if b >= 128 {
1548 crate::gpu::OpClass::MatmatWide
1549 } else {
1550 crate::gpu::OpClass::Matmat
1551 };
1552 if let Self::Mapped { model, idx, .. } = self {
1553 let t0 = std::time::Instant::now();
1554 match crate::gpu::probe_arm(class) {
1555 crate::gpu::ProbeArm::Gpu => {
1556 if crate::gpu::q4t_matmat(
1557 model, *idx, xs_all, b, rows, cols, out,
1558 ) {
1559 let el = t0.elapsed();
1560 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1570 let budget = std::time::Duration::from_secs_f64(
1571 flops / 1.5e12 * 8.0 + 0.020,
1572 );
1573 if el > budget && !crate::gpu::probe_was_cold() {
1574 tracing::warn!(
1575 "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1576 device contended, CPU for the rest of the process"
1577 );
1578 crate::gpu::mm_kill();
1579 }
1580 crate::gpu::probe_record(class, true, el);
1581 return;
1582 }
1583 }
1584 crate::gpu::ProbeArm::CpuTimed => {
1585 q4t_matmat(
1586 self.quant_bytes(),
1587 xs_all,
1588 b,
1589 rows,
1590 cols,
1591 out,
1592 pool,
1593 );
1594 crate::gpu::probe_record(class, false, t0.elapsed());
1595 return;
1596 }
1597 crate::gpu::ProbeArm::Cpu => {}
1598 }
1599 }
1600 }
1601 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1602 return;
1603 }
1604 if *dtype == TensorDtype::Q1 {
1605 if b >= 32
1608 && b * rows * cols >= 128_000_000
1609 && cols % 64 == 0
1610 && crate::gpu::enabled_here()
1611 {
1612 if let Self::Mapped { model, idx, .. } = self {
1613 let t0 = std::time::Instant::now();
1614 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1615 crate::gpu::ProbeArm::Gpu => {
1616 if crate::gpu::q1_matmat(
1617 model, *idx, xs_all, b, rows, cols, out,
1618 ) {
1619 crate::gpu::probe_record(
1620 crate::gpu::OpClass::Matmat,
1621 true,
1622 t0.elapsed(),
1623 );
1624 return;
1625 }
1626 }
1627 crate::gpu::ProbeArm::CpuTimed => {
1628 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1629 crate::gpu::probe_record(
1630 crate::gpu::OpClass::Matmat,
1631 false,
1632 t0.elapsed(),
1633 );
1634 return;
1635 }
1636 crate::gpu::ProbeArm::Cpu => {}
1637 }
1638 }
1639 }
1640 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1641 return;
1642 }
1643 if *dtype == TensorDtype::Q1T {
1644 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1647 if let Self::Mapped { model, idx, .. } = self {
1648 let t0 = std::time::Instant::now();
1649 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1650 crate::gpu::ProbeArm::Gpu => {
1651 if crate::gpu::q1t_matmat(
1652 model, *idx, xs_all, b, rows, cols, out,
1653 ) {
1654 crate::gpu::probe_record(
1655 crate::gpu::OpClass::Matmat,
1656 true,
1657 t0.elapsed(),
1658 );
1659 return;
1660 }
1661 }
1662 crate::gpu::ProbeArm::CpuTimed => {
1663 q1t_matmat(
1664 self.quant_bytes(),
1665 xs_all,
1666 b,
1667 rows,
1668 cols,
1669 out,
1670 pool,
1671 );
1672 crate::gpu::probe_record(
1673 crate::gpu::OpClass::Matmat,
1674 false,
1675 t0.elapsed(),
1676 );
1677 return;
1678 }
1679 crate::gpu::ProbeArm::Cpu => {}
1680 }
1681 }
1682 }
1683 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1684 return;
1685 }
1686 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1687 vbitmatmat(
1688 self.quant_bytes(),
1689 vbit_offsets,
1690 xs_all,
1691 b,
1692 rows,
1693 cols,
1694 out,
1695 pool,
1696 );
1697 return;
1698 }
1699 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1700 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1701 .collect();
1702 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1708 if let Self::Mapped { model, idx, .. } = self {
1709 let t0 = std::time::Instant::now();
1710 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1711 crate::gpu::ProbeArm::Gpu
1712 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1713 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1714 {
1715 let q = self.quant_bytes();
1719 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1720 return;
1721 }
1722 crate::gpu::ProbeArm::Gpu => {
1723 let flat: Vec<f32> =
1724 pre.iter().flat_map(|v| v.iter().copied()).collect();
1725 if crate::gpu::q8_matmat(
1726 model, *idx, row_scale, &flat, b, rows, cols, out,
1727 ) {
1728 crate::gpu::probe_record(
1729 crate::gpu::OpClass::Matmat,
1730 true,
1731 t0.elapsed(),
1732 );
1733 return;
1734 }
1735 }
1736 crate::gpu::ProbeArm::CpuTimed => {
1737 let q = self.quant_bytes();
1738 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1739 crate::gpu::probe_record(
1740 crate::gpu::OpClass::Matmat,
1741 false,
1742 t0.elapsed(),
1743 );
1744 return;
1745 }
1746 crate::gpu::ProbeArm::Cpu => {}
1747 }
1748 }
1749 }
1750 let q = self.quant_bytes();
1751 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1752 }
1753 }
1754 }
1755}
1756
1757impl QTensor {
1758 pub fn matvec_many<const N: usize>(
1765 ts: [&QTensor; N],
1766 x: &[f32],
1767 mut outs: [&mut [f32]; N],
1768 pool: Option<&Pool>,
1769 ) {
1770 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1771 let uniform_q8 = ts.iter().all(|t| {
1772 matches!(
1773 t,
1774 Self::Mapped {
1775 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1776 ..
1777 }
1778 )
1779 });
1780 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1781 let uniform_q4 = ts.iter().all(|t| {
1782 matches!(
1783 t,
1784 Self::Mapped {
1785 dtype: TensorDtype::Q4Block,
1786 ..
1787 }
1788 )
1789 });
1790 let uniform_vbit = ts.iter().all(|t| {
1791 matches!(
1792 t,
1793 Self::Mapped {
1794 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1795 ..
1796 }
1797 )
1798 });
1799 let uniform_q1 = ts.iter().all(|t| {
1800 matches!(
1801 t,
1802 Self::Mapped {
1803 dtype: TensorDtype::Q1,
1804 ..
1805 }
1806 )
1807 });
1808 let uniform_q1t = ts.iter().all(|t| {
1809 matches!(
1810 t,
1811 Self::Mapped {
1812 dtype: TensorDtype::Q1T,
1813 ..
1814 }
1815 )
1816 });
1817 let uniform_q4tp = ts.iter().all(|t| {
1822 matches!(
1823 t,
1824 Self::Mapped {
1825 dtype: TensorDtype::Q4TiledP,
1826 ..
1827 }
1828 )
1829 }) && ts.iter().all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1830 let Some(pool) = pool else {
1831 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1832 t.matvec(x, o, None);
1833 }
1834 return;
1835 };
1836 if total_rows < 256
1837 || !(uniform_q8
1838 || uniform_f32
1839 || uniform_q4
1840 || uniform_vbit
1841 || uniform_q1
1842 || uniform_q1t
1843 || uniform_q4tp)
1844 {
1845 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1846 t.matvec(x, o, Some(pool));
1847 }
1848 return;
1849 }
1850
1851 if uniform_q4tp {
1852 let cols = ts[0].cols();
1858 let gpr = cols / GROUP_SIZE;
1859 let views: Vec<Q4tpView> = ts
1860 .iter()
1861 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1862 .collect();
1863 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1864 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1865 let locate = |flat: usize| -> (usize, usize) {
1867 let mut acc = 0;
1868 for (i, &r) in rows_of.iter().enumerate() {
1869 if flat < acc + r {
1870 return (i, flat - acc);
1871 }
1872 acc += r;
1873 }
1874 (rows_of.len() - 1, 0)
1875 };
1876 let (views, outs_addr) = (&views, &outs_addr);
1877 if a8w8_enabled() {
1878 let act = split_act(x);
1879 let act = &act;
1880 let run = |start: usize, end: usize| {
1881 let mut sc = vec![0f32; gpr];
1882 for flat in start..end {
1883 let (t, r) = locate(flat);
1884 let v = &views[t];
1885 v.scales_into(r, gpr, &mut sc);
1886 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1887 for &(j, xv) in &act.outliers {
1888 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1889 acc += w * s * xv;
1890 }
1891 unsafe { *outs_addr[t].at(r) = acc };
1893 }
1894 };
1895 pool.run_rows(total_rows, &run);
1896 } else {
1897 let run = |start: usize, end: usize| {
1898 let mut sc = vec![0f32; gpr];
1899 for flat in start..end {
1900 let (t, r) = locate(flat);
1901 let v = &views[t];
1902 v.scales_into(r, gpr, &mut sc);
1903 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1905 }
1906 };
1907 pool.run_rows(total_rows, &run);
1908 }
1909 return;
1910 }
1911
1912 if uniform_q1 {
1913 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1916 if a8w8_enabled() {
1917 let act = split_act(x);
1918 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1919 let (act, gsum) = (&act, &gsum);
1920 let closures: [_; N] = std::array::from_fn(|i| {
1921 let (bytes, gpr, out) =
1922 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1923 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1924 });
1925 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1926 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1927 pool.run_many(&parts);
1928 } else {
1929 let closures: [_; N] = std::array::from_fn(|i| {
1930 let (bytes, gpr, out) =
1931 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1932 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1933 });
1934 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1935 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1936 pool.run_many(&parts);
1937 }
1938 return;
1939 }
1940
1941 if uniform_q1t {
1942 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1946 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1947 if a8w8_enabled() {
1948 let act = split_act(x);
1949 let act = &act;
1950 let x_ref = x;
1951 let closures: [_; N] = std::array::from_fn(|i| {
1952 let bytes = ts[i].quant_bytes();
1953 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1954 let gpr = cols / GROUP_SIZE;
1955 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1956 let out = outs_addr[i];
1957 move |s: usize, e: usize| {
1958 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1959 }
1960 });
1961 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1962 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1963 pool.run_many(&parts);
1964 } else {
1965 let x_ref = x;
1966 let closures: [_; N] = std::array::from_fn(|i| {
1967 let bytes = ts[i].quant_bytes();
1968 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1969 let gpr = cols / GROUP_SIZE;
1970 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1971 let out = outs_addr[i];
1972 move |s: usize, e: usize| {
1973 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1974 }
1975 });
1976 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1977 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1978 pool.run_many(&parts);
1979 }
1980 return;
1981 }
1982
1983 if uniform_q4 || uniform_vbit {
1984 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1985 if a8w8_enabled() {
1987 let act = split_act(x);
1988 let act = &act;
1989 if uniform_q4 {
1990 let closures: [_; N] = std::array::from_fn(|i| {
1991 let (packed, scales) =
1992 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1993 let (gpr, cols, out) =
1994 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1995 move |s: usize, e: usize| {
1996 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1997 }
1998 });
1999 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2000 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2001 pool.run_many(&parts);
2002 } else {
2003 let closures: [_; N] = std::array::from_fn(|i| {
2004 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2005 unreachable!()
2006 };
2007 let (bytes, rows, cols, out) = (
2008 ts[i].quant_bytes(),
2009 ts[i].rows(),
2010 ts[i].cols(),
2011 outs_addr[i],
2012 );
2013 move |s: usize, e: usize| {
2014 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2015 }
2016 });
2017 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2018 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2019 pool.run_many(&parts);
2020 }
2021 return;
2022 }
2023 if uniform_q4 {
2024 let closures: [_; N] = std::array::from_fn(|i| {
2025 let (packed, scales) =
2026 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2027 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2028 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2029 });
2030 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2031 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2032 pool.run_many(&parts);
2033 } else {
2034 let closures: [_; N] = std::array::from_fn(|i| {
2035 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2036 unreachable!()
2037 };
2038 let (bytes, rows, cols, out) = (
2039 ts[i].quant_bytes(),
2040 ts[i].rows(),
2041 ts[i].cols(),
2042 outs_addr[i],
2043 );
2044 move |s: usize, e: usize| {
2045 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2046 }
2047 });
2048 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2049 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2050 pool.run_many(&parts);
2051 }
2052 return;
2053 }
2054
2055 if uniform_f32 {
2056 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2057 let closures: [_; N] = std::array::from_fn(|i| {
2058 let Self::F32 { data, cols, .. } = ts[i] else {
2059 unreachable!()
2060 };
2061 let out = outs_addr[i];
2062 move |start: usize, end: usize| {
2063 for o in start..end {
2064 let row = &data[o * cols..(o + 1) * cols];
2065 let mut sum = 0.0f32;
2066 for j in 0..*cols {
2067 sum += row[j] * x[j];
2068 }
2069 unsafe { *out.at(o) = sum };
2071 }
2072 }
2073 });
2074 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2075 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2076 pool.run_many(&parts);
2077 return;
2078 }
2079
2080 struct Ctx<'a> {
2083 bytes: &'a [u8],
2084 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2085 rep: &'a [u8],
2086 row_scale: &'a [f32],
2087 cols: usize,
2088 xs: std::borrow::Cow<'a, [f32]>,
2089 }
2090 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2091 let Self::Mapped {
2092 dtype,
2093 cols,
2094 row_scale,
2095 col_field,
2096 repack,
2097 ..
2098 } = ts[i]
2099 else {
2100 unreachable!()
2101 };
2102 Ctx {
2103 bytes: ts[i].quant_bytes(),
2104 rep: repack,
2105 row_scale,
2106 cols: *cols,
2107 xs: prescale(x, col_field, *dtype),
2108 }
2109 });
2110 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2111 #[cfg(target_arch = "aarch64")]
2112 if sdot_enabled() {
2113 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2114 let closures: [_; N] = std::array::from_fn(|i| {
2115 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2116 move |start: usize, end: usize| {
2117 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2118 }
2119 });
2120 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2121 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2122 pool.run_many(&parts);
2123 return;
2124 }
2125 #[cfg(target_arch = "x86_64")]
2126 if avx2_a8w8_enabled() {
2127 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2128 let closures: [_; N] = std::array::from_fn(|i| {
2129 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2130 move |start: usize, end: usize| {
2131 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2132 }
2133 });
2134 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2135 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2136 pool.run_many(&parts);
2137 return;
2138 }
2139 let closures: [_; N] = std::array::from_fn(|i| {
2140 let (c, out) = (&ctxs[i], outs_addr[i]);
2141 move |start: usize, end: usize| {
2142 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2143 }
2144 });
2145 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2146 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2147 pool.run_many(&parts);
2148 }
2149}
2150
2151impl QTensor {
2152 #[allow(clippy::needless_range_loop)]
2157 pub fn matvec2_many<const N: usize>(
2158 ts: [&QTensor; N],
2159 x1: &[f32],
2160 x2: &[f32],
2161 mut o1s: [&mut [f32]; N],
2162 mut o2s: [&mut [f32]; N],
2163 pool: Option<&Pool>,
2164 ) {
2165 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2166 let uniform_q8 = ts.iter().all(|t| {
2167 matches!(
2168 t,
2169 Self::Mapped {
2170 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2171 ..
2172 }
2173 )
2174 });
2175 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2176 let uniform_q4 = ts.iter().all(|t| {
2177 matches!(
2178 t,
2179 Self::Mapped {
2180 dtype: TensorDtype::Q4Block,
2181 ..
2182 }
2183 )
2184 });
2185 let uniform_vbit = ts.iter().all(|t| {
2186 matches!(
2187 t,
2188 Self::Mapped {
2189 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2190 ..
2191 }
2192 )
2193 });
2194 let fusable = pool.is_some()
2195 && total_rows >= 256
2196 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2197 if !fusable {
2198 for i in 0..N {
2199 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2200 }
2201 return;
2202 }
2203 let pool = pool.unwrap();
2204
2205 if uniform_q4 || uniform_vbit {
2206 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2207 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2208 if a8w8_enabled() {
2210 let a1 = split_act(x1);
2211 let a2 = split_act(x2);
2212 let (a1, a2) = (&a1, &a2);
2213 if uniform_q4 {
2214 let closures: [_; N] = std::array::from_fn(|i| {
2215 let (packed, scales) =
2216 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2217 let (gpr, cols, o1, o2) =
2218 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2219 move |s: usize, e: usize| {
2220 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2221 }
2222 });
2223 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2224 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2225 pool.run_many(&parts);
2226 } else {
2227 let closures: [_; N] = std::array::from_fn(|i| {
2228 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2229 unreachable!()
2230 };
2231 let (bytes, rows, cols, o1, o2) = (
2232 ts[i].quant_bytes(),
2233 ts[i].rows(),
2234 ts[i].cols(),
2235 p1[i],
2236 p2[i],
2237 );
2238 move |s: usize, e: usize| {
2239 vbit_range2_a8w8(
2240 bytes,
2241 vbit_offsets,
2242 x1,
2243 x2,
2244 a1,
2245 a2,
2246 rows,
2247 cols,
2248 o1,
2249 o2,
2250 s,
2251 e,
2252 )
2253 }
2254 });
2255 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2256 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2257 pool.run_many(&parts);
2258 }
2259 return;
2260 }
2261 if uniform_q4 {
2262 let closures: [_; N] = std::array::from_fn(|i| {
2263 let (packed, scales) =
2264 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2265 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2266 move |s: usize, e: usize| {
2267 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2268 }
2269 });
2270 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2271 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2272 pool.run_many(&parts);
2273 } else {
2274 let closures: [_; N] = std::array::from_fn(|i| {
2275 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2276 unreachable!()
2277 };
2278 let (bytes, rows, cols, o1, o2) = (
2279 ts[i].quant_bytes(),
2280 ts[i].rows(),
2281 ts[i].cols(),
2282 p1[i],
2283 p2[i],
2284 );
2285 move |s: usize, e: usize| {
2286 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2287 }
2288 });
2289 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2290 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2291 pool.run_many(&parts);
2292 }
2293 return;
2294 }
2295
2296 if uniform_f32 {
2297 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2298 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2299 let closures: [_; N] = std::array::from_fn(|i| {
2300 let Self::F32 { data, cols, .. } = ts[i] else {
2301 unreachable!()
2302 };
2303 let (o1, o2) = (p1[i], p2[i]);
2304 move |start: usize, end: usize| {
2305 for o in start..end {
2306 let row = &data[o * cols..(o + 1) * cols];
2307 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2308 for j in 0..*cols {
2309 s1 += row[j] * x1[j];
2310 s2 += row[j] * x2[j];
2311 }
2312 unsafe {
2314 *o1.at(o) = s1;
2315 *o2.at(o) = s2;
2316 }
2317 }
2318 }
2319 });
2320 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2321 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2322 pool.run_many(&parts);
2323 return;
2324 }
2325
2326 struct Ctx<'a> {
2327 bytes: &'a [u8],
2328 row_scale: &'a [f32],
2329 cols: usize,
2330 xs1: std::borrow::Cow<'a, [f32]>,
2331 xs2: std::borrow::Cow<'a, [f32]>,
2332 }
2333 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2334 let Self::Mapped {
2335 dtype,
2336 cols,
2337 row_scale,
2338 col_field,
2339 ..
2340 } = ts[i]
2341 else {
2342 unreachable!()
2343 };
2344 Ctx {
2345 bytes: ts[i].quant_bytes(),
2346 row_scale,
2347 cols: *cols,
2348 xs1: prescale(x1, col_field, *dtype),
2349 xs2: prescale(x2, col_field, *dtype),
2350 }
2351 });
2352 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2353 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2354 #[cfg(target_arch = "aarch64")]
2355 if sdot_enabled() {
2356 let acts: [(SplitAct, SplitAct); N] =
2357 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2358 let closures: [_; N] = std::array::from_fn(|i| {
2359 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2360 move |start: usize, end: usize| {
2361 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2362 }
2363 });
2364 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2365 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2366 pool.run_many(&parts);
2367 return;
2368 }
2369 #[cfg(target_arch = "x86_64")]
2370 if avx2_a8w8_enabled() {
2371 let acts: [(SplitAct, SplitAct); N] =
2372 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2373 let closures: [_; N] = std::array::from_fn(|i| {
2374 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2375 move |start: usize, end: usize| {
2376 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2377 }
2378 });
2379 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2380 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2381 pool.run_many(&parts);
2382 return;
2383 }
2384 let closures: [_; N] = std::array::from_fn(|i| {
2385 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2386 move |start: usize, end: usize| {
2387 q8_range2_f32(
2388 c.bytes,
2389 c.row_scale,
2390 &c.xs1,
2391 &c.xs2,
2392 c.cols,
2393 o1,
2394 o2,
2395 start,
2396 end,
2397 )
2398 }
2399 });
2400 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2401 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2402 pool.run_many(&parts);
2403 }
2404
2405 pub fn matvec_silu_mul(
2410 gate: &QTensor,
2411 up: &QTensor,
2412 x: &[f32],
2413 out: &mut [f32],
2414 pool: Option<&Pool>,
2415 ) -> bool {
2416 let inter = gate.rows();
2417 debug_assert_eq!(up.rows(), inter);
2418 debug_assert_eq!(out.len(), inter);
2419 debug_assert_eq!(gate.cols(), up.cols());
2420 if !a8w8_enabled() {
2421 return false;
2422 }
2423 let act = split_act(x);
2424 let act = &act;
2425 let x_ref = x;
2426 let out_addr = SendMut(out.as_mut_ptr());
2427
2428 match (gate, up) {
2429 (
2431 Self::Mapped {
2432 dtype: TensorDtype::Q4Block,
2433 ..
2434 },
2435 Self::Mapped {
2436 dtype: TensorDtype::Q4Block,
2437 ..
2438 },
2439 ) => {
2440 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2441 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2442 let gpr = gate.cols() / GROUP_SIZE;
2443 let cols = gate.cols();
2444 let run = move |start: usize, end: usize| {
2445 for r in start..end {
2446 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2447 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2448 for &(j, xv) in &act.outliers {
2449 let flat = r * cols + j;
2450 let gb = gp[flat / 2];
2451 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2452 let gsc = f16_to_f32(u16::from_le_bytes([
2453 gs[(flat / GROUP_SIZE) * 2],
2454 gs[(flat / GROUP_SIZE) * 2 + 1],
2455 ]));
2456 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2457 let ub = up_p[flat / 2];
2458 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2459 let usc = f16_to_f32(u16::from_le_bytes([
2460 up_s[(flat / GROUP_SIZE) * 2],
2461 up_s[(flat / GROUP_SIZE) * 2 + 1],
2462 ]));
2463 uv += ((un as i32 - 8) as f32) * usc * xv;
2464 }
2465 let silu_g = gv / (1.0 + (-gv).exp());
2466 unsafe { *out_addr.at(r) = silu_g * uv };
2468 }
2469 };
2470 dispatch_rows(pool, inter, &run);
2471 true
2472 }
2473 (
2477 Self::Mapped {
2478 dtype: TensorDtype::Q4Tiled,
2479 ..
2480 },
2481 Self::Mapped {
2482 dtype: TensorDtype::Q4Tiled,
2483 ..
2484 },
2485 ) => {
2486 let g_bytes = gate.quant_bytes();
2487 let u_bytes = up.quant_bytes();
2488 let gpr = gate.cols() / GROUP_SIZE;
2489 let run = move |start: usize, end: usize| {
2490 for r in start..end {
2491 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2492 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2493 for &(j, xv) in &act.outliers {
2494 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2495 gv += w * s * xv;
2496 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2497 uv += w * s * xv;
2498 }
2499 let silu_g = gv / (1.0 + (-gv).exp());
2500 unsafe { *out_addr.at(r) = silu_g * uv };
2502 }
2503 };
2504 dispatch_rows(pool, inter, &run);
2505 true
2506 }
2507 (
2510 Self::Mapped {
2511 dtype: TensorDtype::Q4TiledP,
2512 ..
2513 },
2514 Self::Mapped {
2515 dtype: TensorDtype::Q4TiledP,
2516 ..
2517 },
2518 ) => {
2519 let cols = gate.cols();
2520 let gpr = cols / GROUP_SIZE;
2521 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2522 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2523 let run = |start: usize, end: usize| {
2524 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2525 for r in start..end {
2526 gv_view.scales_into(r, gpr, &mut gsc);
2527 uv_view.scales_into(r, gpr, &mut usc);
2528 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2529 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2530 for &(j, xv) in &act.outliers {
2531 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2532 gv += w * s * xv;
2533 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2534 uv += w * s * xv;
2535 }
2536 let silu_g = gv / (1.0 + (-gv).exp());
2537 unsafe { *out_addr.at(r) = silu_g * uv };
2539 }
2540 };
2541 dispatch_rows(pool, inter, &run);
2542 true
2543 }
2544 (
2550 Self::Mapped {
2551 dtype: TensorDtype::Q1,
2552 ..
2553 },
2554 Self::Mapped {
2555 dtype: TensorDtype::Q1,
2556 ..
2557 },
2558 ) => {
2559 let g_bytes = gate.quant_bytes();
2560 let u_bytes = up.quant_bytes();
2561 let gpr = gate.cols() / GROUP_SIZE;
2562 let gsum = q1_group_sums(&act.xq, gpr);
2563 let gsum = &gsum;
2564 let run = move |start: usize, end: usize| {
2565 for r in start..end {
2566 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2567 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2568 for &(j, xv) in &act.outliers {
2569 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2570 gv += w * s * xv;
2571 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2572 uv += w * s * xv;
2573 }
2574 let silu_g = gv / (1.0 + (-gv).exp());
2575 unsafe { *out_addr.at(r) = silu_g * uv };
2577 }
2578 };
2579 dispatch_rows(pool, inter, &run);
2580 true
2581 }
2582 (
2586 Self::Mapped {
2587 dtype: TensorDtype::Q2TiledP,
2588 ..
2589 },
2590 Self::Mapped {
2591 dtype: TensorDtype::Q2TiledP,
2592 ..
2593 },
2594 ) => {
2595 let cols = gate.cols();
2596 let gpr = cols / GROUP_SIZE;
2597 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2598 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2599 let gsum = q1_group_sums(&act.xq, gpr);
2600 let gsum = &gsum;
2601 let run = move |start: usize, end: usize| {
2602 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2603 for r in start..end {
2604 gv_view.scales_into(r, gpr, &mut gsc);
2605 uv_view.scales_into(r, gpr, &mut usc);
2606 let mut gv =
2607 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2608 let mut uv =
2609 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2610 for &(j, xv) in &act.outliers {
2611 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2612 gv += w * s * xv;
2613 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2614 uv += w * s * xv;
2615 }
2616 let silu_g = gv / (1.0 + (-gv).exp());
2617 unsafe { *out_addr.at(r) = silu_g * uv };
2619 }
2620 };
2621 dispatch_rows(pool, inter, &run);
2622 true
2623 }
2624 (
2629 Self::Mapped {
2630 dtype: TensorDtype::Q8Row,
2631 row_scale: g_rs,
2632 ..
2633 },
2634 Self::Mapped {
2635 dtype: TensorDtype::Q8Row,
2636 row_scale: u_rs,
2637 ..
2638 },
2639 ) => {
2640 let g_bytes = gate.quant_bytes();
2641 let u_bytes = up.quant_bytes();
2642 let cols = gate.cols();
2643 let run = move |start: usize, end: usize| {
2644 for r in start..end {
2645 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2646 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2647 let silu_g = gv / (1.0 + (-gv).exp());
2648 unsafe { *out_addr.at(r) = silu_g * uv };
2650 }
2651 };
2652 dispatch_rows(pool, inter, &run);
2653 true
2654 }
2655 (
2657 Self::Mapped {
2658 dtype: TensorDtype::Q1T,
2659 ..
2660 },
2661 Self::Mapped {
2662 dtype: TensorDtype::Q1T,
2663 ..
2664 },
2665 ) => {
2666 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2667 let g_bytes = gate.quant_bytes();
2668 let u_bytes = up.quant_bytes();
2669 let gpr = gate.cols() / GROUP_SIZE;
2670 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2671 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2672 let run = move |start: usize, end: usize| {
2673 for r in start..end {
2674 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2675 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2676 for &(j, xv) in &act.outliers {
2677 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2678 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2679 }
2680 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2681 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2682 let silu_g = gv / (1.0 + (-gv).exp());
2683 unsafe { *out_addr.at(r) = silu_g * uv };
2685 }
2686 };
2687 dispatch_rows(pool, inter, &run);
2688 true
2689 }
2690 _ => false,
2691 }
2692 }
2693
2694 pub fn moe_gate_up_many(
2708 pairs: &[(&QTensor, &QTensor)],
2709 x: &[f32],
2710 outs: &mut [Vec<f32>],
2711 pool: Option<&Pool>,
2712 ) -> bool {
2713 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2714 return false;
2715 }
2716 let inter = pairs[0].0.rows();
2717 let cols = pairs[0].0.cols();
2718 if cols % GROUP_SIZE != 0 {
2719 return false;
2720 }
2721 let gpr = cols / GROUP_SIZE;
2722 let q2 = matches!(
2725 pairs[0].0,
2726 Self::Mapped {
2727 dtype: TensorDtype::Q2TiledP,
2728 ..
2729 }
2730 );
2731 let want = if q2 {
2732 TensorDtype::Q2TiledP
2733 } else {
2734 TensorDtype::Q4TiledP
2735 };
2736 let mut views = Vec::with_capacity(pairs.len() * 2);
2737 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2738 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2739 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2740 if !both
2741 || g.rows() != inter
2742 || u.rows() != inter
2743 || g.cols() != cols
2744 || u.cols() != cols
2745 || o.len() != inter
2746 {
2747 return false;
2748 }
2749 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2750 views.push(mk(g.quant_bytes(), inter, cols));
2751 views.push(mk(u.quant_bytes(), inter, cols));
2752 }
2753 let act = split_act(x);
2754 let gsum = if q2 {
2755 q1_group_sums(&act.xq, gpr)
2756 } else {
2757 Vec::new()
2758 };
2759 let (act, gsum) = (&act, &gsum);
2760 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2761 let (views, ptrs) = (&views, &ptrs);
2762 let run = |start: usize, end: usize| {
2763 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2764 for flat in start..end {
2765 let (e, r) = (flat / inter, flat % inter);
2766 let gv_view = &views[e * 2];
2767 let uv_view = &views[e * 2 + 1];
2768 gv_view.scales_into(r, gpr, &mut gsc);
2769 uv_view.scales_into(r, gpr, &mut usc);
2770 let (mut gv, mut uv) = if q2 {
2771 (
2772 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2773 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2774 )
2775 } else {
2776 (
2777 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2778 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2779 )
2780 };
2781 for &(j, xv) in &act.outliers {
2782 let (og, ou) = if q2 {
2783 (
2784 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2785 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2786 )
2787 } else {
2788 (
2789 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2790 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2791 )
2792 };
2793 gv += og.0 * og.1 * xv;
2794 uv += ou.0 * ou.1 * xv;
2795 }
2796 let silu_g = gv / (1.0 + (-gv).exp());
2797 unsafe { *ptrs[e].at(r) = silu_g * uv };
2799 }
2800 };
2801 dispatch_rows(pool, pairs.len() * inter, &run);
2802 true
2803 }
2804
2805 pub fn moe_down_many(
2814 downs: &[&QTensor],
2815 gs: &[Vec<f32>],
2816 weights: &[f32],
2817 out: &mut [f32],
2818 pool: Option<&Pool>,
2819 ) -> bool {
2820 if downs.is_empty()
2821 || downs.len() != gs.len()
2822 || downs.len() != weights.len()
2823 || !a8w8_enabled()
2824 {
2825 return false;
2826 }
2827 let rows = out.len();
2828 let cols = downs[0].cols();
2829 if cols % GROUP_SIZE != 0 {
2830 return false;
2831 }
2832 let gpr = cols / GROUP_SIZE;
2833 let mut views = Vec::with_capacity(downs.len());
2834 for (d, g) in downs.iter().zip(gs.iter()) {
2835 if !matches!(
2836 d,
2837 Self::Mapped {
2838 dtype: TensorDtype::Q4TiledP,
2839 ..
2840 }
2841 ) || d.rows() != rows
2842 || d.cols() != cols
2843 || g.len() != cols
2844 {
2845 return false;
2846 }
2847 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2848 }
2849 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2851 let out_addr = SendMut(out.as_mut_ptr());
2859 let (views, acts, weights) = (&views, &acts, &weights);
2860 let run = |start: usize, end: usize| {
2861 let mut sc = vec![0f32; gpr];
2862 for r in start..end {
2863 let mut acc = 0f32;
2864 for (e, v) in views.iter().enumerate() {
2865 v.scales_into(r, gpr, &mut sc);
2866 let a = &acts[e];
2867 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2868 for &(j, xv) in &a.outliers {
2869 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2870 d += w * s * xv;
2871 }
2872 acc += weights[e] * d;
2873 }
2874 unsafe { *out_addr.at(r) = acc };
2876 }
2877 };
2878 dispatch_rows(pool, rows, &run);
2879 true
2880 }
2881}
2882
2883#[cfg(target_os = "macos")]
2888mod accel_blas {
2889 #[link(name = "Accelerate", kind = "framework")]
2890 unsafe extern "C" {
2891 pub fn cblas_sgemm(
2892 order: i32,
2893 trans_a: i32,
2894 trans_b: i32,
2895 m: i32,
2896 n: i32,
2897 k: i32,
2898 alpha: f32,
2899 a: *const f32,
2900 lda: i32,
2901 b: *const f32,
2902 ldb: i32,
2903 beta: f32,
2904 c: *mut f32,
2905 ldc: i32,
2906 );
2907 }
2908}
2909
2910#[cfg(target_os = "macos")]
2911pub(crate) fn accel_gemm_enabled() -> bool {
2912 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2913 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2914}
2915
2916#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2919pub(crate) fn accel_gemm_enabled() -> bool {
2920 true
2921}
2922
2923#[cfg(target_arch = "aarch64")]
2930#[allow(clippy::too_many_arguments)]
2931pub(crate) fn neon_gemm_rm(
2932 m: usize,
2933 n: usize,
2934 k: usize,
2935 alpha: f32,
2936 a: &[f32],
2937 lda: usize,
2938 b_mat: &[f32],
2939 ldb: usize,
2940 b_rows_are_n: bool,
2941 c: &mut [f32],
2942 ldc: usize,
2943) {
2944 debug_assert!(a.len() >= (m - 1) * lda + k);
2945 debug_assert!(c.len() >= (m - 1) * ldc + n);
2946 unsafe {
2948 use core::arch::aarch64::*;
2949 let mut i = 0usize;
2950 while i < m {
2951 let mi = (m - i).min(4);
2952 let mut j = 0usize;
2953 while j < n {
2954 let nj = (n - j).min(8);
2955 if mi == 4 && nj == 8 {
2956 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2957 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2958 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2959 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2960 for p in 0..k {
2961 let (b0, b1) = if b_rows_are_n {
2962 let base = b_mat.as_ptr().add(j * ldb + p);
2965 let g = |o: usize| *base.add(o * ldb);
2966 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2967 } else {
2968 let base = b_mat.as_ptr().add(p * ldb + j);
2969 (
2970 [*base, *base.add(1), *base.add(2), *base.add(3)],
2971 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2972 )
2973 };
2974 let bv0 = vld1q_f32(b0.as_ptr());
2975 let bv1 = vld1q_f32(b1.as_ptr());
2976 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2977 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2978 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2979 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2980 c0a = vfmaq_f32(c0a, a0, bv0);
2981 c0b = vfmaq_f32(c0b, a0, bv1);
2982 c1a = vfmaq_f32(c1a, a1, bv0);
2983 c1b = vfmaq_f32(c1b, a1, bv1);
2984 c2a = vfmaq_f32(c2a, a2, bv0);
2985 c2b = vfmaq_f32(c2b, a2, bv1);
2986 c3a = vfmaq_f32(c3a, a3, bv0);
2987 c3b = vfmaq_f32(c3b, a3, bv1);
2988 }
2989 let al = vdupq_n_f32(alpha);
2990 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2991 .iter()
2992 .enumerate()
2993 {
2994 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2995 vst1q_f32(dst, vmulq_f32(*ca, al));
2996 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2997 }
2998 } else {
2999 for r in 0..mi {
3000 for q in 0..nj {
3001 let mut acc = 0f32;
3002 for p in 0..k {
3003 let bv = if b_rows_are_n {
3004 b_mat[(j + q) * ldb + p]
3005 } else {
3006 b_mat[p * ldb + j + q]
3007 };
3008 acc += a[(i + r) * lda + p] * bv;
3009 }
3010 c[(i + r) * ldc + j + q] = acc * alpha;
3011 }
3012 }
3013 }
3014 j += nj;
3015 }
3016 i += mi;
3017 }
3018 }
3019}
3020
3021#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3023#[allow(clippy::too_many_arguments)]
3024pub(crate) fn sgemm_rm(
3025 m: usize,
3026 n: usize,
3027 k: usize,
3028 alpha: f32,
3029 a: &[f32],
3030 lda: usize,
3031 b_mat: &[f32],
3032 ldb: usize,
3033 b_rows_are_n: bool,
3034 c: &mut [f32],
3035 ldc: usize,
3036) {
3037 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3038}
3039
3040#[allow(clippy::too_many_arguments)]
3044pub fn sgemm_public(
3045 m: usize,
3046 n: usize,
3047 k: usize,
3048 alpha: f32,
3049 a: &[f32],
3050 lda: usize,
3051 b_mat: &[f32],
3052 ldb: usize,
3053 b_rows_are_n: bool,
3054 c: &mut [f32],
3055 ldc: usize,
3056) {
3057 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3058 {
3059 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3060 }
3061 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3066 {
3067 for i in 0..m {
3068 for j in 0..n {
3069 let mut acc = 0f32;
3070 for p in 0..k {
3071 let bv = if b_rows_are_n {
3072 b_mat[j * ldb + p]
3073 } else {
3074 b_mat[p * ldb + j]
3075 };
3076 acc += a[i * lda + p] * bv;
3077 }
3078 c[i * ldc + j] = alpha * acc;
3079 }
3080 }
3081 }
3082}
3083
3084#[cfg(target_os = "macos")]
3087#[allow(clippy::too_many_arguments)]
3088pub(crate) fn sgemm_rm(
3089 m: usize,
3090 n: usize,
3091 k: usize,
3092 alpha: f32,
3093 a: &[f32],
3094 lda: usize,
3095 b_mat: &[f32],
3096 ldb: usize,
3097 b_rows_are_n: bool,
3098 c: &mut [f32],
3099 ldc: usize,
3100) {
3101 debug_assert!(a.len() >= (m - 1) * lda + k);
3102 debug_assert!(c.len() >= (m - 1) * ldc + n);
3103 #[cfg(target_arch = "aarch64")]
3108 if std::env::var("CMF_FORCE_NEON_GEMM")
3109 .map(|v| v == "1")
3110 .unwrap_or(false)
3111 {
3112 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3113 }
3114 unsafe {
3115 accel_blas::cblas_sgemm(
3116 101, 111, if b_rows_are_n { 112 } else { 111 },
3119 m as i32,
3120 n as i32,
3121 k as i32,
3122 alpha,
3123 a.as_ptr(),
3124 lda as i32,
3125 b_mat.as_ptr(),
3126 ldb as i32,
3127 0.0,
3128 c.as_mut_ptr(),
3129 ldc as i32,
3130 );
3131 }
3132}
3133
3134#[cfg(target_os = "macos")]
3141fn qmatmat_accel(
3142 q: &[u8],
3143 row_scale: &[f32],
3144 pre: &[std::borrow::Cow<'_, [f32]>],
3145 rows: usize,
3146 cols: usize,
3147 out: &mut [f32],
3148 pool: Option<&Pool>,
3149) {
3150 const TR: usize = 2048;
3155 let b = pre.len();
3156 thread_local! {
3157 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3158 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3159 }
3160 XPANEL.with(|xp| {
3161 WTILE.with(|wt| {
3162 let mut xpanel = xp.borrow_mut();
3163 xpanel.clear();
3164 for x in pre {
3165 xpanel.extend_from_slice(x);
3166 }
3167 let mut wtile = wt.borrow_mut();
3168 wtile.resize(TR * cols, 0.0);
3169 let mut r0 = 0usize;
3170 while r0 < rows {
3171 let tr = TR.min(rows - r0);
3172 let wt_addr = SendMut(wtile.as_mut_ptr());
3174 let run = |start: usize, end: usize| {
3175 for r in start..end {
3176 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3177 let s = row_scale[r0 + r];
3178 let dst =
3180 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3181 for (d, &v) in dst.iter_mut().zip(row) {
3182 *d = (v as i8) as f32 * s;
3183 }
3184 }
3185 };
3186 dispatch_rows(pool, tr, &run);
3187 unsafe {
3189 accel_blas::cblas_sgemm(
3190 101, 111, 112, b as i32,
3194 tr as i32,
3195 cols as i32,
3196 1.0,
3197 xpanel.as_ptr(),
3198 cols as i32,
3199 wtile.as_ptr(),
3200 cols as i32,
3201 0.0,
3202 out.as_mut_ptr().add(r0),
3203 rows as i32,
3204 );
3205 }
3206 r0 += tr;
3207 }
3208 })
3209 });
3210}
3211
3212fn qmatmat(
3213 q: &[u8],
3214 row_scale: &[f32],
3215 pre: &[std::borrow::Cow<'_, [f32]>],
3216 rows: usize,
3217 cols: usize,
3218 out: &mut [f32],
3219 pool: Option<&Pool>,
3220) {
3221 let b = pre.len();
3222 debug_assert_eq!(out.len(), b * rows);
3223 #[cfg(target_os = "macos")]
3228 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3229 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3230 return;
3231 }
3232 #[cfg(target_arch = "aarch64")]
3233 if sdot_enabled() {
3234 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3235 let out_addr = SendMut(out.as_mut_ptr());
3236 let blocked_ok = blocked_enabled();
3239 let use_i8mm = i8mm_enabled();
3240 if blocked_ok {
3241 let run = |start: usize, end: usize| {
3242 let mut o = start;
3243 while o < end {
3244 if o + 2 <= end {
3245 let r0 = &q[o * cols..(o + 1) * cols];
3246 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3247 let mut bi = 0usize;
3248 while bi + 4 <= acts.len() {
3249 let xs = [
3250 acts[bi].xq.as_slice(),
3251 acts[bi + 1].xq.as_slice(),
3252 acts[bi + 2].xq.as_slice(),
3253 acts[bi + 3].xq.as_slice(),
3254 ];
3255 let d = if use_i8mm {
3256 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3257 } else {
3258 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3259 };
3260 for (r, row) in [r0, r1].into_iter().enumerate() {
3261 for k in 0..4 {
3262 let act = &acts[bi + k];
3263 let mut v = d[r][k] as f32 * act.sx;
3264 for &(j, xv) in &act.outliers {
3265 v += (row[j] as i8) as f32 * xv;
3266 }
3267 unsafe {
3268 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3269 };
3270 }
3271 }
3272 bi += 4;
3273 }
3274 while bi < acts.len() {
3275 for (r, row) in [r0, r1].into_iter().enumerate() {
3276 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3277 unsafe { *out_addr.at(bi * rows + o + r) = v };
3278 }
3279 bi += 1;
3280 }
3281 o += 2;
3282 } else {
3283 let row = &q[o * cols..(o + 1) * cols];
3284 for (bi, act) in acts.iter().enumerate() {
3285 let v = row_dot_sdot(row, act) * row_scale[o];
3286 unsafe { *out_addr.at(bi * rows + o) = v };
3287 }
3288 o += 1;
3289 }
3290 }
3291 };
3292 dispatch_rows(pool, rows, &run);
3293 return;
3294 }
3295 let run = |start: usize, end: usize| {
3296 for o in start..end {
3297 let row = &q[o * cols..(o + 1) * cols];
3298 for (bi, act) in acts.iter().enumerate() {
3299 let v = row_dot_sdot(row, act) * row_scale[o];
3300 unsafe { *out_addr.at(bi * rows + o) = v };
3301 }
3302 }
3303 };
3304 dispatch_rows(pool, rows, &run);
3305 return;
3306 }
3307 #[cfg(target_arch = "x86_64")]
3312 if avx2_a8w8_enabled() {
3313 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3314 let out_addr = SendMut(out.as_mut_ptr());
3315 let blocked_ok = blocked_enabled();
3318 if !avx512vnni_enabled() && blocked_ok {
3319 let run = |start: usize, end: usize| {
3320 let mut o = start;
3321 while o < end {
3322 if o + 2 <= end {
3323 let r0 = &q[o * cols..(o + 1) * cols];
3324 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3325 let mut bi = 0usize;
3326 while bi + 4 <= acts.len() {
3327 let xs = [
3328 acts[bi].xq.as_slice(),
3329 acts[bi + 1].xq.as_slice(),
3330 acts[bi + 2].xq.as_slice(),
3331 acts[bi + 3].xq.as_slice(),
3332 ];
3333 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3334 for (r, row) in [r0, r1].into_iter().enumerate() {
3335 for k in 0..4 {
3336 let act = &acts[bi + k];
3337 let mut v = d[r][k] as f32 * act.sx;
3338 for &(j, xv) in &act.outliers {
3339 v += (row[j] as i8) as f32 * xv;
3340 }
3341 unsafe {
3342 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3343 };
3344 }
3345 }
3346 bi += 4;
3347 }
3348 while bi < acts.len() {
3349 for (r, row) in [r0, r1].into_iter().enumerate() {
3350 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3351 unsafe { *out_addr.at(bi * rows + o + r) = v };
3352 }
3353 bi += 1;
3354 }
3355 o += 2;
3356 } else {
3357 let row = &q[o * cols..(o + 1) * cols];
3358 for (bi, act) in acts.iter().enumerate() {
3359 let v = row_dot_avx2(row, act) * row_scale[o];
3360 unsafe { *out_addr.at(bi * rows + o) = v };
3361 }
3362 o += 1;
3363 }
3364 }
3365 };
3366 dispatch_rows(pool, rows, &run);
3367 return;
3368 }
3369 let run = |start: usize, end: usize| {
3370 for o in start..end {
3371 let row = &q[o * cols..(o + 1) * cols];
3372 for (bi, act) in acts.iter().enumerate() {
3373 let v = row_dot_avx2(row, act) * row_scale[o];
3374 unsafe { *out_addr.at(bi * rows + o) = v };
3375 }
3376 }
3377 };
3378 dispatch_rows(pool, rows, &run);
3379 return;
3380 }
3381 let out_addr = SendMut(out.as_mut_ptr());
3382 let run = |start: usize, end: usize| {
3383 for o in start..end {
3384 let row = &q[o * cols..(o + 1) * cols];
3385 for (bi, x) in pre.iter().enumerate() {
3386 let mut acc = 0f32;
3387 for j in 0..cols {
3388 acc += (row[j] as i8) as f32 * x[j];
3389 }
3390 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3391 }
3392 }
3393 };
3394 dispatch_rows(pool, rows, &run);
3395}
3396
3397fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3400 match pool {
3401 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3402 _ => run(0, rows),
3403 }
3404}
3405
3406fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3408 let groups = rows * cols / GROUP_SIZE;
3409 bytes.split_at(groups * 16)
3410}
3411
3412#[inline]
3417fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3418 #[cfg(target_arch = "aarch64")]
3419 unsafe {
3420 return vbit_fill4_neon(data, buf);
3421 }
3422 #[cfg(target_arch = "x86_64")]
3423 if avx2_enabled() {
3424 return unsafe { vbit_fill4_avx2(data, buf) };
3425 }
3426 #[allow(unreachable_code)]
3427 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3428 let u = unpack8::<4>(&data[blk * 4..]);
3429 for k in 0..8 {
3430 chunk[k] = (u[k] - 7) as i8 as u8;
3431 }
3432 }
3433}
3434
3435#[cfg(target_arch = "aarch64")]
3436#[target_feature(enable = "neon")]
3437unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3438 unsafe {
3441 use core::arch::aarch64::*;
3442 let n = buf.len();
3443 let mask = vdupq_n_u8(0x0F);
3444 let seven = vdupq_n_s8(7);
3445 let mut g = 0usize;
3446 while g * 32 + 32 <= n {
3447 let b = vld1q_u8(data.as_ptr().add(g * 16));
3448 let hi = vshrq_n_u8::<4>(b);
3449 let lo = vandq_u8(b, mask);
3450 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3451 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3452 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3453 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3454 g += 1;
3455 }
3456 }
3457}
3458
3459#[cfg(target_arch = "x86_64")]
3460#[target_feature(enable = "avx2")]
3461unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3462 unsafe {
3464 use core::arch::x86_64::*;
3465 let n = buf.len();
3466 let mask = _mm_set1_epi8(0x0F);
3467 let seven = _mm256_set1_epi8(7);
3468 let mut g = 0usize;
3469 while g * 32 + 32 <= n {
3470 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3471 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3472 let lo = _mm_and_si128(b, mask);
3473 let z = _mm256_sub_epi8(
3474 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3475 seven,
3476 );
3477 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3478 g += 1;
3479 }
3480 }
3481}
3482
3483#[inline(always)]
3488fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3489 let mut acc = 0u64;
3490 for i in 0..B {
3491 acc = (acc << 8) | data[i] as u64;
3492 }
3493 let mask = (1u64 << B) - 1;
3494 let mut out = [0i32; 8];
3495 for (k, o) in out.iter_mut().enumerate() {
3496 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3497 }
3498 out
3499}
3500
3501#[allow(clippy::too_many_arguments)]
3507fn vbitmatvec(
3508 bytes: &[u8],
3509 offsets: &[usize],
3510 x: &[f32],
3511 rows: usize,
3512 cols: usize,
3513 out: &mut [f32],
3514 pool: Option<&Pool>,
3515) {
3516 debug_assert_eq!(out.len(), rows);
3517 debug_assert_eq!(offsets.len(), rows + 1);
3518
3519 if a8w8_enabled() {
3523 let act = split_act(x);
3524 let out_addr = SendMut(out.as_mut_ptr());
3525 let run = move |start: usize, end: usize| {
3526 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3527 };
3528 dispatch_rows(pool, rows, &run);
3529 return;
3530 }
3531
3532 let out_addr = SendMut(out.as_mut_ptr());
3533 let run = move |start: usize, end: usize| {
3534 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3535 };
3536 dispatch_rows(pool, rows, &run);
3537}
3538
3539#[allow(clippy::too_many_arguments)]
3543fn vbit_range_a8w8(
3544 bytes: &[u8],
3545 offsets: &[usize],
3546 x: &[f32],
3547 act: &SplitAct,
3548 rows: usize,
3549 cols: usize,
3550 out: SendMut,
3551 start: usize,
3552 end: usize,
3553) {
3554 let ng = cols / GROUP_SIZE;
3555 let bits = &bytes[..rows];
3556 let sc_off = rows;
3557 let row_dot = |r: usize| -> f32 {
3558 let b = bits[r] as usize;
3559 let l = (1i32 << (b - 1)) - 1;
3560 let mask = (1u64 << b) - 1;
3561 let data = &bytes[offsets[r]..offsets[r + 1]];
3562 if b == 8 {
3563 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3565 let mut dot = 0f32;
3566 for g in 0..ng {
3567 let so = (r * ng + g) * 2;
3568 let sgf = f16_to_f32(u16::from_le_bytes([
3569 bytes[sc_off + so],
3570 bytes[sc_off + so + 1],
3571 ]));
3572 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3573 let mut gd = 0f32;
3574 for &xv in xg.iter() {
3575 if nbits < 8 {
3576 acc = (acc << 8) | data[idx] as u64;
3577 idx += 1;
3578 nbits += 8;
3579 }
3580 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3581 nbits -= 8;
3582 gd += (u - l) as f32 * xv;
3583 }
3584 dot += gd * sgf;
3585 }
3586 return dot;
3587 }
3588 thread_local! {
3592 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3593 const { std::cell::RefCell::new(Vec::new()) };
3594 }
3595 #[inline(always)]
3596 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3597 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3598 let u = unpack8::<B>(&data[blk * B..]);
3599 for k in 0..8 {
3600 chunk[k] = (u[k] - l) as i8 as u8;
3601 }
3602 }
3603 }
3604 let _ = mask;
3605 VBIT_SCRATCH.with(|scratch| {
3606 let mut buf = scratch.borrow_mut();
3607 buf.resize(cols, 0);
3608 match b {
3609 3 => fill::<3>(data, l, &mut buf),
3610 4 => vbit_fill4(data, &mut buf),
3611 5 => fill::<5>(data, l, &mut buf),
3612 6 => fill::<6>(data, l, &mut buf),
3613 _ => unreachable!(),
3614 }
3615 let mut dot = 0f32;
3616 for g in 0..ng {
3617 let so = (r * ng + g) * 2;
3618 let s = f16_to_f32(u16::from_le_bytes([
3619 bytes[sc_off + so],
3620 bytes[sc_off + so + 1],
3621 ]));
3622 let d = dot_i8_i8(
3623 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3624 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3625 ) as f32
3626 * act.sx;
3627 dot += d * s;
3628 }
3629 for &(j, xv) in &act.outliers {
3630 let so = (r * ng + j / GROUP_SIZE) * 2;
3631 let s = f16_to_f32(u16::from_le_bytes([
3632 bytes[sc_off + so],
3633 bytes[sc_off + so + 1],
3634 ]));
3635 dot += (buf[j] as i8) as f32 * s * xv;
3637 }
3638 dot
3639 })
3640 };
3641 for r in start..end {
3642 unsafe { *out.at(r) = row_dot(r) };
3644 }
3645}
3646
3647#[allow(clippy::too_many_arguments)]
3649fn vbit_range_f32(
3650 bytes: &[u8],
3651 offsets: &[usize],
3652 x: &[f32],
3653 rows: usize,
3654 cols: usize,
3655 out: SendMut,
3656 start: usize,
3657 end: usize,
3658) {
3659 let ng = cols / GROUP_SIZE;
3660 let bits = &bytes[..rows];
3661 let sc_off = rows;
3662 #[inline(always)]
3666 fn dot_row<const B: usize>(
3667 data: &[u8],
3668 bytes: &[u8],
3669 sc_off: usize,
3670 r: usize,
3671 ng: usize,
3672 x: &[f32],
3673 ) -> f32 {
3674 let l = ((1i32 << (B - 1)) - 1) as f32;
3675 let gbytes = GROUP_SIZE * B / 8;
3676 let mut dot = 0f32;
3677 for g in 0..ng {
3678 let so = (r * ng + g) * 2;
3679 let s = f16_to_f32(u16::from_le_bytes([
3680 bytes[sc_off + so],
3681 bytes[sc_off + so + 1],
3682 ]));
3683 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3684 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3685 let mut gd = 0f32;
3686 for blk in 0..GROUP_SIZE / 8 {
3687 let u = unpack8::<B>(&gd0[blk * B..]);
3688 let xb = &xg[blk * 8..blk * 8 + 8];
3689 for k in 0..8 {
3690 gd += (u[k] as f32 - l) * xb[k];
3691 }
3692 }
3693 dot += gd * s;
3694 }
3695 dot
3696 }
3697 for r in start..end {
3698 let data = &bytes[offsets[r]..offsets[r + 1]];
3699 let v = match bits[r] {
3700 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3701 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3702 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3703 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3704 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3705 b => unreachable!("vbit bit-width {b} (validated at load)"),
3706 };
3707 unsafe { *out.at(r) = v };
3709 }
3710}
3711
3712#[allow(clippy::too_many_arguments)]
3717fn vbitmatvec2(
3718 bytes: &[u8],
3719 offsets: &[usize],
3720 x1: &[f32],
3721 x2: &[f32],
3722 rows: usize,
3723 cols: usize,
3724 o1: &mut [f32],
3725 o2: &mut [f32],
3726 pool: Option<&Pool>,
3727) {
3728 debug_assert_eq!(o1.len(), rows);
3729 debug_assert_eq!(o2.len(), rows);
3730
3731 if a8w8_enabled() {
3732 let a1 = split_act(x1);
3733 let a2 = split_act(x2);
3734 let p1 = SendMut(o1.as_mut_ptr());
3735 let p2 = SendMut(o2.as_mut_ptr());
3736 let run = move |start: usize, end: usize| {
3737 vbit_range2_a8w8(
3738 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3739 )
3740 };
3741 dispatch_rows(pool, rows, &run);
3742 return;
3743 }
3744
3745 let p1 = SendMut(o1.as_mut_ptr());
3746 let p2 = SendMut(o2.as_mut_ptr());
3747 let run = move |start: usize, end: usize| {
3748 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3749 };
3750 dispatch_rows(pool, rows, &run);
3751}
3752
3753#[allow(clippy::too_many_arguments)]
3757fn vbit_range2_a8w8(
3758 bytes: &[u8],
3759 offsets: &[usize],
3760 x1: &[f32],
3761 x2: &[f32],
3762 a1: &SplitAct,
3763 a2: &SplitAct,
3764 rows: usize,
3765 cols: usize,
3766 p1: SendMut,
3767 p2: SendMut,
3768 start: usize,
3769 end: usize,
3770) {
3771 let ng = cols / GROUP_SIZE;
3772 let bits = &bytes[..rows];
3773 let sc_off = rows;
3774 let row_dots = |r: usize| -> (f32, f32) {
3775 let b = bits[r] as usize;
3776 let l = (1i32 << (b - 1)) - 1;
3777 let data = &bytes[offsets[r]..offsets[r + 1]];
3778 if b == 8 {
3779 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3782 let (mut d1, mut d2) = (0f32, 0f32);
3783 for g in 0..ng {
3784 let so = (r * ng + g) * 2;
3785 let sgf = f16_to_f32(u16::from_le_bytes([
3786 bytes[sc_off + so],
3787 bytes[sc_off + so + 1],
3788 ]));
3789 let (mut g1, mut g2) = (0f32, 0f32);
3790 for k in 0..GROUP_SIZE {
3791 if nbits < 8 {
3792 acc = (acc << 8) | data[idx] as u64;
3793 idx += 1;
3794 nbits += 8;
3795 }
3796 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3797 nbits -= 8;
3798 let w = (u - l) as f32;
3799 g1 += w * x1[g * GROUP_SIZE + k];
3800 g2 += w * x2[g * GROUP_SIZE + k];
3801 }
3802 d1 += g1 * sgf;
3803 d2 += g2 * sgf;
3804 }
3805 return (d1, d2);
3806 }
3807 thread_local! {
3808 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3809 const { std::cell::RefCell::new(Vec::new()) };
3810 }
3811 #[inline(always)]
3812 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3813 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3814 let u = unpack8::<B>(&data[blk * B..]);
3815 for k in 0..8 {
3816 chunk[k] = (u[k] - l) as i8 as u8;
3817 }
3818 }
3819 }
3820 VBIT_SCRATCH2.with(|scratch| {
3821 let mut buf = scratch.borrow_mut();
3822 buf.resize(cols, 0);
3823 match b {
3824 3 => fill::<3>(data, l, &mut buf),
3825 4 => vbit_fill4(data, &mut buf),
3826 5 => fill::<5>(data, l, &mut buf),
3827 6 => fill::<6>(data, l, &mut buf),
3828 _ => unreachable!(),
3829 }
3830 let (mut d1, mut d2) = (0f32, 0f32);
3831 for g in 0..ng {
3832 let so = (r * ng + g) * 2;
3833 let s = f16_to_f32(u16::from_le_bytes([
3834 bytes[sc_off + so],
3835 bytes[sc_off + so + 1],
3836 ]));
3837 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3838 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3839 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3840 d1 += v1 * s;
3841 d2 += v2 * s;
3842 }
3843 for &(j, xv) in &a1.outliers {
3844 let so = (r * ng + j / GROUP_SIZE) * 2;
3845 let s = f16_to_f32(u16::from_le_bytes([
3846 bytes[sc_off + so],
3847 bytes[sc_off + so + 1],
3848 ]));
3849 d1 += (buf[j] as i8) as f32 * s * xv;
3850 }
3851 for &(j, xv) in &a2.outliers {
3852 let so = (r * ng + j / GROUP_SIZE) * 2;
3853 let s = f16_to_f32(u16::from_le_bytes([
3854 bytes[sc_off + so],
3855 bytes[sc_off + so + 1],
3856 ]));
3857 d2 += (buf[j] as i8) as f32 * s * xv;
3858 }
3859 (d1, d2)
3860 })
3861 };
3862 for r in start..end {
3863 let (v1, v2) = row_dots(r);
3864 unsafe {
3866 *p1.at(r) = v1;
3867 *p2.at(r) = v2;
3868 }
3869 }
3870}
3871
3872#[allow(clippy::too_many_arguments)]
3876fn vbit_range2_f32(
3877 bytes: &[u8],
3878 offsets: &[usize],
3879 x1: &[f32],
3880 x2: &[f32],
3881 rows: usize,
3882 cols: usize,
3883 p1: SendMut,
3884 p2: SendMut,
3885 start: usize,
3886 end: usize,
3887) {
3888 let ng = cols / GROUP_SIZE;
3889 let bits = &bytes[..rows];
3890 let sc_off = rows;
3891 #[inline(always)]
3892 #[allow(clippy::too_many_arguments)]
3893 fn dot_row2<const B: usize>(
3894 data: &[u8],
3895 bytes: &[u8],
3896 sc_off: usize,
3897 r: usize,
3898 ng: usize,
3899 x1: &[f32],
3900 x2: &[f32],
3901 ) -> (f32, f32) {
3902 let l = ((1i32 << (B - 1)) - 1) as f32;
3903 let gbytes = GROUP_SIZE * B / 8;
3904 let (mut d1, mut d2) = (0f32, 0f32);
3905 for g in 0..ng {
3906 let so = (r * ng + g) * 2;
3907 let s = f16_to_f32(u16::from_le_bytes([
3908 bytes[sc_off + so],
3909 bytes[sc_off + so + 1],
3910 ]));
3911 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3912 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3913 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3914 let (mut g1, mut g2) = (0f32, 0f32);
3915 for blk in 0..GROUP_SIZE / 8 {
3916 let u = unpack8::<B>(&gd0[blk * B..]);
3917 for k in 0..8 {
3918 let w = u[k] as f32 - l;
3919 g1 += w * x1g[blk * 8 + k];
3920 g2 += w * x2g[blk * 8 + k];
3921 }
3922 }
3923 d1 += g1 * s;
3924 d2 += g2 * s;
3925 }
3926 (d1, d2)
3927 }
3928 for r in start..end {
3929 let data = &bytes[offsets[r]..offsets[r + 1]];
3930 let (v1, v2) = match bits[r] {
3931 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3932 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3933 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3934 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3935 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3936 b => unreachable!("vbit bit-width {b} (validated at load)"),
3937 };
3938 unsafe {
3940 *p1.at(r) = v1;
3941 *p2.at(r) = v2;
3942 }
3943 }
3944}
3945
3946#[inline]
3953#[allow(unreachable_code)]
3954fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3955 #[cfg(target_arch = "aarch64")]
3956 unsafe {
3957 return dot_q4t_row_sdot(bytes, r, gpr, xq);
3958 }
3959 #[cfg(target_arch = "x86_64")]
3960 unsafe {
3961 if vnni_tiles_enabled() {
3962 return dot_q4t_row_vnni(bytes, r, gpr, xq);
3963 }
3964 return dot_q4t_row_avx2(bytes, r, gpr, xq);
3965 }
3966 let mut acc = 0f32;
3967 for gi in 0..gpr {
3968 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3969 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3970 let mut d = 0i32;
3971 for (k, &b) in tile[2..].iter().enumerate() {
3972 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3973 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3974 }
3975 acc += d as f32 * s;
3976 }
3977 acc
3978}
3979
3980#[cfg(target_arch = "aarch64")]
3981#[target_feature(enable = "neon,dotprod")]
3982unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3983 unsafe {
3986 use core::arch::aarch64::*;
3987 use core::arch::asm;
3988 let lomask = vdupq_n_u8(0x0F);
3989 let eight = vdupq_n_s8(8);
3990 let mut acc = 0f32;
3991 for gi in 0..gpr {
3992 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3993 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3994 let b = vld1q_u8(t.add(2));
3995 let lo = vandq_u8(b, lomask);
3996 let hi = vshrq_n_u8::<4>(b);
3997 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3998 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3999 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4000 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4001 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4002 asm!(
4003 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4004 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4005 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4006 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4007 options(pure, nomem, nostack),
4008 );
4009 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4010 }
4011 acc
4012 }
4013}
4014
4015#[cfg(target_arch = "x86_64")]
4016#[target_feature(enable = "avx2")]
4017unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4018 unsafe {
4020 use core::arch::x86_64::*;
4021 let lomask = _mm_set1_epi8(0x0F);
4022 let eight = _mm256_set1_epi8(8);
4023 let ones = _mm256_set1_epi16(1);
4024 let mut acc = 0f32;
4025 for gi in 0..gpr {
4026 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4027 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4028 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4029 let lo = _mm_and_si128(b, lomask);
4030 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4031 let w = _mm256_sub_epi8(
4032 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4033 eight,
4034 );
4035 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4036 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4037 let d = _mm256_madd_epi16(p16, ones);
4038 let hi128 = _mm256_extracti128_si256::<1>(d);
4039 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4040 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4041 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4042 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4043 }
4044 acc
4045 }
4046}
4047
4048#[cfg(target_arch = "x86_64")]
4052#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4053unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4054 unsafe {
4056 use core::arch::x86_64::*;
4057 let lomask = _mm_set1_epi8(0x0F);
4058 let eight = _mm256_set1_epi8(8);
4059 let mut acc = 0f32;
4060 for gi in 0..gpr {
4061 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4062 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4063 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4064 let lo = _mm_and_si128(b, lomask);
4065 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4066 let w = _mm256_sub_epi8(
4067 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4068 eight,
4069 );
4070 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4071 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4072 acc += d as f32 * s;
4073 }
4074 acc
4075 }
4076}
4077
4078#[cfg(target_arch = "x86_64")]
4083#[target_feature(enable = "avx2,fma")]
4088unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4089 unsafe {
4091 use core::arch::x86_64::*;
4092 let lomask = _mm_set1_epi8(0x0F);
4093 let eight = _mm256_set1_epi8(8);
4094 let ones = _mm256_set1_epi16(1);
4095 let mut f0 = _mm256_setzero_ps();
4108 let mut f1 = _mm256_setzero_ps();
4109 let mut f2 = _mm256_setzero_ps();
4110 let mut f3 = _mm256_setzero_ps();
4111 for gi in 0..gpr {
4112 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4113 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4114 let sv = _mm256_set1_ps(s);
4115 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4116 let lo = _mm_and_si128(bb, lomask);
4117 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4118 let w = _mm256_sub_epi8(
4119 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4120 eight,
4121 );
4122 let aw = _mm256_abs_epi8(w);
4123 let off = gi * GROUP_SIZE;
4124 let dot = |xq: &[i8]| {
4125 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4126 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4127 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4128 };
4129 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4130 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4131 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4132 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4133 }
4134 [
4135 hsum256_ps(f0),
4136 hsum256_ps(f1),
4137 hsum256_ps(f2),
4138 hsum256_ps(f3),
4139 ]
4140 }
4141}
4142
4143#[cfg(target_arch = "x86_64")]
4146#[target_feature(enable = "avx2")]
4147#[inline]
4148unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4149 unsafe {
4151 use core::arch::x86_64::*;
4152 let hi = _mm256_extractf128_ps::<1>(v);
4153 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4154 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4155 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4156 _mm_cvtss_f32(s)
4157 }
4158}
4159
4160#[cfg(target_arch = "x86_64")]
4162#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4163unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4164 unsafe {
4166 use core::arch::x86_64::*;
4167 let lomask = _mm_set1_epi8(0x0F);
4168 let eight = _mm256_set1_epi8(8);
4169 let mut f0 = _mm256_setzero_ps();
4172 let mut f1 = _mm256_setzero_ps();
4173 let mut f2 = _mm256_setzero_ps();
4174 let mut f3 = _mm256_setzero_ps();
4175 for gi in 0..gpr {
4176 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4177 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4178 let sv = _mm256_set1_ps(s);
4179 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4180 let lo = _mm_and_si128(bb, lomask);
4181 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4182 let w = _mm256_sub_epi8(
4183 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4184 eight,
4185 );
4186 let aw = _mm256_abs_epi8(w);
4187 let off = gi * GROUP_SIZE;
4188 let dot = |xq: &[i8]| {
4189 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4190 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4191 _mm256_setzero_si256(),
4192 aw,
4193 _mm256_sign_epi8(x, w),
4194 ))
4195 };
4196 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4197 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4198 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4199 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4200 }
4201 let acc = [
4202 hsum256_ps(f0),
4203 hsum256_ps(f1),
4204 hsum256_ps(f2),
4205 hsum256_ps(f3),
4206 ];
4207 acc
4208 }
4209}
4210
4211#[cfg(target_arch = "aarch64")]
4216#[target_feature(enable = "neon,dotprod")]
4217unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4218 unsafe {
4220 use core::arch::aarch64::*;
4221 use core::arch::asm;
4222 let lomask = vdupq_n_u8(0x0F);
4223 let eight = vdupq_n_s8(8);
4224 let mut acc = [0f32; 4];
4225 for gi in 0..gpr {
4226 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4227 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4228 let b = vld1q_u8(t.add(2));
4229 let lo = vandq_u8(b, lomask);
4230 let hi = vshrq_n_u8::<4>(b);
4231 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4232 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4233 for (k, xq) in xs.iter().enumerate() {
4234 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4235 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4236 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4237 asm!(
4238 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4239 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4240 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4241 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4242 options(pure, nomem, nostack),
4243 );
4244 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4245 }
4246 }
4247 acc
4248 }
4249}
4250
4251#[inline]
4253fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4254 let gi = j / GROUP_SIZE;
4255 let k = j % GROUP_SIZE;
4256 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4257 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4258 let byte = tile[2 + k / 2];
4259 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4260 ((nib as i32 - 8) as f32, s)
4261}
4262
4263#[inline]
4266fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4267 let mut acc = 0f32;
4268 for gi in 0..gpr {
4269 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4270 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4271 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4272 let mut ga = 0f32;
4273 for (k, &b) in tile[2..].iter().enumerate() {
4274 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4275 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4276 }
4277 acc += ga * s;
4278 }
4279 acc
4280}
4281
4282struct Q4tpView<'a> {
4286 nib: &'a [u8],
4287 params: &'a [u8],
4288 codes: &'a [u8],
4289 stride: usize,
4290 zero_rung: bool,
4292}
4293
4294impl<'a> Q4tpView<'a> {
4295 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4296 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4297 Self {
4298 nib: &bytes[..params_off],
4299 params: &bytes[params_off..codes_off],
4300 codes: &bytes[codes_off..],
4301 stride,
4302 zero_rung: false,
4303 }
4304 }
4305
4306 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4308 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4309 Self {
4310 nib: &bytes[..params_off],
4311 params: &bytes[params_off..codes_off],
4312 codes: &bytes[codes_off..],
4313 stride,
4314 zero_rung: true,
4315 }
4316 }
4317
4318 #[inline]
4334 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4335 let tab = if self.zero_rung {
4336 q2tp_ladder(self.params, r)
4337 } else {
4338 q4tp_ladder(self.params, r)
4339 };
4340 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4341 let out = &mut out[..gpr];
4342 let mut chunks = out.chunks_exact_mut(8);
4343 let mut ci = 0usize;
4344 for c in &mut chunks {
4345 let w = u64::from(codes[ci])
4346 | u64::from(codes[ci + 1]) << 8
4347 | u64::from(codes[ci + 2]) << 16
4348 | u64::from(codes[ci + 3]) << 24
4349 | u64::from(codes[ci + 4]) << 32;
4350 for (k, o) in c.iter_mut().enumerate() {
4351 *o = tab[((w >> (5 * k)) & 31) as usize];
4352 }
4353 ci += 5;
4354 }
4355 let tail = &codes[ci..];
4358 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4359 *o = tab[q4tp_code(tail, k)];
4360 }
4361 }
4362}
4363
4364#[inline]
4365fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4366 #[cfg(target_arch = "aarch64")]
4367 unsafe {
4368 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4369 }
4370 #[cfg(target_arch = "x86_64")]
4371 unsafe {
4372 if vnni_tiles_enabled() {
4373 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4374 }
4375 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4376 }
4377 #[allow(unreachable_code)]
4378 {
4379 let mut acc = 0f32;
4380 for gi in 0..gpr {
4381 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4382 let s = scales[gi];
4383 let mut d = 0i32;
4384 for (k, &b) in tile.iter().enumerate() {
4385 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4386 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4387 }
4388 acc += d as f32 * s;
4389 }
4390 acc
4391 }
4392}
4393
4394#[cfg(target_arch = "aarch64")]
4397#[target_feature(enable = "neon,dotprod")]
4398unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4399 unsafe {
4402 use core::arch::aarch64::*;
4403 use core::arch::asm;
4404 let lomask = vdupq_n_u8(0x0F);
4405 let eight = vdupq_n_s8(8);
4406 let mut acc = 0f32;
4407 for gi in 0..gpr {
4408 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4409 let s = *scales.get_unchecked(gi);
4410 let b = vld1q_u8(t);
4411 let lo = vandq_u8(b, lomask);
4412 let hi = vshrq_n_u8::<4>(b);
4413 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4414 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4415 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4416 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4417 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4418 asm!(
4419 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4420 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4421 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4422 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4423 options(pure, nomem, nostack),
4424 );
4425 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4426 }
4427 acc
4428 }
4429}
4430
4431#[cfg(target_arch = "x86_64")]
4432#[target_feature(enable = "avx2")]
4433unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4434 unsafe {
4436 use core::arch::x86_64::*;
4437 let lomask = _mm_set1_epi8(0x0F);
4438 let eight = _mm256_set1_epi8(8);
4439 let ones = _mm256_set1_epi16(1);
4440 let mut acc = 0f32;
4441 for gi in 0..gpr {
4442 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4443 let s = *scales.get_unchecked(gi);
4444 let b = _mm_loadu_si128(t as *const __m128i);
4445 let lo = _mm_and_si128(b, lomask);
4446 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4447 let w = _mm256_sub_epi8(
4448 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4449 eight,
4450 );
4451 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4452 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4453 let d = _mm256_madd_epi16(p16, ones);
4454 let hi128 = _mm256_extracti128_si256::<1>(d);
4455 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4456 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4457 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4458 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4459 }
4460 acc
4461 }
4462}
4463
4464
4465#[cfg(target_arch = "x86_64")]
4468#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4469unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4470 unsafe {
4472 use core::arch::x86_64::*;
4473 let lomask = _mm_set1_epi8(0x0F);
4474 let eight = _mm256_set1_epi8(8);
4475 let mut acc = 0f32;
4476 for gi in 0..gpr {
4477 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4478 let s = *scales.get_unchecked(gi);
4479 let b = _mm_loadu_si128(t as *const __m128i);
4480 let lo = _mm_and_si128(b, lomask);
4481 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4482 let w = _mm256_sub_epi8(
4483 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4484 eight,
4485 );
4486 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4487 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4488 }
4489 acc
4490 }
4491}
4492
4493#[inline]
4496fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4497 let mut acc = 0f32;
4498 for gi in 0..gpr {
4499 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4500 let s = scales[gi];
4501 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4502 let mut ga = 0f32;
4503 for (k, &b) in tile.iter().enumerate() {
4504 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4505 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4506 }
4507 acc += ga * s;
4508 }
4509 acc
4510}
4511
4512#[inline]
4515fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4516 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4517 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4518 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4519 ((n as i32 - 8) as f32, scales[gi])
4520}
4521
4522fn q4tp_matvec(
4524 bytes: &[u8],
4525 x: &[f32],
4526 rows: usize,
4527 cols: usize,
4528 out: &mut [f32],
4529 pool: Option<&Pool>,
4530) {
4531 debug_assert_eq!(out.len(), rows);
4532 let gpr = cols / GROUP_SIZE;
4533 let v = Q4tpView::new(bytes, rows, cols);
4534 let out_addr = SendMut(out.as_mut_ptr());
4535 if a8w8_enabled() {
4536 let act = split_act(x);
4537 let run = |start: usize, end: usize| {
4538 with_krow(gpr, |sc| {
4540 for r in start..end {
4541 v.scales_into(r, gpr, sc);
4542 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4543 for &(j, xv) in &act.outliers {
4544 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4545 acc += w * s * xv;
4546 }
4547 unsafe { *out_addr.at(r) = acc };
4549 }
4550 })
4551 };
4552 dispatch_rows(pool, rows, &run);
4553 return;
4554 }
4555 let run = |start: usize, end: usize| {
4556 with_krow(gpr, |sc| {
4557 for r in start..end {
4558 v.scales_into(r, gpr, sc);
4559 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4561 }
4562 })
4563 };
4564 dispatch_rows(pool, rows, &run);
4565}
4566
4567#[allow(clippy::too_many_arguments)]
4570fn q4tp_matvec2(
4571 bytes: &[u8],
4572 x1: &[f32],
4573 x2: &[f32],
4574 rows: usize,
4575 cols: usize,
4576 o1: &mut [f32],
4577 o2: &mut [f32],
4578 pool: Option<&Pool>,
4579) {
4580 let gpr = cols / GROUP_SIZE;
4581 let v = Q4tpView::new(bytes, rows, cols);
4582 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4583 let run = |start: usize, end: usize| {
4584 let mut sc = vec![0f32; gpr];
4585 for r in start..end {
4586 v.scales_into(r, gpr, &mut sc);
4587 unsafe {
4589 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4590 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4591 }
4592 }
4593 };
4594 dispatch_rows(pool, rows, &run);
4595}
4596
4597#[inline]
4600fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4601 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4602 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4603 let c = (byte >> (2 * (k % 4))) & 3;
4604 (c as f32 - 1.5, scales[gi])
4605}
4606
4607#[inline]
4614fn dot_q2tp_row_i8(
4615 chunks: &[u8],
4616 r: usize,
4617 gpr: usize,
4618 xq: &[i8],
4619 gsum: &[i32],
4620 scales: &[f32],
4621) -> f32 {
4622 let mut acc = 0f32;
4623 let base = r * gpr * Q2TP_CHUNK;
4624 #[cfg(not(target_arch = "aarch64"))]
4625 let mut codes = [0i8; GROUP_SIZE];
4626 for gi in 0..gpr {
4627 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4628 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4629 #[cfg(target_arch = "aarch64")]
4630 let dot = unsafe {
4636 use core::arch::aarch64::*;
4637 let b = vld1_u8(ch.as_ptr());
4638 let three = vdup_n_u8(3);
4639 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4640 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4641 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4642 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4643 let x4 = vld4_s8(xg.as_ptr());
4644 let mut acc4 = vdupq_n_s32(0);
4645 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4646 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4647 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4648 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4649 vaddvq_s32(acc4)
4650 };
4651 #[cfg(not(target_arch = "aarch64"))]
4652 let dot: i32 = {
4653 for (k, &b) in ch.iter().enumerate() {
4654 codes[k * 4] = (b & 3) as i8;
4655 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4656 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4657 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4658 }
4659 codes
4660 .iter()
4661 .zip(xg)
4662 .map(|(&c, &x)| c as i32 * x as i32)
4663 .sum()
4664 };
4665 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4666 }
4667 acc
4668}
4669
4670fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4674 let mut acc = 0f32;
4675 for gi in 0..gpr {
4676 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4677 let s = scales[gi];
4678 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4679 let mut g = 0f32;
4680 for (k, &b) in ch.iter().enumerate() {
4681 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4682 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4683 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4684 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4685 }
4686 acc += s * g;
4687 }
4688 acc
4689}
4690
4691fn q2tp_matvec(
4692 bytes: &[u8],
4693 x: &[f32],
4694 rows: usize,
4695 cols: usize,
4696 out: &mut [f32],
4697 pool: Option<&Pool>,
4698) {
4699 debug_assert_eq!(out.len(), rows);
4700 let gpr = cols / GROUP_SIZE;
4701 let v = Q4tpView::new_q2(bytes, rows, cols);
4702 let out_addr = SendMut(out.as_mut_ptr());
4703 if a8w8_enabled() {
4708 let act = split_act(x);
4709 let gsum = q1_group_sums(&act.xq, gpr);
4710 let (act, gsum) = (&act, &gsum);
4711 let run = move |start: usize, end: usize| {
4712 with_krow(gpr, |sc| {
4713 for r in start..end {
4714 v.scales_into(r, gpr, sc);
4715 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4716 for &(j, xv) in &act.outliers {
4717 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4718 acc += w * s * xv;
4719 }
4720 unsafe { *out_addr.at(r) = acc };
4722 }
4723 })
4724 };
4725 dispatch_rows(pool, rows, &run);
4726 return;
4727 }
4728 let run = |start: usize, end: usize| {
4729 with_krow(gpr, |sc| {
4730 for r in start..end {
4731 v.scales_into(r, gpr, sc);
4732 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4734 }
4735 })
4736 };
4737 dispatch_rows(pool, rows, &run);
4738}
4739
4740#[allow(clippy::too_many_arguments)]
4742fn q2tp_matvec2(
4743 bytes: &[u8],
4744 x1: &[f32],
4745 x2: &[f32],
4746 rows: usize,
4747 cols: usize,
4748 o1: &mut [f32],
4749 o2: &mut [f32],
4750 pool: Option<&Pool>,
4751) {
4752 let gpr = cols / GROUP_SIZE;
4753 let v = Q4tpView::new_q2(bytes, rows, cols);
4754 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4755 let run = |start: usize, end: usize| {
4756 let mut sc = vec![0f32; gpr];
4757 for r in start..end {
4758 v.scales_into(r, gpr, &mut sc);
4759 unsafe {
4761 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4762 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4763 }
4764 }
4765 };
4766 dispatch_rows(pool, rows, &run);
4767}
4768
4769pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4776 let gpr = cols / GROUP_SIZE;
4781 let v = Q4tpView::new_q2(bytes, rows, cols);
4782 with_krow(gpr, |sc| {
4783 for r in 0..rows {
4784 v.scales_into(r, gpr, sc);
4785 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4786 }
4787 });
4788}
4789
4790pub fn q2tp_matmat_for_test(
4791 bytes: &[u8],
4792 xs_all: &[f32],
4793 b: usize,
4794 rows: usize,
4795 cols: usize,
4796 out: &mut [f32],
4797) {
4798 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4799}
4800
4801fn q2tp_matmat(
4802 bytes: &[u8],
4803 xs_all: &[f32],
4804 b: usize,
4805 rows: usize,
4806 cols: usize,
4807 out: &mut [f32],
4808 pool: Option<&Pool>,
4809) {
4810 debug_assert_eq!(out.len(), b * rows);
4811 let gpr = cols / GROUP_SIZE;
4812 let v = Q4tpView::new_q2(bytes, rows, cols);
4813 let out_addr = SendMut(out.as_mut_ptr());
4814 let run = |start: usize, end: usize| {
4815 let mut sc = vec![0f32; gpr];
4816 for r in start..end {
4817 v.scales_into(r, gpr, &mut sc);
4818 for bi in 0..b {
4819 let x = &xs_all[bi * cols..(bi + 1) * cols];
4820 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4822 }
4823 }
4824 };
4825 dispatch_rows(pool, rows, &run);
4826}
4827
4828#[cfg(target_arch = "aarch64")]
4832#[target_feature(enable = "neon,dotprod")]
4833unsafe fn dot_q4tp_row_1x4_sdot_v1(
4834 nib: &[u8],
4835 r: usize,
4836 gpr: usize,
4837 xs: [&[i8]; 4],
4838 scales: &[f32],
4839) -> [f32; 4] {
4840 unsafe {
4841 use core::arch::aarch64::*;
4842 use core::arch::asm;
4843 let lomask = vdupq_n_u8(0x0F);
4844 let eight = vdupq_n_s8(8);
4845 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4846 for gi in 0..gpr {
4847 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4848 let s = *scales.get_unchecked(gi);
4849 let bb = vld1q_u8(t);
4850 let lo = vandq_u8(bb, lomask);
4851 let hi = vshrq_n_u8::<4>(bb);
4852 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4853 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4854 let mut d = [0f32; 4];
4855 for (k, dk) in d.iter_mut().enumerate() {
4856 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4857 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4858 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4859 asm!(
4860 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4861 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4862 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4863 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4864 options(pure, nomem, nostack),
4865 );
4866 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4867 }
4868 f0 += d[0];
4869 f1 += d[1];
4870 f2 += d[2];
4871 f3 += d[3];
4872 }
4873 [f0, f1, f2, f3]
4874 }
4875}
4876
4877#[allow(dead_code)]
4884static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4885
4886#[cfg(target_arch = "x86_64")]
4892fn q4tp_blocked_x86() -> bool {
4893 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4894 1 => false,
4895 2 => avx512vnni_enabled(),
4900 _ => blocked_enabled() && avx512vnni_enabled(),
4904 }
4905}
4906
4907#[cfg(target_arch = "aarch64")]
4909#[allow(dead_code)]
4910fn q4tp_v1() -> bool {
4911 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4912 1 => true,
4913 2 => false,
4914 _ => {
4915 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4916 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4917 }
4918 }
4919}
4920
4921#[cfg(target_arch = "x86_64")]
4933#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4934unsafe fn dot_q4tp_2x8_avx512(
4935 nib: &[u8],
4936 r0: usize,
4937 gpr: usize,
4938 xs: [&[i8]; 8],
4939 sc0: &[f32],
4940 sc1: &[f32],
4941) -> [[f32; 8]; 2] {
4942 unsafe {
4945 use core::arch::x86_64::*;
4946 let lomask = _mm256_set1_epi8(0x0F);
4947 let eight = _mm256_set1_epi8(8);
4948 let zero = _mm512_setzero_si512();
4949 let mut v0 = [_mm512_setzero_ps(); 8];
4950 let mut v1 = [_mm512_setzero_ps(); 8];
4951 let pairs = gpr / 2;
4952 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4953 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4954 let bb = _mm256_loadu_si256(t as *const __m256i);
4955 let lo = _mm256_and_si256(bb, lomask);
4956 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4957 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4958 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4959 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4960 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4961 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
4962 };
4963 for gp in 0..pairs {
4964 let gi = gp * 2;
4965 let (wa0, neg0) = unpack(r0, gi);
4966 let (wa1, neg1) = unpack(r0 + 1, gi);
4967 let off = gi * GROUP_SIZE;
4968 let sv = |sc: &[f32]| {
4969 _mm512_insertf32x8::<1>(
4970 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
4971 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
4972 )
4973 };
4974 let s0 = sv(sc0);
4975 let s1 = sv(sc1);
4976 for k in 0..8 {
4977 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
4978 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4979 zero,
4980 wa0,
4981 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
4982 ));
4983 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4984 zero,
4985 wa1,
4986 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
4987 ));
4988 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
4989 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
4990 }
4991 }
4992 let mut acc = [[0f32; 8]; 2];
4993 for k in 0..8 {
4994 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
4995 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
4996 }
4997 if gpr % 2 == 1 {
4998 let off = (gpr - 1) * GROUP_SIZE;
4999 for j in off..off + GROUP_SIZE {
5000 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5001 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5002 for k in 0..8 {
5003 let x = *xs[k].get_unchecked(j) as f32;
5004 acc[0][k] += w0 * sa * x;
5005 acc[1][k] += w1 * sb * x;
5006 }
5007 }
5008 }
5009 acc
5010 }
5011}
5012
5013#[cfg(target_arch = "x86_64")]
5018#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5019unsafe fn dot_q4tp_row_1x8_avx512(
5020 nib: &[u8],
5021 r: usize,
5022 gpr: usize,
5023 xs: [&[i8]; 8],
5024 scales: &[f32],
5025) -> [f32; 8] {
5026 unsafe {
5028 use core::arch::x86_64::*;
5029 let lomask = _mm256_set1_epi8(0x0F);
5030 let eight = _mm256_set1_epi8(8);
5031 let zero = _mm512_setzero_si512();
5032 let (mut v0, mut v1, mut v2, mut v3) = (
5033 _mm512_setzero_ps(),
5034 _mm512_setzero_ps(),
5035 _mm512_setzero_ps(),
5036 _mm512_setzero_ps(),
5037 );
5038 let (mut v4, mut v5, mut v6, mut v7) = (
5039 _mm512_setzero_ps(),
5040 _mm512_setzero_ps(),
5041 _mm512_setzero_ps(),
5042 _mm512_setzero_ps(),
5043 );
5044 let pairs = gpr / 2;
5045 for gp in 0..pairs {
5046 let gi = gp * 2;
5047 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5048 let bb = _mm256_loadu_si256(t as *const __m256i);
5049 let lo = _mm256_and_si256(bb, lomask);
5050 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5051 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5056 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5057 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5058 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5059 let wabs = _mm512_abs_epi8(w);
5060 let neg = _mm512_movepi8_mask(w);
5061 let off = gi * GROUP_SIZE;
5062 let sv = _mm512_insertf32x8::<1>(
5063 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5064 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5065 );
5066 let dot = |x: &[i8]| -> __m512 {
5067 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5068 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5069 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5070 };
5071 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5072 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5073 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5074 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5075 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5076 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5077 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5078 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5079 }
5080 let mut acc = [
5081 _mm512_reduce_add_ps(v0),
5082 _mm512_reduce_add_ps(v1),
5083 _mm512_reduce_add_ps(v2),
5084 _mm512_reduce_add_ps(v3),
5085 _mm512_reduce_add_ps(v4),
5086 _mm512_reduce_add_ps(v5),
5087 _mm512_reduce_add_ps(v6),
5088 _mm512_reduce_add_ps(v7),
5089 ];
5090 if gpr % 2 == 1 {
5093 let off = (gpr - 1) * GROUP_SIZE;
5094 for j in off..off + GROUP_SIZE {
5095 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5096 let ws = w * s;
5097 for k in 0..8 {
5098 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5099 }
5100 }
5101 }
5102 acc
5103 }
5104}
5105
5106#[cfg(target_arch = "x86_64")]
5119#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5120unsafe fn dot_q4tp_row_1x4_avx512(
5121 nib: &[u8],
5122 r: usize,
5123 gpr: usize,
5124 xs: [&[i8]; 4],
5125 scales: &[f32],
5126) -> [f32; 4] {
5127 unsafe {
5129 use core::arch::x86_64::*;
5130 let lomask = _mm256_set1_epi8(0x0F);
5131 let eight = _mm256_set1_epi8(8);
5132 let zero = _mm512_setzero_si512();
5133 let (mut v0, mut v1, mut v2, mut v3) = (
5134 _mm512_setzero_ps(),
5135 _mm512_setzero_ps(),
5136 _mm512_setzero_ps(),
5137 _mm512_setzero_ps(),
5138 );
5139 let pairs = gpr / 2;
5140 for gp in 0..pairs {
5141 let gi = gp * 2;
5142 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5143 let bb = _mm256_loadu_si256(t as *const __m256i);
5144 let lo = _mm256_and_si256(bb, lomask);
5145 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5146 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5151 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5152 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5153 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5154 let wabs = _mm512_abs_epi8(w);
5155 let neg = _mm512_movepi8_mask(w);
5156 let off = gi * GROUP_SIZE;
5157 let sv = _mm512_insertf32x8::<1>(
5158 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5159 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5160 );
5161 let dot = |x: &[i8]| -> __m512 {
5162 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5163 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5164 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5165 };
5166 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5167 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5168 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5169 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5170 }
5171 let mut acc = [
5172 _mm512_reduce_add_ps(v0),
5173 _mm512_reduce_add_ps(v1),
5174 _mm512_reduce_add_ps(v2),
5175 _mm512_reduce_add_ps(v3),
5176 ];
5177 if gpr % 2 == 1 {
5180 let off = (gpr - 1) * GROUP_SIZE;
5181 for j in off..off + GROUP_SIZE {
5182 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5183 let ws = w * s;
5184 for k in 0..4 {
5185 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5186 }
5187 }
5188 }
5189 acc
5190 }
5191}
5192
5193#[cfg(target_arch = "aarch64")]
5197#[target_feature(enable = "neon,dotprod")]
5198unsafe fn dot_q4tp_row_1x4_sdot(
5199 nib: &[u8],
5200 r: usize,
5201 gpr: usize,
5202 xs: [&[i8]; 4],
5203 scales: &[f32],
5204) -> [f32; 4] {
5205 unsafe {
5207 use core::arch::aarch64::*;
5208 use core::arch::asm;
5209 let lomask = vdupq_n_u8(0x0F);
5210 let eight = vdupq_n_s8(8);
5211 let (mut v0, mut v1, mut v2, mut v3) = (
5227 vdupq_n_f32(0.0),
5228 vdupq_n_f32(0.0),
5229 vdupq_n_f32(0.0),
5230 vdupq_n_f32(0.0),
5231 );
5232 for gi in 0..gpr {
5233 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5234 let s = *scales.get_unchecked(gi);
5235 let bb = vld1q_u8(t);
5236 let lo = vandq_u8(bb, lomask);
5237 let hi = vshrq_n_u8::<4>(bb);
5238 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5239 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5240 let off = gi * GROUP_SIZE;
5241 let dot4 = |x: &[i8]| -> int32x4_t {
5242 let x0 = vld1q_s8(x.as_ptr().add(off));
5243 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5244 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5245 asm!(
5246 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5247 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5248 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5249 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5250 options(pure, nomem, nostack),
5251 );
5252 vaddq_s32(a0, a1)
5253 };
5254 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5255 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5256 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5257 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5258 }
5259 [
5260 vaddvq_f32(v0),
5261 vaddvq_f32(v1),
5262 vaddvq_f32(v2),
5263 vaddvq_f32(v3),
5264 ]
5265 }
5266}
5267
5268fn q4tp_matmat(
5272 bytes: &[u8],
5273 xs_all: &[f32],
5274 b: usize,
5275 rows: usize,
5276 cols: usize,
5277 out: &mut [f32],
5278 pool: Option<&Pool>,
5279) {
5280 debug_assert_eq!(out.len(), b * rows);
5281 let gpr = cols / GROUP_SIZE;
5282 let v = Q4tpView::new(bytes, rows, cols);
5283
5284 #[cfg(target_os = "macos")]
5286 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5287 dequant_matmat_accel(
5288 &|r, dst| {
5289 let mut sc = [0f32; 32];
5290 let mut scv;
5291 let s: &[f32] = if gpr <= 32 {
5292 v.scales_into(r, gpr, &mut sc);
5293 &sc[..gpr]
5294 } else {
5295 scv = vec![0f32; gpr];
5296 v.scales_into(r, gpr, &mut scv);
5297 &scv
5298 };
5299 for gi in 0..gpr {
5300 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5301 for (k, &bb) in tile.iter().enumerate() {
5302 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5303 dst[gi * GROUP_SIZE + k * 2 + 1] =
5304 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5305 }
5306 }
5307 },
5308 xs_all,
5309 b,
5310 rows,
5311 cols,
5312 out,
5313 pool,
5314 );
5315 return;
5316 }
5317
5318 let out_addr = SendMut(out.as_mut_ptr());
5319 if a8w8_enabled() {
5320 let acts: Vec<SplitAct> = (0..b)
5321 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5322 .collect();
5323 let acts = &acts;
5324 #[cfg(target_arch = "aarch64")]
5325 let blocked_ok = sdot_enabled() && blocked_enabled();
5326 #[cfg(target_arch = "x86_64")]
5333 let blocked_ok = q4tp_blocked_x86();
5334 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5335 let blocked_ok = false;
5336 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5346 .ok()
5347 .and_then(|v| v.parse().ok())
5348 .filter(|v| *v > 0)
5349 .unwrap_or(256);
5350 let run = |start: usize, end: usize| {
5351 for abase in (0..acts.len()).step_by(panel_cols) {
5352 let alen = (acts.len() - abase).min(panel_cols);
5353 let mut sc = vec![0f32; gpr];
5354 #[cfg(target_arch = "x86_64")]
5355 let mut r_lo = start;
5356 #[cfg(target_arch = "x86_64")]
5357 if blocked_ok && alen >= 8 {
5358 let mut sc1 = vec![0f32; gpr];
5359 while r_lo + 2 <= end {
5360 v.scales_into(r_lo, gpr, &mut sc);
5361 v.scales_into(r_lo + 1, gpr, &mut sc1);
5362 let mut bi = 0usize;
5363 while bi + 8 <= alen {
5364 let xs = [
5365 acts[abase + bi].xq.as_slice(),
5366 acts[abase + bi + 1].xq.as_slice(),
5367 acts[abase + bi + 2].xq.as_slice(),
5368 acts[abase + bi + 3].xq.as_slice(),
5369 acts[abase + bi + 4].xq.as_slice(),
5370 acts[abase + bi + 5].xq.as_slice(),
5371 acts[abase + bi + 6].xq.as_slice(),
5372 acts[abase + bi + 7].xq.as_slice(),
5373 ];
5374 let d =
5375 unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5376 for (row, dr, scr) in
5377 [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)]
5378 {
5379 for k in 0..8 {
5380 let act = &acts[abase + bi + k];
5381 let mut acc = dr[k] * act.sx;
5382 for &(j, xv) in &act.outliers {
5383 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5384 acc += w * s * xv;
5385 }
5386 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5388 }
5389 }
5390 bi += 8;
5391 }
5392 for row in [r_lo, r_lo + 1] {
5395 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5396 for b2 in bi..alen {
5397 let act = &acts[abase + b2];
5398 let xs4 = [
5399 act.xq.as_slice(),
5400 act.xq.as_slice(),
5401 act.xq.as_slice(),
5402 act.xq.as_slice(),
5403 ];
5404 let d =
5405 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5406 let mut acc = d[0] * act.sx;
5407 for &(j, xv) in &act.outliers {
5408 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5409 acc += w * s * xv;
5410 }
5411 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5413 }
5414 }
5415 r_lo += 2;
5416 }
5417 }
5418 #[cfg(target_arch = "x86_64")]
5419 let row_start = r_lo;
5420 #[cfg(not(target_arch = "x86_64"))]
5421 let row_start = start;
5422 for r in row_start..end {
5423 v.scales_into(r, gpr, &mut sc);
5424 let mut bi = 0usize;
5425 #[cfg(target_arch = "x86_64")]
5426 if blocked_ok {
5427 while bi + 8 <= alen {
5428 let xs = [
5429 acts[abase + bi].xq.as_slice(),
5430 acts[abase + bi + 1].xq.as_slice(),
5431 acts[abase + bi + 2].xq.as_slice(),
5432 acts[abase + bi + 3].xq.as_slice(),
5433 acts[abase + bi + 4].xq.as_slice(),
5434 acts[abase + bi + 5].xq.as_slice(),
5435 acts[abase + bi + 6].xq.as_slice(),
5436 acts[abase + bi + 7].xq.as_slice(),
5437 ];
5438 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5439 for k in 0..8 {
5440 let act = &acts[abase + bi + k];
5441 let mut acc = d[k] * act.sx;
5442 for &(j, xv) in &act.outliers {
5443 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5444 acc += w * s * xv;
5445 }
5446 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5448 }
5449 bi += 8;
5450 }
5451 while bi + 4 <= alen {
5452 let xs = [
5453 acts[abase + bi].xq.as_slice(),
5454 acts[abase + bi + 1].xq.as_slice(),
5455 acts[abase + bi + 2].xq.as_slice(),
5456 acts[abase + bi + 3].xq.as_slice(),
5457 ];
5458 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5459 for k in 0..4 {
5460 let act = &acts[abase + bi + k];
5461 let mut acc = d[k] * act.sx;
5462 for &(j, xv) in &act.outliers {
5463 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5464 acc += w * s * xv;
5465 }
5466 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5468 }
5469 bi += 4;
5470 }
5471 }
5472 #[cfg(target_arch = "aarch64")]
5473 if blocked_ok {
5474 while bi + 4 <= alen {
5475 let xs = [
5476 acts[abase + bi].xq.as_slice(),
5477 acts[abase + bi + 1].xq.as_slice(),
5478 acts[abase + bi + 2].xq.as_slice(),
5479 acts[abase + bi + 3].xq.as_slice(),
5480 ];
5481 let d = unsafe {
5482 if q4tp_v1() {
5483 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5484 } else {
5485 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5486 }
5487 };
5488 for k in 0..4 {
5489 let act = &acts[abase + bi + k];
5490 let mut acc = d[k] * act.sx;
5491 for &(j, xv) in &act.outliers {
5492 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5493 acc += w * s * xv;
5494 }
5495 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5497 }
5498 bi += 4;
5499 }
5500 }
5501 let _ = blocked_ok;
5502 while bi < alen {
5503 let act = &acts[abase + bi];
5504 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5505 for &(j, xv) in &act.outliers {
5506 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5507 acc += w * s * xv;
5508 }
5509 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5511 bi += 1;
5512 }
5513 }
5514
5515 }
5516 };
5517 dispatch_rows(pool, rows, &run);
5518 return;
5519 }
5520
5521 let run = |start: usize, end: usize| {
5522 let mut sc = vec![0f32; gpr];
5523 for r in start..end {
5524 v.scales_into(r, gpr, &mut sc);
5525 for bi in 0..b {
5526 let x = &xs_all[bi * cols..(bi + 1) * cols];
5527 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5529 }
5530 }
5531 };
5532 dispatch_rows(pool, rows, &run);
5533}
5534
5535fn q4t_matvec(
5537 bytes: &[u8],
5538 x: &[f32],
5539 rows: usize,
5540 cols: usize,
5541 out: &mut [f32],
5542 pool: Option<&Pool>,
5543) {
5544 debug_assert_eq!(out.len(), rows);
5545 let gpr = cols / GROUP_SIZE;
5546 let out_addr = SendMut(out.as_mut_ptr());
5547 if a8w8_enabled() {
5548 let act = split_act(x);
5549 let run = move |start: usize, end: usize| {
5550 for r in start..end {
5551 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5552 for &(j, xv) in &act.outliers {
5553 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5554 acc += w * s * xv;
5555 }
5556 unsafe { *out_addr.at(r) = acc };
5558 }
5559 };
5560 dispatch_rows(pool, rows, &run);
5561 return;
5562 }
5563 let run = move |start: usize, end: usize| {
5564 for r in start..end {
5565 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5567 }
5568 };
5569 dispatch_rows(pool, rows, &run);
5570}
5571
5572#[allow(clippy::too_many_arguments)]
5574fn q4t_matvec2(
5575 bytes: &[u8],
5576 x1: &[f32],
5577 x2: &[f32],
5578 rows: usize,
5579 cols: usize,
5580 o1: &mut [f32],
5581 o2: &mut [f32],
5582 pool: Option<&Pool>,
5583) {
5584 let gpr = cols / GROUP_SIZE;
5585 let p1 = SendMut(o1.as_mut_ptr());
5586 let p2 = SendMut(o2.as_mut_ptr());
5587 if a8w8_enabled() {
5588 let a1 = split_act(x1);
5589 let a2 = split_act(x2);
5590 let run = move |start: usize, end: usize| {
5591 for r in start..end {
5592 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5593 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5594 for &(j, xv) in &a1.outliers {
5595 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5596 v1 += w * s * xv;
5597 }
5598 for &(j, xv) in &a2.outliers {
5599 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5600 v2 += w * s * xv;
5601 }
5602 unsafe {
5604 *p1.at(r) = v1;
5605 *p2.at(r) = v2;
5606 }
5607 }
5608 };
5609 dispatch_rows(pool, rows, &run);
5610 return;
5611 }
5612 let run = move |start: usize, end: usize| {
5613 for r in start..end {
5614 unsafe {
5616 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5617 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5618 }
5619 }
5620 };
5621 dispatch_rows(pool, rows, &run);
5622}
5623
5624#[allow(clippy::too_many_arguments)]
5626#[cfg(target_os = "macos")]
5632fn dequant_matmat_accel(
5633 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5634 xs_all: &[f32],
5635 b: usize,
5636 rows: usize,
5637 cols: usize,
5638 out: &mut [f32],
5639 pool: Option<&Pool>,
5640) {
5641 const TR: usize = 2048;
5642 thread_local! {
5643 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5644 }
5645 WTILE.with(|wt| {
5646 let mut wtile = wt.borrow_mut();
5647 wtile.resize(TR * cols, 0.0);
5648 let mut r0 = 0usize;
5649 while r0 < rows {
5650 let tr = TR.min(rows - r0);
5651 let wt_addr = SendMut(wtile.as_mut_ptr());
5652 let run = |start: usize, end: usize| {
5653 for r in start..end {
5654 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5656 dequant_row(r0 + r, dst);
5657 }
5658 };
5659 dispatch_rows(pool, tr, &run);
5660 unsafe {
5661 accel_blas::cblas_sgemm(
5662 101, 111, 112, b as i32,
5666 tr as i32,
5667 cols as i32,
5668 1.0,
5669 xs_all.as_ptr(),
5670 cols as i32,
5671 wtile.as_ptr(),
5672 cols as i32,
5673 0.0,
5674 out.as_mut_ptr().add(r0),
5675 rows as i32,
5676 );
5677 }
5678 r0 += tr;
5679 }
5680 });
5681}
5682
5683fn q4t_matmat(
5684 bytes: &[u8],
5685 xs_all: &[f32],
5686 b: usize,
5687 rows: usize,
5688 cols: usize,
5689 out: &mut [f32],
5690 pool: Option<&Pool>,
5691) {
5692 debug_assert_eq!(out.len(), b * rows);
5693 let gpr = cols / GROUP_SIZE;
5694 #[cfg(target_os = "macos")]
5698 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5699 dequant_matmat_accel(
5700 &|r, dst| {
5701 for gi in 0..gpr {
5702 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5703 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5704 for (k, &bb) in tile[2..].iter().enumerate() {
5705 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5706 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5707 }
5708 }
5709 },
5710 xs_all,
5711 b,
5712 rows,
5713 cols,
5714 out,
5715 pool,
5716 );
5717 return;
5718 }
5719 let out_addr = SendMut(out.as_mut_ptr());
5720 if a8w8_enabled() {
5721 let acts: Vec<SplitAct> = (0..b)
5722 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5723 .collect();
5724 let acts = &acts;
5725 #[cfg(target_arch = "x86_64")]
5726 let blocked_ok = avx2_enabled()
5727 && blocked_enabled();
5728 #[cfg(target_arch = "aarch64")]
5729 let blocked_ok = sdot_enabled()
5730 && blocked_enabled();
5731 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5732 let blocked_ok = false;
5733 let run = move |start: usize, end: usize| {
5734 for r in start..end {
5735 let mut bi = 0usize;
5736 #[cfg(target_arch = "aarch64")]
5737 if blocked_ok {
5738 while bi + 4 <= acts.len() {
5739 let xs = [
5740 acts[bi].xq.as_slice(),
5741 acts[bi + 1].xq.as_slice(),
5742 acts[bi + 2].xq.as_slice(),
5743 acts[bi + 3].xq.as_slice(),
5744 ];
5745 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5746 for k in 0..4 {
5747 let act = &acts[bi + k];
5748 let mut acc = d[k] * act.sx;
5749 for &(j, xv) in &act.outliers {
5750 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5751 acc += w * sc * xv;
5752 }
5753 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5755 }
5756 bi += 4;
5757 }
5758 }
5759 #[cfg(target_arch = "x86_64")]
5760 if blocked_ok {
5761 while bi + 4 <= acts.len() {
5762 let xs = [
5763 acts[bi].xq.as_slice(),
5764 acts[bi + 1].xq.as_slice(),
5765 acts[bi + 2].xq.as_slice(),
5766 acts[bi + 3].xq.as_slice(),
5767 ];
5768 let d = unsafe {
5769 if vnni_tiles_enabled() {
5770 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5771 } else {
5772 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5773 }
5774 };
5775 for k in 0..4 {
5776 let act = &acts[bi + k];
5777 let mut acc = d[k] * act.sx;
5778 for &(j, xv) in &act.outliers {
5779 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5780 acc += w * sc * xv;
5781 }
5782 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5784 }
5785 bi += 4;
5786 }
5787 }
5788 let _ = blocked_ok;
5789 while bi < acts.len() {
5790 let act = &acts[bi];
5791 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5792 for &(j, xv) in &act.outliers {
5793 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5794 acc += w * s * xv;
5795 }
5796 unsafe { *out_addr.at(bi * rows + r) = acc };
5798 bi += 1;
5799 }
5800 }
5801 };
5802 dispatch_rows(pool, rows, &run);
5803 return;
5804 }
5805 let run = move |start: usize, end: usize| {
5806 for r in start..end {
5807 for bi in 0..b {
5808 let x = &xs_all[bi * cols..(bi + 1) * cols];
5809 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5811 }
5812 }
5813 };
5814 dispatch_rows(pool, rows, &run);
5815}
5816
5817fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5826 (0..gpr)
5827 .map(|gi| {
5828 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5829 .iter()
5830 .map(|&v| v as i32)
5831 .sum()
5832 })
5833 .collect()
5834}
5835
5836#[inline]
5840#[allow(unreachable_code)]
5841#[cfg(target_arch = "x86_64")]
5846#[target_feature(enable = "avx2")]
5847unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5848 unsafe {
5850 use core::arch::x86_64::*;
5851 let expand = _mm256_setr_epi8(
5853 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
5854 3, 3, 3,
5855 );
5856 let bitsel = _mm256_setr_epi8(
5857 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5858 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5859 );
5860 let ones8 = _mm256_set1_epi8(1);
5861 let ones16 = _mm256_set1_epi16(1);
5862 let mut acc = 0f32;
5863 for gi in 0..gpr {
5864 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5865 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5866 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5867 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5868 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5869 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5870 let sel = _mm256_and_si256(x, mask);
5871 let p16 = _mm256_maddubs_epi16(ones8, sel);
5873 let d32 = _mm256_madd_epi16(p16, ones16);
5874 let hi128 = _mm256_extracti128_si256::<1>(d32);
5875 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5876 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5877 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5878 let msum = _mm_cvtsi128_si32(s32);
5879 let d = 2 * msum - gsum[gi];
5882 acc += d as f32 * s;
5883 }
5884 acc
5885 }
5886}
5887
5888#[cfg(target_arch = "x86_64")]
5891#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5892unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5893 unsafe {
5895 use core::arch::x86_64::*;
5896 let expand = _mm256_setr_epi8(
5897 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
5898 3, 3, 3,
5899 );
5900 let bitsel = _mm256_setr_epi8(
5901 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5902 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5903 );
5904 let ones8 = _mm256_set1_epi8(1);
5905 let mut acc = 0f32;
5906 for gi in 0..gpr {
5907 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5908 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5909 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5910 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5911 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5912 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5913 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5914 let d = 2 * msum - gsum[gi];
5915 acc += d as f32 * s;
5916 }
5917 acc
5918 }
5919}
5920
5921#[cfg(target_arch = "x86_64")]
5923#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5924unsafe fn dot_q1_row_1x4_vnni(
5925 bytes: &[u8],
5926 r: usize,
5927 gpr: usize,
5928 xs: [&[i8]; 4],
5929 gsums: [&[i32]; 4],
5930) -> [f32; 4] {
5931 unsafe {
5933 use core::arch::x86_64::*;
5934 let expand = _mm256_setr_epi8(
5935 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
5936 3, 3, 3,
5937 );
5938 let bitsel = _mm256_setr_epi8(
5939 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5940 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5941 );
5942 let ones8 = _mm256_set1_epi8(1);
5943 let mut acc = [0f32; 4];
5944 for gi in 0..gpr {
5945 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5946 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5947 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5948 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5949 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5950 for (k, xq) in xs.iter().enumerate() {
5951 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5952 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5953 let d = 2 * msum - gsums[k][gi];
5954 acc[k] += d as f32 * s;
5955 }
5956 }
5957 acc
5958 }
5959}
5960
5961#[cfg(target_arch = "x86_64")]
5964#[target_feature(enable = "avx2")]
5965unsafe fn dot_q1_row_1x4_avx2(
5966 bytes: &[u8],
5967 r: usize,
5968 gpr: usize,
5969 xs: [&[i8]; 4],
5970 gsums: [&[i32]; 4],
5971) -> [f32; 4] {
5972 unsafe {
5974 use core::arch::x86_64::*;
5975 let expand = _mm256_setr_epi8(
5976 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3,
5977 3, 3, 3,
5978 );
5979 let bitsel = _mm256_setr_epi8(
5980 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5981 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5982 );
5983 let ones8 = _mm256_set1_epi8(1);
5984 let ones16 = _mm256_set1_epi16(1);
5985 let mut acc = [0f32; 4];
5986 for gi in 0..gpr {
5987 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5988 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5989 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5990 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5991 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5992 for (k, xq) in xs.iter().enumerate() {
5993 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5994 let sel = _mm256_and_si256(x, mask);
5995 let p16 = _mm256_maddubs_epi16(ones8, sel);
5996 let d32 = _mm256_madd_epi16(p16, ones16);
5997 let hi128 = _mm256_extracti128_si256::<1>(d32);
5998 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5999 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6000 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6001 let msum = _mm_cvtsi128_si32(s32);
6002 let d = 2 * msum - gsums[k][gi];
6003 acc[k] += d as f32 * s;
6004 }
6005 }
6006 acc
6007 }
6008}
6009
6010#[allow(unreachable_code)]
6011fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6012 #[cfg(target_arch = "aarch64")]
6013 unsafe {
6014 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6015 }
6016 #[cfg(target_arch = "x86_64")]
6017 if avx2_enabled() {
6018 unsafe {
6019 if vnni_tiles_enabled() {
6020 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6021 }
6022 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6023 }
6024 }
6025 let _ = gsum;
6026 let mut acc = 0f32;
6027 for gi in 0..gpr {
6028 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6029 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6030 let mut d = 0i32;
6031 for (j, &b) in tile[2..].iter().enumerate() {
6032 for k in 0..8 {
6033 let w = ((b >> k) & 1) as i32 * 2 - 1;
6034 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6035 }
6036 }
6037 acc += d as f32 * s;
6038 }
6039 acc
6040}
6041
6042#[cfg(target_arch = "aarch64")]
6051#[target_feature(enable = "neon,dotprod")]
6052unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6053 unsafe {
6056 use core::arch::aarch64::*;
6057 use core::arch::asm;
6058 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6059 let m = vld1q_u8(MASKS.as_ptr());
6060 macro_rules! tile_dot {
6062 ($t:expr, $x:expr) => {{
6063 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6064 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6065 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6066 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6067 let x0 = vld1q_s8($x);
6068 let x1 = vld1q_s8($x.add(16));
6069 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6070 asm!(
6071 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6072 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6073 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6074 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6075 options(pure, nomem, nostack),
6076 );
6077 vaddq_s32(a0, a1)
6078 }};
6079 }
6080 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6089 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6090 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6091 const IW11: [u8; 16] = [
6092 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6093 ];
6094 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6095 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6096 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6097 let isc = vld1_u8(ISC.as_ptr());
6098 macro_rules! tile_dot_tbl {
6100 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6101 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6102 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6103 let x0 = vld1q_s8($x);
6104 let x1 = vld1q_s8($x.add(16));
6105 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6106 asm!(
6107 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6108 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6109 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6110 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6111 options(pure, nomem, nostack),
6112 );
6113 vaddq_s32(a0, a1)
6114 }};
6115 }
6116 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6117 let row_base = r * gpr * Q1_TILE;
6118 let abs_end = bytes.len();
6119 let xp = xq.as_ptr();
6120 let gp = gsum.as_ptr();
6121 let mut accv = vdupq_n_f32(0.0);
6122 let mut gi = 0;
6123 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6126 let t0 = base.add(gi * Q1_TILE);
6127 let ld_a = vld1q_u8(t0);
6128 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6129 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6130 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6131 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6132 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6133 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6135 let g = vld1q_s32(gp.add(gi));
6136 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6137 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6138 let scf: float32x4_t;
6139 asm!(
6140 "fcvtl {o:v}.4s, {i:v}.4h",
6141 o = out(vreg) scf, i = in(vreg) sc16,
6142 options(pure, nomem, nostack),
6143 );
6144 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6145 gi += 4;
6146 }
6147 let mut acc = vaddvq_f32(accv);
6148 while gi < gpr {
6149 let t = base.add(gi * Q1_TILE);
6150 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6151 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6152 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6153 gi += 1;
6154 }
6155 acc
6156 }
6157}
6158
6159#[cfg(target_arch = "aarch64")]
6164#[target_feature(enable = "neon,dotprod")]
6165unsafe fn dot_q1_row_1x4_sdot(
6166 bytes: &[u8],
6167 r: usize,
6168 gpr: usize,
6169 xs: [&[i8]; 4],
6170 gs: [&[i32]; 4],
6171) -> [f32; 4] {
6172 unsafe {
6174 use core::arch::aarch64::*;
6175 use core::arch::asm;
6176 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6177 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6178 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6179 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6180 const IW11: [u8; 16] = [
6181 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6182 ];
6183 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6184 let m = vld1q_u8(MASKS.as_ptr());
6185 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6186 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6187 let isc = vld1_u8(ISC.as_ptr());
6188 macro_rules! sdot2 {
6189 ($w0:expr, $w1:expr, $x:expr) => {{
6190 let x0 = vld1q_s8($x);
6191 let x1 = vld1q_s8($x.add(16));
6192 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6193 asm!(
6194 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6195 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6196 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6197 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6198 options(pure, nomem, nostack),
6199 );
6200 vaddq_s32(a0, a1)
6201 }};
6202 }
6203 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6204 let row_base = r * gpr * Q1_TILE;
6205 let abs_end = bytes.len();
6206 let mut accv = [vdupq_n_f32(0.0); 4];
6207 let mut gi = 0;
6208 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6209 let t0 = base.add(gi * Q1_TILE);
6210 let ld_a = vld1q_u8(t0);
6211 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6212 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6214 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6215 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6216 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6217 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6218 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6219 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6220 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6221 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6222 let scf: float32x4_t;
6223 asm!(
6224 "fcvtl {o:v}.4s, {i:v}.4h",
6225 o = out(vreg) scf, i = in(vreg) sc16,
6226 options(pure, nomem, nostack),
6227 );
6228 for k in 0..4 {
6229 let xp = xs[k].as_ptr();
6230 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6231 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6232 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6233 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6234 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6235 let g = vld1q_s32(gs[k].as_ptr().add(gi));
6236 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6237 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6238 }
6239 gi += 4;
6240 }
6241 let mut acc = [
6242 vaddvq_f32(accv[0]),
6243 vaddvq_f32(accv[1]),
6244 vaddvq_f32(accv[2]),
6245 vaddvq_f32(accv[3]),
6246 ];
6247 while gi < gpr {
6248 let t = base.add(gi * Q1_TILE);
6249 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6250 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6251 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6252 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6253 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6254 for k in 0..4 {
6255 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6256 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6257 }
6258 gi += 1;
6259 }
6260 acc
6261 }
6262}
6263
6264#[inline]
6266fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6267 let gi = j / GROUP_SIZE;
6268 let k = j % GROUP_SIZE;
6269 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6270 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6271 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6272 ((bit as i32 * 2 - 1) as f32, s)
6273}
6274
6275#[inline]
6277fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6278 let mut acc = 0f32;
6279 for gi in 0..gpr {
6280 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6281 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6282 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6283 let mut ga = 0f32;
6284 for (j, &b) in tile[2..].iter().enumerate() {
6285 for k in 0..8 {
6286 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6287 }
6288 }
6289 acc += ga * s;
6290 }
6291 acc
6292}
6293
6294#[allow(clippy::too_many_arguments)]
6297fn q1_range_a8w8(
6298 bytes: &[u8],
6299 gpr: usize,
6300 act: &SplitAct,
6301 gsum: &[i32],
6302 out: SendMut,
6303 start: usize,
6304 end: usize,
6305) {
6306 for r in start..end {
6307 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6308 for &(j, xv) in &act.outliers {
6309 let (w, s) = q1_outlier(bytes, r, gpr, j);
6310 acc += w * s * xv;
6311 }
6312 unsafe { *out.at(r) = acc };
6314 }
6315}
6316
6317fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6319 for r in start..end {
6320 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6322 }
6323}
6324
6325fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6330 let entries = base_len + (rows + 1) * 4;
6331 (base_len, entries, entries <= bytes.len())
6332}
6333
6334#[inline]
6336fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6337 let o = rp_off + r * 4;
6338 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6339}
6340
6341const SIGN5: [[f32; 5]; 256] = {
6345 let mut lut = [[0.0f32; 5]; 256];
6346 let pow3 = [1u16, 3, 9, 27, 81];
6347 let mut byte = 0usize;
6348 while byte < 256 {
6349 let mut i = 0usize;
6350 while i < 5 {
6351 let code = (byte as u16 / pow3[i]) % 3;
6352 lut[byte][i] = if code == 1 {
6353 1.0
6354 } else if code == 2 {
6355 -1.0
6356 } else {
6357 0.0
6358 };
6359 i += 1;
6360 }
6361 byte += 1;
6362 }
6363 lut
6364};
6365
6366const SIGN5_I8: [[i8; 5]; 256] = {
6368 let mut lut = [[0i8; 5]; 256];
6369 let pow3 = [1u16, 3, 9, 27, 81];
6370 let mut byte = 0usize;
6371 while byte < 256 {
6372 let mut i = 0usize;
6373 while i < 5 {
6374 let code = (byte as u16 / pow3[i]) % 3;
6375 lut[byte][i] = if code == 1 {
6376 1
6377 } else if code == 2 {
6378 -1
6379 } else {
6380 0
6381 };
6382 i += 1;
6383 }
6384 byte += 1;
6385 }
6386 lut
6387};
6388
6389const SIGN5_U64: [u64; 256] = {
6395 let mut lut = [0u64; 256];
6396 let pow3 = [1u16, 3, 9, 27, 81];
6397 let mut byte = 0usize;
6398 while byte < 256 {
6399 let mut v = 0u64;
6400 let mut i = 0usize;
6401 while i < 5 {
6402 let code = (byte as u16 / pow3[i]) % 3;
6403 let s: u8 = if code == 1 {
6404 1
6405 } else if code == 2 {
6406 0xFF
6407 } else {
6408 0
6409 };
6410 v |= (s as u64) << (i * 8);
6411 i += 1;
6412 }
6413 lut[byte] = v;
6414 byte += 1;
6415 }
6416 lut
6417};
6418
6419#[inline]
6424fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6425 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6426 let off = (r * gpr + j / GROUP_SIZE) * TILE;
6427 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6428 let within = j % GROUP_SIZE;
6429 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6430}
6431
6432#[cfg(target_arch = "aarch64")]
6435#[target_feature(enable = "neon,dotprod")]
6436#[inline]
6437unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6438 unsafe {
6440 use core::arch::aarch64::*;
6441 use core::arch::asm;
6442 let w0 = vld1q_s8(w);
6443 let w1 = vld1q_s8(w.add(16));
6444 let x0 = vld1q_s8(x);
6445 let x1 = vld1q_s8(x.add(16));
6446 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6447 asm!(
6448 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6449 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6450 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6451 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6452 options(pure, nomem, nostack),
6453 );
6454 vaddvq_s32(vaddq_s32(a0, a1))
6455 }
6456}
6457
6458#[cfg(target_arch = "x86_64")]
6461#[target_feature(enable = "avx2")]
6462#[inline]
6463unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6464 unsafe {
6466 use core::arch::x86_64::*;
6467 let wv = _mm256_loadu_si256(w as *const __m256i);
6468 let xv = _mm256_loadu_si256(x as *const __m256i);
6469 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6470 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6471 let hi128 = _mm256_extracti128_si256::<1>(d);
6472 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6473 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6474 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6475 _mm_cvtsi128_si32(s32)
6476 }
6477}
6478
6479#[inline]
6484fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6485 debug_assert!(dst.len() >= 40);
6486 unsafe {
6489 let p = dst.as_mut_ptr();
6490 for bi in 0..7 {
6491 core::ptr::write_unaligned(
6492 p.add(bi * 5) as *mut u64,
6493 SIGN5_U64[*codes.add(bi) as usize],
6494 );
6495 }
6496 }
6497}
6498
6499#[inline]
6504fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6505 #[cfg(target_arch = "aarch64")]
6506 unsafe {
6507 return sdot32_i8(w, x);
6508 }
6509 #[cfg(target_arch = "x86_64")]
6510 unsafe {
6511 return i8dot32_avx2(w, x);
6512 }
6513 #[allow(unreachable_code)]
6514 unsafe {
6515 let mut s = 0i32;
6516 for k in 0..GROUP_SIZE {
6517 s += *w.add(k) as i32 * *x.add(k) as i32;
6518 }
6519 s
6520 }
6521}
6522
6523#[inline]
6524unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6525 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6526 (
6527 SIGN5_U64[*codes as usize],
6528 SIGN5_U64[*codes.add(1) as usize],
6529 SIGN5_U64[*codes.add(2) as usize],
6530 SIGN5_U64[*codes.add(3) as usize],
6531 SIGN5_U64[*codes.add(4) as usize],
6532 SIGN5_U64[*codes.add(5) as usize],
6533 SIGN5_U64[*codes.add(6) as usize],
6534 )
6535 };
6536
6537 let u0 = s0 | (s1 << 40);
6538 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6539 let u2 = (s3 >> 8) | (s4 << 32);
6540 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6541
6542 (u0, u1, u2, u3)
6543}
6544
6545#[cfg(target_arch = "aarch64")]
6549#[target_feature(enable = "neon,dotprod")]
6550unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6551 use core::arch::aarch64::*;
6552 use core::arch::asm;
6553 unsafe {
6554 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6555 let mut acc = 0f32;
6556 let bytes_ptr = bytes.as_ptr();
6557 let xq_ptr = xq.as_ptr();
6558 let row_off = r * gpr * TILE;
6559
6560 let gpr2 = gpr & !1;
6561 let mut gi = 0;
6562 while gi < gpr2 {
6563 let off0 = row_off + gi * TILE;
6564 let off1 = off0 + TILE;
6565 let s0 = f16_to_f32(u16::from_le_bytes([
6566 *bytes_ptr.add(off0),
6567 *bytes_ptr.add(off0 + 1),
6568 ]));
6569 let s1 = f16_to_f32(u16::from_le_bytes([
6570 *bytes_ptr.add(off1),
6571 *bytes_ptr.add(off1 + 1),
6572 ]));
6573
6574 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6575 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6576
6577 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6578 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6579 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6580 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6581
6582 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6583 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6584 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6585 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6586
6587 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6588 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6589 asm!(
6590 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6591 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6592 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6593 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6594 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6595 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6596 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6597 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6598 options(pure, nomem, nostack),
6599 );
6600 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6601 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6602 acc += d0 as f32 * s0 + d1 as f32 * s1;
6603 gi += 2;
6604 }
6605
6606 if gi < gpr {
6607 let off = row_off + gi * TILE;
6608 let s = f16_to_f32(u16::from_le_bytes([
6609 *bytes_ptr.add(off),
6610 *bytes_ptr.add(off + 1),
6611 ]));
6612 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6613 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6614 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6615 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6616 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6617 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6618 asm!(
6619 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6620 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6621 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6622 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6623 options(pure, nomem, nostack),
6624 );
6625 let d = vaddvq_s32(vaddq_s32(a0, a1));
6626 acc += d as f32 * s;
6627 }
6628 acc
6629 }
6630}
6631
6632#[cfg(target_arch = "x86_64")]
6634#[target_feature(enable = "avx2")]
6635unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6636 use core::arch::x86_64::*;
6637 unsafe {
6638 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6639 let mut acc = 0f32;
6640 let bytes_ptr = bytes.as_ptr();
6641 let xq_ptr = xq.as_ptr();
6642 let row_off = r * gpr * TILE;
6643
6644 let ones = _mm256_set1_epi16(1);
6645 for gi in 0..gpr {
6646 let off = row_off + gi * TILE;
6647 let s = f16_to_f32(u16::from_le_bytes([
6648 *bytes_ptr.add(off),
6649 *bytes_ptr.add(off + 1),
6650 ]));
6651 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6652 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6653 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6654 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6655 let d256 = _mm256_madd_epi16(p16, ones);
6656 let d128 = _mm_add_epi32(
6657 _mm256_castsi256_si128(d256),
6658 _mm256_extracti128_si256(d256, 1),
6659 );
6660 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6661 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6662 acc += d32 as f32 * s;
6663 }
6664 acc
6665 }
6666}
6667
6668#[cfg(target_arch = "x86_64")]
6670#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6671unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6672 use core::arch::x86_64::*;
6673 unsafe {
6675 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6676 let mut acc = 0f32;
6677 let bytes_ptr = bytes.as_ptr();
6678 let xq_ptr = xq.as_ptr();
6679 let row_off = r * gpr * TILE;
6680 for gi in 0..gpr {
6681 let off = row_off + gi * TILE;
6682 let s = f16_to_f32(u16::from_le_bytes([
6683 *bytes_ptr.add(off),
6684 *bytes_ptr.add(off + 1),
6685 ]));
6686 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6687 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6688 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6689 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6690 acc += d as f32 * s;
6691 }
6692 acc
6693 }
6694}
6695
6696#[inline]
6700fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6701 #[cfg(target_arch = "aarch64")]
6702 unsafe {
6703 return q1t_dot_row_sdot(bytes, r, gpr, xq);
6704 }
6705 #[cfg(target_arch = "x86_64")]
6706 unsafe {
6707 if vnni_tiles_enabled() {
6708 return q1t_dot_row_vnni(bytes, r, gpr, xq);
6709 }
6710 return q1t_dot_row_avx2(bytes, r, gpr, xq);
6711 }
6712 #[allow(unreachable_code)]
6713 {
6714 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6715 let mut acc = 0f32;
6716 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
6718 let off = (r * gpr + gi) * TILE;
6719 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6720 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6721 let mut d = 0i32;
6722 for k in 0..GROUP_SIZE {
6723 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6724 }
6725 acc += d as f32 * s;
6726 }
6727 acc
6728 }
6729}
6730
6731fn q1t_row_outlier_correction(
6738 bytes: &[u8],
6739 r: usize,
6740 rp_off: usize,
6741 entries_off: usize,
6742 has_ov: bool,
6743 x: &[f32],
6744) -> f32 {
6745 if !has_ov {
6746 return 0.0;
6747 }
6748 let (c0, c1) = (
6749 q1t_rowptr(bytes, rp_off, r),
6750 q1t_rowptr(bytes, rp_off, r + 1),
6751 );
6752 let mut corr = 0f32;
6753 for p in c0..c1 {
6754 let e = entries_off + p * 4;
6755 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6756 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6757 corr += val * x[col];
6758 }
6759 corr
6760}
6761
6762fn q1t_dequant_row(
6766 bytes: &[u8],
6767 r: usize,
6768 gpr: usize,
6769 rp_off: usize,
6770 entries_off: usize,
6771 has_ov: bool,
6772 buf: &mut [f32],
6773) {
6774 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6775 for g in 0..gpr {
6776 let off = (r * gpr + g) * TILE;
6777 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6778 let codes = &bytes[off + 2..off + TILE];
6779 let bc = g * GROUP_SIZE;
6780 for bi in 0..6 {
6782 let lut = &SIGN5[codes[bi] as usize];
6783 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6784 for i in 0..5 {
6785 d[i] = lut[i] * s;
6786 }
6787 }
6788 let lut = &SIGN5[codes[6] as usize];
6789 buf[bc + 30] = lut[0] * s;
6790 buf[bc + 31] = lut[1] * s;
6791 }
6792 if !has_ov {
6793 return;
6794 }
6795 let (c0, c1) = (
6796 q1t_rowptr(bytes, rp_off, r),
6797 q1t_rowptr(bytes, rp_off, r + 1),
6798 );
6799 for p in c0..c1 {
6800 let e = entries_off + p * 4;
6801 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6802 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6803 }
6804}
6805
6806fn q1t_add_overlay(
6810 bytes: &[u8],
6811 x: &[f32],
6812 rows: usize,
6813 cols: usize,
6814 out: &mut [f32],
6815 pool: Option<&Pool>,
6816) {
6817 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6818 let gpr = cols / GROUP_SIZE;
6819 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6820 if !has_ov {
6821 return;
6822 }
6823 let out_addr = SendMut(out.as_mut_ptr());
6824 let run = move |start: usize, end: usize| {
6825 for r in start..end {
6826 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6827 unsafe { *out_addr.at(r) += corr };
6829 }
6830 };
6831 dispatch_rows(pool, rows, &run);
6832}
6833
6834#[allow(clippy::too_many_arguments)]
6837fn q1t_range_a8w8(
6838 bytes: &[u8],
6839 gpr: usize,
6840 rp_off: usize,
6841 ent_off: usize,
6842 has_ov: bool,
6843 act: &SplitAct,
6844 x: &[f32],
6845 out: SendMut,
6846 start: usize,
6847 end: usize,
6848) {
6849 for r in start..end {
6850 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6851 for &(j, xv) in &act.outliers {
6852 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6853 }
6854 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6855 unsafe { *out.at(r) = acc };
6857 }
6858}
6859
6860#[allow(clippy::too_many_arguments)]
6863fn q1t_range_f32_batch(
6864 bytes: &[u8],
6865 gpr: usize,
6866 rp_off: usize,
6867 ent_off: usize,
6868 has_ov: bool,
6869 x: &[f32],
6870 out: SendMut,
6871 start: usize,
6872 end: usize,
6873) {
6874 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6875 let mut sg = [0f32; GROUP_SIZE];
6876 for r in start..end {
6877 let mut acc = 0f32;
6878 for g in 0..gpr {
6879 let off = (r * gpr + g) * TILE;
6880 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6881 let codes = &bytes[off + 2..off + TILE];
6882 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6883 for bi in 0..6 {
6884 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6885 }
6886 let lut = &SIGN5[codes[6] as usize];
6887 sg[30] = lut[0];
6888 sg[31] = lut[1];
6889 let mut gsum = 0f32;
6890 for k in 0..GROUP_SIZE {
6891 gsum += sg[k] * xg[k];
6892 }
6893 acc += s * gsum;
6894 }
6895 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6896 unsafe { *out.at(r) = acc };
6898 }
6899}
6900
6901fn q1t_matvec(
6905 bytes: &[u8],
6906 x: &[f32],
6907 rows: usize,
6908 cols: usize,
6909 out: &mut [f32],
6910 pool: Option<&Pool>,
6911) {
6912 debug_assert_eq!(out.len(), rows);
6913 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6914 let gpr = cols / GROUP_SIZE;
6915 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6916 let out_addr = SendMut(out.as_mut_ptr());
6917 if a8w8_enabled() {
6921 let act = split_act(x);
6922 let act = &act;
6923 let run = move |start: usize, end: usize| {
6924 for r in start..end {
6925 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6926 for &(j, xv) in &act.outliers {
6927 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6928 }
6929 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6930 unsafe { *out_addr.at(r) = acc };
6932 }
6933 };
6934 dispatch_rows(pool, rows, &run);
6935 return;
6936 }
6937 let run = move |start: usize, end: usize| {
6938 let mut sg = [0f32; GROUP_SIZE];
6942 for r in start..end {
6943 let mut acc = 0f32;
6944 for g in 0..gpr {
6945 let off = (r * gpr + g) * TILE;
6946 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6947 let codes = &bytes[off + 2..off + TILE];
6948 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6949 for bi in 0..6 {
6950 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6951 }
6952 let lut = &SIGN5[codes[6] as usize];
6953 sg[30] = lut[0];
6954 sg[31] = lut[1];
6955 let mut gsum = 0f32;
6956 for k in 0..GROUP_SIZE {
6957 gsum += sg[k] * xg[k];
6958 }
6959 acc += s * gsum;
6960 }
6961 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6962 unsafe { *out_addr.at(r) = acc };
6963 }
6964 };
6965 dispatch_rows(pool, rows, &run);
6966}
6967
6968#[cfg(target_arch = "aarch64")]
6974#[target_feature(enable = "neon,dotprod")]
6975unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6976 use core::arch::aarch64::*;
6977 use core::arch::asm;
6978 unsafe {
6980 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6981 let bytes_ptr = bytes.as_ptr();
6982 let row_off = r * gpr * TILE;
6983 let xp = [xa.as_ptr(), xb.as_ptr()];
6984 let mut acc = [0f32; 2];
6985 macro_rules! sdot2 {
6986 ($w0:expr, $w1:expr, $x:expr) => {{
6987 let x0 = vld1q_s8($x);
6988 let x1 = vld1q_s8($x.add(16));
6989 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6990 asm!(
6991 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6992 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6993 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6994 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6995 options(pure, nomem, nostack),
6996 );
6997 vaddvq_s32(vaddq_s32(a0, a1))
6998 }};
6999 }
7000 let gpr2 = gpr & !1;
7001 let mut gi = 0;
7002 while gi < gpr2 {
7003 let off0 = row_off + gi * TILE;
7004 let off1 = off0 + TILE;
7005 let s0 = f16_to_f32(u16::from_le_bytes([
7006 *bytes_ptr.add(off0),
7007 *bytes_ptr.add(off0 + 1),
7008 ]));
7009 let s1 = f16_to_f32(u16::from_le_bytes([
7010 *bytes_ptr.add(off1),
7011 *bytes_ptr.add(off1 + 1),
7012 ]));
7013 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7014 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7015 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7016 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7017 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7018 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7019 for k in 0..2 {
7020 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7021 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7022 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7023 }
7024 gi += 2;
7025 }
7026 if gi < gpr {
7027 let off = row_off + gi * TILE;
7028 let s = f16_to_f32(u16::from_le_bytes([
7029 *bytes_ptr.add(off),
7030 *bytes_ptr.add(off + 1),
7031 ]));
7032 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7033 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7034 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7035 for k in 0..2 {
7036 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7037 acc[k] += d as f32 * s;
7038 }
7039 }
7040 acc
7041 }
7042}
7043
7044fn q1t_matvec2(
7050 bytes: &[u8],
7051 x1: &[f32],
7052 x2: &[f32],
7053 rows: usize,
7054 cols: usize,
7055 o1: &mut [f32],
7056 o2: &mut [f32],
7057 pool: Option<&Pool>,
7058) {
7059 debug_assert_eq!(o1.len(), rows);
7060 debug_assert_eq!(o2.len(), rows);
7061 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7062 let gpr = cols / GROUP_SIZE;
7063 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7064 let out1 = SendMut(o1.as_mut_ptr());
7065 let out2 = SendMut(o2.as_mut_ptr());
7066 if a8w8_enabled() {
7067 let a1 = split_act(x1);
7068 let a2 = split_act(x2);
7069 let (a1, a2) = (&a1, &a2);
7070 let run = move |start: usize, end: usize| {
7071 for r in start..end {
7072 #[cfg(target_arch = "aarch64")]
7073 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7076 #[cfg(not(target_arch = "aarch64"))]
7077 let ds = [
7078 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7079 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7080 ];
7081 let mut acc1 = ds[0] * a1.sx;
7082 for &(j, xv) in &a1.outliers {
7083 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7084 }
7085 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7086 let mut acc2 = ds[1] * a2.sx;
7087 for &(j, xv) in &a2.outliers {
7088 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7089 }
7090 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7091 unsafe {
7093 *out1.at(r) = acc1;
7094 *out2.at(r) = acc2;
7095 }
7096 }
7097 };
7098 dispatch_rows(pool, rows, &run);
7099 return;
7100 }
7101 let run = move |start: usize, end: usize| {
7102 let mut sg = [0f32; GROUP_SIZE];
7105 for r in start..end {
7106 let mut acc1 = 0f32;
7107 let mut acc2 = 0f32;
7108 for g in 0..gpr {
7109 let off = (r * gpr + g) * TILE;
7110 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7111 let codes = &bytes[off + 2..off + TILE];
7112 for bi in 0..6 {
7113 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7114 }
7115 let lut = &SIGN5[codes[6] as usize];
7116 sg[30] = lut[0];
7117 sg[31] = lut[1];
7118 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7119 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7120 let mut gsum1 = 0f32;
7121 for k in 0..GROUP_SIZE {
7122 gsum1 += sg[k] * xg1[k];
7123 }
7124 acc1 += s * gsum1;
7125 let mut gsum2 = 0f32;
7126 for k in 0..GROUP_SIZE {
7127 gsum2 += sg[k] * xg2[k];
7128 }
7129 acc2 += s * gsum2;
7130 }
7131 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7132 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7133 unsafe {
7135 *out1.at(r) = acc1;
7136 *out2.at(r) = acc2;
7137 }
7138 }
7139 };
7140 dispatch_rows(pool, rows, &run);
7141}
7142
7143fn q1t_matmat(
7146 bytes: &[u8],
7147 xs: &[f32],
7148 b: usize,
7149 rows: usize,
7150 cols: usize,
7151 out: &mut [f32],
7152 pool: Option<&Pool>,
7153) {
7154 debug_assert_eq!(out.len(), b * rows);
7155 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7156 let gpr = cols / GROUP_SIZE;
7157 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7158 let out_addr = SendMut(out.as_mut_ptr());
7159 if a8w8_enabled() {
7163 let acts: Vec<SplitAct> = (0..b)
7164 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7165 .collect();
7166 let acts = &acts;
7167 let run = move |start: usize, end: usize| {
7168 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
7172 for g in 0..gpr {
7173 let off = (r * gpr + g) * TILE;
7174 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7175 q1t_unpack_group_i8(
7176 bytes.as_ptr().wrapping_add(off + 2),
7177 &mut sg[g * GROUP_SIZE..],
7178 );
7179 }
7180 for bi in 0..b {
7181 let act = &acts[bi];
7182 let mut isum = 0f32;
7183 for g in 0..gpr {
7184 let d = q1t_i8dot32(
7185 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7186 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7187 );
7188 isum += d as f32 * sc[g];
7189 }
7190 let mut acc = isum * act.sx;
7191 for &(j, xv) in &act.outliers {
7192 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7193 }
7194 accs[bi] = acc;
7195 }
7196 if has_ov {
7200 let (c0, c1) = (
7201 q1t_rowptr(bytes, rp_off, r),
7202 q1t_rowptr(bytes, rp_off, r + 1),
7203 );
7204 for p in c0..c1 {
7205 let e = ent_off + p * 4;
7206 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7207 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7208 for bi in 0..b {
7209 accs[bi] += val * xs[bi * cols + col];
7210 }
7211 }
7212 }
7213 for bi in 0..b {
7214 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7215 }
7216 }
7217 };
7218 dispatch_rows(pool, rows, &run);
7219 return;
7220 }
7221 let run = move |start: usize, end: usize| {
7222 let mut buf = vec![0f32; cols];
7223 for r in start..end {
7224 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7225 for bi in 0..b {
7226 let xr = &xs[bi * cols..(bi + 1) * cols];
7227 let mut acc = 0f32;
7228 for j in 0..cols {
7229 acc += buf[j] * xr[j];
7230 }
7231 unsafe { *out_addr.at(bi * rows + r) = acc };
7232 }
7233 }
7234 };
7235 dispatch_rows(pool, rows, &run);
7236}
7237
7238fn q1_matvec(
7239 bytes: &[u8],
7240 x: &[f32],
7241 rows: usize,
7242 cols: usize,
7243 out: &mut [f32],
7244 pool: Option<&Pool>,
7245) {
7246 debug_assert_eq!(out.len(), rows);
7247 let gpr = cols / GROUP_SIZE;
7248 let out_addr = SendMut(out.as_mut_ptr());
7249 if a8w8_enabled() {
7250 let act = split_act(x);
7251 let gsum = q1_group_sums(&act.xq, gpr);
7252 let (act, gsum) = (&act, &gsum);
7253 let run = move |start: usize, end: usize| {
7254 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7255 };
7256 dispatch_rows(pool, rows, &run);
7257 return;
7258 }
7259 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7260 dispatch_rows(pool, rows, &run);
7261}
7262
7263#[allow(clippy::too_many_arguments)]
7265fn q1_matvec2(
7266 bytes: &[u8],
7267 x1: &[f32],
7268 x2: &[f32],
7269 rows: usize,
7270 cols: usize,
7271 o1: &mut [f32],
7272 o2: &mut [f32],
7273 pool: Option<&Pool>,
7274) {
7275 let gpr = cols / GROUP_SIZE;
7276 let p1 = SendMut(o1.as_mut_ptr());
7277 let p2 = SendMut(o2.as_mut_ptr());
7278 if a8w8_enabled() {
7279 let a1 = split_act(x1);
7280 let a2 = split_act(x2);
7281 let g1 = q1_group_sums(&a1.xq, gpr);
7282 let g2 = q1_group_sums(&a2.xq, gpr);
7283 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7284 let run = move |start: usize, end: usize| {
7285 for r in start..end {
7286 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7287 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7288 for &(j, xv) in &a1.outliers {
7289 let (w, s) = q1_outlier(bytes, r, gpr, j);
7290 v1 += w * s * xv;
7291 }
7292 for &(j, xv) in &a2.outliers {
7293 let (w, s) = q1_outlier(bytes, r, gpr, j);
7294 v2 += w * s * xv;
7295 }
7296 unsafe {
7298 *p1.at(r) = v1;
7299 *p2.at(r) = v2;
7300 }
7301 }
7302 };
7303 dispatch_rows(pool, rows, &run);
7304 return;
7305 }
7306 let run = move |start: usize, end: usize| {
7307 for r in start..end {
7308 unsafe {
7310 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7311 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7312 }
7313 }
7314 };
7315 dispatch_rows(pool, rows, &run);
7316}
7317
7318#[allow(clippy::too_many_arguments)]
7320fn q1_matmat(
7321 bytes: &[u8],
7322 xs_all: &[f32],
7323 b: usize,
7324 rows: usize,
7325 cols: usize,
7326 out: &mut [f32],
7327 pool: Option<&Pool>,
7328) {
7329 debug_assert_eq!(out.len(), b * rows);
7330 let gpr = cols / GROUP_SIZE;
7331 let out_addr = SendMut(out.as_mut_ptr());
7332 if a8w8_enabled() {
7333 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7334 .map(|bi| {
7335 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7336 let gsum = q1_group_sums(&act.xq, gpr);
7337 (act, gsum)
7338 })
7339 .collect();
7340 let acts = &acts;
7341 #[cfg(target_arch = "x86_64")]
7342 let blocked_ok = avx2_enabled()
7343 && blocked_enabled();
7344 #[cfg(target_arch = "aarch64")]
7345 let blocked_ok = sdot_enabled()
7346 && blocked_enabled();
7347 let run = move |start: usize, end: usize| {
7348 for r in start..end {
7349 let mut bi = 0usize;
7350 #[cfg(target_arch = "aarch64")]
7353 if blocked_ok {
7354 while bi + 4 <= acts.len() {
7355 let xs = [
7356 acts[bi].0.xq.as_slice(),
7357 acts[bi + 1].0.xq.as_slice(),
7358 acts[bi + 2].0.xq.as_slice(),
7359 acts[bi + 3].0.xq.as_slice(),
7360 ];
7361 let gs = [
7362 acts[bi].1.as_slice(),
7363 acts[bi + 1].1.as_slice(),
7364 acts[bi + 2].1.as_slice(),
7365 acts[bi + 3].1.as_slice(),
7366 ];
7367 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7368 for k in 0..4 {
7369 let (act, _) = &acts[bi + k];
7370 let mut acc = d[k] * act.sx;
7371 for &(j, xv) in &act.outliers {
7372 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7373 acc += w * sc * xv;
7374 }
7375 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7377 }
7378 bi += 4;
7379 }
7380 }
7381 #[cfg(target_arch = "x86_64")]
7382 if blocked_ok {
7383 while bi + 4 <= acts.len() {
7384 let xs = [
7385 acts[bi].0.xq.as_slice(),
7386 acts[bi + 1].0.xq.as_slice(),
7387 acts[bi + 2].0.xq.as_slice(),
7388 acts[bi + 3].0.xq.as_slice(),
7389 ];
7390 let gs = [
7391 acts[bi].1.as_slice(),
7392 acts[bi + 1].1.as_slice(),
7393 acts[bi + 2].1.as_slice(),
7394 acts[bi + 3].1.as_slice(),
7395 ];
7396 let d = unsafe {
7397 if vnni_tiles_enabled() {
7398 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7399 } else {
7400 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7401 }
7402 };
7403 for k in 0..4 {
7404 let (act, _) = &acts[bi + k];
7405 let mut acc = d[k] * act.sx;
7406 for &(j, xv) in &act.outliers {
7407 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7408 acc += w * sc * xv;
7409 }
7410 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7412 }
7413 bi += 4;
7414 }
7415 }
7416 while bi < acts.len() {
7417 let (act, gsum) = &acts[bi];
7418 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7419 for &(j, xv) in &act.outliers {
7420 let (w, s) = q1_outlier(bytes, r, gpr, j);
7421 acc += w * s * xv;
7422 }
7423 unsafe { *out_addr.at(bi * rows + r) = acc };
7425 bi += 1;
7426 }
7427 }
7428 };
7429 dispatch_rows(pool, rows, &run);
7430 return;
7431 }
7432 let run = move |start: usize, end: usize| {
7433 for r in start..end {
7434 for bi in 0..b {
7435 let x = &xs_all[bi * cols..(bi + 1) * cols];
7436 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7438 }
7439 }
7440 };
7441 dispatch_rows(pool, rows, &run);
7442}
7443
7444fn q4matvec(
7450 bytes: &[u8],
7451 x: &[f32],
7452 rows: usize,
7453 cols: usize,
7454 out: &mut [f32],
7455 pool: Option<&Pool>,
7456) {
7457 debug_assert_eq!(out.len(), rows);
7458 let (packed, scales) = q4_split(bytes, rows, cols);
7459 let gpr = cols / GROUP_SIZE;
7460 let out_addr = SendMut(out.as_mut_ptr());
7461
7462 if a8w8_enabled() {
7463 let act = split_act(x);
7464 let run = move |start: usize, end: usize| {
7465 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7466 };
7467 dispatch_rows(pool, rows, &run);
7468 return;
7469 }
7470
7471 let run =
7472 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7473 dispatch_rows(pool, rows, &run);
7474}
7475
7476#[inline]
7479#[allow(unreachable_code)]
7480#[cfg(target_arch = "x86_64")]
7485#[target_feature(enable = "avx2")]
7486unsafe fn dot_q4b_row_1x4_avx2(
7487 buf: &[u8],
7488 scales: &[u8],
7489 g0: usize,
7490 gpr: usize,
7491 xs: [&[i8]; 4],
7492) -> [f32; 4] {
7493 unsafe {
7495 use core::arch::x86_64::*;
7496 let ones = _mm256_set1_epi16(1);
7497 let mut acc = [0f32; 4];
7498 for gi in 0..gpr {
7499 let s = f16_to_f32(u16::from_le_bytes([
7500 scales[(g0 + gi) * 2],
7501 scales[(g0 + gi) * 2 + 1],
7502 ]));
7503 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7504 let aw = _mm256_abs_epi8(w);
7505 for (k, xq) in xs.iter().enumerate() {
7506 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7507 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7508 let d = _mm256_madd_epi16(p16, ones);
7509 let hi128 = _mm256_extracti128_si256::<1>(d);
7510 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7511 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7512 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7513 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7514 }
7515 }
7516 acc
7517 }
7518}
7519
7520#[cfg(target_arch = "x86_64")]
7522#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7523unsafe fn dot_q4b_row_1x4_vnni(
7524 buf: &[u8],
7525 scales: &[u8],
7526 g0: usize,
7527 gpr: usize,
7528 xs: [&[i8]; 4],
7529) -> [f32; 4] {
7530 unsafe {
7532 use core::arch::x86_64::*;
7533 let mut acc = [0f32; 4];
7534 for gi in 0..gpr {
7535 let s = f16_to_f32(u16::from_le_bytes([
7536 scales[(g0 + gi) * 2],
7537 scales[(g0 + gi) * 2 + 1],
7538 ]));
7539 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7540 let aw = _mm256_abs_epi8(w);
7541 for (k, xq) in xs.iter().enumerate() {
7542 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7543 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7544 acc[k] += d as f32 * s;
7545 }
7546 }
7547 acc
7548 }
7549}
7550
7551#[cfg(target_arch = "x86_64")]
7557#[target_feature(enable = "avx2")]
7558unsafe fn dot_q4b_row_1x4_sx_avx2(
7559 buf: &[u8],
7560 scales: &[u8],
7561 g0: usize,
7562 gpr: usize,
7563 xs: [&[i8]; 4],
7564 sxs: [f32; 4],
7565) -> [f32; 4] {
7566 unsafe {
7568 use core::arch::x86_64::*;
7569 let ones = _mm256_set1_epi16(1);
7570 let mut acc = [0f32; 4];
7571 for gi in 0..gpr {
7572 let s = f16_to_f32(u16::from_le_bytes([
7573 scales[(g0 + gi) * 2],
7574 scales[(g0 + gi) * 2 + 1],
7575 ]));
7576 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7577 let aw = _mm256_abs_epi8(w);
7578 for (k, xq) in xs.iter().enumerate() {
7579 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7580 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7581 let d = _mm256_madd_epi16(p16, ones);
7582 let hi128 = _mm256_extracti128_si256::<1>(d);
7583 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7584 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7585 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7586 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7587 }
7588 }
7589 acc
7590 }
7591}
7592
7593#[cfg(target_arch = "x86_64")]
7596#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7597unsafe fn dot_q4b_row_1x4_sx_vnni(
7598 buf: &[u8],
7599 scales: &[u8],
7600 g0: usize,
7601 gpr: usize,
7602 xs: [&[i8]; 4],
7603 sxs: [f32; 4],
7604) -> [f32; 4] {
7605 unsafe {
7607 use core::arch::x86_64::*;
7608 let mut acc = [0f32; 4];
7609 for gi in 0..gpr {
7610 let s = f16_to_f32(u16::from_le_bytes([
7611 scales[(g0 + gi) * 2],
7612 scales[(g0 + gi) * 2 + 1],
7613 ]));
7614 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7615 let aw = _mm256_abs_epi8(w);
7616 for (k, xq) in xs.iter().enumerate() {
7617 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7618 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7619 acc[k] += (d as f32 * sxs[k]) * s;
7620 }
7621 }
7622 acc
7623 }
7624}
7625
7626#[allow(unreachable_code)]
7627fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7628 #[cfg(target_arch = "aarch64")]
7629 unsafe {
7630 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7631 }
7632 #[cfg(target_arch = "x86_64")]
7633 unsafe {
7634 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7635 }
7636 let mut acc = 0f32;
7637 for gi in 0..gpr {
7638 let g = g0 + gi;
7639 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7640 let mut d = 0i32;
7641 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7642 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7643 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7644 }
7645 acc += d as f32 * s;
7646 }
7647 acc
7648}
7649
7650#[inline]
7652#[allow(unreachable_code)]
7653fn dot_q4_row_i8_2(
7654 packed: &[u8],
7655 scales: &[u8],
7656 g0: usize,
7657 gpr: usize,
7658 xq1: &[i8],
7659 xq2: &[i8],
7660) -> (f32, f32) {
7661 #[cfg(target_arch = "aarch64")]
7662 unsafe {
7663 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7664 }
7665 #[cfg(target_arch = "x86_64")]
7666 unsafe {
7667 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7668 }
7669 (
7670 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7671 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7672 )
7673}
7674
7675#[allow(clippy::too_many_arguments)]
7678fn q4_range_a8w8(
7679 packed: &[u8],
7680 scales: &[u8],
7681 gpr: usize,
7682 cols: usize,
7683 act: &SplitAct,
7684 out: SendMut,
7685 start: usize,
7686 end: usize,
7687) {
7688 for r in start..end {
7689 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7690 for &(j, xv) in &act.outliers {
7692 let flat = r * cols + j;
7693 let byte = packed[flat / 2];
7694 let nib = if flat & 1 == 0 {
7695 byte & 0x0F
7696 } else {
7697 byte >> 4
7698 };
7699 let s = f16_to_f32(u16::from_le_bytes([
7700 scales[(flat / GROUP_SIZE) * 2],
7701 scales[(flat / GROUP_SIZE) * 2 + 1],
7702 ]));
7703 acc += ((nib as i32 - 8) as f32) * s * xv;
7704 }
7705 unsafe { *out.at(r) = acc };
7707 }
7708}
7709
7710#[allow(clippy::too_many_arguments)]
7713fn q4_range2_a8w8(
7714 packed: &[u8],
7715 scales: &[u8],
7716 gpr: usize,
7717 cols: usize,
7718 a1: &SplitAct,
7719 a2: &SplitAct,
7720 p1: SendMut,
7721 p2: SendMut,
7722 start: usize,
7723 end: usize,
7724) {
7725 for r in start..end {
7726 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7727 let mut acc1 = s1 * a1.sx;
7728 let mut acc2 = s2 * a2.sx;
7729 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7731 for &(j, xv) in outliers {
7732 let flat = r * cols + j;
7733 let byte = packed[flat / 2];
7734 let nib = if flat & 1 == 0 {
7735 byte & 0x0F
7736 } else {
7737 byte >> 4
7738 };
7739 let s = f16_to_f32(u16::from_le_bytes([
7740 scales[(flat / GROUP_SIZE) * 2],
7741 scales[(flat / GROUP_SIZE) * 2 + 1],
7742 ]));
7743 *acc += ((nib as i32 - 8) as f32) * s * xv;
7744 }
7745 };
7746 fix(&a1.outliers, &mut acc1);
7747 fix(&a2.outliers, &mut acc2);
7748 unsafe {
7750 *p1.at(r) = acc1;
7751 *p2.at(r) = acc2;
7752 }
7753 }
7754}
7755
7756fn q4_range_f32(
7758 packed: &[u8],
7759 scales: &[u8],
7760 gpr: usize,
7761 x: &[f32],
7762 out: SendMut,
7763 start: usize,
7764 end: usize,
7765) {
7766 for r in start..end {
7767 let mut acc = 0f32;
7768 for gi in 0..gpr {
7769 let g = r * gpr + gi;
7770 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7771 let pk = &packed[g * 16..(g + 1) * 16];
7772 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7773 let mut ga = 0f32;
7774 for (k, &b) in pk.iter().enumerate() {
7775 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7776 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7777 }
7778 acc += ga * s;
7779 }
7780 unsafe { *out.at(r) = acc };
7782 }
7783}
7784
7785#[allow(clippy::too_many_arguments)]
7789fn q4matvec2(
7790 bytes: &[u8],
7791 x1: &[f32],
7792 x2: &[f32],
7793 rows: usize,
7794 cols: usize,
7795 o1: &mut [f32],
7796 o2: &mut [f32],
7797 pool: Option<&Pool>,
7798) {
7799 debug_assert_eq!(o1.len(), rows);
7800 debug_assert_eq!(o2.len(), rows);
7801 let (packed, scales) = q4_split(bytes, rows, cols);
7802 let gpr = cols / GROUP_SIZE;
7803
7804 if a8w8_enabled() {
7805 let a1 = split_act(x1);
7806 let a2 = split_act(x2);
7807 let p1 = SendMut(o1.as_mut_ptr());
7808 let p2 = SendMut(o2.as_mut_ptr());
7809 let run = move |start: usize, end: usize| {
7810 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7811 };
7812 dispatch_rows(pool, rows, &run);
7813 return;
7814 }
7815
7816 let p1 = SendMut(o1.as_mut_ptr());
7817 let p2 = SendMut(o2.as_mut_ptr());
7818 let run = move |start: usize, end: usize| {
7819 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7820 };
7821 dispatch_rows(pool, rows, &run);
7822}
7823
7824#[allow(clippy::too_many_arguments)]
7826fn q4_range2_f32(
7827 packed: &[u8],
7828 scales: &[u8],
7829 gpr: usize,
7830 x1: &[f32],
7831 x2: &[f32],
7832 p1: SendMut,
7833 p2: SendMut,
7834 start: usize,
7835 end: usize,
7836) {
7837 for r in start..end {
7838 let (mut acc1, mut acc2) = (0f32, 0f32);
7839 for gi in 0..gpr {
7840 let g = r * gpr + gi;
7841 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7842 let pk = &packed[g * 16..(g + 1) * 16];
7843 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7844 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7845 let (mut g1, mut g2) = (0f32, 0f32);
7846 for (k, &b) in pk.iter().enumerate() {
7847 let wl = (b & 0x0F) as f32 - 8.0;
7848 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7849 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7850 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7851 }
7852 acc1 += g1 * s;
7853 acc2 += g2 * s;
7854 }
7855 unsafe {
7857 *p1.at(r) = acc1;
7858 *p2.at(r) = acc2;
7859 }
7860 }
7861}
7862
7863thread_local! {
7864 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7867 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7868}
7869
7870#[allow(clippy::too_many_arguments)]
7876fn q4matmat(
7877 bytes: &[u8],
7878 xs_all: &[f32],
7879 b: usize,
7880 rows: usize,
7881 cols: usize,
7882 out: &mut [f32],
7883 pool: Option<&Pool>,
7884) {
7885 debug_assert_eq!(xs_all.len(), b * cols);
7886 debug_assert_eq!(out.len(), b * rows);
7887 let (packed, scales) = q4_split(bytes, rows, cols);
7888 let gpr = cols / GROUP_SIZE;
7889 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7890
7891 if a8w8_enabled() {
7892 let acts: Vec<SplitAct> = (0..b)
7893 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7894 .collect();
7895 let acts = &acts;
7896 let out_addr = SendMut(out.as_mut_ptr());
7897 let run = move |start: usize, end: usize| {
7898 ROW_I8.with(|rb| {
7899 let mut buf = rb.borrow_mut();
7900 buf.resize(cols, 0);
7901 for r in start..end {
7902 for gi in 0..gpr {
7906 let g = r * gpr + gi;
7907 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7908 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7909 buf[gi * GROUP_SIZE + k * 2 + 1] =
7910 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7911 }
7912 }
7913 let mut bi = 0usize;
7914 #[cfg(target_arch = "x86_64")]
7915 if avx2_enabled()
7916 && blocked_enabled()
7917 {
7918 while bi + 4 <= acts.len() {
7919 let xs = [
7920 acts[bi].xq.as_slice(),
7921 acts[bi + 1].xq.as_slice(),
7922 acts[bi + 2].xq.as_slice(),
7923 acts[bi + 3].xq.as_slice(),
7924 ];
7925 let d = unsafe {
7926 if vnni_tiles_enabled() {
7927 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7928 } else {
7929 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7930 }
7931 };
7932 for k in 0..4 {
7933 let act = &acts[bi + k];
7934 let mut acc = d[k] * act.sx;
7935 for &(j, xv) in &act.outliers {
7936 acc += (buf[j] as i8) as f32
7937 * gscale((r * cols + j) / GROUP_SIZE)
7938 * xv;
7939 }
7940 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7942 }
7943 bi += 4;
7944 }
7945 }
7946 while bi < acts.len() {
7947 let act = &acts[bi];
7948 let mut acc = 0f32;
7949 for gi in 0..gpr {
7950 let d = dot_i8_i8(
7951 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7952 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7953 );
7954 acc += d as f32 * gscale(r * gpr + gi);
7955 }
7956 acc *= act.sx;
7957 for &(j, xv) in &act.outliers {
7959 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7960 }
7961 unsafe { *out_addr.at(bi * rows + r) = acc };
7963 bi += 1;
7964 }
7965 }
7966 })
7967 };
7968 dispatch_rows(pool, rows, &run);
7969 return;
7970 }
7971
7972 let out_addr = SendMut(out.as_mut_ptr());
7973 let run = move |start: usize, end: usize| {
7974 ROW_F32.with(|rb| {
7975 let mut buf = rb.borrow_mut();
7976 buf.resize(cols, 0.0);
7977 for r in start..end {
7978 for gi in 0..gpr {
7981 let g = r * gpr + gi;
7982 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7983 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7984 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7985 }
7986 }
7987 for bi in 0..b {
7988 let x = &xs_all[bi * cols..(bi + 1) * cols];
7989 let mut acc = 0f32;
7990 for gi in 0..gpr {
7991 let mut ga = 0f32;
7992 for k in 0..GROUP_SIZE / 2 {
7997 let e = gi * GROUP_SIZE + k * 2;
7998 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
7999 }
8000 acc += ga * gscale(r * gpr + gi);
8001 }
8002 unsafe { *out_addr.at(bi * rows + r) = acc };
8004 }
8005 }
8006 })
8007 };
8008 dispatch_rows(pool, rows, &run);
8009}
8010
8011#[allow(clippy::too_many_arguments)]
8016fn vbitmatmat(
8017 bytes: &[u8],
8018 offsets: &[usize],
8019 xs_all: &[f32],
8020 b: usize,
8021 rows: usize,
8022 cols: usize,
8023 out: &mut [f32],
8024 pool: Option<&Pool>,
8025) {
8026 debug_assert_eq!(xs_all.len(), b * cols);
8027 debug_assert_eq!(out.len(), b * rows);
8028 debug_assert_eq!(offsets.len(), rows + 1);
8029 let ng = cols / GROUP_SIZE;
8030 let bits = &bytes[..rows];
8031 let sc_off = rows;
8032 let gscale = |r: usize, g: usize| {
8033 let so = (r * ng + g) * 2;
8034 f16_to_f32(u16::from_le_bytes([
8035 bytes[sc_off + so],
8036 bytes[sc_off + so + 1],
8037 ]))
8038 };
8039
8040 let decode_f32 = |r: usize, dst: &mut [f32]| {
8042 let bw = bits[r] as usize;
8043 let l = ((1i32 << (bw - 1)) - 1) as f32;
8044 let data = &bytes[offsets[r]..offsets[r + 1]];
8045 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8046 for d in dst.iter_mut() {
8047 while nbits < bw {
8048 acc = (acc << 8) | data[idx] as u64;
8049 idx += 1;
8050 nbits += 8;
8051 }
8052 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8053 nbits -= bw;
8054 *d = u - l;
8055 }
8056 };
8057
8058 if a8w8_enabled() {
8059 let acts: Vec<SplitAct> = (0..b)
8060 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8061 .collect();
8062 let acts = &acts;
8063 let out_addr = SendMut(out.as_mut_ptr());
8064 let run = move |start: usize, end: usize| {
8065 for r in start..end {
8066 let bw = bits[r] as usize;
8067 if bw == 8 {
8068 ROW_F32.with(|rb| {
8071 let mut buf = rb.borrow_mut();
8072 buf.resize(cols, 0.0);
8073 decode_f32(r, &mut buf);
8074 for bi in 0..b {
8075 let x = &xs_all[bi * cols..(bi + 1) * cols];
8076 let mut dot = 0f32;
8077 for g in 0..ng {
8078 let mut gd = 0f32;
8079 for k in 0..GROUP_SIZE {
8080 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8081 }
8082 dot += gd * gscale(r, g);
8083 }
8084 unsafe { *out_addr.at(bi * rows + r) = dot };
8086 }
8087 });
8088 continue;
8089 }
8090 let l = (1i32 << (bw - 1)) - 1;
8091 let data = &bytes[offsets[r]..offsets[r + 1]];
8092 ROW_I8.with(|rb| {
8093 let mut buf = rb.borrow_mut();
8094 buf.resize(cols, 0);
8095 #[inline(always)]
8096 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8097 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8098 let u = unpack8::<B>(&data[blk * B..]);
8099 for k in 0..8 {
8100 chunk[k] = (u[k] - l) as i8 as u8;
8101 }
8102 }
8103 }
8104 match bw {
8105 3 => fill::<3>(data, l, &mut buf),
8106 4 => vbit_fill4(data, &mut buf),
8107 5 => fill::<5>(data, l, &mut buf),
8108 6 => fill::<6>(data, l, &mut buf),
8109 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8110 }
8111 let mut bi = 0usize;
8112 #[cfg(target_arch = "x86_64")]
8116 if avx2_enabled()
8117 && blocked_enabled()
8118 {
8119 while bi + 4 <= acts.len() {
8120 let xs = [
8121 acts[bi].xq.as_slice(),
8122 acts[bi + 1].xq.as_slice(),
8123 acts[bi + 2].xq.as_slice(),
8124 acts[bi + 3].xq.as_slice(),
8125 ];
8126 let sxs = [
8127 acts[bi].sx,
8128 acts[bi + 1].sx,
8129 acts[bi + 2].sx,
8130 acts[bi + 3].sx,
8131 ];
8132 let d = unsafe {
8133 if vnni_tiles_enabled() {
8134 dot_q4b_row_1x4_sx_vnni(
8135 &buf,
8136 &bytes[sc_off..],
8137 r * ng,
8138 ng,
8139 xs,
8140 sxs,
8141 )
8142 } else {
8143 dot_q4b_row_1x4_sx_avx2(
8144 &buf,
8145 &bytes[sc_off..],
8146 r * ng,
8147 ng,
8148 xs,
8149 sxs,
8150 )
8151 }
8152 };
8153 for k in 0..4 {
8154 let act = &acts[bi + k];
8155 let mut dot = d[k];
8156 for &(j, xv) in &act.outliers {
8157 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8158 }
8159 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8161 }
8162 bi += 4;
8163 }
8164 }
8165 while bi < acts.len() {
8166 let act = &acts[bi];
8167 let mut dot = 0f32;
8168 for g in 0..ng {
8169 let d = dot_i8_i8(
8170 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8171 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8172 ) as f32
8173 * act.sx;
8174 dot += d * gscale(r, g);
8175 }
8176 for &(j, xv) in &act.outliers {
8177 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8178 }
8179 unsafe { *out_addr.at(bi * rows + r) = dot };
8181 bi += 1;
8182 }
8183 });
8184 }
8185 };
8186 dispatch_rows(pool, rows, &run);
8187 return;
8188 }
8189
8190 let out_addr = SendMut(out.as_mut_ptr());
8191 let run = move |start: usize, end: usize| {
8192 ROW_F32.with(|rb| {
8193 let mut buf = rb.borrow_mut();
8194 buf.resize(cols, 0.0);
8195 for r in start..end {
8196 decode_f32(r, &mut buf);
8197 for bi in 0..b {
8198 let x = &xs_all[bi * cols..(bi + 1) * cols];
8199 let mut dot = 0f32;
8200 for g in 0..ng {
8201 let mut gd = 0f32;
8202 for k in 0..GROUP_SIZE {
8203 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8204 }
8205 dot += gd * gscale(r, g);
8206 }
8207 unsafe { *out_addr.at(bi * rows + r) = dot };
8209 }
8210 }
8211 })
8212 };
8213 dispatch_rows(pool, rows, &run);
8214}
8215
8216pub(crate) fn gpu_batch_job<'a>(
8220 t: &'a QTensor,
8221 x: &[f32],
8222) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8223 match t {
8224 QTensor::Mapped {
8225 model,
8226 idx,
8227 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8228 rows,
8229 cols,
8230 row_scale,
8231 col_field,
8232 ..
8233 } => Some((
8234 model.clone(),
8235 crate::gpu::BatchJob {
8236 idx: *idx,
8237 rows: *rows,
8238 cols: *cols,
8239 row_scale,
8240 xs: prescale(x, col_field, *dt).into_owned(),
8241 layout: crate::gpu::BatchLayout::Q8,
8242 },
8243 )),
8244 QTensor::Mapped {
8246 model,
8247 idx,
8248 dtype: TensorDtype::Q1,
8249 rows,
8250 cols,
8251 ..
8252 } => Some((
8253 model.clone(),
8254 crate::gpu::BatchJob {
8255 idx: *idx,
8256 rows: *rows,
8257 cols: *cols,
8258 row_scale: &[],
8259 xs: x.to_vec(),
8260 layout: crate::gpu::BatchLayout::Q1,
8261 },
8262 )),
8263 QTensor::Mapped {
8268 model,
8269 idx,
8270 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8271 rows,
8272 cols,
8273 ..
8274 } => Some((
8275 model.clone(),
8276 crate::gpu::BatchJob {
8277 idx: *idx,
8278 rows: *rows,
8279 cols: *cols,
8280 row_scale: &[],
8281 xs: x.to_vec(),
8282 layout: if *dt == TensorDtype::Q4Tiled {
8283 crate::gpu::BatchLayout::Q4t
8284 } else {
8285 crate::gpu::BatchLayout::Q4tp
8286 },
8287 },
8288 )),
8289 _ => None,
8290 }
8291}
8292
8293thread_local! {
8294 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8295 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8296}
8297
8298pub(crate) fn prescale<'a>(
8299 x: &'a [f32],
8300 col_field: &[f32],
8301 dtype: TensorDtype,
8302) -> std::borrow::Cow<'a, [f32]> {
8303 if dtype == TensorDtype::Q8_2f {
8304 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8305 } else {
8306 std::borrow::Cow::Borrowed(x)
8307 }
8308}
8309
8310pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8313 x: &[f32],
8314 col_field: &[f32],
8315 dtype: TensorDtype,
8316 buf_id: u8,
8317 f: F,
8318) -> R {
8319 if dtype == TensorDtype::Q8_2f {
8320 if buf_id == 1 {
8321 PRESCALE_BUF1.with(|b| {
8322 let mut buf = b.borrow_mut();
8323 buf.clear();
8324 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8325 f(&buf)
8326 })
8327 } else {
8328 PRESCALE_BUF2.with(|b| {
8329 let mut buf = b.borrow_mut();
8330 buf.clear();
8331 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8332 f(&buf)
8333 })
8334 }
8335 } else {
8336 f(x)
8337 }
8338}
8339
8340#[cfg(target_arch = "x86_64")]
8345pub(crate) fn avx2_enabled() -> bool {
8346 use std::sync::OnceLock;
8347 static ON: OnceLock<bool> = OnceLock::new();
8348 *ON.get_or_init(|| {
8349 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8350 && std::arch::is_x86_feature_detected!("avx2")
8351 && std::arch::is_x86_feature_detected!("fma")
8352 })
8353}
8354
8355#[cfg(target_arch = "x86_64")]
8360fn avx2_a8w8_enabled() -> bool {
8361 use std::sync::OnceLock;
8362 static ON: OnceLock<bool> = OnceLock::new();
8363 *ON.get_or_init(|| {
8364 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8365 })
8366}
8367
8368#[inline]
8372pub(crate) fn a8w8_enabled() -> bool {
8373 #[cfg(target_arch = "aarch64")]
8374 {
8375 sdot_enabled()
8376 }
8377 #[cfg(target_arch = "x86_64")]
8378 {
8379 avx2_a8w8_enabled()
8380 }
8381 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8382 {
8383 false
8384 }
8385}
8386
8387#[inline]
8390#[allow(unreachable_code)]
8391fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8392 #[cfg(target_arch = "aarch64")]
8393 unsafe {
8394 return dot_i8_sdot(w, xq);
8395 }
8396 #[cfg(target_arch = "x86_64")]
8397 unsafe {
8398 if avx512vnni_enabled() {
8399 return dot_i8_i8_vnni(w, xq);
8400 }
8401 return dot_i8_i8_avx2(w, xq);
8402 }
8403 w.iter()
8404 .zip(xq)
8405 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8406 .sum()
8407}
8408
8409#[cfg(target_arch = "x86_64")]
8413fn avx512vnni_enabled() -> bool {
8414 use std::sync::OnceLock;
8415 static ON: OnceLock<bool> = OnceLock::new();
8416 *ON.get_or_init(|| {
8417 std::env::var("CMF_AVX512")
8418 .map(|v| v != "0")
8419 .unwrap_or(true)
8420 && std::arch::is_x86_feature_detected!("avx512f")
8421 && std::arch::is_x86_feature_detected!("avx512bw")
8422 && std::arch::is_x86_feature_detected!("avx512vl")
8423 && std::arch::is_x86_feature_detected!("avx512vnni")
8424 })
8425}
8426
8427#[cfg(target_arch = "x86_64")]
8435fn vnni_tiles_enabled() -> bool {
8436 use std::sync::OnceLock;
8437 static ON: OnceLock<bool> = OnceLock::new();
8438 *ON.get_or_init(|| {
8439 std::env::var("CMF_VNNI_TILES")
8440 .map(|v| v != "0")
8441 .unwrap_or(true)
8442 && avx512vnni_enabled()
8443 })
8444}
8445
8446#[cfg(target_arch = "x86_64")]
8451#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8452#[inline]
8453unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8454 unsafe {
8456 use core::arch::x86_64::*;
8457 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8458 let hi128 = _mm256_extracti128_si256::<1>(d);
8459 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8460 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8461 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8462 _mm_cvtsi128_si32(s32)
8463 }
8464}
8465
8466#[cfg(target_arch = "x86_64")]
8471#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8472unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8473 unsafe {
8475 use core::arch::x86_64::*;
8476 let n = w.len();
8477 let mut j = 0usize;
8478 let mut total: i32;
8479 {
8484 #[inline(always)]
8485 unsafe fn step(
8486 w: *const u8,
8487 x: *const i8,
8488 acc: core::arch::x86_64::__m512i,
8489 ) -> core::arch::x86_64::__m512i {
8490 unsafe {
8491 use core::arch::x86_64::*;
8492 let wv = _mm512_loadu_si512(w as *const _);
8493 let xv = _mm512_loadu_si512(x as *const _);
8494 let aw = _mm512_abs_epi8(wv);
8495 let neg = _mm512_movepi8_mask(wv);
8496 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8497 _mm512_dpbusd_epi32(acc, aw, sx)
8498 }
8499 }
8500 let (mut a0, mut a1, mut a2, mut a3) = (
8501 _mm512_setzero_si512(),
8502 _mm512_setzero_si512(),
8503 _mm512_setzero_si512(),
8504 _mm512_setzero_si512(),
8505 );
8506 while j + 256 <= n {
8507 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8508 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8509 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8510 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8511 j += 256;
8512 }
8513 while j + 64 <= n {
8514 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8515 j += 64;
8516 }
8517 let s01 = _mm512_add_epi32(a0, a1);
8518 let s23 = _mm512_add_epi32(a2, a3);
8519 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8520 }
8521 if j + 32 <= n {
8523 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8524 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8525 let d = _mm256_dpbusd_epi32(
8526 _mm256_setzero_si256(),
8527 _mm256_abs_epi8(wv),
8528 _mm256_sign_epi8(xv, wv),
8529 );
8530 let hi128 = _mm256_extracti128_si256::<1>(d);
8531 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8532 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8533 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8534 total += _mm_cvtsi128_si32(s32);
8535 j += 32;
8536 }
8537 while j < n {
8538 total += (w[j] as i8) as i32 * xq[j] as i32;
8539 j += 1;
8540 }
8541 total
8542 }
8543}
8544
8545#[cfg(target_arch = "x86_64")]
8547#[target_feature(enable = "avx2,fma")]
8548unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8549 unsafe {
8551 use core::arch::x86_64::*;
8552 let n = x.len();
8553 let wp = w.as_ptr();
8554 let xp = x.as_ptr();
8555 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8556 let mut j = 0usize;
8557 while j + 16 <= n {
8558 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8559 let lo = _mm256_cvtepi8_epi32(wb);
8560 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8561 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8562 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8563 j += 16;
8564 }
8565 let acc = _mm256_add_ps(a0, a1);
8566 let hi128 = _mm256_extractf128_ps::<1>(acc);
8567 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8568 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8569 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8570 let mut sum = _mm_cvtss_f32(s32);
8571 while j < n {
8572 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8573 j += 1;
8574 }
8575 sum
8576 }
8577}
8578
8579#[cfg(target_arch = "x86_64")]
8584#[target_feature(enable = "avx2")]
8585unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8586 unsafe {
8588 use core::arch::x86_64::*;
8589 let n = w.len();
8590 let ones = _mm256_set1_epi16(1);
8591 let mut acc = _mm256_setzero_si256();
8592 let mut j = 0usize;
8593 while j + 32 <= n {
8594 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8595 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8596 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8597 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8598 j += 32;
8599 }
8600 let hi128 = _mm256_extracti128_si256::<1>(acc);
8601 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8602 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8603 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8604 let mut s = _mm_cvtsi128_si32(s32);
8605 while j < n {
8606 s += (w[j] as i8) as i32 * xq[j] as i32;
8607 j += 1;
8608 }
8609 s
8610 }
8611}
8612
8613#[cfg(target_arch = "aarch64")]
8617#[target_feature(enable = "neon,i8mm")]
8618unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8619 unsafe {
8621 use core::arch::aarch64::*;
8622 use core::arch::asm;
8623 let n = w0.len();
8624 let w0p = w0.as_ptr() as *const i8;
8625 let w1p = w1.as_ptr() as *const i8;
8626 let mut acc01 = vdupq_n_s32(0);
8629 let mut acc23 = vdupq_n_s32(0);
8630 let mut i = 0usize;
8631 while i + 8 <= n {
8632 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8633 let xb01 = vcombine_s8(
8634 vld1_s8(xs[0].as_ptr().add(i)),
8635 vld1_s8(xs[1].as_ptr().add(i)),
8636 );
8637 let xb23 = vcombine_s8(
8638 vld1_s8(xs[2].as_ptr().add(i)),
8639 vld1_s8(xs[3].as_ptr().add(i)),
8640 );
8641 asm!(
8642 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8643 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8644 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8645 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8646 options(pure, nomem, nostack),
8647 );
8648 i += 8;
8649 }
8650 let mut out = [[0i32; 4]; 2];
8651 let a01: [i32; 4] = core::mem::transmute(acc01);
8652 let a23: [i32; 4] = core::mem::transmute(acc23);
8653 out[0][0] = a01[0];
8654 out[0][1] = a01[1];
8655 out[1][0] = a01[2];
8656 out[1][1] = a01[3];
8657 out[0][2] = a23[0];
8658 out[0][3] = a23[1];
8659 out[1][2] = a23[2];
8660 out[1][3] = a23[3];
8661 if i < n {
8662 for (k, x) in xs.iter().enumerate() {
8663 for j in i..n {
8664 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8665 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8666 }
8667 }
8668 }
8669 out
8670 }
8671}
8672
8673#[cfg(target_arch = "aarch64")]
8677#[target_feature(enable = "neon,dotprod")]
8678unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8679 unsafe {
8681 use core::arch::aarch64::*;
8682 use core::arch::asm;
8683 let n = w0.len();
8684 let w0p = w0.as_ptr() as *const i8;
8685 let w1p = w1.as_ptr() as *const i8;
8686 let mut acc = [[vdupq_n_s32(0); 4]; 2];
8687 let mut i = 0usize;
8688 while i + 16 <= n {
8689 let wv0 = vld1q_s8(w0p.add(i));
8690 let wv1 = vld1q_s8(w1p.add(i));
8691 for (k, x) in xs.iter().enumerate() {
8692 let xv = vld1q_s8(x.as_ptr().add(i));
8693 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8694 asm!(
8695 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8696 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8697 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8698 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8699 options(pure, nomem, nostack),
8700 );
8701 acc[0][k] = a0;
8702 acc[1][k] = a1;
8703 }
8704 i += 16;
8705 }
8706 let mut out = [[0i32; 4]; 2];
8707 for r in 0..2 {
8708 for k in 0..4 {
8709 out[r][k] = vaddvq_s32(acc[r][k]);
8710 }
8711 }
8712 if i < n {
8713 for (k, x) in xs.iter().enumerate() {
8714 for j in i..n {
8715 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8716 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8717 }
8718 }
8719 }
8720 out
8721 }
8722}
8723
8724#[cfg(target_arch = "x86_64")]
8730#[target_feature(enable = "avx2")]
8731unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8732 unsafe {
8734 use core::arch::x86_64::*;
8735 let n = w0.len();
8736 let ones = _mm256_set1_epi16(1);
8737 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8738 let mut j = 0usize;
8739 while j + 32 <= n {
8740 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8741 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8742 let aw0 = _mm256_abs_epi8(wv0);
8743 let aw1 = _mm256_abs_epi8(wv1);
8744 for (k, x) in xs.iter().enumerate() {
8745 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8746 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8747 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8748 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8749 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8750 }
8751 j += 32;
8752 }
8753 let mut out = [[0i32; 4]; 2];
8754 for r in 0..2 {
8755 for k in 0..4 {
8756 let a = acc[r][k];
8757 let hi128 = _mm256_extracti128_si256::<1>(a);
8758 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8759 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8760 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8761 out[r][k] = _mm_cvtsi128_si32(s32);
8762 }
8763 }
8764 if j < n {
8765 for (k, x) in xs.iter().enumerate() {
8766 for i in j..n {
8767 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8768 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8769 }
8770 }
8771 }
8772 out
8773 }
8774}
8775
8776#[cfg(target_arch = "x86_64")]
8781#[inline]
8782fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8783 let dot = if avx512vnni_enabled() && row.len() >= 64 {
8784 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8785 } else {
8786 unsafe { dot_i8_i8_avx2(row, &act.xq) }
8787 };
8788 let mut acc = dot as f32 * act.sx;
8789 for &(j, xv) in &act.outliers {
8790 acc += (row[j] as i8) as f32 * xv;
8791 }
8792 acc
8793}
8794
8795#[cfg(target_arch = "x86_64")]
8799#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8800unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8801 unsafe {
8803 use core::arch::x86_64::*;
8804 let n = w.len();
8805 let flip = _mm512_set1_epi8(-128); #[inline(always)]
8807 unsafe fn step(
8808 w: *const u8,
8809 x: *const i8,
8810 flip: core::arch::x86_64::__m512i,
8811 acc: core::arch::x86_64::__m512i,
8812 ) -> core::arch::x86_64::__m512i {
8813 unsafe {
8814 use core::arch::x86_64::*;
8815 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8816 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8817 }
8818 }
8819 let (mut a0, mut a1, mut a2, mut a3) = (
8820 _mm512_setzero_si512(),
8821 _mm512_setzero_si512(),
8822 _mm512_setzero_si512(),
8823 _mm512_setzero_si512(),
8824 );
8825 let mut j = 0usize;
8826 while j + 256 <= n {
8827 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8828 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8829 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8830 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8831 j += 256;
8832 }
8833 while j + 64 <= n {
8834 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8835 j += 64;
8836 }
8837 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8838 _mm512_add_epi32(a0, a1),
8839 _mm512_add_epi32(a2, a3),
8840 ));
8841 while j < n {
8843 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8844 j += 1;
8845 }
8846 total
8847 }
8848}
8849
8850#[cfg(target_arch = "x86_64")]
8856#[target_feature(enable = "avx2")]
8857unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8858 unsafe {
8861 use core::arch::x86_64::*;
8862 let lomask = _mm_set1_epi8(0x0F);
8863 let eight = _mm256_set1_epi8(8);
8864 let ones = _mm256_set1_epi16(1);
8865 let mut acc = 0f32;
8866 for gi in 0..gpr {
8867 let g = g0 + gi;
8868 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8869 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8870 let lo = _mm_and_si128(b, lomask);
8871 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8872 let w = _mm256_sub_epi8(
8873 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8874 eight,
8875 );
8876 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8877 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8878 let d = _mm256_madd_epi16(p16, ones);
8879 let hi128 = _mm256_extracti128_si256::<1>(d);
8880 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8881 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8882 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8883 acc += _mm_cvtsi128_si32(s32) as f32 * s;
8884 }
8885 acc
8886 }
8887}
8888
8889#[cfg(target_arch = "x86_64")]
8892#[target_feature(enable = "avx2")]
8893unsafe fn dot_q4_row_avx2_2(
8894 packed: &[u8],
8895 scales: &[u8],
8896 g0: usize,
8897 gpr: usize,
8898 xq1: &[i8],
8899 xq2: &[i8],
8900) -> (f32, f32) {
8901 unsafe {
8903 use core::arch::x86_64::*;
8904 let lomask = _mm_set1_epi8(0x0F);
8905 let eight = _mm256_set1_epi8(8);
8906 let ones = _mm256_set1_epi16(1);
8907 let (mut acc1, mut acc2) = (0f32, 0f32);
8908 #[inline(always)]
8909 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8910 unsafe {
8911 use core::arch::x86_64::*;
8912 let hi128 = _mm256_extracti128_si256::<1>(d);
8913 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8914 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8915 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8916 _mm_cvtsi128_si32(s32)
8917 }
8918 }
8919 for gi in 0..gpr {
8920 let g = g0 + gi;
8921 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8922 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8923 let lo = _mm_and_si128(b, lomask);
8924 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8925 let w = _mm256_sub_epi8(
8926 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8927 eight,
8928 );
8929 let aw = _mm256_abs_epi8(w);
8930 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8931 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8932 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8933 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8934 acc1 += hsum(d1) as f32 * s;
8935 acc2 += hsum(d2) as f32 * s;
8936 }
8937 (acc1, acc2)
8938 }
8939}
8940
8941#[cfg(target_arch = "x86_64")]
8943fn q8_range_avx2(
8944 q: &[u8],
8945 row_scale: &[f32],
8946 act: &SplitAct,
8947 cols: usize,
8948 out_addr: SendMut,
8949 start: usize,
8950 end: usize,
8951) {
8952 for o in start..end {
8953 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8954 unsafe { *out_addr.at(o) = v };
8956 }
8957}
8958
8959#[cfg(target_arch = "x86_64")]
8961#[allow(clippy::too_many_arguments)]
8962fn q8_range2_avx2(
8963 q: &[u8],
8964 row_scale: &[f32],
8965 a1: &SplitAct,
8966 a2: &SplitAct,
8967 cols: usize,
8968 p1: SendMut,
8969 p2: SendMut,
8970 start: usize,
8971 end: usize,
8972) {
8973 for o in start..end {
8974 let row = &q[o * cols..(o + 1) * cols];
8975 unsafe {
8977 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8978 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8979 }
8980 }
8981}
8982
8983#[cfg(target_arch = "aarch64")]
8994fn i8mm_enabled() -> bool {
8995 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8996 *ON.get_or_init(|| {
8997 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8998 && std::arch::is_aarch64_feature_detected!("i8mm")
8999 })
9000}
9001
9002#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9006fn sdot_enabled() -> bool {
9007 use std::sync::OnceLock;
9008 static ON: OnceLock<bool> = OnceLock::new();
9009 *ON.get_or_init(|| {
9010 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9011 if !want {
9012 return false;
9013 }
9014
9015 #[cfg(target_arch = "aarch64")]
9016 {
9017 if std::arch::is_aarch64_feature_detected!("dotprod") {
9018 return true;
9019 }
9020 #[cfg(target_os = "android")]
9021 {
9022 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9023 if cpuinfo.lines().any(|l| {
9024 (l.starts_with("Features") || l.starts_with("features"))
9025 && l.contains("asimddp")
9026 }) {
9027 return true;
9028 }
9029 }
9030 }
9031 false
9032 }
9033 #[cfg(not(target_arch = "aarch64"))]
9034 {
9035 false
9036 }
9037 })
9038}
9039
9040struct SplitAct {
9045 xq: Vec<i8>,
9046 sx: f32,
9047 outliers: Vec<(usize, f32)>,
9048 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9051 xsum: i32,
9052}
9053
9054thread_local! {
9055 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9058 const { std::cell::RefCell::new(Vec::new()) };
9059}
9060
9061impl Drop for SplitAct {
9062 fn drop(&mut self) {
9063 let buf = std::mem::take(&mut self.xq);
9064 if buf.capacity() > 0 {
9065 XQ_FREE.with(|f| {
9066 let mut f = f.borrow_mut();
9067 if f.len() < 16 {
9068 f.push(buf);
9069 }
9070 });
9071 }
9072 }
9073}
9074
9075thread_local! {
9076 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9083}
9084
9085#[inline]
9088fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9089 KROW.with(|s| {
9090 let mut b = s.borrow_mut();
9091 if b.len() < n {
9092 b.resize(n, 0.0);
9093 }
9094 f(&mut b[..n])
9095 })
9096}
9097
9098fn split_act(x: &[f32]) -> SplitAct {
9099 let n = x.len();
9100 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9101 let thr = 8.0 * rms;
9102 let mut outliers: Vec<(usize, f32)> = Vec::new();
9106 let mut amax = 0f32;
9107 for (j, &v) in x.iter().enumerate() {
9108 let a = v.abs();
9109 if a > thr {
9110 outliers.push((j, v));
9111 } else if a > amax {
9112 amax = a;
9113 }
9114 }
9115 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9116 let inv = 1.0 / sx;
9117 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9118 xq.clear();
9119 xq.reserve(n);
9120 if outliers.is_empty() {
9121 xq.extend(
9122 x.iter()
9123 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9124 );
9125 } else {
9126 xq.extend(x.iter().map(|&v| {
9128 if v.abs() > thr {
9129 0
9130 } else {
9131 (v * inv).round().clamp(-127.0, 127.0) as i8
9132 }
9133 }));
9134 }
9135 let xsum = xq.iter().map(|&v| v as i32).sum();
9136 SplitAct {
9137 xq,
9138 sx,
9139 outliers,
9140 xsum,
9141 }
9142}
9143
9144fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9145 let n = x.len();
9146 let rms = (x
9147 .iter()
9148 .zip(col)
9149 .map(|(&a, &c)| {
9150 let v = a * c;
9151 (v * v) as f64
9152 })
9153 .sum::<f64>()
9154 / n.max(1) as f64)
9155 .sqrt() as f32;
9156 let thr = 8.0 * rms;
9157
9158 let mut outliers = Vec::new();
9159 let mut amax = 0f32;
9160 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9161 let v = a * c;
9162 let s = v.abs();
9163 if s > thr {
9164 outliers.push((j, v));
9165 } else if s > amax {
9166 amax = s;
9167 }
9168 }
9169
9170 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9171 let inv = 1.0 / sx;
9172 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9173 xq.clear();
9174 xq.reserve(n);
9175 if outliers.is_empty() {
9176 xq.extend(
9177 x.iter()
9178 .zip(col)
9179 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9180 );
9181 } else {
9182 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9183 let v = a * c;
9184 if v.abs() > thr {
9185 0
9186 } else {
9187 (v * inv).round().clamp(-127.0, 127.0) as i8
9188 }
9189 }));
9190 }
9191 let xsum = xq.iter().map(|&v| v as i32).sum();
9192 SplitAct {
9193 xq,
9194 sx,
9195 outliers,
9196 xsum,
9197 }
9198}
9199
9200#[cfg(target_arch = "aarch64")]
9203#[target_feature(enable = "neon,dotprod")]
9204unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9205 unsafe {
9207 use core::arch::aarch64::*;
9208 use core::arch::asm;
9209 let wp = w.as_ptr() as *const i8;
9210 let n = w.len();
9211 let (mut a0, mut a1, mut a2, mut a3) = (
9212 vdupq_n_s32(0),
9213 vdupq_n_s32(0),
9214 vdupq_n_s32(0),
9215 vdupq_n_s32(0),
9216 );
9217 let mut i = 0;
9218 while i + 64 <= n {
9219 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9220 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9221 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9222 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9223 asm!(
9224 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9225 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9226 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9227 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9228 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9229 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9230 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9231 options(pure, nomem, nostack),
9232 );
9233 i += 64;
9234 }
9235 while i + 16 <= n {
9236 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9237 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9238 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9239 i += 16;
9240 }
9241 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9242 while i < n {
9243 s += (*wp.add(i)) as i32 * xq[i] as i32;
9244 i += 1;
9245 }
9246 s
9247 }
9248}
9249
9250#[cfg(target_arch = "aarch64")]
9254#[target_feature(enable = "neon,dotprod")]
9255unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9256 unsafe {
9258 use core::arch::aarch64::*;
9259 use core::arch::asm;
9260 let n = xq.len();
9261 let px = xq.as_ptr();
9262 let (p0, p1, p2, p3) = (
9263 w0.as_ptr() as *const i8,
9264 w1.as_ptr() as *const i8,
9265 w2.as_ptr() as *const i8,
9266 w3.as_ptr() as *const i8,
9267 );
9268 let (mut a0, mut a1, mut a2, mut a3) = (
9269 vdupq_n_s32(0),
9270 vdupq_n_s32(0),
9271 vdupq_n_s32(0),
9272 vdupq_n_s32(0),
9273 );
9274 let mut i = 0;
9275 while i + 16 <= n {
9276 let x = vld1q_s8(px.add(i));
9277 let v0 = vld1q_s8(p0.add(i));
9278 let v1 = vld1q_s8(p1.add(i));
9279 let v2 = vld1q_s8(p2.add(i));
9280 let v3 = vld1q_s8(p3.add(i));
9281 asm!(
9282 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9283 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9284 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9285 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9286 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9287 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9288 options(pure, nomem, nostack),
9289 );
9290 i += 16;
9291 }
9292 let mut r = [
9293 vaddvq_s32(a0),
9294 vaddvq_s32(a1),
9295 vaddvq_s32(a2),
9296 vaddvq_s32(a3),
9297 ];
9298 while i < n {
9299 let xi = *px.add(i) as i32;
9300 r[0] += (*p0.add(i)) as i32 * xi;
9301 r[1] += (*p1.add(i)) as i32 * xi;
9302 r[2] += (*p2.add(i)) as i32 * xi;
9303 r[3] += (*p3.add(i)) as i32 * xi;
9304 i += 1;
9305 }
9306 r
9307 }
9308}
9309
9310#[cfg(target_arch = "aarch64")]
9317#[target_feature(enable = "neon,dotprod")]
9318unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9319 unsafe {
9322 use core::arch::aarch64::*;
9323 use core::arch::asm;
9324 let n = xq.len();
9325 let px = xq.as_ptr();
9326 let pg = g.as_ptr() as *const i8;
9327 let (mut a0, mut a1, mut a2, mut a3) = (
9328 vdupq_n_s32(0),
9329 vdupq_n_s32(0),
9330 vdupq_n_s32(0),
9331 vdupq_n_s32(0),
9332 );
9333 let mut i = 0;
9334 while i + 16 <= n {
9335 let x = vld1q_s8(px.add(i));
9336 let base = pg.add(4 * i);
9337 let v0 = vld1q_s8(base);
9338 let v1 = vld1q_s8(base.add(16));
9339 let v2 = vld1q_s8(base.add(32));
9340 let v3 = vld1q_s8(base.add(48));
9341 asm!(
9342 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9343 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9344 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9345 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9346 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9347 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9348 options(pure, nomem, nostack),
9349 );
9350 i += 16;
9351 }
9352 [
9353 vaddvq_s32(a0),
9354 vaddvq_s32(a1),
9355 vaddvq_s32(a2),
9356 vaddvq_s32(a3),
9357 ]
9358 }
9359}
9360
9361#[cfg(target_arch = "aarch64")]
9367fn q8_range_sdot(
9368 q: &[u8],
9369 rep: &[u8],
9370 row_scale: &[f32],
9371 act: &SplitAct,
9372 cols: usize,
9373 out_addr: SendMut,
9374 start: usize,
9375 end: usize,
9376) {
9377 let mut o = start;
9378 if !rep.is_empty() {
9381 while o < end && o % 4 != 0 {
9382 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9383 unsafe { *out_addr.at(o) = v };
9384 o += 1;
9385 }
9386 }
9387 while o + 4 <= end {
9388 let r = if rep.is_empty() {
9389 unsafe {
9390 dot_i8_sdot_4rows(
9391 &q[o * cols..(o + 1) * cols],
9392 &q[(o + 1) * cols..(o + 2) * cols],
9393 &q[(o + 2) * cols..(o + 3) * cols],
9394 &q[(o + 3) * cols..(o + 4) * cols],
9395 &act.xq,
9396 )
9397 }
9398 } else {
9399 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9400 };
9401 for k in 0..4 {
9402 let mut acc = r[k] as f32 * act.sx;
9403 for &(j, xv) in &act.outliers {
9404 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9405 }
9406 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9408 }
9409 o += 4;
9410 }
9411 while o < end {
9412 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9413 unsafe { *out_addr.at(o) = v };
9414 o += 1;
9415 }
9416}
9417
9418#[cfg(target_arch = "aarch64")]
9421#[allow(clippy::too_many_arguments)]
9422fn q8_range2_sdot(
9423 q: &[u8],
9424 row_scale: &[f32],
9425 a1: &SplitAct,
9426 a2: &SplitAct,
9427 cols: usize,
9428 p1: SendMut,
9429 p2: SendMut,
9430 start: usize,
9431 end: usize,
9432) {
9433 for o in start..end {
9434 let row = &q[o * cols..(o + 1) * cols];
9435 unsafe {
9437 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9438 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9439 }
9440 }
9441}
9442
9443#[allow(clippy::too_many_arguments)]
9445fn q8_range2_f32(
9446 q: &[u8],
9447 row_scale: &[f32],
9448 x1: &[f32],
9449 x2: &[f32],
9450 cols: usize,
9451 p1: SendMut,
9452 p2: SendMut,
9453 start: usize,
9454 end: usize,
9455) {
9456 for o in start..end {
9457 let row = &q[o * cols..(o + 1) * cols];
9458 unsafe {
9460 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9461 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9462 }
9463 }
9464}
9465
9466fn q8_range_f32(
9468 q: &[u8],
9469 row_scale: &[f32],
9470 xs: &[f32],
9471 cols: usize,
9472 out_addr: SendMut,
9473 start: usize,
9474 end: usize,
9475) {
9476 for o in start..end {
9477 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9478 unsafe { *out_addr.at(o) = v };
9480 }
9481}
9482
9483#[inline]
9487fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9488 #[cfg(target_arch = "aarch64")]
9489 return row_dot_sdot(row, act);
9490 #[cfg(target_arch = "x86_64")]
9491 return row_dot_avx2(row, act);
9492 #[allow(unreachable_code)]
9493 q8_row_dot_scalar(row, act)
9494}
9495
9496#[allow(dead_code)]
9497fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9498 let mut acc = 0i32;
9499 for (k, &b) in row.iter().enumerate() {
9500 acc += (b as i8) as i32 * act.xq[k] as i32;
9501 }
9502 let mut acc = acc as f32 * act.sx;
9503 for &(j, xv) in &act.outliers {
9504 acc += (row[j] as i8) as f32 * xv;
9505 }
9506 acc
9507}
9508
9509#[cfg(target_arch = "aarch64")]
9512#[inline]
9513fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9514 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9515 for &(j, xv) in &act.outliers {
9516 acc += (row[j] as i8) as f32 * xv;
9517 }
9518 acc
9519}
9520
9521#[cfg(target_arch = "aarch64")]
9529#[target_feature(enable = "neon,dotprod")]
9530unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9531 unsafe {
9534 use core::arch::aarch64::*;
9535 use core::arch::asm;
9536 let lomask = vdupq_n_u8(0x0F);
9537 let eight = vdupq_n_s8(8);
9538 let mut acc = 0f32;
9539 for gi in 0..gpr {
9540 let g = g0 + gi;
9541 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9542 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9543 let lo = vandq_u8(b, lomask);
9544 let hi = vshrq_n_u8::<4>(b);
9545 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9546 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9547 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9548 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9549 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9550 asm!(
9551 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9552 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9553 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9554 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9555 options(pure, nomem, nostack),
9556 );
9557 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9558 }
9559 acc
9560 }
9561}
9562
9563#[cfg(target_arch = "aarch64")]
9568#[target_feature(enable = "neon,dotprod")]
9569unsafe fn dot_q4_row_sdot2(
9570 packed: &[u8],
9571 scales: &[u8],
9572 g0: usize,
9573 gpr: usize,
9574 xq1: &[i8],
9575 xq2: &[i8],
9576) -> (f32, f32) {
9577 unsafe {
9580 use core::arch::aarch64::*;
9581 use core::arch::asm;
9582 let lomask = vdupq_n_u8(0x0F);
9583 let eight = vdupq_n_s8(8);
9584 let (mut acc1, mut acc2) = (0f32, 0f32);
9585 for gi in 0..gpr {
9586 let g = g0 + gi;
9587 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9588 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9589 let lo = vandq_u8(b, lomask);
9590 let hi = vshrq_n_u8::<4>(b);
9591 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9592 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9593 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9594 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9595 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9596 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9597 let (mut a0, mut a1, mut b0, mut b1) = (
9598 vdupq_n_s32(0),
9599 vdupq_n_s32(0),
9600 vdupq_n_s32(0),
9601 vdupq_n_s32(0),
9602 );
9603 asm!(
9604 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9605 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9606 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9607 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9608 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9609 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9610 e0 = in(vreg) e0, e1 = in(vreg) e1,
9611 x10 = in(vreg) x10, x11 = in(vreg) x11,
9612 x20 = in(vreg) x20, x21 = in(vreg) x21,
9613 options(pure, nomem, nostack),
9614 );
9615 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9616 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9617 }
9618 (acc1, acc2)
9619 }
9620}
9621
9622#[inline]
9627pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9628 #[cfg(target_arch = "aarch64")]
9629 unsafe {
9630 return axpy_i8_f32_neon(acc, row, w);
9631 }
9632 #[cfg(target_arch = "x86_64")]
9633 if avx2_enabled() {
9634 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9635 }
9636 #[allow(unreachable_code)]
9637 {
9638 for (a, &b) in acc.iter_mut().zip(row) {
9639 *a += w * b as f32;
9640 }
9641 }
9642}
9643
9644#[cfg(target_arch = "x86_64")]
9646#[target_feature(enable = "avx2,fma")]
9647unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9648 unsafe {
9650 use core::arch::x86_64::*;
9651 let n = acc.len().min(row.len());
9652 let ap = acc.as_mut_ptr();
9653 let rp = row.as_ptr();
9654 let wv = _mm256_set1_ps(w);
9655 let mut j = 0usize;
9656 while j + 16 <= n {
9657 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9658 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9659 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9660 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9661 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9662 _mm256_storeu_ps(ap.add(j), v0);
9663 _mm256_storeu_ps(ap.add(j + 8), v1);
9664 j += 16;
9665 }
9666 while j < n {
9667 *ap.add(j) += w * (*rp.add(j)) as f32;
9668 j += 1;
9669 }
9670 }
9671}
9672
9673#[cfg(target_arch = "aarch64")]
9674#[target_feature(enable = "neon")]
9675unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9676 unsafe {
9678 use core::arch::aarch64::*;
9679 let n = acc.len().min(row.len());
9680 let ap = acc.as_mut_ptr();
9681 let rp = row.as_ptr();
9682 let wv = vdupq_n_f32(w);
9683 let mut j = 0usize;
9684 while j + 16 <= n {
9685 let rb = vld1q_s8(rp.add(j));
9686 let lo = vmovl_s8(vget_low_s8(rb));
9687 let hi = vmovl_s8(vget_high_s8(rb));
9688 for (off, half) in [(0, lo), (8, hi)] {
9689 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9690 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9691 let o = j + off;
9692 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9693 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9694 }
9695 j += 16;
9696 }
9697 while j < n {
9698 *ap.add(j) += w * (*rp.add(j)) as f32;
9699 j += 1;
9700 }
9701 }
9702}
9703
9704#[inline]
9707pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9708 #[cfg(target_arch = "aarch64")]
9709 unsafe {
9710 return dot_i8_f32_neon(w, x);
9711 }
9712 #[cfg(target_arch = "x86_64")]
9713 if avx2_enabled() {
9714 return unsafe { dot_i8_f32_avx2(w, x) };
9715 }
9716 #[allow(unreachable_code)]
9717 {
9718 let mut sum = 0.0f32;
9719 for (j, &b) in w.iter().enumerate() {
9720 sum += (b as i8) as f32 * x[j];
9721 }
9722 sum
9723 }
9724}
9725
9726#[inline]
9730fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9731 #[cfg(target_arch = "aarch64")]
9732 unsafe {
9733 return dot_i8_col_f32_neon(w, x, col);
9734 }
9735 #[allow(unreachable_code)]
9736 {
9737 let mut sum = 0.0f32;
9738 for (j, &b) in w.iter().enumerate() {
9739 sum += (b as i8) as f32 * x[j] * col[j];
9740 }
9741 sum
9742 }
9743}
9744
9745#[cfg(target_arch = "aarch64")]
9746#[target_feature(enable = "neon")]
9747unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9748 unsafe {
9750 use core::arch::aarch64::*;
9751 let n = x.len();
9752 let wp = w.as_ptr() as *const i8;
9753 let xp = x.as_ptr();
9754 let cp = col.as_ptr();
9755 let (mut a0, mut a1, mut a2, mut a3) = (
9756 vdupq_n_f32(0.0),
9757 vdupq_n_f32(0.0),
9758 vdupq_n_f32(0.0),
9759 vdupq_n_f32(0.0),
9760 );
9761 let mut j = 0usize;
9762 while j + 16 <= n {
9763 let wb = vld1q_s8(wp.add(j));
9764 let lo = vmovl_s8(vget_low_s8(wb));
9765 let hi = vmovl_s8(vget_high_s8(wb));
9766 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9767 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9768 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9769 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9770 a0 = vfmaq_f32(
9771 a0,
9772 w0,
9773 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9774 );
9775 a1 = vfmaq_f32(
9776 a1,
9777 w1,
9778 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9779 );
9780 a2 = vfmaq_f32(
9781 a2,
9782 w2,
9783 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9784 );
9785 a3 = vfmaq_f32(
9786 a3,
9787 w3,
9788 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9789 );
9790 j += 16;
9791 }
9792 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9793 while j < n {
9794 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9795 j += 1;
9796 }
9797 sum
9798 }
9799}
9800
9801#[cfg(target_arch = "aarch64")]
9802#[target_feature(enable = "neon")]
9803unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9804 unsafe {
9806 use core::arch::aarch64::*;
9807 let n = x.len();
9808 let wp = w.as_ptr() as *const i8;
9809 let xp = x.as_ptr();
9810 let (mut a0, mut a1, mut a2, mut a3) = (
9811 vdupq_n_f32(0.0),
9812 vdupq_n_f32(0.0),
9813 vdupq_n_f32(0.0),
9814 vdupq_n_f32(0.0),
9815 );
9816 let mut j = 0usize;
9817 while j + 16 <= n {
9818 let wb = vld1q_s8(wp.add(j));
9819 let lo = vmovl_s8(vget_low_s8(wb));
9820 let hi = vmovl_s8(vget_high_s8(wb));
9821 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9822 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9823 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9824 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9825 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9826 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9827 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9828 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9829 j += 16;
9830 }
9831 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9832 while j < n {
9833 sum += (*wp.add(j)) as f32 * *xp.add(j);
9834 j += 1;
9835 }
9836 sum
9837 }
9838}
9839
9840#[allow(clippy::too_many_arguments)]
9841fn qmatvec(
9842 q: &[u8],
9843 rep: &[u8],
9844 row_scale: &[f32],
9845 x: &[f32],
9846 col_field: &[f32],
9847 dtype: TensorDtype,
9848 rows: usize,
9849 cols: usize,
9850 out: &mut [f32],
9851 pool: Option<&Pool>,
9852) {
9853 debug_assert_eq!(out.len(), rows);
9854 #[cfg(not(target_arch = "aarch64"))]
9855 let _ = rep;
9856
9857 #[cfg(target_arch = "aarch64")]
9858 if sdot_enabled() {
9859 let act = if dtype == TensorDtype::Q8_2f {
9860 split_act_q8_2f(x, col_field)
9861 } else {
9862 split_act(x)
9863 };
9864 let out_addr = SendMut(out.as_mut_ptr());
9865 let run_range = |start: usize, end: usize| {
9866 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9867 };
9868 match pool {
9869 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9870 _ => run_range(0, rows),
9871 }
9872 return;
9873 }
9874 #[cfg(target_arch = "x86_64")]
9877 if avx2_a8w8_enabled() {
9878 let act = if dtype == TensorDtype::Q8_2f {
9879 split_act_q8_2f(x, col_field)
9880 } else {
9881 split_act(x)
9882 };
9883 let out_addr = SendMut(out.as_mut_ptr());
9884 let run_range = |start: usize, end: usize| {
9885 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9886 };
9887 match pool {
9888 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9889 _ => run_range(0, rows),
9890 }
9891 return;
9892 }
9893
9894 prescale_with(x, col_field, dtype, 1, |xs| {
9895 let out_addr = SendMut(out.as_mut_ptr());
9896 let run_range = move |start: usize, end: usize| {
9897 for o in start..end {
9898 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9899 unsafe { *out_addr.at(o) = v };
9901 }
9902 };
9903 match pool {
9904 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9905 _ => run_range(0, rows),
9906 }
9907 });
9908}
9909
9910#[allow(clippy::too_many_arguments)]
9911fn qmatvec2(
9912 q: &[u8],
9913 row_scale: &[f32],
9914 x1: &[f32],
9915 x2: &[f32],
9916 col_field: &[f32],
9917 dtype: TensorDtype,
9918 rows: usize,
9919 cols: usize,
9920 o1: &mut [f32],
9921 o2: &mut [f32],
9922 pool: Option<&Pool>,
9923) {
9924 #[cfg(target_arch = "aarch64")]
9925 if sdot_enabled() {
9926 let a1s = if dtype == TensorDtype::Q8_2f {
9927 split_act_q8_2f(x1, col_field)
9928 } else {
9929 split_act(x1)
9930 };
9931 let a2s = if dtype == TensorDtype::Q8_2f {
9932 split_act_q8_2f(x2, col_field)
9933 } else {
9934 split_act(x2)
9935 };
9936 let p1 = SendMut(o1.as_mut_ptr());
9937 let p2 = SendMut(o2.as_mut_ptr());
9938 let run_range = |start: usize, end: usize| {
9939 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9940 };
9941 match pool {
9942 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9943 _ => run_range(0, rows),
9944 }
9945 return;
9946 }
9947 #[cfg(target_arch = "x86_64")]
9948 if avx2_a8w8_enabled() {
9949 let a1s = if dtype == TensorDtype::Q8_2f {
9950 split_act_q8_2f(x1, col_field)
9951 } else {
9952 split_act(x1)
9953 };
9954 let a2s = if dtype == TensorDtype::Q8_2f {
9955 split_act_q8_2f(x2, col_field)
9956 } else {
9957 split_act(x2)
9958 };
9959 let p1 = SendMut(o1.as_mut_ptr());
9960 let p2 = SendMut(o2.as_mut_ptr());
9961 let run_range = |start: usize, end: usize| {
9962 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9963 };
9964 match pool {
9965 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9966 _ => run_range(0, rows),
9967 }
9968 return;
9969 }
9970
9971 prescale_with(x1, col_field, dtype, 1, |x1s| {
9972 prescale_with(x2, col_field, dtype, 2, |x2s| {
9973 let p1 = SendMut(o1.as_mut_ptr());
9974 let p2 = SendMut(o2.as_mut_ptr());
9975 let run_range = move |start: usize, end: usize| {
9976 for o in start..end {
9977 let row = &q[o * cols..(o + 1) * cols];
9978 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
9979 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
9980 unsafe {
9982 *p1.at(o) = s1;
9983 *p2.at(o) = s2;
9984 }
9985 }
9986 };
9987 match pool {
9988 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9989 _ => run_range(0, rows),
9990 }
9991 });
9992 });
9993}
9994
9995#[derive(Clone, Copy)]
9996struct SendMut(*mut f32);
9997unsafe impl Send for SendMut {}
9998unsafe impl Sync for SendMut {}
9999
10000impl SendMut {
10001 #[inline]
10002 fn at(self, i: usize) -> *mut f32 {
10003 unsafe { self.0.add(i) }
10004 }
10005}
10006
10007#[cfg(test)]
10008mod tests {
10009 use super::*;
10010
10011 #[test]
10012 fn q2tp_i8_dot_matches_exact_on_grid() {
10013 let (rows, cols) = (5, 64);
10017 let gpr = cols / GROUP_SIZE;
10018 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10022 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10023 .collect();
10024 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10025 let x: Vec<f32> = (0..cols).map(|i| if i % 3 == 0 { -1.0 } else { 1.0 }).collect();
10026 let act = split_act(&x);
10027 assert!(act.outliers.is_empty(), "on-grid input must have no outliers");
10028 let gsum = q1_group_sums(&act.xq, gpr);
10029 for r in 0..rows {
10030 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10031 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10032 assert!(
10033 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10034 "row {r}: exact {exact} vs i8 {fast}"
10035 );
10036 }
10037 }
10038
10039 #[test]
10040 fn q8_row_dot_fast_matches_scalar() {
10041 let cols = 96;
10044 let row: Vec<u8> = (0..cols)
10045 .map(|i| ((i as i32 * 37 % 251) - 125) as i8 as u8)
10046 .collect();
10047 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10048 let act = split_act(&x);
10049 let fast = q8_row_dot(&row, &act);
10050 let scalar = q8_row_dot_scalar(&row, &act);
10051 assert!(
10052 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10053 "fast {fast} vs scalar {scalar}"
10054 );
10055 }
10056
10057 #[test]
10058 fn f32_matvec_matches_matvec_rows_bitexact() {
10059 let (rows, cols) = (300, 40);
10060 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10061 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10062 let qt = QTensor::from_f32(w.clone(), rows, cols);
10063
10064 let mut a = vec![0.0f32; rows];
10065 matvec_rows(None, &w, &x, &mut a);
10066 let mut b = vec![0.0f32; rows];
10067 qt.matvec(&x, &mut b, None);
10068 assert_eq!(a, b);
10069 }
10070
10071 #[test]
10072 fn sdot_kernel_exact_on_grid() {
10073 eprintln!("sdot_enabled = {}", sdot_enabled());
10078 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
10080 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10081 .collect();
10082 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10083 let x: Vec<f32> = (0..cols)
10084 .map(|i| match i % 3 {
10085 0 => 1.0,
10086 1 => -1.0,
10087 _ => 0.0,
10088 })
10089 .collect();
10090 let mut a = vec![0.0f32; rows];
10091 qmatvec(
10092 &w,
10093 &[],
10094 &scales,
10095 &x,
10096 &[],
10097 TensorDtype::Q8Row,
10098 rows,
10099 cols,
10100 &mut a,
10101 None,
10102 );
10103 for o in 0..rows {
10104 let mut acc = 0.0f32;
10105 for j in 0..cols {
10106 acc += (w[o * cols + j] as i8) as f32 * x[j];
10107 }
10108 let expect = acc * scales[o];
10109 assert!(
10110 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10111 "row {o}: {} vs {expect}",
10112 a[o]
10113 );
10114 }
10115 }
10116
10117 #[test]
10118 fn q1_tbl_fast_path_matches_reference() {
10119 let (rows, cols) = (5, 256);
10124 let gpr = cols / GROUP_SIZE;
10125 let mut bytes = Vec::new();
10126 for t in 0..rows * gpr {
10127 let s = 0.007 + (t % 11) as f32 * 0.004;
10128 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10129 for j in 0..4 {
10130 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10131 }
10132 }
10133 let x: Vec<f32> = (0..cols)
10134 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10135 .collect();
10136 let mut w = vec![0.0f32; rows * cols];
10137 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10138 let mut got = vec![0.0f32; rows];
10139 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10140 for o in 0..rows {
10141 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10142 assert!(
10143 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10144 "row {o}: {} vs {expect}",
10145 got[o]
10146 );
10147 }
10148 let b = 5usize;
10151 let mut xs_all = Vec::new();
10152 for bi in 0..b {
10153 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10154 }
10155 let mut mm = vec![0.0f32; b * rows];
10156 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10157 for bi in 0..b {
10158 let mut single = vec![0.0f32; rows];
10159 q1_matvec(
10160 &bytes,
10161 &xs_all[bi * cols..(bi + 1) * cols],
10162 rows,
10163 cols,
10164 &mut single,
10165 None,
10166 );
10167 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10168 }
10169 }
10170
10171 #[test]
10172 fn q1_kernels_match_exact_reference() {
10173 let (rows, cols) = (7, 96);
10175 let gpr = cols / GROUP_SIZE;
10176 let mut bytes = Vec::new();
10177 for t in 0..rows * gpr {
10178 let s = 0.01 + (t % 13) as f32 * 0.003;
10179 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10180 for j in 0..4 {
10181 bytes.push(((t * 31 + j * 97) % 251) as u8);
10182 }
10183 }
10184 let x: Vec<f32> = (0..cols)
10186 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10187 .collect();
10188 let mut w = vec![0.0f32; rows * cols];
10190 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10191 let mut expect = vec![0.0f32; rows];
10192 for o in 0..rows {
10193 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10194 }
10195 let mut got = vec![0.0f32; rows];
10196 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10197 for o in 0..rows {
10198 assert!(
10199 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10200 "row {o}: {} vs {}",
10201 got[o],
10202 expect[o]
10203 );
10204 }
10205 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10207 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10208 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10209 assert_eq!(a1, got);
10210 let mut xs = x.clone();
10211 xs.extend_from_slice(&x2);
10212 let mut mm = vec![0.0f32; 2 * rows];
10213 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10214 assert_eq!(&mm[..rows], got.as_slice());
10215 assert_eq!(&mm[rows..], a2.as_slice());
10216 }
10217
10218 #[test]
10219 fn repack_is_bit_identical() {
10220 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
10226 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10227 .collect();
10228 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10229 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10230 let rep = q8_repack_layout(&w, rows, cols);
10231 for g in 0..rows / 4 {
10233 for c in 0..cols / 16 {
10234 for lane in 0..4 {
10235 assert_eq!(
10236 &rep[g * 4 * cols + c * 64 + lane * 16
10237 ..g * 4 * cols + c * 64 + lane * 16 + 16],
10238 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10239 );
10240 }
10241 }
10242 }
10243 let mut a = vec![0.0f32; rows];
10244 qmatvec(
10245 &w,
10246 &[],
10247 &scales,
10248 &x,
10249 &[],
10250 TensorDtype::Q8Row,
10251 rows,
10252 cols,
10253 &mut a,
10254 None,
10255 );
10256 let mut b = vec![0.0f32; rows];
10257 qmatvec(
10258 &w,
10259 &rep,
10260 &scales,
10261 &x,
10262 &[],
10263 TensorDtype::Q8Row,
10264 rows,
10265 cols,
10266 &mut b,
10267 None,
10268 );
10269 assert_eq!(a, b, "full-range repack output diverged");
10270
10271 #[cfg(target_arch = "aarch64")]
10272 if sdot_enabled() {
10273 let act = split_act(&x);
10275 let mut c1 = vec![0.0f32; rows];
10276 let mut c2 = vec![0.0f32; rows];
10277 q8_range_sdot(
10278 &w,
10279 &[],
10280 &scales,
10281 &act,
10282 cols,
10283 SendMut(c1.as_mut_ptr()),
10284 3,
10285 rows - 2,
10286 );
10287 q8_range_sdot(
10288 &w,
10289 &rep,
10290 &scales,
10291 &act,
10292 cols,
10293 SendMut(c2.as_mut_ptr()),
10294 3,
10295 rows - 2,
10296 );
10297 assert_eq!(c1, c2, "unaligned-range repack output diverged");
10298 }
10299 }
10300
10301 #[test]
10302 fn sdot_a8w8_noise_is_bounded() {
10303 let (rows, cols) = (16, 512);
10307 let w: Vec<u8> = (0..rows * cols)
10308 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10309 .collect();
10310 let scales = vec![0.01f32; rows];
10311 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10312 let mut a = vec![0.0f32; rows];
10313 qmatvec(
10314 &w,
10315 &[],
10316 &scales,
10317 &x,
10318 &[],
10319 TensorDtype::Q8Row,
10320 rows,
10321 cols,
10322 &mut a,
10323 None,
10324 );
10325 let (mut num, mut den) = (0f64, 0f64);
10326 for o in 0..rows {
10327 let mut acc = 0.0f32;
10328 for j in 0..cols {
10329 acc += (w[o * cols + j] as i8) as f32 * x[j];
10330 }
10331 let expect = acc * scales[o];
10332 num += ((a[o] - expect) as f64).powi(2);
10333 den += (expect as f64).powi(2);
10334 }
10335 let rel = (num / den.max(1e-12)).sqrt();
10336 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10337 }
10338
10339 #[test]
10340 fn i8_dot_neon_matches_scalar() {
10341 let n = 100;
10342 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10343 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10344 let mut scalar = 0.0f32;
10345 for j in 0..n {
10346 scalar += (w[j] as i8) as f32 * x[j];
10347 }
10348 let fast = dot_i8_f32(&w, &x);
10349 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10350 }
10351
10352 #[test]
10354 fn vbitmatvec_matches_full_dequant() {
10355 let (rows, cols) = (6, 64);
10356 let ng = cols / GROUP_SIZE;
10357 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10359 let mut bytes = bits.clone();
10360 for g in 0..rows * ng {
10361 let s = 0.02 + 0.001 * g as f32;
10362 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10363 }
10364 for r in 0..rows {
10365 let b = bits[r] as usize;
10366 let (mut acc, mut nb) = (0u64, 0usize);
10367 let mut rowbytes = Vec::new();
10368 for i in 0..cols {
10369 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10370 acc = (acc << b) | v;
10371 nb += b;
10372 while nb >= 8 {
10373 nb -= 8;
10374 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10375 }
10376 }
10377 if nb > 0 {
10378 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10379 }
10380 bytes.extend_from_slice(&rowbytes);
10381 }
10382 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10383
10384 let mut reference = vec![0f32; rows * cols];
10385 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10386 let mut expect = vec![0f32; rows];
10387 for r in 0..rows {
10388 expect[r] = reference[r * cols..(r + 1) * cols]
10389 .iter()
10390 .zip(&x)
10391 .map(|(w, xv)| w * xv)
10392 .sum();
10393 }
10394 let mut got = vec![0f32; rows];
10395 let offsets = vbit_row_offsets(&bytes, rows, cols);
10396 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10397 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10401 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10402 for r in 0..rows {
10403 assert!(
10404 (got[r] - expect[r]).abs() < tol * scale,
10405 "row {r}: {} vs {}",
10406 got[r],
10407 expect[r]
10408 );
10409 }
10410 }
10411
10412 #[test]
10417 #[cfg(target_arch = "x86_64")]
10418 fn vbit_matmat_blocked_matches_per_row() {
10419 let (rows, cols, b) = (64usize, 128usize, 9usize);
10420 let ng = cols / GROUP_SIZE;
10421 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10422 let mut bytes = bits.clone();
10423 for g in 0..rows * ng {
10424 let sc = 0.02 + 0.0005 * g as f32;
10425 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10426 }
10427 for r in 0..rows {
10428 let bw = bits[r] as usize;
10429 let (mut acc, mut nb) = (0u64, 0usize);
10430 let mut rowbytes = Vec::new();
10431 for i in 0..cols {
10432 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10433 acc = (acc << bw) | v;
10434 nb += bw;
10435 while nb >= 8 {
10436 nb -= 8;
10437 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10438 }
10439 }
10440 if nb > 0 {
10441 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10442 }
10443 bytes.extend_from_slice(&rowbytes);
10444 }
10445 let x: Vec<f32> = (0..b * cols)
10446 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10447 .collect();
10448 let offsets = vbit_row_offsets(&bytes, rows, cols);
10449 let mut y_a = vec![0f32; b * rows];
10450 let mut y_b = vec![0f32; b * rows];
10451 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10452 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10453 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10454 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10455 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10456 let max_d = y_a
10457 .iter()
10458 .zip(&y_b)
10459 .map(|(p, q)| (p - q).abs())
10460 .fold(0.0f32, f32::max);
10461 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10462 }
10463
10464 #[test]
10472 fn q4t_matmat_blocked_matches_per_row() {
10473 let (rows, cols, b) = (16usize, 64usize, 9usize);
10474 let gpr = cols / GROUP_SIZE;
10475 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10476 for r in 0..rows {
10477 for g in 0..gpr {
10478 let t = (r * gpr + g) * Q4_TILE;
10479 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10480 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10481 for k in 0..16 {
10482 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10483 }
10484 }
10485 }
10486 let x: Vec<f32> = (0..b * cols)
10487 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10488 .collect();
10489 let mut y_blk = vec![0f32; b * rows];
10490 let mut y_row = vec![0f32; b * rows];
10491 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10492 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10493 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10494 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10495 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10496 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10497 }
10498
10499 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10506 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10507 let gpr = cols / GROUP_SIZE;
10508 let stride = q4tp_code_stride(gpr);
10509 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10510 let mut b = vec![0u8; codes_off + rows * stride];
10511 for r in 0..rows {
10512 for g in 0..gpr {
10513 let t = (r * gpr + g) * Q4TP_NIB;
10514 for k in 0..16 {
10515 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10516 }
10517 }
10518 let lo = -6.0 - 0.03 * (r % 17) as f32;
10519 let step = 0.01 + 0.004 * (r % 11) as f32;
10520 let p = params_off + r * 4;
10521 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10522 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10523 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10524 for g in 0..gpr {
10525 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10526 }
10527 }
10528 b
10529 }
10530
10531 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10535 let gpr = cols / GROUP_SIZE;
10536 let v = Q4tpView::new(bytes, rows, cols);
10537 let mut out = vec![0u8; rows * gpr * Q4_TILE];
10538 let mut sc = vec![0f32; gpr];
10539 for r in 0..rows {
10540 v.scales_into(r, gpr, &mut sc);
10541 for g in 0..gpr {
10542 let t = (r * gpr + g) * Q4_TILE;
10543 let s = sc[g];
10544 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10545 let src = (r * gpr + g) * Q4TP_NIB;
10546 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10547 }
10548 }
10549 out
10550 }
10551
10552 #[test]
10558 fn q4tp_exact_path_matches_dequant_reference() {
10559 let (rows, cols) = (256usize, 512usize);
10560 let gpr = cols / GROUP_SIZE;
10561 let bytes = synth_q4tp(rows, cols);
10562 let mut w = vec![0f32; rows * cols];
10563 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10564
10565 let x: Vec<f32> = (0..cols)
10566 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10567 .collect();
10568 let v = Q4tpView::new(&bytes, rows, cols);
10569 let mut sc = vec![0f32; gpr];
10570 for r in 0..rows {
10571 v.scales_into(r, gpr, &mut sc);
10572 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10573 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10574 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10578 assert!(
10579 (got - want).abs() <= 1e-5 * mag,
10580 "row {r}: kernel {got} vs dequant {want}"
10581 );
10582 }
10583 }
10584
10585 #[test]
10591 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10592 let (rows, cols) = (256usize, 512usize);
10593 let bytes = synth_q4tp(rows, cols);
10594 let twin = q4tp_as_q4t(&bytes, rows, cols);
10595 let x: Vec<f32> = (0..cols)
10596 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10597 .collect();
10598
10599 let mut got = vec![0f32; rows];
10600 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10601 let mut want = vec![0f32; rows];
10602 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10603
10604 let mut w = vec![0f32; rows * cols];
10607 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10608 for r in 0..rows {
10609 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10610 assert!(
10611 (got[r] - want[r]).abs() <= 1e-3 * mag,
10612 "row {r}: q4tp {} vs q4t {}",
10613 got[r],
10614 want[r]
10615 );
10616 }
10617 }
10618
10619 #[test]
10624 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10625 let (rows, cols, b) = (256usize, 512usize, 5usize);
10626 let bytes = synth_q4tp(rows, cols);
10627 let twin = q4tp_as_q4t(&bytes, rows, cols);
10628 let xs: Vec<f32> = (0..b * cols)
10629 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10630 .collect();
10631
10632 let mut got = vec![0f32; b * rows];
10633 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10634 let mut want = vec![0f32; b * rows];
10635 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10636
10637 let mut w = vec![0f32; rows * cols];
10638 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10639 for t in 0..b {
10640 for r in 0..rows {
10641 let mag: f32 = (0..cols)
10642 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10643 .sum();
10644 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10645 assert!(
10646 (g - wa).abs() <= 1e-3 * mag,
10647 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10648 );
10649 }
10650 }
10651 }
10652
10653 #[test]
10654 fn q4tp_matvec2_matches_the_single_stream_kernel() {
10655 let (rows, cols) = (128usize, 256usize);
10656 let gpr = cols / GROUP_SIZE;
10657 let bytes = synth_q4tp(rows, cols);
10658 let xs: Vec<f32> = (0..2 * cols)
10659 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10660 .collect();
10661
10662 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10663 q4tp_matvec2(
10664 &bytes,
10665 &xs[..cols],
10666 &xs[cols..],
10667 rows,
10668 cols,
10669 &mut o1,
10670 &mut o2,
10671 None,
10672 );
10673
10674 let v = Q4tpView::new(&bytes, rows, cols);
10677 let mut sc = vec![0f32; gpr];
10678 for r in 0..rows {
10679 v.scales_into(r, gpr, &mut sc);
10680 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10681 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10682 }
10683 }
10684
10685 #[test]
10692 fn q4tp_matvec_keeps_pace_with_q4t() {
10693 let (rows, cols) = (4096usize, 3072usize);
10694 let bytes = synth_q4tp(rows, cols);
10695 let twin = q4tp_as_q4t(&bytes, rows, cols);
10696 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10697 let mut o = vec![0f32; rows];
10698 let n = 12;
10699 let mut best = (f64::MAX, f64::MAX);
10700 for _ in 0..3 {
10703 let t0 = std::time::Instant::now();
10704 for _ in 0..n {
10705 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10706 }
10707 best.0 = best.0.min(t0.elapsed().as_secs_f64());
10708 let t0 = std::time::Instant::now();
10709 for _ in 0..n {
10710 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10711 }
10712 best.1 = best.1.min(t0.elapsed().as_secs_f64());
10713 }
10714 let ratio = best.1 / best.0;
10715 println!(
10716 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10717 best.0 * 1e3 / n as f64,
10718 best.1 * 1e3 / n as f64
10719 );
10720 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10721 }
10722
10723 #[cfg(target_os = "macos")]
10724 #[test]
10725 fn q4t_matmat_accel_matches_dequant_reference() {
10726 if !accel_gemm_enabled() {
10727 return; }
10729 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
10731 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10732 for r in 0..rows {
10733 for g in 0..gpr {
10734 let t = (r * gpr + g) * Q4_TILE;
10735 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10736 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10737 for k in 0..16 {
10738 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10739 }
10740 }
10741 }
10742 let x: Vec<f32> = (0..b * cols)
10743 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10744 .collect();
10745 let mut got = vec![0f32; b * rows];
10746 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10747 let mut w = vec![0f32; rows * cols];
10749 for r in 0..rows {
10750 for g in 0..gpr {
10751 let t = (r * gpr + g) * Q4_TILE;
10752 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10753 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10754 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10755 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10756 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10757 }
10758 }
10759 }
10760 for bi in 0..b {
10761 for r in 0..rows {
10762 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10763 let d = (got[bi * rows + r] - want).abs();
10764 assert!(
10765 d <= want.abs().max(1.0) * 1e-4,
10766 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10767 got[bi * rows + r]
10768 );
10769 }
10770 }
10771 }
10772
10773 #[test]
10774 fn q4matvec_matches_full_dequant() {
10775 let (rows, cols) = (8, 64);
10776 let groups = rows * cols / GROUP_SIZE;
10777 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10779 for i in 0..groups * 16 {
10780 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10781 }
10782 for g in 0..groups {
10783 let s = 0.01 + 0.003 * g as f32;
10784 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10785 }
10786 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10787
10788 let mut reference = vec![0.0f32; rows * cols];
10789 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10790 let mut expect = vec![0.0f32; rows];
10791 for r in 0..rows {
10792 expect[r] = reference[r * cols..(r + 1) * cols]
10793 .iter()
10794 .zip(&x)
10795 .map(|(w, xv)| w * xv)
10796 .sum();
10797 }
10798
10799 let mut got = vec![0.0f32; rows];
10800 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10801 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10805 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10806 for r in 0..rows {
10807 assert!(
10808 (got[r] - expect[r]).abs() < tol * scale,
10809 "row {r}: {} vs {}",
10810 got[r],
10811 expect[r]
10812 );
10813 }
10814 }
10815
10816 #[test]
10819 fn vbitmatvec2_equals_two_singles() {
10820 let (rows, cols) = (6, 64);
10821 let ng = cols / GROUP_SIZE;
10822 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10823 let mut bytes = bits.clone();
10824 for g in 0..rows * ng {
10825 let s = 0.02 + 0.001 * g as f32;
10826 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10827 }
10828 for r in 0..rows {
10829 let b = bits[r] as usize;
10830 let (mut acc, mut nb) = (0u64, 0usize);
10831 let mut rowbytes = Vec::new();
10832 for i in 0..cols {
10833 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10834 acc = (acc << b) | v;
10835 nb += b;
10836 while nb >= 8 {
10837 nb -= 8;
10838 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10839 }
10840 }
10841 if nb > 0 {
10842 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10843 }
10844 bytes.extend_from_slice(&rowbytes);
10845 }
10846 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10847 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10848 let offsets = vbit_row_offsets(&bytes, rows, cols);
10849
10850 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10851 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10852 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10853 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10854 vbitmatvec2(
10855 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10856 );
10857 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10858 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10859 }
10860
10861 #[test]
10863 fn q4matvec2_equals_two_singles() {
10864 let (rows, cols) = (8, 128);
10865 let groups = rows * cols / GROUP_SIZE;
10866 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10867 for i in 0..groups * 16 {
10868 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10869 }
10870 for g in 0..groups {
10871 let s = 0.01 + 0.003 * g as f32;
10872 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10873 }
10874 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10877 x1[9] = 250.0;
10878 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10879
10880 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10881 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10882 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10883 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10884 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10885 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10886 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10887 }
10888
10889 #[test]
10892 fn matvec_many_equals_separate_matvecs() {
10893 use crate::pool::Pool;
10894 let (r1, r2, cols) = (300, 200, 64);
10895 let mk = |salt: usize, rows: usize| {
10896 QTensor::from_f32(
10897 (0..rows * cols)
10898 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10899 .collect(),
10900 rows,
10901 cols,
10902 )
10903 };
10904 let (a, b) = (mk(1, r1), mk(5, r2));
10905 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10906 let pool = Pool::new(3);
10907
10908 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10909 a.matvec(&x, &mut ea, Some(&pool));
10910 b.matvec(&x, &mut eb, Some(&pool));
10911 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10912 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10913 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10914 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10915 }
10916
10917 #[test]
10920 fn batched_matmat_equals_per_position_matvec() {
10921 let (rows, cols, b) = (8, 64, 5);
10922 let groups = rows * cols / GROUP_SIZE;
10924 let mut q4 = Vec::new();
10925 for i in 0..groups * 16 {
10926 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10927 }
10928 for g in 0..groups {
10929 q4.extend_from_slice(
10930 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10931 );
10932 }
10933 let ng = cols / GROUP_SIZE;
10935 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10936 let mut vb = bits.clone();
10937 for g in 0..rows * ng {
10938 vb.extend_from_slice(
10939 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10940 );
10941 }
10942 for r in 0..rows {
10943 let bw = bits[r] as usize;
10944 let (mut acc, mut nb) = (0u64, 0usize);
10945 let mut rowbytes = Vec::new();
10946 for i in 0..cols {
10947 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10948 acc = (acc << bw) | v;
10949 nb += bw;
10950 while nb >= 8 {
10951 nb -= 8;
10952 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10953 }
10954 }
10955 if nb > 0 {
10956 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10957 }
10958 vb.extend_from_slice(&rowbytes);
10959 }
10960 let offsets = vbit_row_offsets(&vb, rows, cols);
10961
10962 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10963
10964 let mut got = vec![0f32; b * rows];
10966 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
10967 for bi in 0..b {
10968 let mut expect = vec![0f32; rows];
10969 q4matvec(
10970 &q4,
10971 &xs[bi * cols..(bi + 1) * cols],
10972 rows,
10973 cols,
10974 &mut expect,
10975 None,
10976 );
10977 assert_eq!(
10978 &got[bi * rows..(bi + 1) * rows],
10979 &expect[..],
10980 "q4 batch pos {bi}"
10981 );
10982 }
10983
10984 let mut got = vec![0f32; b * rows];
10986 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
10987 for bi in 0..b {
10988 let mut expect = vec![0f32; rows];
10989 vbitmatvec(
10990 &vb,
10991 &offsets,
10992 &xs[bi * cols..(bi + 1) * cols],
10993 rows,
10994 cols,
10995 &mut expect,
10996 None,
10997 );
10998 assert_eq!(
10999 &got[bi * rows..(bi + 1) * rows],
11000 &expect[..],
11001 "vbit batch pos {bi}"
11002 );
11003 }
11004 }
11005
11006 #[test]
11010 fn q4_tiled_matches_q4_block_bitexact() {
11011 let (rows, cols, b) = (8usize, 128usize, 3usize);
11012 let groups = rows * cols / GROUP_SIZE;
11013 let mut split = Vec::with_capacity(groups * 18);
11014 for i in 0..groups * 16 {
11015 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11016 }
11017 for g in 0..groups {
11018 split.extend_from_slice(
11019 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11020 );
11021 }
11022 let (packed, scales) = split.split_at(groups * 16);
11024 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
11025 for g in 0..groups {
11026 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
11027 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
11028 }
11029
11030 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11031 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11033
11034 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
11035 q4matvec(&split, &x1, rows, cols, &mut a, None);
11036 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
11037 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11038
11039 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11040 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11041 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11042 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11043 assert_eq!(a1, t1);
11044 assert_eq!(a2, t2);
11045
11046 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11047 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11048 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11049 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11050 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11051 }
11052
11053 #[test]
11060 fn q4matvec_sdot_outlier_exact() {
11061 let (rows, cols) = (4, 128);
11062 let groups = rows * cols / GROUP_SIZE;
11063 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11064 for i in 0..groups * 16 {
11065 bytes.push(((i * 11 + 5) % 256) as u8);
11066 }
11067 for g in 0..groups {
11068 let s = 0.02 + 0.002 * g as f32;
11069 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11070 }
11071 let mut x: Vec<f32> = (0..cols)
11072 .map(|i| match i % 3 {
11073 0 => 1.0,
11074 1 => -1.0,
11075 _ => 0.0,
11076 })
11077 .collect();
11078 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
11081 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11082 let mut expect = vec![0.0f32; rows];
11083 for r in 0..rows {
11084 expect[r] = reference[r * cols..(r + 1) * cols]
11085 .iter()
11086 .zip(&x)
11087 .map(|(w, xv)| w * xv)
11088 .sum();
11089 }
11090 let mut got = vec![0.0f32; rows];
11091 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11092 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11093 for r in 0..rows {
11094 assert!(
11095 (got[r] - expect[r]).abs() < 2e-3 * scale,
11096 "row {r}: {} vs {} (outlier term must be exact)",
11097 got[r],
11098 expect[r]
11099 );
11100 }
11101 }
11102
11103 #[test]
11107 fn q1t_matvec_matches_reference() {
11108 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11109 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
11111 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11112 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11114 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11115 let mut bytes = Vec::new();
11116 for r in 0..rows {
11117 for g in 0..gpr {
11118 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11119 let mut c = [0u8; 7];
11120 for k in 0..GROUP_SIZE {
11121 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11123 0
11124 } else {
11125 ((k + r * 3 + g) % 3) as u8 };
11127 cortiq_core::quant::q1t_pack(&mut c, k, code);
11128 }
11129 bytes.extend_from_slice(&c);
11130 }
11131 }
11132 let mut row_ptr = vec![0u32; rows + 1];
11135 for &(idx, _) in &outliers {
11136 row_ptr[idx as usize / cols + 1] += 1;
11137 }
11138 for r in 0..rows {
11139 row_ptr[r + 1] += row_ptr[r];
11140 }
11141 for &p in &row_ptr {
11142 bytes.extend_from_slice(&p.to_le_bytes());
11143 }
11144 for &(idx, v) in &outliers {
11145 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11146 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11147 }
11148
11149 let mut refw = vec![0f32; rows * cols];
11150 dequant_q1t(&bytes, rows, cols, &mut refw);
11151 let x: Vec<f32> = (0..cols)
11154 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11155 .collect();
11156 let mut expect = vec![0f32; rows];
11157 for r in 0..rows {
11158 let mut a = 0.0f32;
11159 for j in 0..cols {
11160 a += refw[r * cols + j] * x[j];
11161 }
11162 expect[r] = a;
11163 }
11164 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11165 let mut got = vec![0f32; rows];
11166 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11167 for r in 0..rows {
11168 assert!(
11169 (got[r] - expect[r]).abs() < tol(expect[r]),
11170 "row {r}: {} vs {}",
11171 got[r],
11172 expect[r]
11173 );
11174 }
11175 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
11177 let mut gm = vec![0f32; 2 * rows];
11178 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11179 for r in 0..rows {
11180 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11181 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11182 }
11183 let xb: Vec<f32> = (0..cols)
11187 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11188 .collect();
11189 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11190 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11191 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11192 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11193 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11194 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11195 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11196 }
11197
11198 #[test]
11201 fn q1t_matvec2_odd_gpr_matches_singles() {
11202 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11203 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
11205 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11206 for r in 0..rows {
11207 for g in 0..gpr {
11208 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11209 let mut c = [0u8; 7];
11210 for k in 0..GROUP_SIZE {
11211 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11212 }
11213 bytes.extend_from_slice(&c);
11214 }
11215 }
11216 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11217 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11218 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11219 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11220 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11221 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11222 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11223 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11224 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11225 }
11226
11227 #[test]
11231 #[ignore]
11232 fn q1t_matvec2_speed() {
11233 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11234 use std::time::Instant;
11235 let (rows, cols) = (8192usize, 4096usize);
11236 let gpr = cols / GROUP_SIZE;
11237 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11238 for r in 0..rows {
11239 for g in 0..gpr {
11240 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11241 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11242 let mut c = [0u8; 7];
11243 for k in 0..GROUP_SIZE {
11244 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11245 }
11246 bytes.extend_from_slice(&c);
11247 }
11248 }
11249 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11250 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11251 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11252 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11253 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11255 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11256 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11257 for _ in 0..8 {
11258 let t0 = Instant::now();
11259 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11260 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11261 let t1 = Instant::now();
11262 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11263 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11264 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11265 }
11266 assert_eq!(p1, s1);
11267 assert_eq!(p2, s2);
11268 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11269 }
11270
11271 #[test]
11275 #[ignore]
11276 fn q1t_matvec_speed() {
11277 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11278 use std::time::Instant;
11279 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
11281 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11282 for r in 0..rows {
11283 for g in 0..gpr {
11284 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11285 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11286 let mut c = [0u8; 7];
11287 for k in 0..GROUP_SIZE {
11288 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11289 }
11290 bytes.extend_from_slice(&c);
11291 }
11292 }
11293 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
11295 let mut idx = 0usize;
11296 while idx < n {
11297 row_ptr[idx / cols + 1] += 1;
11298 idx += stride;
11299 }
11300 for r in 0..rows {
11301 row_ptr[r + 1] += row_ptr[r];
11302 }
11303 for &p in &row_ptr {
11304 bytes.extend_from_slice(&p.to_le_bytes());
11305 }
11306 let mut idx = 0usize;
11307 while idx < n {
11308 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11309 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11310 idx += stride;
11311 }
11312 let x: Vec<f32> = (0..cols)
11315 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11316 .collect();
11317 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11318
11319 let slow = |out: &mut [f32]| {
11321 let mut buf = vec![0f32; cols];
11322 for r in 0..rows {
11323 for g in 0..gpr {
11324 let off = (r * gpr + g) * Q1T_TILE;
11325 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11326 let codes = &bytes[off + 2..off + Q1T_TILE];
11327 for k in 0..GROUP_SIZE {
11328 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11329 1 => s,
11330 2 => -s,
11331 _ => 0.0,
11332 };
11333 }
11334 }
11335 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11336 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11337 }
11338 };
11339 let iters = 5;
11340 let mut a = vec![0f32; rows];
11341 slow(&mut a); let t = Instant::now();
11343 for _ in 0..iters {
11344 slow(&mut a);
11345 }
11346 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11347
11348 let mut b = vec![0f32; rows];
11349 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
11351 for _ in 0..iters {
11352 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11353 }
11354 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11355
11356 for r in 0..rows {
11357 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11358 }
11359 println!(
11360 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
11361 slow_ms / fast_ms
11362 );
11363 }
11364}
11365
11366
11367#[cfg(test)]
11368mod gemm_bench {
11369 #[test]
11380 #[ignore]
11381 fn q4tp_matmat_throughput() {
11382 let b: usize = std::env::var("CMF_BENCH_B")
11386 .ok()
11387 .and_then(|v| v.parse().ok())
11388 .unwrap_or(296);
11389 let (rows, cols) = (9216usize, 2304usize);
11390 let (_, _, _) = (rows, cols, b);
11391 let total = cortiq_core::quant::expected_nbytes(
11392 cortiq_core::TensorDtype::Q4TiledP,
11393 &[rows, cols],
11394 )
11395 .unwrap();
11396 let (params_off, codes_off, _) =
11401 cortiq_core::quant::q4tp_sections(rows, cols);
11402 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11403 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11404 let step = cortiq_core::quant::f32_to_f16(0.1);
11405 for r in 0..rows {
11406 let o = params_off + r * 4;
11407 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11408 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11409 }
11410 let _ = codes_off;
11411 let xs: Vec<f32> = (0..b * cols)
11412 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11413 .collect();
11414 let mut out = vec![0f32; b * rows];
11415 let pool = crate::pool::Pool::from_env();
11416 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11422 let reps: usize = std::env::var("CMF_BENCH_REPS")
11423 .ok()
11424 .and_then(|v| v.parse().ok())
11425 .unwrap_or(10);
11426 let mut best = [f64::MAX; 2];
11427 let mut sums = [0f32; 2];
11428 for _ in 0..reps {
11429 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11430 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11431 let t = std::time::Instant::now();
11432 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11433 best[k] = best[k].min(t.elapsed().as_secs_f64());
11434 sums[k] = out.iter().take(64).sum::<f32>();
11435 }
11436 }
11437 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11438 for (k, name) in ["previous", "tuned "].iter().enumerate() {
11439 println!(
11440 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11441 best[k] * 1e3,
11442 flops / best[k] / 1e9,
11443 sums[k]
11444 );
11445 }
11446 assert!(
11447 (sums[0] - sums[1]).abs() < 1e-2,
11448 "the tuned kernel changed the result: {} vs {}",
11449 sums[0],
11450 sums[1]
11451 );
11452 }
11453
11454 #[test]
11460 fn q4tp_matmat_blocked_matches_scalar() {
11461 use std::sync::atomic::Ordering::Relaxed;
11462 for &(rows, cols, b) in &[
11470 (64usize, 128usize, 7usize),
11471 (33, 96, 4),
11472 (16, 256, 9),
11473 (192, 2304, 37),
11474 ] {
11475 let total = cortiq_core::quant::expected_nbytes(
11476 cortiq_core::TensorDtype::Q4TiledP,
11477 &[rows, cols],
11478 )
11479 .unwrap();
11480 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11481 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11482 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11483 let step = cortiq_core::quant::f32_to_f16(0.1);
11484 for r in 0..rows {
11485 let o = params_off + r * 4;
11486 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11487 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11488 }
11489 let xs: Vec<f32> = (0..b * cols)
11490 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11491 .collect();
11492 let mut got = vec![0f32; b * rows];
11493 let mut want = vec![0f32; b * rows];
11494 let gpr = cols / 32;
11495 let view = super::Q4tpView::new(&bytes, rows, cols);
11496 let pool = crate::pool::Pool::from_env();
11497 super::Q4TP_ALT.store(2, Relaxed);
11498 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11499 super::Q4TP_ALT.store(1, Relaxed);
11500 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11501 super::Q4TP_ALT.store(0, Relaxed);
11502 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11509 let (mut worst, mut at) = (0f32, 0usize);
11510 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11511 if (g - w).abs() > worst {
11512 worst = (g - w).abs();
11513 at = i;
11514 }
11515 }
11516 assert!(
11517 worst <= 1e-4 * scale,
11518 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11519 (scale {scale:.3e}) at cell {at}: {} vs {}",
11520 got[at],
11521 want[at]
11522 );
11523
11524 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11531 for bi in 0..b {
11532 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11533 for r in 0..rows {
11534 let mut sc = vec![0f32; gpr];
11535 view.scales_into(r, gpr, &mut sc);
11536 let mut exact = 0f64;
11537 for j in 0..cols {
11538 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11539 exact += w as f64 * sq as f64 * act.xq[j] as f64;
11540 }
11541 exact *= act.sx as f64;
11542 for &(j, xv) in &act.outliers {
11543 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11544 exact += w as f64 * sq as f64 * xv as f64;
11545 }
11546 let i = bi * rows + r;
11547 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11548 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11549 }
11550 }
11551 println!(
11552 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11553 per-column {e_scalar:.3e}"
11554 );
11555 assert!(
11560 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11561 "{rows}x{cols} b={b}: error against f64 too large — blocked \
11562 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11563 );
11564 }
11565 }
11566
11567 #[test]
11571 #[ignore]
11572 fn q4t_matmat_throughput() {
11573 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11574 let total = cortiq_core::quant::expected_nbytes(
11575 cortiq_core::TensorDtype::Q4Tiled,
11576 &[rows, cols],
11577 )
11578 .unwrap();
11579 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11582 let sc = cortiq_core::quant::f32_to_f16(0.02);
11583 for t in bytes.chunks_mut(super::Q4_TILE) {
11584 t[..2].copy_from_slice(&sc.to_le_bytes());
11585 }
11586 let xs: Vec<f32> = (0..b * cols)
11587 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11588 .collect();
11589 let mut out = vec![0f32; b * rows];
11590 let pool = crate::pool::Pool::from_env();
11591 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11592 let reps: usize = std::env::var("CMF_BENCH_REPS")
11593 .ok()
11594 .and_then(|v| v.parse().ok())
11595 .unwrap_or(10);
11596 let mut best = f64::MAX;
11597 for _ in 0..reps {
11598 let t = std::time::Instant::now();
11599 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11600 best = best.min(t.elapsed().as_secs_f64());
11601 }
11602 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11603 println!(
11604 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11605 best * 1e3,
11606 flops / best / 1e9,
11607 out.iter().take(64).sum::<f32>()
11608 );
11609 }
11610
11611}