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::cell::UnsafeCell;
21use std::sync::Arc;
22
23pub enum QTensor {
24 F32 {
25 data: Vec<f32>,
26 rows: usize,
27 cols: usize,
28 },
29 Mapped {
30 model: Arc<CmfModel>,
31 idx: usize,
33 dtype: TensorDtype,
34 rows: usize,
35 cols: usize,
36 row_scale: Vec<f32>,
38 col_field: Vec<f32>,
40 vbit_offsets: Vec<usize>,
44 repack: Vec<u8>,
52 },
53}
54
55fn repack_enabled() -> bool {
62 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
63 *ON.get_or_init(|| {
64 std::env::var("CMF_REPACK")
65 .map(|v| v == "1")
66 .unwrap_or(cfg!(target_os = "android"))
67 })
68}
69
70fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
74 #[cfg(target_arch = "aarch64")]
75 let arch_ok = sdot_enabled();
76 #[cfg(not(target_arch = "aarch64"))]
77 let arch_ok = false;
78 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
79 return Vec::new();
80 }
81 q8_repack_layout(bytes, rows, cols)
82}
83
84fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
87 let groups = rows / 4;
88 let mut rep = vec![0u8; groups * 4 * cols];
89 for g in 0..groups {
90 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
91 for c in 0..cols / 16 {
92 for lane in 0..4 {
93 let src = (g * 4 + lane) * cols + c * 16;
94 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
95 .copy_from_slice(&bytes[src..src + 16]);
96 }
97 }
98 }
99 rep
100}
101
102fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
105 let ng = cols / GROUP_SIZE;
106 let bits = &bytes[..rows];
107 let mut offsets = Vec::with_capacity(rows + 1);
108 let mut off = rows + rows * ng * 2;
109 for r in 0..rows {
110 offsets.push(off);
111 off += (cols * bits[r] as usize).div_ceil(8);
112 }
113 offsets.push(off);
114 offsets
115}
116
117fn blocked_enabled() -> bool {
123 use std::sync::atomic::Ordering::Relaxed;
124 match BLOCKED_OVERRIDE.load(Relaxed) {
125 1 => false,
126 2 => true,
127 _ => {
128 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
129 *ON.get_or_init(|| {
130 std::env::var("CMF_X86_BLOCKED")
131 .map(|v| v != "0")
132 .unwrap_or(true)
133 })
134 }
135 }
136}
137
138static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
139
140pub fn set_blocked_override(on: Option<bool>) {
147 let v = match on {
148 None => 0,
149 Some(false) => 1,
150 Some(true) => 2,
151 };
152 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
153}
154
155fn gpu_lmhead_enabled() -> bool {
156 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
157 *ON.get_or_init(|| {
158 std::env::var("CMF_GPU_LMHEAD")
159 .map(|v| v != "0")
160 .unwrap_or(true)
161 })
162}
163
164fn gpu_split_frac() -> f32 {
165 if FULL_GPU_Q8.get() {
168 return 1.0;
169 }
170 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
171 *F.get_or_init(|| {
172 std::env::var("CMF_GPU_SPLIT")
173 .ok()
174 .and_then(|v| v.parse::<f32>().ok())
175 .unwrap_or(0.5)
176 .clamp(0.0, 1.0)
177 })
178}
179
180impl QTensor {
181 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
182 debug_assert_eq!(data.len(), rows * cols);
183 Self::F32 { data, rows, cols }
184 }
185
186 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
189 let idx = model
192 .tensor_index(name)
193 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
194 let entry = &model.tensors[idx];
195 if entry.shape.len() != 2 {
196 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
197 }
198 let (rows, cols) = (entry.shape[0], entry.shape[1]);
199 let bytes = model.entry_bytes(entry);
200
201 match entry.dtype {
202 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
203 let n = rows * cols;
204 let scales_off = n;
205 let row_scale: Vec<f32> = (0..rows)
206 .map(|o| {
207 f16_to_f32(u16::from_le_bytes([
208 bytes[scales_off + o * 2],
209 bytes[scales_off + o * 2 + 1],
210 ]))
211 })
212 .collect();
213 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
214 let col_off = n + rows * 2;
215 (0..cols)
216 .map(|i| {
217 f16_to_f32(u16::from_le_bytes([
218 bytes[col_off + i * 2],
219 bytes[col_off + i * 2 + 1],
220 ]))
221 })
222 .collect()
223 } else {
224 Vec::new()
225 };
226 Ok(Self::Mapped {
227 model: model.clone(),
228 idx,
229 dtype: entry.dtype,
230 rows,
231 cols,
232 row_scale,
233 col_field,
234 vbit_offsets: Vec::new(),
235 repack: q8_repack(bytes, rows, cols),
236 })
237 }
238 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
240 model: model.clone(),
241 idx,
242 dtype: entry.dtype,
243 rows,
244 cols,
245 row_scale: Vec::new(),
246 col_field: Vec::new(),
247 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
248 repack: Vec::new(),
249 }),
250 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
254 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
255 let offsets: Vec<usize> = (0..=rows)
256 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
257 .collect();
258 Ok(Self::Mapped {
259 model: model.clone(),
260 idx,
261 dtype: entry.dtype,
262 rows,
263 cols,
264 row_scale: Vec::new(),
265 col_field: Vec::new(),
266 vbit_offsets: offsets,
267 repack: Vec::new(),
268 })
269 }
270 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
276 model: model.clone(),
277 idx,
278 dtype: entry.dtype,
279 rows,
280 cols,
281 row_scale: Vec::new(),
282 col_field: Vec::new(),
283 vbit_offsets: Vec::new(),
284 repack: Vec::new(),
285 }),
286 TensorDtype::Q4TiledP 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::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
301 model: model.clone(),
302 idx,
303 dtype: entry.dtype,
304 rows,
305 cols,
306 row_scale: Vec::new(),
307 col_field: Vec::new(),
308 vbit_offsets: Vec::new(),
309 repack: Vec::new(),
310 }),
311 TensorDtype::Q4Block 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::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
324 model: model.clone(),
325 idx,
326 dtype: entry.dtype,
327 rows,
328 cols,
329 row_scale: Vec::new(),
330 col_field: Vec::new(),
331 vbit_offsets: Vec::new(),
332 repack: Vec::new(),
333 }),
334 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
338 model: model.clone(),
339 idx,
340 dtype: entry.dtype,
341 rows,
342 cols,
343 row_scale: Vec::new(),
344 col_field: Vec::new(),
345 vbit_offsets: Vec::new(),
346 repack: Vec::new(),
347 }),
348 _ => {
350 let mut data = vec![0.0f32; rows * cols];
351 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
352 Ok(Self::from_f32(data, rows, cols))
353 }
354 }
355 }
356
357 pub(crate) fn is_q1(&self) -> bool {
360 matches!(
361 self,
362 Self::Mapped {
363 dtype: TensorDtype::Q1,
364 ..
365 }
366 )
367 }
368
369 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
372 match self {
373 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
374 _ => None,
375 }
376 }
377
378 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
385 if self.has_prism_contract() {
386 return None;
387 }
388 match self {
389 #[cfg(target_os = "macos")]
390 Self::Mapped {
391 dtype: TensorDtype::Q1T,
392 ..
393 } if !crate::gpu::metal_q1t_enabled() => None,
394 Self::Mapped {
395 idx,
396 dtype:
397 TensorDtype::Q1
398 | TensorDtype::Q1T
399 | TensorDtype::Q4Block
400 | TensorDtype::Q4Tiled
401 | TensorDtype::Q4TiledP
405 | TensorDtype::Q8Row
406 | TensorDtype::Q8_2f,
407 rows,
408 cols,
409 ..
410 } => Some((*idx, *rows, *cols)),
411 _ => None,
412 }
413 }
414
415 #[cfg(target_os = "macos")]
422 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
423 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
424 let name = &model.tensors[idx].name;
425 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
426 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
427 if forward && affine {
428 let e = model.tensors.get(idx)?;
429 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
430 }
431 }
432 self.q1_parts()
433 }
434
435 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
441 if self.has_prism_contract() {
442 return None;
443 }
444 match self {
445 Self::Mapped {
446 idx,
447 dtype: TensorDtype::Q4Tiled,
448 rows,
449 cols,
450 ..
451 } => Some((*idx, *rows, *cols)),
452 _ => None,
453 }
454 }
455
456 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
460 if self.has_prism_contract() {
461 return None;
462 }
463 match self {
464 Self::Mapped {
465 idx,
466 dtype: TensorDtype::Q4TiledP,
467 rows,
468 cols,
469 ..
470 } => Some((*idx, *rows, *cols)),
471 _ => None,
472 }
473 }
474
475 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
480 if self.has_prism_contract() {
481 return None;
482 }
483 match self {
484 Self::Mapped {
485 idx,
486 dtype: TensorDtype::Q8Row,
487 rows,
488 cols,
489 row_scale,
490 col_field,
491 ..
492 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
493 _ => None,
494 }
495 }
496
497 pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
500 if self.has_prism_contract() {
501 return None;
502 }
503 match self {
504 Self::Mapped {
505 idx,
506 dtype: TensorDtype::Q8_2f,
507 rows,
508 cols,
509 row_scale,
510 col_field,
511 ..
512 } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
513 _ => None,
514 }
515 }
516
517 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
521 match self {
522 Self::Mapped { dtype, .. } => Some(*dtype),
523 _ => None,
524 }
525 }
526
527 pub fn model_idx(&self) -> Option<usize> {
532 match self {
533 Self::Mapped { idx, .. } => Some(*idx),
534 _ => None,
535 }
536 }
537
538 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
543 match self {
544 Self::Mapped { model, .. } => Some(model.clone()),
545 _ => None,
546 }
547 }
548
549 pub(crate) fn has_prism_contract(&self) -> bool {
554 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
555 }
556
557 pub fn rows(&self) -> usize {
558 match self {
559 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
560 }
561 }
562
563 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
566 if self.has_prism_contract() {
567 return None;
568 }
569 match self {
570 Self::Mapped {
571 model,
572 idx,
573 dtype: TensorDtype::Q4Tiled,
574 ..
575 } => Some((model, *idx)),
576 _ => None,
577 }
578 }
579
580 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
583 if self.has_prism_contract() {
584 return None;
585 }
586 match self {
587 Self::Mapped {
588 model,
589 idx,
590 dtype: TensorDtype::Q4TiledP,
591 ..
592 } => Some((model, *idx)),
593 _ => None,
594 }
595 }
596
597 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
605 if self.has_prism_contract() {
606 return None;
607 }
608 match self {
609 Self::Mapped {
610 model,
611 idx,
612 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
613 ..
614 } => Some((model, *idx)),
615 _ => None,
616 }
617 }
618
619 pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
622 if self.has_prism_contract() {
623 return None;
624 }
625 match self {
626 Self::Mapped {
627 model,
628 idx,
629 dtype: TensorDtype::Q8_2f,
630 ..
631 } => Some((model, *idx)),
632 _ => None,
633 }
634 }
635
636 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
639 if self.has_prism_contract() {
640 return None;
641 }
642 match self {
643 Self::Mapped {
644 model,
645 idx,
646 dtype: TensorDtype::Q2TiledP,
647 ..
648 } => Some((model, *idx)),
649 _ => None,
650 }
651 }
652
653 pub fn cols(&self) -> usize {
654 match self {
655 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
656 }
657 }
658
659 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
662 if self.has_prism_contract() {
663 return None;
664 }
665 match self {
666 Self::Mapped {
667 model,
668 idx,
669 dtype: TensorDtype::Q1,
670 ..
671 } => Some((model, *idx)),
672 _ => None,
673 }
674 }
675
676 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
686 if self.has_prism_contract() {
687 return None;
688 }
689 self.graph_weight_descriptor()
690 }
691
692 pub(crate) fn graph_weight_descriptor(
697 &self,
698 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
699 match self {
700 Self::Mapped {
701 model,
702 idx,
703 dtype: TensorDtype::Q8Row,
704 row_scale,
705 ..
706 } => Some((model, *idx, 0, row_scale.as_slice())),
707 Self::Mapped {
708 model,
709 idx,
710 dtype: TensorDtype::Q1,
711 ..
712 } => Some((model, *idx, 1, &[])),
713 Self::Mapped {
718 model,
719 idx,
720 dtype: TensorDtype::Q4Tiled,
721 ..
722 } => Some((model, *idx, 5, &[])),
723 Self::Mapped {
727 model,
728 idx,
729 dtype: TensorDtype::Q4TiledP,
730 ..
731 } => Some((model, *idx, 6, &[])),
732 Self::Mapped {
733 model,
734 idx,
735 dtype: TensorDtype::Q4Block,
736 ..
737 } => Some((model, *idx, 2, &[])),
738 Self::Mapped {
743 model,
744 idx,
745 dtype: TensorDtype::Q8_2f,
746 ..
747 } => Some((model, *idx, 7, &[])),
748 Self::Mapped {
749 model,
750 idx,
751 dtype: TensorDtype::Q1T,
752 ..
753 } => Some((model, *idx, 3, &[])),
754 Self::Mapped {
758 model,
759 idx,
760 dtype: TensorDtype::Q2TiledP,
761 ..
762 } => Some((model, *idx, 9, &[])),
763 _ => None,
764 }
765 }
766
767 pub fn as_f32(&self) -> Option<&[f32]> {
770 match self {
771 Self::F32 { data, .. } => Some(data),
772 Self::Mapped { .. } => None,
773 }
774 }
775
776 fn quant_bytes(&self) -> &[u8] {
777 match self {
778 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
779 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
780 }
781 }
782
783 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
785 let cols = self.cols();
786 debug_assert_eq!(dst.len(), cols);
787 match self {
788 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
789 Self::Mapped {
790 model,
791 idx,
792 dtype,
793 row_scale,
794 col_field,
795 vbit_offsets,
796 ..
797 } => {
798 if *dtype == TensorDtype::Q4Tiled {
799 let bytes = self.quant_bytes();
800 let gpr = cols / GROUP_SIZE;
801 for gi in 0..gpr {
802 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
803 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
804 for (k, &b) in tile[2..].iter().enumerate() {
805 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
806 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
807 }
808 }
809 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
810 crate::prism::inverse_embedding(model, dst);
811 }
812 return;
813 }
814 if *dtype == TensorDtype::Q4TiledP {
815 let bytes = self.quant_bytes();
816 let gpr = cols / GROUP_SIZE;
817 let v = Q4tpView::new(bytes, self.rows(), cols);
818 let mut sc = vec![0f32; gpr];
819 v.scales_into(r, gpr, &mut sc);
820 for gi in 0..gpr {
821 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
822 let s = sc[gi];
823 for (k, &b) in tile.iter().enumerate() {
824 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
825 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
826 }
827 }
828 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
829 crate::prism::inverse_embedding(model, dst);
830 }
831 return;
832 }
833 if *dtype == TensorDtype::Q2TiledP {
834 let bytes = self.quant_bytes();
835 let gpr = cols / GROUP_SIZE;
836 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
837 let mut sc = vec![0f32; gpr];
838 v.scales_into(r, gpr, &mut sc);
839 for gi in 0..gpr {
840 let ch =
841 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
842 let s = sc[gi];
843 for (k, &b) in ch.iter().enumerate() {
844 for j in 0..4 {
845 let center = if crate::prism::is_affine_target(
846 model,
847 &model.tensors[*idx].name,
848 ) {
849 1.0
850 } else {
851 1.5
852 };
853 dst[gi * GROUP_SIZE + k * 4 + j] =
854 (((b >> (2 * j)) & 3) as f32 - center) * s;
855 }
856 }
857 }
858 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
859 crate::prism::inverse_embedding(model, dst);
860 }
861 return;
862 }
863 if *dtype == TensorDtype::Q4Block {
864 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
865 let gpr = cols / GROUP_SIZE;
866 for gi in 0..gpr {
867 let g = r * gpr + gi;
868 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
869 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
870 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
871 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
872 }
873 }
874 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
875 crate::prism::inverse_embedding(model, dst);
876 }
877 return;
878 }
879 if *dtype == TensorDtype::Q1 {
880 let bytes = self.quant_bytes();
881 let gpr = cols / GROUP_SIZE;
882 for gi in 0..gpr {
883 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
884 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
885 for (j, &b) in tile[2..].iter().enumerate() {
886 for k in 0..8 {
887 dst[gi * GROUP_SIZE + j * 8 + k] =
888 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
889 }
890 }
891 }
892 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
893 crate::prism::inverse_embedding(model, dst);
894 }
895 return;
896 }
897 if *dtype == TensorDtype::Q1T {
898 let bytes = self.quant_bytes();
899 let gpr = cols / GROUP_SIZE;
900 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
901 for gi in 0..gpr {
902 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
903 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
904 bytes[off],
905 bytes[off + 1],
906 ]));
907 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
908 for k in 0..GROUP_SIZE {
909 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
910 {
911 1 => s,
912 2 => -s,
913 _ => 0.0,
914 };
915 }
916 }
917 let rows = self.rows();
919 let entries = base_len + (rows + 1) * 4;
920 if entries <= bytes.len() {
921 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
922 let r0 = u32::from_le_bytes([
923 ptrs[r * 4],
924 ptrs[r * 4 + 1],
925 ptrs[r * 4 + 2],
926 ptrs[r * 4 + 3],
927 ]) as usize;
928 let r1 = u32::from_le_bytes([
929 ptrs[(r + 1) * 4],
930 ptrs[(r + 1) * 4 + 1],
931 ptrs[(r + 1) * 4 + 2],
932 ptrs[(r + 1) * 4 + 3],
933 ]) as usize;
934 let off = entries + r0 * 4;
935 for i in 0..r1 - r0 {
936 let item = &bytes[off + i * 4..off + i * 4 + 4];
937 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
938 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
939 item[2], item[3],
940 ]));
941 if c < cols {
942 dst[c] = v;
943 }
944 }
945 }
946 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
947 crate::prism::inverse_embedding(model, dst);
948 }
949 return;
950 }
951 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
952 let bytes = self.quant_bytes();
953 let rows = self.rows();
954 let ng = cols / GROUP_SIZE;
955 let bits = &bytes[..rows];
956 let sc_off = rows;
957 let off = vbit_offsets[r];
960 let b = bits[r] as usize;
961 let l = ((1usize << (b - 1)) - 1) as f32;
962 let data = &bytes[off..];
963 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
964 for (i, d) in dst.iter_mut().enumerate() {
965 while nbits < b {
966 acc = (acc << 8) | data[byte_idx] as u64;
967 byte_idx += 1;
968 nbits += 8;
969 }
970 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
971 nbits -= b;
972 let so = (r * ng + i / GROUP_SIZE) * 2;
973 let sv = f16_to_f32(u16::from_le_bytes([
974 bytes[sc_off + so],
975 bytes[sc_off + so + 1],
976 ]));
977 *d = (u - l) * sv;
978 }
979 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
980 crate::prism::inverse_embedding(model, dst);
981 }
982 return;
983 }
984 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
985 let s = row_scale[r];
986 match dtype {
987 TensorDtype::Q8Row => {
988 for (d, &b) in dst.iter_mut().zip(q) {
989 *d = (b as i8) as f32 * s;
990 }
991 }
992 TensorDtype::Q8_2f => {
993 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
994 *d = (b as i8) as f32 * s * col_field[i];
995 }
996 }
997 _ => unreachable!(),
998 }
999 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1000 crate::prism::inverse_embedding(model, dst);
1001 }
1002 }
1003 }
1004 }
1005
1006 pub fn sparse_col_ok(&self) -> bool {
1011 match self {
1012 Self::F32 { .. } => true,
1013 Self::Mapped { dtype, .. } => {
1014 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1015 }
1016 }
1017 }
1018
1019 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1023 let inter = self.cols();
1024 let hidden = self.rows();
1025 debug_assert_eq!(out.len(), hidden);
1026 match self {
1027 Self::F32 { data, .. } => {
1028 for (k, o) in out.iter_mut().enumerate() {
1029 *o += w * data[k * inter + c];
1030 }
1031 }
1032 Self::Mapped {
1033 dtype,
1034 row_scale,
1035 col_field,
1036 ..
1037 } => {
1038 let q = self.quant_bytes();
1039 let colf = if *dtype == TensorDtype::Q8_2f {
1040 col_field[c]
1041 } else {
1042 1.0
1043 };
1044 let wc = w * colf;
1045 for (k, o) in out.iter_mut().enumerate() {
1046 let b = q[k * inter + c] as i8 as f32;
1047 *o += wc * b * row_scale[k];
1048 }
1049 }
1050 }
1051 }
1052
1053 #[inline]
1062 pub fn prefetch_row(&self, r: usize) {
1063 let Self::Mapped { dtype, .. } = self else {
1064 return;
1065 };
1066 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1067 return;
1068 }
1069 let cols = self.cols();
1070 let q = self.quant_bytes();
1071 let (a, b) = (r * cols, (r + 1) * cols);
1072 if b > q.len() {
1073 return;
1074 }
1075 let mut j = a;
1076 while j < b {
1077 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1078 j += 512;
1079 }
1080 }
1081
1082 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1091 let cols = self.cols();
1092 debug_assert_eq!(out.len(), cols);
1093 match self {
1094 Self::F32 { data, .. } => {
1095 let row = &data[r * cols..(r + 1) * cols];
1096 for (o, v) in out.iter_mut().zip(row) {
1097 *o += w * v;
1098 }
1099 }
1100 Self::Mapped {
1101 dtype,
1102 row_scale,
1103 col_field,
1104 ..
1105 } => match dtype {
1106 TensorDtype::Q8Row => {
1107 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1108 let ws = w * row_scale[r];
1109 let row: &[i8] =
1110 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1111 axpy_i8_f32(out, row, ws);
1112 }
1113 TensorDtype::Q8_2f => {
1114 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1115 let ws = w * row_scale[r];
1116 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1117 *o += ws * c * (*b as i8 as f32);
1118 }
1119 }
1120 _ => {
1121 self.row_f32(r, scratch);
1122 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1123 *o += w * v;
1124 }
1125 }
1126 },
1127 }
1128 }
1129
1130 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1134 let cols = self.cols();
1135 match self {
1136 Self::F32 { data, .. } => {
1137 let row = &data[r * cols..(r + 1) * cols];
1138 row.iter().zip(x).map(|(w, v)| w * v).sum()
1139 }
1140 Self::Mapped {
1141 model,
1142 idx,
1143 dtype,
1144 row_scale,
1145 col_field,
1146 ..
1147 } => {
1148 let prism_forward =
1149 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1150 if prism_forward {
1151 let transformed = crate::prism::forward(model, &x[..cols]);
1152 let gpr = cols / GROUP_SIZE;
1153 match dtype {
1154 TensorDtype::Q2TiledP => {
1155 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1156 let mut sc = vec![0f32; gpr];
1157 v.scales_into(r, gpr, &mut sc);
1158 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1159 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1160 }
1161 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1162 }
1163 _ => {
1164 self.row_f32(r, scratch);
1165 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1166 }
1167 }
1168 }
1169 match dtype {
1170 TensorDtype::Q8Row => {
1171 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1172 dot_i8_f32(q, x) * row_scale[r]
1173 }
1174 TensorDtype::Q8_2f => {
1175 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1176 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1177 }
1178 _ => {
1179 self.row_f32(r, scratch);
1180 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1181 }
1182 }
1183 }
1184 }
1185 }
1186
1187 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1190 match self {
1191 Self::F32 { data, .. } => {
1195 if !crate::f32_backend::matvec(data, x, out) {
1196 matvec_rows(pool, data, x, out);
1197 }
1198 }
1199 Self::Mapped {
1200 model,
1201 idx,
1202 dtype,
1203 rows,
1204 cols,
1205 row_scale,
1206 col_field,
1207 vbit_offsets,
1208 repack,
1209 } => {
1210 let _ = (model, idx);
1211 assert!(
1221 out.len() >= *rows && x.len() >= *cols,
1222 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1223 out.len(),
1224 x.len(),
1225 );
1226 let prism_forward =
1227 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1228 if *dtype == TensorDtype::Q2TiledP
1229 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1230 {
1231 use std::sync::atomic::{AtomicUsize, Ordering};
1232 static N: AtomicUsize = AtomicUsize::new(0);
1233 let n = N.fetch_add(1, Ordering::Relaxed);
1234 if n < 128 {
1235 eprintln!(
1236 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1237 model.tensors[*idx].name,
1238 prism_forward,
1239 rows,
1240 cols,
1241 crate::gpu::enabled_here(),
1242 crate::gpu::q2tp_gpu_opt_in(),
1243 crate::gpu::cur_layer(),
1244 );
1245 }
1246 }
1247 if prism_forward {
1252 let transformed = crate::prism::forward(model, &x[..*cols]);
1253 match dtype {
1254 TensorDtype::Q4Block => {
1255 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1256 }
1257 TensorDtype::Q4Tiled => {
1258 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1259 }
1260 TensorDtype::Q4TiledP => {
1261 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1262 }
1263 TensorDtype::Q2TiledP => {
1264 let affine =
1265 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1266 if *rows * *cols >= 8_388_608
1267 && crate::gpu::enabled_here()
1268 && crate::gpu::q2tp_gpu_opt_in()
1269 {
1270 let gpu_ok = if affine {
1271 crate::gpu::q2tp_affine_matvec(
1272 model,
1273 *idx,
1274 &transformed,
1275 *rows,
1276 *cols,
1277 out,
1278 )
1279 } else {
1280 crate::gpu::q2tp_matvec(
1281 model,
1282 *idx,
1283 &transformed,
1284 *rows,
1285 *cols,
1286 out,
1287 )
1288 };
1289 if gpu_ok {
1290 return;
1291 }
1292 }
1293 if affine {
1294 q2tp_affine_matvec(
1295 self.quant_bytes(),
1296 &transformed,
1297 *rows,
1298 *cols,
1299 out,
1300 pool,
1301 )
1302 } else {
1303 q2tp_matvec(
1304 self.quant_bytes(),
1305 &transformed,
1306 *rows,
1307 *cols,
1308 out,
1309 pool,
1310 )
1311 }
1312 }
1313 TensorDtype::Q1 => {
1314 q1_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1315 }
1316 TensorDtype::Q1T => {
1317 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1318 }
1319 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1320 self.quant_bytes(),
1321 vbit_offsets,
1322 &transformed,
1323 *rows,
1324 *cols,
1325 out,
1326 pool,
1327 ),
1328 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1329 self.quant_bytes(),
1330 repack,
1331 row_scale,
1332 &transformed,
1333 col_field,
1334 *dtype,
1335 *rows,
1336 *cols,
1337 out,
1338 pool,
1339 ),
1340 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1341 }
1342 return;
1343 }
1344 if *dtype == TensorDtype::Q4Block {
1345 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1349 let t0 = std::time::Instant::now();
1350 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1351 crate::gpu::ProbeArm::Gpu => {
1352 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1353 crate::gpu::probe_record(
1354 crate::gpu::OpClass::Matvec,
1355 true,
1356 t0.elapsed(),
1357 );
1358 return;
1359 }
1360 }
1361 crate::gpu::ProbeArm::CpuTimed => {
1362 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1363 crate::gpu::probe_record(
1364 crate::gpu::OpClass::Matvec,
1365 false,
1366 t0.elapsed(),
1367 );
1368 return;
1369 }
1370 crate::gpu::ProbeArm::Cpu => {}
1371 }
1372 }
1373 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1374 return;
1375 }
1376 if *dtype == TensorDtype::Q4Tiled {
1377 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1382 let t0 = std::time::Instant::now();
1383 let cls = crate::gpu::matvec_class(*rows, *cols);
1384 match crate::gpu::probe_arm(cls) {
1385 crate::gpu::ProbeArm::Gpu => {
1386 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1387 crate::gpu::probe_record(cls, true, t0.elapsed());
1388 return;
1389 }
1390 }
1391 crate::gpu::ProbeArm::CpuTimed => {
1392 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1393 crate::gpu::probe_record(cls, false, t0.elapsed());
1394 return;
1395 }
1396 crate::gpu::ProbeArm::Cpu => {}
1397 }
1398 }
1399 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1400 return;
1401 }
1402 if *dtype == TensorDtype::Q4TiledP {
1403 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1409 let t0 = std::time::Instant::now();
1410 let cls = crate::gpu::matvec_class(*rows, *cols);
1411 match crate::gpu::probe_arm(cls) {
1412 crate::gpu::ProbeArm::Gpu => {
1413 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1414 crate::gpu::probe_record(cls, true, t0.elapsed());
1415 return;
1416 }
1417 }
1418 crate::gpu::ProbeArm::CpuTimed => {
1419 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1420 crate::gpu::probe_record(cls, false, t0.elapsed());
1421 return;
1422 }
1423 crate::gpu::ProbeArm::Cpu => {}
1424 }
1425 }
1426 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1427 return;
1428 }
1429 if *dtype == TensorDtype::Q2TiledP {
1430 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1431 return;
1432 }
1433 if *dtype == TensorDtype::Q1 {
1434 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1439 let t0 = std::time::Instant::now();
1440 let arm = if crate::gpu::q1_force() {
1441 crate::gpu::ProbeArm::Gpu
1442 } else {
1443 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1444 };
1445 match arm {
1446 crate::gpu::ProbeArm::Gpu => {
1447 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1448 crate::gpu::probe_record(
1449 crate::gpu::OpClass::Matvec,
1450 true,
1451 t0.elapsed(),
1452 );
1453 return;
1454 }
1455 }
1456 crate::gpu::ProbeArm::CpuTimed => {
1457 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1458 crate::gpu::probe_record(
1459 crate::gpu::OpClass::Matvec,
1460 false,
1461 t0.elapsed(),
1462 );
1463 return;
1464 }
1465 crate::gpu::ProbeArm::Cpu => {}
1466 }
1467 }
1468 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1469 return;
1470 }
1471 if *dtype == TensorDtype::Q1T {
1472 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1476 let t0 = std::time::Instant::now();
1477 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1478 crate::gpu::ProbeArm::Gpu => {
1479 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1480 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1481 crate::gpu::probe_record(
1482 crate::gpu::OpClass::Matvec,
1483 true,
1484 t0.elapsed(),
1485 );
1486 return;
1487 }
1488 }
1489 crate::gpu::ProbeArm::CpuTimed => {
1490 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1491 crate::gpu::probe_record(
1492 crate::gpu::OpClass::Matvec,
1493 false,
1494 t0.elapsed(),
1495 );
1496 return;
1497 }
1498 crate::gpu::ProbeArm::Cpu => {}
1499 }
1500 }
1501 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1502 return;
1503 }
1504 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1505 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1506 return;
1507 }
1508 let xs = prescale(x, col_field, *dtype);
1509 if *rows >= crate::gpu::min_rows()
1514 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1515 && gpu_lmhead_enabled()
1516 && crate::gpu::enabled_here()
1517 {
1518 let t0 = std::time::Instant::now();
1521 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1522 crate::gpu::ProbeArm::Gpu => {}
1523 crate::gpu::ProbeArm::CpuTimed => {
1524 qmatvec(
1525 self.quant_bytes(),
1526 repack,
1527 row_scale,
1528 x,
1529 col_field,
1530 *dtype,
1531 *rows,
1532 *cols,
1533 out,
1534 pool,
1535 );
1536 crate::gpu::probe_record(
1537 crate::gpu::OpClass::Matvec,
1538 false,
1539 t0.elapsed(),
1540 );
1541 return;
1542 }
1543 crate::gpu::ProbeArm::Cpu => {
1544 qmatvec(
1545 self.quant_bytes(),
1546 repack,
1547 row_scale,
1548 x,
1549 col_field,
1550 *dtype,
1551 *rows,
1552 *cols,
1553 out,
1554 pool,
1555 );
1556 return;
1557 }
1558 }
1559 let frac = gpu_split_frac();
1560 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1561 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1562 let bytes = self.quant_bytes();
1563 let ok = std::thread::scope(|sc| {
1564 let g = sc.spawn(|| {
1565 crate::gpu::q8_matvec_range(
1566 model,
1567 *idx,
1568 cpu_rows,
1569 &row_scale[cpu_rows..],
1570 &xs,
1571 *rows - cpu_rows,
1572 *cols,
1573 out_gpu,
1574 )
1575 });
1576 if cpu_rows > 0 {
1577 let rep_cpu = if repack.is_empty() {
1580 &[][..]
1581 } else {
1582 &repack[..(cpu_rows / 4) * 4 * *cols]
1583 };
1584 qmatvec(
1585 &bytes[..cpu_rows * *cols],
1586 rep_cpu,
1587 &row_scale[..cpu_rows],
1588 x,
1589 col_field,
1590 *dtype,
1591 cpu_rows,
1592 *cols,
1593 out_cpu,
1594 pool,
1595 );
1596 }
1597 g.join().unwrap_or(false)
1598 });
1599 if ok {
1600 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1601 return;
1602 }
1603 qmatvec(
1606 &bytes[cpu_rows * *cols..(*rows) * *cols],
1607 &[],
1608 &row_scale[cpu_rows..],
1609 x,
1610 col_field,
1611 *dtype,
1612 *rows - cpu_rows,
1613 *cols,
1614 out_gpu,
1615 pool,
1616 );
1617 return;
1618 }
1619 qmatvec(
1620 self.quant_bytes(),
1621 repack,
1622 row_scale,
1623 x,
1624 col_field,
1625 *dtype,
1626 *rows,
1627 *cols,
1628 out,
1629 pool,
1630 );
1631 }
1632 }
1633 }
1634
1635 pub fn matvec2(
1637 &self,
1638 x1: &[f32],
1639 x2: &[f32],
1640 o1: &mut [f32],
1641 o2: &mut [f32],
1642 pool: Option<&Pool>,
1643 ) {
1644 match self {
1645 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1646 Self::Mapped {
1647 model,
1648 idx,
1649 dtype,
1650 rows,
1651 cols,
1652 row_scale,
1653 col_field,
1654 vbit_offsets,
1655 ..
1656 } => {
1657 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1658 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1659 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1660 match dtype {
1661 TensorDtype::Q4Block => {
1662 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1663 }
1664 TensorDtype::Q4Tiled => {
1665 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1666 }
1667 TensorDtype::Q4TiledP => {
1668 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1669 }
1670 TensorDtype::Q2TiledP => {
1671 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1672 q2tp_affine_matvec2(
1673 self.quant_bytes(),
1674 &tx1,
1675 &tx2,
1676 *rows,
1677 *cols,
1678 o1,
1679 o2,
1680 pool,
1681 )
1682 } else {
1683 q2tp_matvec2(
1684 self.quant_bytes(),
1685 &tx1,
1686 &tx2,
1687 *rows,
1688 *cols,
1689 o1,
1690 o2,
1691 pool,
1692 )
1693 }
1694 }
1695 TensorDtype::Q1 => {
1696 q1_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1697 }
1698 TensorDtype::Q1T => {
1699 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1700 }
1701 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1702 self.quant_bytes(),
1703 vbit_offsets,
1704 &tx1,
1705 &tx2,
1706 *rows,
1707 *cols,
1708 o1,
1709 o2,
1710 pool,
1711 ),
1712 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1713 self.quant_bytes(),
1714 row_scale,
1715 &tx1,
1716 &tx2,
1717 col_field,
1718 *dtype,
1719 *rows,
1720 *cols,
1721 o1,
1722 o2,
1723 pool,
1724 ),
1725 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1726 }
1727 return;
1728 }
1729 if *dtype == TensorDtype::Q4Block {
1730 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1731 return;
1732 }
1733 if *dtype == TensorDtype::Q4Tiled {
1734 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1735 return;
1736 }
1737 if *dtype == TensorDtype::Q4TiledP {
1738 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1739 return;
1740 }
1741 if *dtype == TensorDtype::Q2TiledP {
1742 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1743 return;
1744 }
1745 if *dtype == TensorDtype::Q1 {
1746 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1747 return;
1748 }
1749 if *dtype == TensorDtype::Q1T {
1750 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1756 return;
1757 }
1758 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1759 vbitmatvec2(
1760 self.quant_bytes(),
1761 vbit_offsets,
1762 x1,
1763 x2,
1764 *rows,
1765 *cols,
1766 o1,
1767 o2,
1768 pool,
1769 );
1770 return;
1771 }
1772 qmatvec2(
1773 self.quant_bytes(),
1774 row_scale,
1775 x1,
1776 x2,
1777 col_field,
1778 *dtype,
1779 *rows,
1780 *cols,
1781 o1,
1782 o2,
1783 pool,
1784 );
1785 }
1786 }
1787 }
1788}
1789
1790impl QTensor {
1791 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1799 if self.has_prism_contract() {
1800 return None;
1801 }
1802 match self {
1803 Self::Mapped {
1804 model, idx, dtype, ..
1805 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1806 _ => None,
1807 }
1808 }
1809
1810 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1811 let cols = self.cols();
1812 let rows = self.rows();
1813 debug_assert_eq!(xs_all.len(), b * cols);
1814 debug_assert_eq!(out.len(), b * rows);
1815 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1816 if crate::gptq_capture::capturing() {
1820 if let Self::Mapped { model, idx, .. } = self {
1821 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1822 }
1823 }
1824 match self {
1825 Self::F32 { data, .. } => {
1826 if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
1827 return;
1828 }
1829 let out_addr = SendMut(out.as_mut_ptr());
1830 let run = |start: usize, end: usize| {
1837 let mut bi = start;
1838 while bi < end {
1839 let n = (end - bi).min(4);
1840 for o in 0..rows {
1841 let row = &data[o * cols..(o + 1) * cols];
1842 if n == 4 {
1843 let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
1844 let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
1845 let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
1846 for j in 0..cols {
1847 let w = row[j];
1848 a0 += w * x0[j];
1849 a1 += w * x1[j];
1850 a2 += w * x2[j];
1851 a3 += w * x3[j];
1852 }
1853 for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
1854 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
1855 }
1856 } else {
1857 for k in 0..n {
1858 let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
1859 let mut acc = 0f32;
1860 for j in 0..cols {
1861 acc += row[j] * x[j];
1862 }
1863 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
1864 }
1865 }
1866 }
1867 bi += n;
1868 }
1869 };
1870 dispatch_rows(pool, b, &run);
1871 }
1872 Self::Mapped {
1873 model,
1874 idx,
1875 dtype,
1876 row_scale,
1877 col_field,
1878 vbit_offsets,
1879 ..
1880 } => {
1881 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1882 let mut transformed = Vec::with_capacity(xs_all.len());
1883 for bi in 0..b {
1884 transformed.extend_from_slice(&crate::prism::forward(
1885 model,
1886 &xs_all[bi * cols..(bi + 1) * cols],
1887 ));
1888 }
1889 match dtype {
1890 TensorDtype::Q4Block => {
1891 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1892 }
1893 TensorDtype::Q4Tiled => {
1894 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1895 }
1896 TensorDtype::Q4TiledP => {
1897 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1898 }
1899 TensorDtype::Q2TiledP => {
1900 let affine =
1901 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1902 let gpu_batch_ok = if affine {
1908 b >= 2
1909 } else {
1910 b >= 32 && b * rows * cols >= 128_000_000
1911 };
1912 if gpu_batch_ok
1913 && cols % 32 == 0
1914 && crate::gpu::enabled_here()
1915 && crate::gpu::q2tp_gpu_opt_in()
1916 {
1917 let gpu_ok = if affine {
1918 crate::gpu::q2tp_affine_matmat(
1919 model,
1920 *idx,
1921 &transformed,
1922 b,
1923 rows,
1924 cols,
1925 out,
1926 )
1927 } else {
1928 crate::gpu::q2tp_matmat(
1929 model,
1930 *idx,
1931 &transformed,
1932 b,
1933 rows,
1934 cols,
1935 out,
1936 )
1937 };
1938 if gpu_ok {
1939 return;
1940 }
1941 }
1942 if affine
1946 && b == 1
1947 && cols % 32 == 0
1948 && crate::gpu::enabled_here()
1949 && crate::gpu::q2tp_gpu_opt_in()
1950 && crate::gpu::q2tp_affine_matvec(
1951 model,
1952 *idx,
1953 &transformed[..cols],
1954 rows,
1955 cols,
1956 &mut out[..rows],
1957 )
1958 {
1959 return;
1960 }
1961 if affine {
1962 q2tp_affine_matmat(
1963 self.quant_bytes(),
1964 &transformed,
1965 b,
1966 rows,
1967 cols,
1968 out,
1969 pool,
1970 )
1971 } else {
1972 q2tp_matmat(
1973 self.quant_bytes(),
1974 &transformed,
1975 b,
1976 rows,
1977 cols,
1978 out,
1979 pool,
1980 )
1981 }
1982 }
1983 TensorDtype::Q1 => {
1984 q1_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1985 }
1986 TensorDtype::Q1T => {
1987 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1988 }
1989 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
1990 self.quant_bytes(),
1991 vbit_offsets,
1992 &transformed,
1993 b,
1994 rows,
1995 cols,
1996 out,
1997 pool,
1998 ),
1999 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2000 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2001 .map(|bi| {
2002 prescale(
2003 &transformed[bi * cols..(bi + 1) * cols],
2004 col_field,
2005 *dtype,
2006 )
2007 })
2008 .collect();
2009 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2010 }
2011 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2012 }
2013 return;
2014 }
2015 if *dtype == TensorDtype::Q4Block {
2016 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2017 return;
2018 }
2019 if *dtype == TensorDtype::Q4TiledP {
2020 if b >= 32
2034 && b * rows * cols >= 128_000_000
2035 && cols % 32 == 0
2036 && !row_exact()
2037 && !crate::gpu::mm_killed()
2038 && crate::gpu::enabled_here()
2039 {
2040 let class = if b >= 128 {
2041 crate::gpu::OpClass::MatmatWide
2042 } else {
2043 crate::gpu::OpClass::Matmat
2044 };
2045 if let Self::Mapped { model, idx, .. } = self {
2046 if crate::mm_ab::on() {
2058 let mut g = vec![0f32; b * rows];
2059 let t = std::time::Instant::now();
2060 let took = crate::gpu::q4tp_matmat(
2061 model, *idx, xs_all, b, rows, cols, &mut g,
2062 );
2063 let dg = t.elapsed();
2064 let t = std::time::Instant::now();
2065 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2066 let dc = t.elapsed();
2067 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2068 return;
2069 }
2070 let t0 = std::time::Instant::now();
2071 let resident = crate::gpu::weight_is_resident(model, *idx);
2075 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2076 crate::gpu::ProbeArm::Gpu => {
2077 if crate::gpu::q4tp_matmat(
2078 model, *idx, xs_all, b, rows, cols, out,
2079 ) {
2080 let el = t0.elapsed();
2081 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2091 let budget = std::time::Duration::from_secs_f64(
2092 flops / 1.5e12 * 8.0 + 0.020,
2093 );
2094 crate::gpu::mm_budget_check(
2095 "q4tp matmat",
2096 el,
2097 budget,
2098 crate::gpu::probe_was_cold() || !resident,
2099 );
2100 crate::gpu::probe_record(class, true, el);
2101 return;
2102 }
2103 }
2104 crate::gpu::ProbeArm::CpuTimed => {
2105 q4tp_matmat(
2106 self.quant_bytes(),
2107 xs_all,
2108 b,
2109 rows,
2110 cols,
2111 out,
2112 pool,
2113 );
2114 crate::gpu::probe_record(class, false, t0.elapsed());
2115 return;
2116 }
2117 crate::gpu::ProbeArm::Cpu => {}
2118 }
2119 }
2120 }
2121 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2122 return;
2123 }
2124 if *dtype == TensorDtype::Q2TiledP {
2125 if b >= 32
2131 && b * rows * cols >= 128_000_000
2132 && cols % 32 == 0
2133 && !crate::gpu::mm_killed()
2134 && crate::gpu::enabled_here()
2135 {
2136 let class = if b >= 128 {
2137 crate::gpu::OpClass::MatmatWide
2138 } else {
2139 crate::gpu::OpClass::Matmat
2140 };
2141 if let Self::Mapped { model, idx, .. } = self {
2142 let t0 = std::time::Instant::now();
2143 match crate::gpu::probe_arm(class) {
2144 crate::gpu::ProbeArm::Gpu => {
2145 if crate::gpu::q2tp_matmat(
2146 model, *idx, xs_all, b, rows, cols, out,
2147 ) {
2148 crate::gpu::probe_record(class, true, t0.elapsed());
2149 return;
2150 }
2151 }
2152 crate::gpu::ProbeArm::CpuTimed => {
2153 q2tp_matmat(
2154 self.quant_bytes(),
2155 xs_all,
2156 b,
2157 rows,
2158 cols,
2159 out,
2160 pool,
2161 );
2162 crate::gpu::probe_record(class, false, t0.elapsed());
2163 return;
2164 }
2165 crate::gpu::ProbeArm::Cpu => {}
2166 }
2167 }
2168 }
2169 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2174 return;
2175 }
2176 if *dtype == TensorDtype::Q4Tiled {
2177 if b >= 32
2189 && b * rows * cols >= 128_000_000
2190 && cols % 32 == 0
2191 && !crate::gpu::mm_killed()
2192 && crate::gpu::enabled_here()
2193 {
2194 let class = if b >= 128 {
2195 crate::gpu::OpClass::MatmatWide
2196 } else {
2197 crate::gpu::OpClass::Matmat
2198 };
2199 if let Self::Mapped { model, idx, .. } = self {
2200 let t0 = std::time::Instant::now();
2201 match crate::gpu::probe_arm(class) {
2202 crate::gpu::ProbeArm::Gpu => {
2203 if crate::gpu::q4t_matmat(
2204 model, *idx, xs_all, b, rows, cols, out,
2205 ) {
2206 let el = t0.elapsed();
2207 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2217 let budget = std::time::Duration::from_secs_f64(
2218 flops / 1.5e12 * 8.0 + 0.020,
2219 );
2220 crate::gpu::mm_budget_check(
2221 "q4t matmat",
2222 el,
2223 budget,
2224 crate::gpu::probe_was_cold(),
2225 );
2226 crate::gpu::probe_record(class, true, el);
2227 return;
2228 }
2229 }
2230 crate::gpu::ProbeArm::CpuTimed => {
2231 q4t_matmat(
2232 self.quant_bytes(),
2233 xs_all,
2234 b,
2235 rows,
2236 cols,
2237 out,
2238 pool,
2239 );
2240 crate::gpu::probe_record(class, false, t0.elapsed());
2241 return;
2242 }
2243 crate::gpu::ProbeArm::Cpu => {}
2244 }
2245 }
2246 }
2247 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2248 return;
2249 }
2250 if *dtype == TensorDtype::Q1 {
2251 if b >= 32
2254 && b * rows * cols >= 128_000_000
2255 && cols % 64 == 0
2256 && crate::gpu::enabled_here()
2257 {
2258 if let Self::Mapped { model, idx, .. } = self {
2259 let t0 = std::time::Instant::now();
2260 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2261 crate::gpu::ProbeArm::Gpu => {
2262 if crate::gpu::q1_matmat(
2263 model, *idx, xs_all, b, rows, cols, out,
2264 ) {
2265 crate::gpu::probe_record(
2266 crate::gpu::OpClass::Matmat,
2267 true,
2268 t0.elapsed(),
2269 );
2270 return;
2271 }
2272 }
2273 crate::gpu::ProbeArm::CpuTimed => {
2274 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2275 crate::gpu::probe_record(
2276 crate::gpu::OpClass::Matmat,
2277 false,
2278 t0.elapsed(),
2279 );
2280 return;
2281 }
2282 crate::gpu::ProbeArm::Cpu => {}
2283 }
2284 }
2285 }
2286 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2287 return;
2288 }
2289 if *dtype == TensorDtype::Q1T {
2290 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2293 if let Self::Mapped { model, idx, .. } = self {
2294 let t0 = std::time::Instant::now();
2295 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2296 crate::gpu::ProbeArm::Gpu => {
2297 if crate::gpu::q1t_matmat(
2298 model, *idx, xs_all, b, rows, cols, out,
2299 ) {
2300 crate::gpu::probe_record(
2301 crate::gpu::OpClass::Matmat,
2302 true,
2303 t0.elapsed(),
2304 );
2305 return;
2306 }
2307 }
2308 crate::gpu::ProbeArm::CpuTimed => {
2309 q1t_matmat(
2310 self.quant_bytes(),
2311 xs_all,
2312 b,
2313 rows,
2314 cols,
2315 out,
2316 pool,
2317 );
2318 crate::gpu::probe_record(
2319 crate::gpu::OpClass::Matmat,
2320 false,
2321 t0.elapsed(),
2322 );
2323 return;
2324 }
2325 crate::gpu::ProbeArm::Cpu => {}
2326 }
2327 }
2328 }
2329 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2330 return;
2331 }
2332 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2333 vbitmatmat(
2334 self.quant_bytes(),
2335 vbit_offsets,
2336 xs_all,
2337 b,
2338 rows,
2339 cols,
2340 out,
2341 pool,
2342 );
2343 return;
2344 }
2345 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2346 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2347 .collect();
2348 if row_exact()
2354 && (1..=4).contains(&b)
2355 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2356 && crate::gpu::enabled_here()
2357 && crate::gpu::wgpu_active()
2358 {
2359 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2360 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2361 return;
2362 }
2363 }
2364 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2370 if let Self::Mapped { model, idx, .. } = self {
2371 let t0 = std::time::Instant::now();
2372 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2373 crate::gpu::ProbeArm::Gpu
2374 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2375 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2376 {
2377 let q = self.quant_bytes();
2381 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2382 return;
2383 }
2384 crate::gpu::ProbeArm::Gpu => {
2385 if *dtype == TensorDtype::Q8_2f
2393 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2394 && crate::gpu::q8_matmat_2f(
2395 model, *idx, row_scale, col_field, xs_all, b, rows,
2396 cols, out,
2397 )
2398 {
2399 crate::gpu::probe_record(
2400 crate::gpu::OpClass::Matmat,
2401 true,
2402 t0.elapsed(),
2403 );
2404 return;
2405 }
2406 let flat: Vec<f32> =
2407 pre.iter().flat_map(|v| v.iter().copied()).collect();
2408 if crate::gpu::q8_matmat(
2409 model, *idx, row_scale, &flat, b, rows, cols, out,
2410 ) {
2411 crate::gpu::probe_record(
2412 crate::gpu::OpClass::Matmat,
2413 true,
2414 t0.elapsed(),
2415 );
2416 return;
2417 }
2418 }
2419 crate::gpu::ProbeArm::CpuTimed => {
2420 let q = self.quant_bytes();
2421 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2422 crate::gpu::probe_record(
2423 crate::gpu::OpClass::Matmat,
2424 false,
2425 t0.elapsed(),
2426 );
2427 return;
2428 }
2429 crate::gpu::ProbeArm::Cpu => {}
2430 }
2431 }
2432 }
2433 let q = self.quant_bytes();
2434 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2435 }
2436 }
2437 }
2438}
2439
2440impl QTensor {
2441 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2451 let (rows, cols) = (self.rows(), self.cols());
2452 let Self::Mapped {
2453 model,
2454 idx,
2455 dtype,
2456 row_scale,
2457 col_field,
2458 ..
2459 } = self
2460 else {
2461 return false;
2462 };
2463 if crate::prism::has_contract(model) {
2464 return false;
2465 }
2466 match *dtype {
2467 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2468 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2472 if *dtype == TensorDtype::Q8_2f
2475 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2476 && crate::gpu::q8_matmat_2f(
2477 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2478 )
2479 {
2480 return true;
2481 }
2482 let flat: Vec<f32> = (0..b)
2483 .flat_map(|bi| {
2484 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2485 })
2486 .collect();
2487 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2488 }
2489 _ => false,
2490 }
2491 }
2492
2493 pub fn matvec_many<const N: usize>(
2500 ts: [&QTensor; N],
2501 x: &[f32],
2502 mut outs: [&mut [f32]; N],
2503 pool: Option<&Pool>,
2504 ) {
2505 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2506 if ts.iter().any(|t| t.has_prism_contract()) {
2507 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2511 t.matvec(x, o, pool);
2512 }
2513 return;
2514 }
2515 let uniform_q8 = ts.iter().all(|t| {
2516 matches!(
2517 t,
2518 Self::Mapped {
2519 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2520 ..
2521 }
2522 )
2523 });
2524 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2525 if uniform_f32 && crate::f32_backend::active() {
2526 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2527 t.matvec(x, o, pool);
2528 }
2529 return;
2530 }
2531 let uniform_q4 = ts.iter().all(|t| {
2532 matches!(
2533 t,
2534 Self::Mapped {
2535 dtype: TensorDtype::Q4Block,
2536 ..
2537 }
2538 )
2539 });
2540 let uniform_vbit = ts.iter().all(|t| {
2541 matches!(
2542 t,
2543 Self::Mapped {
2544 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2545 ..
2546 }
2547 )
2548 });
2549 let uniform_q1 = ts.iter().all(|t| {
2550 matches!(
2551 t,
2552 Self::Mapped {
2553 dtype: TensorDtype::Q1,
2554 ..
2555 }
2556 )
2557 });
2558 let uniform_q1t = ts.iter().all(|t| {
2559 matches!(
2560 t,
2561 Self::Mapped {
2562 dtype: TensorDtype::Q1T,
2563 ..
2564 }
2565 )
2566 });
2567 let uniform_q4tp = ts.iter().all(|t| {
2572 matches!(
2573 t,
2574 Self::Mapped {
2575 dtype: TensorDtype::Q4TiledP,
2576 ..
2577 }
2578 )
2579 }) && ts
2580 .iter()
2581 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2582 let Some(pool) = pool else {
2583 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2584 t.matvec(x, o, None);
2585 }
2586 return;
2587 };
2588 if total_rows < 256
2589 || !(uniform_q8
2590 || uniform_f32
2591 || uniform_q4
2592 || uniform_vbit
2593 || uniform_q1
2594 || uniform_q1t
2595 || uniform_q4tp)
2596 {
2597 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2598 t.matvec(x, o, Some(pool));
2599 }
2600 return;
2601 }
2602
2603 if uniform_q4tp {
2604 let cols = ts[0].cols();
2610 let gpr = cols / GROUP_SIZE;
2611 let views: [Q4tpView; N] =
2612 std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2613 let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2614 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2615 let locate = |flat: usize| -> (usize, usize) {
2617 let mut acc = 0;
2618 for (i, &r) in rows_of.iter().enumerate() {
2619 if flat < acc + r {
2620 return (i, flat - acc);
2621 }
2622 acc += r;
2623 }
2624 (rows_of.len() - 1, 0)
2625 };
2626 let (views, outs_addr) = (&views, &outs_addr);
2627 if a8w8_enabled() {
2628 let act = split_act(x);
2629 let act = &act;
2630 let run = |start: usize, end: usize| {
2631 with_krow(gpr, |sc| {
2632 for flat in start..end {
2633 let (t, r) = locate(flat);
2634 let v = &views[t];
2635 v.scales_into(r, gpr, sc);
2636 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2637 for &(j, xv) in &act.outliers {
2638 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2639 acc += w * s * xv;
2640 }
2641 unsafe { *outs_addr[t].at(r) = acc };
2643 }
2644 });
2645 };
2646 pool.run_rows(total_rows, &run);
2647 } else {
2648 let run = |start: usize, end: usize| {
2649 with_krow(gpr, |sc| {
2650 for flat in start..end {
2651 let (t, r) = locate(flat);
2652 let v = &views[t];
2653 v.scales_into(r, gpr, sc);
2654 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2656 }
2657 });
2658 };
2659 pool.run_rows(total_rows, &run);
2660 }
2661 return;
2662 }
2663
2664 if uniform_q1 {
2665 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2668 if a8w8_enabled() {
2669 let act = split_act(x);
2670 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2671 let (act, gsum) = (&act, &gsum);
2672 let closures: [_; N] = std::array::from_fn(|i| {
2673 let (bytes, gpr, out) =
2674 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2675 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2676 });
2677 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2678 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2679 pool.run_many(&parts);
2680 } else {
2681 let closures: [_; N] = std::array::from_fn(|i| {
2682 let (bytes, gpr, out) =
2683 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2684 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2685 });
2686 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2687 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2688 pool.run_many(&parts);
2689 }
2690 return;
2691 }
2692
2693 if uniform_q1t {
2694 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2698 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2699 if a8w8_enabled() {
2700 let act = split_act(x);
2701 let act = &act;
2702 let x_ref = x;
2703 let closures: [_; N] = std::array::from_fn(|i| {
2704 let bytes = ts[i].quant_bytes();
2705 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2706 let gpr = cols / GROUP_SIZE;
2707 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2708 let out = outs_addr[i];
2709 move |s: usize, e: usize| {
2710 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2711 }
2712 });
2713 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2714 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2715 pool.run_many(&parts);
2716 } else {
2717 let x_ref = x;
2718 let closures: [_; N] = std::array::from_fn(|i| {
2719 let bytes = ts[i].quant_bytes();
2720 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2721 let gpr = cols / GROUP_SIZE;
2722 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2723 let out = outs_addr[i];
2724 move |s: usize, e: usize| {
2725 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2726 }
2727 });
2728 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2729 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2730 pool.run_many(&parts);
2731 }
2732 return;
2733 }
2734
2735 if uniform_q4 || uniform_vbit {
2736 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2737 if a8w8_enabled() {
2739 let act = split_act(x);
2740 let act = &act;
2741 if uniform_q4 {
2742 let closures: [_; N] = std::array::from_fn(|i| {
2743 let (packed, scales) =
2744 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2745 let (gpr, cols, out) =
2746 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2747 move |s: usize, e: usize| {
2748 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2749 }
2750 });
2751 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2752 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2753 pool.run_many(&parts);
2754 } else {
2755 let closures: [_; N] = std::array::from_fn(|i| {
2756 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2757 unreachable!()
2758 };
2759 let (bytes, rows, cols, out) = (
2760 ts[i].quant_bytes(),
2761 ts[i].rows(),
2762 ts[i].cols(),
2763 outs_addr[i],
2764 );
2765 move |s: usize, e: usize| {
2766 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2767 }
2768 });
2769 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2770 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2771 pool.run_many(&parts);
2772 }
2773 return;
2774 }
2775 if uniform_q4 {
2776 let closures: [_; N] = std::array::from_fn(|i| {
2777 let (packed, scales) =
2778 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2779 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2780 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2781 });
2782 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2783 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2784 pool.run_many(&parts);
2785 } else {
2786 let closures: [_; N] = std::array::from_fn(|i| {
2787 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2788 unreachable!()
2789 };
2790 let (bytes, rows, cols, out) = (
2791 ts[i].quant_bytes(),
2792 ts[i].rows(),
2793 ts[i].cols(),
2794 outs_addr[i],
2795 );
2796 move |s: usize, e: usize| {
2797 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2798 }
2799 });
2800 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2801 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2802 pool.run_many(&parts);
2803 }
2804 return;
2805 }
2806
2807 if uniform_f32 {
2808 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2809 let closures: [_; N] = std::array::from_fn(|i| {
2810 let Self::F32 { data, cols, .. } = ts[i] else {
2811 unreachable!()
2812 };
2813 let out = outs_addr[i];
2814 move |start: usize, end: usize| {
2815 for o in start..end {
2816 let row = &data[o * cols..(o + 1) * cols];
2817 let mut sum = 0.0f32;
2818 for j in 0..*cols {
2819 sum += row[j] * x[j];
2820 }
2821 unsafe { *out.at(o) = sum };
2823 }
2824 }
2825 });
2826 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2827 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2828 pool.run_many(&parts);
2829 return;
2830 }
2831
2832 struct Ctx<'a> {
2835 bytes: &'a [u8],
2836 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2837 rep: &'a [u8],
2838 row_scale: &'a [f32],
2839 cols: usize,
2840 xs: std::borrow::Cow<'a, [f32]>,
2841 }
2842 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2843 let Self::Mapped {
2844 dtype,
2845 cols,
2846 row_scale,
2847 col_field,
2848 repack,
2849 ..
2850 } = ts[i]
2851 else {
2852 unreachable!()
2853 };
2854 Ctx {
2855 bytes: ts[i].quant_bytes(),
2856 rep: repack,
2857 row_scale,
2858 cols: *cols,
2859 xs: prescale(x, col_field, *dtype),
2860 }
2861 });
2862 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2863 #[cfg(target_arch = "aarch64")]
2864 if sdot_enabled() {
2865 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2866 let closures: [_; N] = std::array::from_fn(|i| {
2867 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2868 move |start: usize, end: usize| {
2869 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2870 }
2871 });
2872 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2873 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2874 pool.run_many(&parts);
2875 return;
2876 }
2877 #[cfg(target_arch = "x86_64")]
2878 if avx2_a8w8_enabled() {
2879 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2880 let closures: [_; N] = std::array::from_fn(|i| {
2881 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2882 move |start: usize, end: usize| {
2883 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2884 }
2885 });
2886 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2887 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2888 pool.run_many(&parts);
2889 return;
2890 }
2891 let closures: [_; N] = std::array::from_fn(|i| {
2892 let (c, out) = (&ctxs[i], outs_addr[i]);
2893 move |start: usize, end: usize| {
2894 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2895 }
2896 });
2897 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2898 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2899 pool.run_many(&parts);
2900 }
2901}
2902
2903impl QTensor {
2904 #[allow(clippy::needless_range_loop)]
2909 pub fn matvec2_many<const N: usize>(
2910 ts: [&QTensor; N],
2911 x1: &[f32],
2912 x2: &[f32],
2913 mut o1s: [&mut [f32]; N],
2914 mut o2s: [&mut [f32]; N],
2915 pool: Option<&Pool>,
2916 ) {
2917 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2918 if ts.iter().any(|t| t.has_prism_contract()) {
2919 for i in 0..N {
2920 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2921 }
2922 return;
2923 }
2924 let uniform_q8 = ts.iter().all(|t| {
2925 matches!(
2926 t,
2927 Self::Mapped {
2928 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2929 ..
2930 }
2931 )
2932 });
2933 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2934 let uniform_q4 = ts.iter().all(|t| {
2935 matches!(
2936 t,
2937 Self::Mapped {
2938 dtype: TensorDtype::Q4Block,
2939 ..
2940 }
2941 )
2942 });
2943 let uniform_vbit = ts.iter().all(|t| {
2944 matches!(
2945 t,
2946 Self::Mapped {
2947 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2948 ..
2949 }
2950 )
2951 });
2952 let fusable = pool.is_some()
2953 && total_rows >= 256
2954 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2955 if !fusable {
2956 for i in 0..N {
2957 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2958 }
2959 return;
2960 }
2961 let pool = pool.unwrap();
2962
2963 if uniform_q4 || uniform_vbit {
2964 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2965 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2966 if a8w8_enabled() {
2968 let a1 = split_act(x1);
2969 let a2 = split_act(x2);
2970 let (a1, a2) = (&a1, &a2);
2971 if uniform_q4 {
2972 let closures: [_; N] = std::array::from_fn(|i| {
2973 let (packed, scales) =
2974 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2975 let (gpr, cols, o1, o2) =
2976 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2977 move |s: usize, e: usize| {
2978 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2979 }
2980 });
2981 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2982 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2983 pool.run_many(&parts);
2984 } else {
2985 let closures: [_; N] = std::array::from_fn(|i| {
2986 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2987 unreachable!()
2988 };
2989 let (bytes, rows, cols, o1, o2) = (
2990 ts[i].quant_bytes(),
2991 ts[i].rows(),
2992 ts[i].cols(),
2993 p1[i],
2994 p2[i],
2995 );
2996 move |s: usize, e: usize| {
2997 vbit_range2_a8w8(
2998 bytes,
2999 vbit_offsets,
3000 x1,
3001 x2,
3002 a1,
3003 a2,
3004 rows,
3005 cols,
3006 o1,
3007 o2,
3008 s,
3009 e,
3010 )
3011 }
3012 });
3013 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3014 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3015 pool.run_many(&parts);
3016 }
3017 return;
3018 }
3019 if uniform_q4 {
3020 let closures: [_; N] = std::array::from_fn(|i| {
3021 let (packed, scales) =
3022 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3023 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3024 move |s: usize, e: usize| {
3025 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3026 }
3027 });
3028 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3029 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3030 pool.run_many(&parts);
3031 } else {
3032 let closures: [_; N] = std::array::from_fn(|i| {
3033 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3034 unreachable!()
3035 };
3036 let (bytes, rows, cols, o1, o2) = (
3037 ts[i].quant_bytes(),
3038 ts[i].rows(),
3039 ts[i].cols(),
3040 p1[i],
3041 p2[i],
3042 );
3043 move |s: usize, e: usize| {
3044 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3045 }
3046 });
3047 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3048 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3049 pool.run_many(&parts);
3050 }
3051 return;
3052 }
3053
3054 if uniform_f32 {
3055 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3056 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3057 let closures: [_; N] = std::array::from_fn(|i| {
3058 let Self::F32 { data, cols, .. } = ts[i] else {
3059 unreachable!()
3060 };
3061 let (o1, o2) = (p1[i], p2[i]);
3062 move |start: usize, end: usize| {
3063 for o in start..end {
3064 let row = &data[o * cols..(o + 1) * cols];
3065 let (mut s1, mut s2) = (0.0f32, 0.0f32);
3066 for j in 0..*cols {
3067 s1 += row[j] * x1[j];
3068 s2 += row[j] * x2[j];
3069 }
3070 unsafe {
3072 *o1.at(o) = s1;
3073 *o2.at(o) = s2;
3074 }
3075 }
3076 }
3077 });
3078 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3079 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3080 pool.run_many(&parts);
3081 return;
3082 }
3083
3084 struct Ctx<'a> {
3085 bytes: &'a [u8],
3086 row_scale: &'a [f32],
3087 cols: usize,
3088 xs1: std::borrow::Cow<'a, [f32]>,
3089 xs2: std::borrow::Cow<'a, [f32]>,
3090 }
3091 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3092 let Self::Mapped {
3093 dtype,
3094 cols,
3095 row_scale,
3096 col_field,
3097 ..
3098 } = ts[i]
3099 else {
3100 unreachable!()
3101 };
3102 Ctx {
3103 bytes: ts[i].quant_bytes(),
3104 row_scale,
3105 cols: *cols,
3106 xs1: prescale(x1, col_field, *dtype),
3107 xs2: prescale(x2, col_field, *dtype),
3108 }
3109 });
3110 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3111 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3112 #[cfg(target_arch = "aarch64")]
3113 if sdot_enabled() {
3114 let acts: [(SplitAct, SplitAct); N] =
3115 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3116 let closures: [_; N] = std::array::from_fn(|i| {
3117 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3118 move |start: usize, end: usize| {
3119 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3120 }
3121 });
3122 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3123 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3124 pool.run_many(&parts);
3125 return;
3126 }
3127 #[cfg(target_arch = "x86_64")]
3128 if avx2_a8w8_enabled() {
3129 let acts: [(SplitAct, SplitAct); N] =
3130 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3131 let closures: [_; N] = std::array::from_fn(|i| {
3132 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3133 move |start: usize, end: usize| {
3134 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3135 }
3136 });
3137 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3138 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3139 pool.run_many(&parts);
3140 return;
3141 }
3142 let closures: [_; N] = std::array::from_fn(|i| {
3143 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3144 move |start: usize, end: usize| {
3145 q8_range2_f32(
3146 c.bytes,
3147 c.row_scale,
3148 &c.xs1,
3149 &c.xs2,
3150 c.cols,
3151 o1,
3152 o2,
3153 start,
3154 end,
3155 )
3156 }
3157 });
3158 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3159 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3160 pool.run_many(&parts);
3161 }
3162
3163 pub fn matvec_silu_mul(
3168 gate: &QTensor,
3169 up: &QTensor,
3170 x: &[f32],
3171 out: &mut [f32],
3172 pool: Option<&Pool>,
3173 ) -> bool {
3174 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3175 }
3176
3177 pub fn matvec_silu_mul_limited(
3183 gate: &QTensor,
3184 up: &QTensor,
3185 x: &[f32],
3186 out: &mut [f32],
3187 limit: f32,
3188 pool: Option<&Pool>,
3189 ) -> bool {
3190 if gate.has_prism_contract() || up.has_prism_contract() {
3191 return false;
3195 }
3196 let inter = gate.rows();
3197 debug_assert_eq!(up.rows(), inter);
3198 debug_assert_eq!(out.len(), inter);
3199 debug_assert_eq!(gate.cols(), up.cols());
3200 if !a8w8_enabled() {
3201 return false;
3202 }
3203 let act = split_act(x);
3204 let act = &act;
3205 let x_ref = x;
3206 let out_addr = SendMut(out.as_mut_ptr());
3207
3208 match (gate, up) {
3209 (
3211 Self::Mapped {
3212 dtype: TensorDtype::Q4Block,
3213 ..
3214 },
3215 Self::Mapped {
3216 dtype: TensorDtype::Q4Block,
3217 ..
3218 },
3219 ) => {
3220 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3221 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3222 let gpr = gate.cols() / GROUP_SIZE;
3223 let cols = gate.cols();
3224 let run = move |start: usize, end: usize| {
3225 for r in start..end {
3226 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3227 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3228 for &(j, xv) in &act.outliers {
3229 let flat = r * cols + j;
3230 let gb = gp[flat / 2];
3231 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3232 let gsc = f16_to_f32(u16::from_le_bytes([
3233 gs[(flat / GROUP_SIZE) * 2],
3234 gs[(flat / GROUP_SIZE) * 2 + 1],
3235 ]));
3236 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3237 let ub = up_p[flat / 2];
3238 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3239 let usc = f16_to_f32(u16::from_le_bytes([
3240 up_s[(flat / GROUP_SIZE) * 2],
3241 up_s[(flat / GROUP_SIZE) * 2 + 1],
3242 ]));
3243 uv += ((un as i32 - 8) as f32) * usc * xv;
3244 }
3245 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3247 }
3248 };
3249 dispatch_rows(pool, inter, &run);
3250 true
3251 }
3252 (
3256 Self::Mapped {
3257 dtype: TensorDtype::Q4Tiled,
3258 ..
3259 },
3260 Self::Mapped {
3261 dtype: TensorDtype::Q4Tiled,
3262 ..
3263 },
3264 ) => {
3265 let g_bytes = gate.quant_bytes();
3266 let u_bytes = up.quant_bytes();
3267 let gpr = gate.cols() / GROUP_SIZE;
3268 let run = move |start: usize, end: usize| {
3269 for r in start..end {
3270 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3271 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3272 for &(j, xv) in &act.outliers {
3273 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3274 gv += w * s * xv;
3275 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3276 uv += w * s * xv;
3277 }
3278 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3280 }
3281 };
3282 dispatch_rows(pool, inter, &run);
3283 true
3284 }
3285 (
3288 Self::Mapped {
3289 dtype: TensorDtype::Q4TiledP,
3290 ..
3291 },
3292 Self::Mapped {
3293 dtype: TensorDtype::Q4TiledP,
3294 ..
3295 },
3296 ) => {
3297 let cols = gate.cols();
3298 let gpr = cols / GROUP_SIZE;
3299 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3300 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3301 let run = |start: usize, end: usize| {
3302 with_krows(gpr, |gsc, usc| {
3303 for r in start..end {
3304 gv_view.scales_into(r, gpr, gsc);
3305 uv_view.scales_into(r, gpr, usc);
3306 let mut gv =
3307 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3308 let mut uv =
3309 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3310 for &(j, xv) in &act.outliers {
3311 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3312 gv += w * s * xv;
3313 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3314 uv += w * s * xv;
3315 }
3316 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3318 }
3319 });
3320 };
3321 dispatch_rows(pool, inter, &run);
3322 true
3323 }
3324 (
3330 Self::Mapped {
3331 dtype: TensorDtype::Q1,
3332 ..
3333 },
3334 Self::Mapped {
3335 dtype: TensorDtype::Q1,
3336 ..
3337 },
3338 ) => {
3339 let g_bytes = gate.quant_bytes();
3340 let u_bytes = up.quant_bytes();
3341 let gpr = gate.cols() / GROUP_SIZE;
3342 let gsum = q1_group_sums(&act.xq, gpr);
3343 let gsum = &gsum;
3344 let run = move |start: usize, end: usize| {
3345 for r in start..end {
3346 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
3347 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
3348 for &(j, xv) in &act.outliers {
3349 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
3350 gv += w * s * xv;
3351 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
3352 uv += w * s * xv;
3353 }
3354 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3356 }
3357 };
3358 dispatch_rows(pool, inter, &run);
3359 true
3360 }
3361 (
3365 Self::Mapped {
3366 dtype: TensorDtype::Q2TiledP,
3367 ..
3368 },
3369 Self::Mapped {
3370 dtype: TensorDtype::Q2TiledP,
3371 ..
3372 },
3373 ) => {
3374 let cols = gate.cols();
3375 let gpr = cols / GROUP_SIZE;
3376 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3377 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3378 let gsum = q1_group_sums(&act.xq, gpr);
3379 let gsum = &gsum;
3380 let run = move |start: usize, end: usize| {
3381 with_krows(gpr, |gsc, usc| {
3382 for r in start..end {
3383 gv_view.scales_into(r, gpr, gsc);
3384 uv_view.scales_into(r, gpr, usc);
3385 let mut gv =
3386 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3387 let mut uv =
3388 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3389 for &(j, xv) in &act.outliers {
3390 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3391 gv += w * s * xv;
3392 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3393 uv += w * s * xv;
3394 }
3395 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3397 }
3398 });
3399 };
3400 dispatch_rows(pool, inter, &run);
3401 true
3402 }
3403 (
3408 Self::Mapped {
3409 dtype: TensorDtype::Q8Row,
3410 row_scale: g_rs,
3411 ..
3412 },
3413 Self::Mapped {
3414 dtype: TensorDtype::Q8Row,
3415 row_scale: u_rs,
3416 ..
3417 },
3418 ) => {
3419 let g_bytes = gate.quant_bytes();
3420 let u_bytes = up.quant_bytes();
3421 let cols = gate.cols();
3422 let run = move |start: usize, end: usize| {
3423 for r in start..end {
3424 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3425 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3426 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3428 }
3429 };
3430 dispatch_rows(pool, inter, &run);
3431 true
3432 }
3433 (
3435 Self::Mapped {
3436 dtype: TensorDtype::Q1T,
3437 ..
3438 },
3439 Self::Mapped {
3440 dtype: TensorDtype::Q1T,
3441 ..
3442 },
3443 ) => {
3444 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3445 let g_bytes = gate.quant_bytes();
3446 let u_bytes = up.quant_bytes();
3447 let gpr = gate.cols() / GROUP_SIZE;
3448 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3449 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3450 let run = move |start: usize, end: usize| {
3451 for r in start..end {
3452 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3453 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3454 for &(j, xv) in &act.outliers {
3455 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3456 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3457 }
3458 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3459 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3460 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3462 }
3463 };
3464 dispatch_rows(pool, inter, &run);
3465 true
3466 }
3467 _ => false,
3468 }
3469 }
3470
3471 pub fn moe_gate_up_many(
3486 pairs: &[(&QTensor, &QTensor)],
3487 x: &[f32],
3488 outs: &mut [Vec<f32>],
3489 pool: Option<&Pool>,
3490 ) -> bool {
3491 Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3492 }
3493
3494 pub fn moe_gate_up_many_limited(
3498 pairs: &[(&QTensor, &QTensor)],
3499 x: &[f32],
3500 outs: &mut [Vec<f32>],
3501 limit: f32,
3502 pool: Option<&Pool>,
3503 ) -> bool {
3504 if pairs.is_empty() || pairs.len() != outs.len() {
3505 return false;
3506 }
3507 if !a8w8_enabled() {
3508 if limit > 0.0 {
3509 return false;
3512 }
3513 let groups = vec![vec![0]; pairs.len()];
3514 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3515 }
3516 let inter = pairs[0].0.rows();
3517 let cols = pairs[0].0.cols();
3518 if cols % GROUP_SIZE != 0 {
3519 return false;
3520 }
3521 let gpr = cols / GROUP_SIZE;
3522 let q2 = matches!(
3525 pairs[0].0,
3526 Self::Mapped {
3527 dtype: TensorDtype::Q2TiledP,
3528 ..
3529 }
3530 );
3531 let want = if q2 {
3532 TensorDtype::Q2TiledP
3533 } else {
3534 TensorDtype::Q4TiledP
3535 };
3536 let mut views = Vec::with_capacity(pairs.len() * 2);
3537 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3538 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3539 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3540 if !both
3541 || g.rows() != inter
3542 || u.rows() != inter
3543 || g.cols() != cols
3544 || u.cols() != cols
3545 || o.len() != inter
3546 {
3547 return false;
3548 }
3549 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3550 views.push(mk(g.quant_bytes(), inter, cols));
3551 views.push(mk(u.quant_bytes(), inter, cols));
3552 }
3553 let act = split_act(x);
3554 let gsum = if q2 {
3555 q1_group_sums(&act.xq, gpr)
3556 } else {
3557 Vec::new()
3558 };
3559 let (act, gsum) = (&act, &gsum);
3560 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3561 let (views, ptrs) = (&views, &ptrs);
3562 let run = |start: usize, end: usize| {
3563 with_krows(gpr, |gsc, usc| {
3564 for flat in start..end {
3565 let (e, r) = (flat / inter, flat % inter);
3566 let gv_view = &views[e * 2];
3567 let uv_view = &views[e * 2 + 1];
3568 gv_view.scales_into(r, gpr, gsc);
3569 uv_view.scales_into(r, gpr, usc);
3570 let (mut gv, mut uv) = if q2 {
3571 (
3572 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3573 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3574 )
3575 } else {
3576 (
3577 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3578 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3579 )
3580 };
3581 for &(j, xv) in &act.outliers {
3582 let (og, ou) = if q2 {
3583 (
3584 q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3585 q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3586 )
3587 } else {
3588 (
3589 q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3590 q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3591 )
3592 };
3593 gv += og.0 * og.1 * xv;
3594 uv += ou.0 * ou.1 * xv;
3595 }
3596 unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3598 }
3599 });
3600 };
3601 dispatch_rows(pool, pairs.len() * inter, &run);
3602 true
3603 }
3604
3605 pub fn moe_down_many(
3614 downs: &[&QTensor],
3615 gs: &[Vec<f32>],
3616 weights: &[f32],
3617 out: &mut [f32],
3618 pool: Option<&Pool>,
3619 ) -> bool {
3620 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3621 return false;
3622 }
3623 if !a8w8_enabled() {
3624 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3625 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3626 return false;
3627 }
3628 out.fill(0.0);
3629 for (row, &w) in terms.iter().zip(weights) {
3630 for (o, &v) in out.iter_mut().zip(row) {
3631 *o += w * v;
3632 }
3633 }
3634 return true;
3635 }
3636 let rows = out.len();
3637 let cols = downs[0].cols();
3638 if cols % GROUP_SIZE != 0 {
3639 return false;
3640 }
3641 let gpr = cols / GROUP_SIZE;
3642 let mut views = Vec::with_capacity(downs.len());
3643 for (d, g) in downs.iter().zip(gs.iter()) {
3644 if !matches!(
3645 d,
3646 Self::Mapped {
3647 dtype: TensorDtype::Q4TiledP,
3648 ..
3649 }
3650 ) || d.rows() != rows
3651 || d.cols() != cols
3652 || g.len() != cols
3653 {
3654 return false;
3655 }
3656 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3657 }
3658 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3660 let out_addr = SendMut(out.as_mut_ptr());
3668 let (views, acts, weights) = (&views, &acts, &weights);
3669 let run = |start: usize, end: usize| {
3670 with_krow(gpr, |sc| {
3671 for r in start..end {
3672 let mut acc = 0f32;
3673 for (e, v) in views.iter().enumerate() {
3674 v.scales_into(r, gpr, sc);
3675 let a = &acts[e];
3676 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
3677 for &(j, xv) in &a.outliers {
3678 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3679 d += w * s * xv;
3680 }
3681 acc += weights[e] * d;
3682 }
3683 unsafe { *out_addr.at(r) = acc };
3685 }
3686 });
3687 };
3688 dispatch_rows(pool, rows, &run);
3689 true
3690 }
3691
3692 pub fn moe_gate_up_rows(
3704 pairs: &[(&QTensor, &QTensor)],
3705 groups: &[Vec<usize>],
3706 xs: &[f32],
3707 outs: &mut [Vec<f32>],
3708 pool: Option<&Pool>,
3709 ) -> bool {
3710 if pairs.is_empty() || pairs.len() != groups.len() {
3711 return false;
3712 }
3713 let inter = pairs[0].0.rows();
3714 let cols = pairs[0].0.cols();
3715 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
3716 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
3717 return false;
3718 }
3719 let b = xs.len() / cols;
3720 let gpr = cols / GROUP_SIZE;
3721 let mut views = Vec::with_capacity(pairs.len() * 2);
3722 for (g, u) in pairs {
3723 let q4tp = |t: &QTensor| {
3724 matches!(
3725 t,
3726 Self::Mapped {
3727 dtype: TensorDtype::Q4TiledP,
3728 ..
3729 }
3730 )
3731 };
3732 if g.has_prism_contract()
3733 || u.has_prism_contract()
3734 || !q4tp(g)
3735 || !q4tp(u)
3736 || g.rows() != inter
3737 || u.rows() != inter
3738 || g.cols() != cols
3739 || u.cols() != cols
3740 {
3741 return false;
3742 }
3743 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
3744 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
3745 }
3746 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
3747 return false;
3748 }
3749 let quantized = a8w8_enabled();
3750 let acts: Vec<SplitAct> = if quantized {
3751 (0..b)
3752 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
3753 .collect()
3754 } else {
3755 Vec::new()
3756 };
3757 let mut offs = Vec::with_capacity(groups.len());
3758 let mut o = 0usize;
3759 for g in groups {
3760 offs.push(o);
3761 o += g.len();
3762 }
3763 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3764 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3765 let run = |start: usize, end: usize| {
3766 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3767 for flat in start..end {
3768 let (e, r) = (flat / inter, flat % inter);
3769 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
3770 gv_view.scales_into(r, gpr, &mut gsc);
3771 uv_view.scales_into(r, gpr, &mut usc);
3772 for (k, &t) in groups[e].iter().enumerate() {
3773 if !quantized {
3774 let x = &xs[t * cols..(t + 1) * cols];
3775 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
3776 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
3777 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
3778 continue;
3779 }
3780 let act = &acts[t];
3781 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3782 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3783 for &(j, xv) in &act.outliers {
3784 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3785 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3786 gv += og.0 * og.1 * xv;
3787 uv += ou.0 * ou.1 * xv;
3788 }
3789 let silu_g = gv / (1.0 + (-gv).exp());
3790 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
3793 }
3794 }
3795 };
3796 dispatch_rows(pool, pairs.len() * inter, &run);
3797 true
3798 }
3799
3800 pub fn moe_down_rows(
3807 downs: &[&QTensor],
3808 group_lens: &[usize],
3809 gs: &[Vec<f32>],
3810 outs: &mut [Vec<f32>],
3811 pool: Option<&Pool>,
3812 ) -> bool {
3813 if downs.is_empty() || downs.len() != group_lens.len() {
3814 return false;
3815 }
3816 let rows = downs[0].rows();
3817 let cols = downs[0].cols();
3818 let n_pairs: usize = group_lens.iter().sum();
3819 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
3820 return false;
3821 }
3822 let gpr = cols / GROUP_SIZE;
3823 let mut views = Vec::with_capacity(downs.len());
3824 for d in downs {
3825 if d.has_prism_contract()
3826 || !matches!(
3827 d,
3828 Self::Mapped {
3829 dtype: TensorDtype::Q4TiledP,
3830 ..
3831 }
3832 )
3833 || d.rows() != rows
3834 || d.cols() != cols
3835 {
3836 return false;
3837 }
3838 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3839 }
3840 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
3841 return false;
3842 }
3843 let quantized = a8w8_enabled();
3844 let acts: Vec<SplitAct> = if quantized {
3845 gs.iter().map(|g| split_act(g)).collect()
3846 } else {
3847 Vec::new()
3848 };
3849 let mut offs = Vec::with_capacity(group_lens.len());
3850 let mut o = 0usize;
3851 for &l in group_lens {
3852 offs.push(o);
3853 o += l;
3854 }
3855 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3856 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3857 let run = |start: usize, end: usize| {
3858 let mut sc = vec![0f32; gpr];
3859 for flat in start..end {
3860 let (e, r) = (flat / rows, flat % rows);
3861 let v = &views[e];
3862 v.scales_into(r, gpr, &mut sc);
3863 for k in 0..group_lens[e] {
3864 if !quantized {
3865 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
3866 unsafe { *ptrs[offs[e] + k].at(r) = d };
3867 continue;
3868 }
3869 let a = &acts[offs[e] + k];
3870 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3871 for &(j, xv) in &a.outliers {
3872 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3873 d += w * s * xv;
3874 }
3875 unsafe { *ptrs[offs[e] + k].at(r) = d };
3877 }
3878 }
3879 };
3880 dispatch_rows(pool, downs.len() * rows, &run);
3881 true
3882 }
3883}
3884
3885#[cfg(target_os = "macos")]
3890mod accel_blas {
3891 #[link(name = "Accelerate", kind = "framework")]
3892 unsafe extern "C" {
3893 pub fn cblas_sgemm(
3894 order: i32,
3895 trans_a: i32,
3896 trans_b: i32,
3897 m: i32,
3898 n: i32,
3899 k: i32,
3900 alpha: f32,
3901 a: *const f32,
3902 lda: i32,
3903 b: *const f32,
3904 ldb: i32,
3905 beta: f32,
3906 c: *mut f32,
3907 ldc: i32,
3908 );
3909 }
3910}
3911
3912#[cfg(target_os = "macos")]
3913pub(crate) fn accel_gemm_enabled() -> bool {
3914 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3915 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3916}
3917
3918#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3921pub(crate) fn accel_gemm_enabled() -> bool {
3922 true
3923}
3924
3925#[cfg(target_arch = "aarch64")]
3932#[allow(clippy::too_many_arguments)]
3933pub(crate) fn neon_gemm_rm(
3934 m: usize,
3935 n: usize,
3936 k: usize,
3937 alpha: f32,
3938 a: &[f32],
3939 lda: usize,
3940 b_mat: &[f32],
3941 ldb: usize,
3942 b_rows_are_n: bool,
3943 c: &mut [f32],
3944 ldc: usize,
3945) {
3946 debug_assert!(a.len() >= (m - 1) * lda + k);
3947 debug_assert!(c.len() >= (m - 1) * ldc + n);
3948 unsafe {
3950 use core::arch::aarch64::*;
3951 let mut i = 0usize;
3952 while i < m {
3953 let mi = (m - i).min(4);
3954 let mut j = 0usize;
3955 while j < n {
3956 let nj = (n - j).min(8);
3957 if mi == 4 && nj == 8 {
3958 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3959 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3960 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3961 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3962 for p in 0..k {
3963 let (b0, b1) = if b_rows_are_n {
3964 let base = b_mat.as_ptr().add(j * ldb + p);
3967 let g = |o: usize| *base.add(o * ldb);
3968 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3969 } else {
3970 let base = b_mat.as_ptr().add(p * ldb + j);
3971 (
3972 [*base, *base.add(1), *base.add(2), *base.add(3)],
3973 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3974 )
3975 };
3976 let bv0 = vld1q_f32(b0.as_ptr());
3977 let bv1 = vld1q_f32(b1.as_ptr());
3978 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3979 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3980 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3981 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3982 c0a = vfmaq_f32(c0a, a0, bv0);
3983 c0b = vfmaq_f32(c0b, a0, bv1);
3984 c1a = vfmaq_f32(c1a, a1, bv0);
3985 c1b = vfmaq_f32(c1b, a1, bv1);
3986 c2a = vfmaq_f32(c2a, a2, bv0);
3987 c2b = vfmaq_f32(c2b, a2, bv1);
3988 c3a = vfmaq_f32(c3a, a3, bv0);
3989 c3b = vfmaq_f32(c3b, a3, bv1);
3990 }
3991 let al = vdupq_n_f32(alpha);
3992 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3993 .iter()
3994 .enumerate()
3995 {
3996 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3997 vst1q_f32(dst, vmulq_f32(*ca, al));
3998 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3999 }
4000 } else {
4001 for r in 0..mi {
4002 for q in 0..nj {
4003 let mut acc = 0f32;
4004 for p in 0..k {
4005 let bv = if b_rows_are_n {
4006 b_mat[(j + q) * ldb + p]
4007 } else {
4008 b_mat[p * ldb + j + q]
4009 };
4010 acc += a[(i + r) * lda + p] * bv;
4011 }
4012 c[(i + r) * ldc + j + q] = acc * alpha;
4013 }
4014 }
4015 }
4016 j += nj;
4017 }
4018 i += mi;
4019 }
4020 }
4021}
4022
4023#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4025#[allow(clippy::too_many_arguments)]
4026pub(crate) fn sgemm_rm(
4027 m: usize,
4028 n: usize,
4029 k: usize,
4030 alpha: f32,
4031 a: &[f32],
4032 lda: usize,
4033 b_mat: &[f32],
4034 ldb: usize,
4035 b_rows_are_n: bool,
4036 c: &mut [f32],
4037 ldc: usize,
4038) {
4039 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4040}
4041
4042#[allow(clippy::too_many_arguments)]
4046pub fn sgemm_public(
4047 m: usize,
4048 n: usize,
4049 k: usize,
4050 alpha: f32,
4051 a: &[f32],
4052 lda: usize,
4053 b_mat: &[f32],
4054 ldb: usize,
4055 b_rows_are_n: bool,
4056 c: &mut [f32],
4057 ldc: usize,
4058) {
4059 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4060 {
4061 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4062 }
4063 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4068 {
4069 for i in 0..m {
4070 for j in 0..n {
4071 let mut acc = 0f32;
4072 for p in 0..k {
4073 let bv = if b_rows_are_n {
4074 b_mat[j * ldb + p]
4075 } else {
4076 b_mat[p * ldb + j]
4077 };
4078 acc += a[i * lda + p] * bv;
4079 }
4080 c[i * ldc + j] = alpha * acc;
4081 }
4082 }
4083 }
4084}
4085
4086#[cfg(target_os = "macos")]
4089#[allow(clippy::too_many_arguments)]
4090pub(crate) fn sgemm_rm(
4091 m: usize,
4092 n: usize,
4093 k: usize,
4094 alpha: f32,
4095 a: &[f32],
4096 lda: usize,
4097 b_mat: &[f32],
4098 ldb: usize,
4099 b_rows_are_n: bool,
4100 c: &mut [f32],
4101 ldc: usize,
4102) {
4103 debug_assert!(a.len() >= (m - 1) * lda + k);
4104 debug_assert!(c.len() >= (m - 1) * ldc + n);
4105 #[cfg(target_arch = "aarch64")]
4110 if std::env::var("CMF_FORCE_NEON_GEMM")
4111 .map(|v| v == "1")
4112 .unwrap_or(false)
4113 {
4114 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4115 }
4116 unsafe {
4117 accel_blas::cblas_sgemm(
4118 101, 111, if b_rows_are_n { 112 } else { 111 },
4121 m as i32,
4122 n as i32,
4123 k as i32,
4124 alpha,
4125 a.as_ptr(),
4126 lda as i32,
4127 b_mat.as_ptr(),
4128 ldb as i32,
4129 0.0,
4130 c.as_mut_ptr(),
4131 ldc as i32,
4132 );
4133 }
4134}
4135
4136#[cfg(target_os = "macos")]
4143fn qmatmat_accel(
4144 q: &[u8],
4145 row_scale: &[f32],
4146 pre: &[std::borrow::Cow<'_, [f32]>],
4147 rows: usize,
4148 cols: usize,
4149 out: &mut [f32],
4150 pool: Option<&Pool>,
4151) {
4152 const TR: usize = 2048;
4157 let b = pre.len();
4158 thread_local! {
4159 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4160 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4161 }
4162 XPANEL.with(|xp| {
4163 WTILE.with(|wt| {
4164 let mut xpanel = xp.borrow_mut();
4165 xpanel.clear();
4166 for x in pre {
4167 xpanel.extend_from_slice(x);
4168 }
4169 let mut wtile = wt.borrow_mut();
4170 wtile.resize(TR * cols, 0.0);
4171 let mut r0 = 0usize;
4172 while r0 < rows {
4173 let tr = TR.min(rows - r0);
4174 let wt_addr = SendMut(wtile.as_mut_ptr());
4176 let run = |start: usize, end: usize| {
4177 for r in start..end {
4178 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4179 let s = row_scale[r0 + r];
4180 let dst =
4182 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4183 for (d, &v) in dst.iter_mut().zip(row) {
4184 *d = (v as i8) as f32 * s;
4185 }
4186 }
4187 };
4188 dispatch_rows(pool, tr, &run);
4189 unsafe {
4191 accel_blas::cblas_sgemm(
4192 101, 111, 112, b as i32,
4196 tr as i32,
4197 cols as i32,
4198 1.0,
4199 xpanel.as_ptr(),
4200 cols as i32,
4201 wtile.as_ptr(),
4202 cols as i32,
4203 0.0,
4204 out.as_mut_ptr().add(r0),
4205 rows as i32,
4206 );
4207 }
4208 r0 += tr;
4209 }
4210 })
4211 });
4212}
4213
4214fn qmatmat(
4215 q: &[u8],
4216 row_scale: &[f32],
4217 pre: &[std::borrow::Cow<'_, [f32]>],
4218 rows: usize,
4219 cols: usize,
4220 out: &mut [f32],
4221 pool: Option<&Pool>,
4222) {
4223 let b = pre.len();
4224 debug_assert_eq!(out.len(), b * rows);
4225 #[cfg(target_os = "macos")]
4230 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4231 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4232 return;
4233 }
4234 #[cfg(target_arch = "aarch64")]
4235 if sdot_enabled() {
4236 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4237 let out_addr = SendMut(out.as_mut_ptr());
4238 let blocked_ok = blocked_enabled();
4241 let use_i8mm = i8mm_enabled();
4242 if blocked_ok {
4243 let run = |start: usize, end: usize| {
4244 let mut o = start;
4245 while o < end {
4246 if o + 2 <= end {
4247 let r0 = &q[o * cols..(o + 1) * cols];
4248 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4249 let mut bi = 0usize;
4250 while bi + 4 <= acts.len() {
4251 let xs = [
4252 acts[bi].xq.as_slice(),
4253 acts[bi + 1].xq.as_slice(),
4254 acts[bi + 2].xq.as_slice(),
4255 acts[bi + 3].xq.as_slice(),
4256 ];
4257 let d = if use_i8mm {
4258 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4259 } else {
4260 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4261 };
4262 for (r, row) in [r0, r1].into_iter().enumerate() {
4263 for k in 0..4 {
4264 let act = &acts[bi + k];
4265 let mut v = d[r][k] as f32 * act.sx;
4266 for &(j, xv) in &act.outliers {
4267 v += (row[j] as i8) as f32 * xv;
4268 }
4269 unsafe {
4270 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4271 };
4272 }
4273 }
4274 bi += 4;
4275 }
4276 while bi < acts.len() {
4277 for (r, row) in [r0, r1].into_iter().enumerate() {
4278 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4279 unsafe { *out_addr.at(bi * rows + o + r) = v };
4280 }
4281 bi += 1;
4282 }
4283 o += 2;
4284 } else {
4285 let row = &q[o * cols..(o + 1) * cols];
4286 for (bi, act) in acts.iter().enumerate() {
4287 let v = row_dot_sdot(row, act) * row_scale[o];
4288 unsafe { *out_addr.at(bi * rows + o) = v };
4289 }
4290 o += 1;
4291 }
4292 }
4293 };
4294 dispatch_rows(pool, rows, &run);
4295 return;
4296 }
4297 let run = |start: usize, end: usize| {
4298 for o in start..end {
4299 let row = &q[o * cols..(o + 1) * cols];
4300 for (bi, act) in acts.iter().enumerate() {
4301 let v = row_dot_sdot(row, act) * row_scale[o];
4302 unsafe { *out_addr.at(bi * rows + o) = v };
4303 }
4304 }
4305 };
4306 dispatch_rows(pool, rows, &run);
4307 return;
4308 }
4309 #[cfg(target_arch = "x86_64")]
4314 if avx2_a8w8_enabled() {
4315 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4316 let out_addr = SendMut(out.as_mut_ptr());
4317 let blocked_ok = blocked_enabled();
4320 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4321 let run = |start: usize, end: usize| {
4322 let mut o = start;
4323 while o < end {
4324 if o + 2 <= end {
4325 let r0 = &q[o * cols..(o + 1) * cols];
4326 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4327 let mut bi = 0usize;
4328 while bi + 4 <= acts.len() {
4329 let xs = [
4330 acts[bi].xq.as_slice(),
4331 acts[bi + 1].xq.as_slice(),
4332 acts[bi + 2].xq.as_slice(),
4333 acts[bi + 3].xq.as_slice(),
4334 ];
4335 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4336 for (r, row) in [r0, r1].into_iter().enumerate() {
4337 for k in 0..4 {
4338 let act = &acts[bi + k];
4339 let mut v = d[r][k] as f32 * act.sx;
4340 for &(j, xv) in &act.outliers {
4341 v += (row[j] as i8) as f32 * xv;
4342 }
4343 unsafe {
4344 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4345 };
4346 }
4347 }
4348 bi += 4;
4349 }
4350 while bi < acts.len() {
4351 for (r, row) in [r0, r1].into_iter().enumerate() {
4352 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4353 unsafe { *out_addr.at(bi * rows + o + r) = v };
4354 }
4355 bi += 1;
4356 }
4357 o += 2;
4358 } else {
4359 let row = &q[o * cols..(o + 1) * cols];
4360 for (bi, act) in acts.iter().enumerate() {
4361 let v = row_dot_avx2(row, act) * row_scale[o];
4362 unsafe { *out_addr.at(bi * rows + o) = v };
4363 }
4364 o += 1;
4365 }
4366 }
4367 };
4368 dispatch_rows(pool, rows, &run);
4369 return;
4370 }
4371 let run = |start: usize, end: usize| {
4372 for o in start..end {
4373 let row = &q[o * cols..(o + 1) * cols];
4374 for (bi, act) in acts.iter().enumerate() {
4375 let v = row_dot_avx2(row, act) * row_scale[o];
4376 unsafe { *out_addr.at(bi * rows + o) = v };
4377 }
4378 }
4379 };
4380 dispatch_rows(pool, rows, &run);
4381 return;
4382 }
4383 let out_addr = SendMut(out.as_mut_ptr());
4384 let run = |start: usize, end: usize| {
4385 for o in start..end {
4386 let row = &q[o * cols..(o + 1) * cols];
4387 for (bi, x) in pre.iter().enumerate() {
4388 let mut acc = 0f32;
4389 for j in 0..cols {
4390 acc += (row[j] as i8) as f32 * x[j];
4391 }
4392 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4393 }
4394 }
4395 };
4396 dispatch_rows(pool, rows, &run);
4397}
4398
4399fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4402 match pool {
4403 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4404 _ => run(0, rows),
4405 }
4406}
4407
4408fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4410 let groups = rows * cols / GROUP_SIZE;
4411 bytes.split_at(groups * 16)
4412}
4413
4414#[inline]
4419fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4420 #[cfg(target_arch = "aarch64")]
4421 unsafe {
4422 return vbit_fill4_neon(data, buf);
4423 }
4424 #[cfg(target_arch = "x86_64")]
4425 if avx2_enabled() {
4426 return unsafe { vbit_fill4_avx2(data, buf) };
4427 }
4428 #[allow(unreachable_code)]
4429 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4430 let u = unpack8::<4>(&data[blk * 4..]);
4431 for k in 0..8 {
4432 chunk[k] = (u[k] - 7) as i8 as u8;
4433 }
4434 }
4435}
4436
4437#[cfg(target_arch = "aarch64")]
4438#[target_feature(enable = "neon")]
4439unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4440 unsafe {
4443 use core::arch::aarch64::*;
4444 let n = buf.len();
4445 let mask = vdupq_n_u8(0x0F);
4446 let seven = vdupq_n_s8(7);
4447 let mut g = 0usize;
4448 while g * 32 + 32 <= n {
4449 let b = vld1q_u8(data.as_ptr().add(g * 16));
4450 let hi = vshrq_n_u8::<4>(b);
4451 let lo = vandq_u8(b, mask);
4452 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4453 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4454 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4455 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4456 g += 1;
4457 }
4458 }
4459}
4460
4461#[cfg(target_arch = "x86_64")]
4462#[target_feature(enable = "avx2")]
4463unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4464 unsafe {
4466 use core::arch::x86_64::*;
4467 let n = buf.len();
4468 let mask = _mm_set1_epi8(0x0F);
4469 let seven = _mm256_set1_epi8(7);
4470 let mut g = 0usize;
4471 while g * 32 + 32 <= n {
4472 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4473 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4474 let lo = _mm_and_si128(b, mask);
4475 let z = _mm256_sub_epi8(
4476 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4477 seven,
4478 );
4479 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4480 g += 1;
4481 }
4482 }
4483}
4484
4485#[inline(always)]
4490fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4491 let mut acc = 0u64;
4492 for i in 0..B {
4493 acc = (acc << 8) | data[i] as u64;
4494 }
4495 let mask = (1u64 << B) - 1;
4496 let mut out = [0i32; 8];
4497 for (k, o) in out.iter_mut().enumerate() {
4498 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4499 }
4500 out
4501}
4502
4503#[allow(clippy::too_many_arguments)]
4509fn vbitmatvec(
4510 bytes: &[u8],
4511 offsets: &[usize],
4512 x: &[f32],
4513 rows: usize,
4514 cols: usize,
4515 out: &mut [f32],
4516 pool: Option<&Pool>,
4517) {
4518 debug_assert_eq!(out.len(), rows);
4519 debug_assert_eq!(offsets.len(), rows + 1);
4520
4521 if a8w8_enabled() {
4525 let act = split_act(x);
4526 let out_addr = SendMut(out.as_mut_ptr());
4527 let run = move |start: usize, end: usize| {
4528 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4529 };
4530 dispatch_rows(pool, rows, &run);
4531 return;
4532 }
4533
4534 let out_addr = SendMut(out.as_mut_ptr());
4535 let run = move |start: usize, end: usize| {
4536 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4537 };
4538 dispatch_rows(pool, rows, &run);
4539}
4540
4541#[allow(clippy::too_many_arguments)]
4545fn vbit_range_a8w8(
4546 bytes: &[u8],
4547 offsets: &[usize],
4548 x: &[f32],
4549 act: &SplitAct,
4550 rows: usize,
4551 cols: usize,
4552 out: SendMut,
4553 start: usize,
4554 end: usize,
4555) {
4556 let ng = cols / GROUP_SIZE;
4557 let bits = &bytes[..rows];
4558 let sc_off = rows;
4559 let row_dot = |r: usize| -> f32 {
4560 let b = bits[r] as usize;
4561 let l = (1i32 << (b - 1)) - 1;
4562 let mask = (1u64 << b) - 1;
4563 let data = &bytes[offsets[r]..offsets[r + 1]];
4564 if b == 8 {
4565 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4567 let mut dot = 0f32;
4568 for g in 0..ng {
4569 let so = (r * ng + g) * 2;
4570 let sgf = f16_to_f32(u16::from_le_bytes([
4571 bytes[sc_off + so],
4572 bytes[sc_off + so + 1],
4573 ]));
4574 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4575 let mut gd = 0f32;
4576 for &xv in xg.iter() {
4577 if nbits < 8 {
4578 acc = (acc << 8) | data[idx] as u64;
4579 idx += 1;
4580 nbits += 8;
4581 }
4582 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4583 nbits -= 8;
4584 gd += (u - l) as f32 * xv;
4585 }
4586 dot += gd * sgf;
4587 }
4588 return dot;
4589 }
4590 thread_local! {
4594 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4595 const { std::cell::RefCell::new(Vec::new()) };
4596 }
4597 #[inline(always)]
4598 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4599 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4600 let u = unpack8::<B>(&data[blk * B..]);
4601 for k in 0..8 {
4602 chunk[k] = (u[k] - l) as i8 as u8;
4603 }
4604 }
4605 }
4606 let _ = mask;
4607 VBIT_SCRATCH.with(|scratch| {
4608 let mut buf = scratch.borrow_mut();
4609 buf.resize(cols, 0);
4610 match b {
4611 3 => fill::<3>(data, l, &mut buf),
4612 4 => vbit_fill4(data, &mut buf),
4613 5 => fill::<5>(data, l, &mut buf),
4614 6 => fill::<6>(data, l, &mut buf),
4615 _ => unreachable!(),
4616 }
4617 let mut dot = 0f32;
4618 for g in 0..ng {
4619 let so = (r * ng + g) * 2;
4620 let s = f16_to_f32(u16::from_le_bytes([
4621 bytes[sc_off + so],
4622 bytes[sc_off + so + 1],
4623 ]));
4624 let d = dot_i8_i8(
4625 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4626 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4627 ) as f32
4628 * act.sx;
4629 dot += d * s;
4630 }
4631 for &(j, xv) in &act.outliers {
4632 let so = (r * ng + j / GROUP_SIZE) * 2;
4633 let s = f16_to_f32(u16::from_le_bytes([
4634 bytes[sc_off + so],
4635 bytes[sc_off + so + 1],
4636 ]));
4637 dot += (buf[j] as i8) as f32 * s * xv;
4639 }
4640 dot
4641 })
4642 };
4643 for r in start..end {
4644 unsafe { *out.at(r) = row_dot(r) };
4646 }
4647}
4648
4649#[allow(clippy::too_many_arguments)]
4651fn vbit_range_f32(
4652 bytes: &[u8],
4653 offsets: &[usize],
4654 x: &[f32],
4655 rows: usize,
4656 cols: usize,
4657 out: SendMut,
4658 start: usize,
4659 end: usize,
4660) {
4661 let ng = cols / GROUP_SIZE;
4662 let bits = &bytes[..rows];
4663 let sc_off = rows;
4664 #[inline(always)]
4668 fn dot_row<const B: usize>(
4669 data: &[u8],
4670 bytes: &[u8],
4671 sc_off: usize,
4672 r: usize,
4673 ng: usize,
4674 x: &[f32],
4675 ) -> f32 {
4676 let l = ((1i32 << (B - 1)) - 1) as f32;
4677 let gbytes = GROUP_SIZE * B / 8;
4678 let mut dot = 0f32;
4679 for g in 0..ng {
4680 let so = (r * ng + g) * 2;
4681 let s = f16_to_f32(u16::from_le_bytes([
4682 bytes[sc_off + so],
4683 bytes[sc_off + so + 1],
4684 ]));
4685 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4686 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4687 let mut gd = 0f32;
4688 for blk in 0..GROUP_SIZE / 8 {
4689 let u = unpack8::<B>(&gd0[blk * B..]);
4690 let xb = &xg[blk * 8..blk * 8 + 8];
4691 for k in 0..8 {
4692 gd += (u[k] as f32 - l) * xb[k];
4693 }
4694 }
4695 dot += gd * s;
4696 }
4697 dot
4698 }
4699 for r in start..end {
4700 let data = &bytes[offsets[r]..offsets[r + 1]];
4701 let v = match bits[r] {
4702 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4703 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4704 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4705 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4706 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4707 b => unreachable!("vbit bit-width {b} (validated at load)"),
4708 };
4709 unsafe { *out.at(r) = v };
4711 }
4712}
4713
4714#[allow(clippy::too_many_arguments)]
4719fn vbitmatvec2(
4720 bytes: &[u8],
4721 offsets: &[usize],
4722 x1: &[f32],
4723 x2: &[f32],
4724 rows: usize,
4725 cols: usize,
4726 o1: &mut [f32],
4727 o2: &mut [f32],
4728 pool: Option<&Pool>,
4729) {
4730 debug_assert_eq!(o1.len(), rows);
4731 debug_assert_eq!(o2.len(), rows);
4732
4733 if a8w8_enabled() {
4734 let a1 = split_act(x1);
4735 let a2 = split_act(x2);
4736 let p1 = SendMut(o1.as_mut_ptr());
4737 let p2 = SendMut(o2.as_mut_ptr());
4738 let run = move |start: usize, end: usize| {
4739 vbit_range2_a8w8(
4740 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4741 )
4742 };
4743 dispatch_rows(pool, rows, &run);
4744 return;
4745 }
4746
4747 let p1 = SendMut(o1.as_mut_ptr());
4748 let p2 = SendMut(o2.as_mut_ptr());
4749 let run = move |start: usize, end: usize| {
4750 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4751 };
4752 dispatch_rows(pool, rows, &run);
4753}
4754
4755#[allow(clippy::too_many_arguments)]
4759fn vbit_range2_a8w8(
4760 bytes: &[u8],
4761 offsets: &[usize],
4762 x1: &[f32],
4763 x2: &[f32],
4764 a1: &SplitAct,
4765 a2: &SplitAct,
4766 rows: usize,
4767 cols: usize,
4768 p1: SendMut,
4769 p2: SendMut,
4770 start: usize,
4771 end: usize,
4772) {
4773 let ng = cols / GROUP_SIZE;
4774 let bits = &bytes[..rows];
4775 let sc_off = rows;
4776 let row_dots = |r: usize| -> (f32, f32) {
4777 let b = bits[r] as usize;
4778 let l = (1i32 << (b - 1)) - 1;
4779 let data = &bytes[offsets[r]..offsets[r + 1]];
4780 if b == 8 {
4781 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4784 let (mut d1, mut d2) = (0f32, 0f32);
4785 for g in 0..ng {
4786 let so = (r * ng + g) * 2;
4787 let sgf = f16_to_f32(u16::from_le_bytes([
4788 bytes[sc_off + so],
4789 bytes[sc_off + so + 1],
4790 ]));
4791 let (mut g1, mut g2) = (0f32, 0f32);
4792 for k in 0..GROUP_SIZE {
4793 if nbits < 8 {
4794 acc = (acc << 8) | data[idx] as u64;
4795 idx += 1;
4796 nbits += 8;
4797 }
4798 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4799 nbits -= 8;
4800 let w = (u - l) as f32;
4801 g1 += w * x1[g * GROUP_SIZE + k];
4802 g2 += w * x2[g * GROUP_SIZE + k];
4803 }
4804 d1 += g1 * sgf;
4805 d2 += g2 * sgf;
4806 }
4807 return (d1, d2);
4808 }
4809 thread_local! {
4810 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4811 const { std::cell::RefCell::new(Vec::new()) };
4812 }
4813 #[inline(always)]
4814 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4815 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4816 let u = unpack8::<B>(&data[blk * B..]);
4817 for k in 0..8 {
4818 chunk[k] = (u[k] - l) as i8 as u8;
4819 }
4820 }
4821 }
4822 VBIT_SCRATCH2.with(|scratch| {
4823 let mut buf = scratch.borrow_mut();
4824 buf.resize(cols, 0);
4825 match b {
4826 3 => fill::<3>(data, l, &mut buf),
4827 4 => vbit_fill4(data, &mut buf),
4828 5 => fill::<5>(data, l, &mut buf),
4829 6 => fill::<6>(data, l, &mut buf),
4830 _ => unreachable!(),
4831 }
4832 let (mut d1, mut d2) = (0f32, 0f32);
4833 for g in 0..ng {
4834 let so = (r * ng + g) * 2;
4835 let s = f16_to_f32(u16::from_le_bytes([
4836 bytes[sc_off + so],
4837 bytes[sc_off + so + 1],
4838 ]));
4839 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4840 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
4841 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
4842 d1 += v1 * s;
4843 d2 += v2 * s;
4844 }
4845 for &(j, xv) in &a1.outliers {
4846 let so = (r * ng + j / GROUP_SIZE) * 2;
4847 let s = f16_to_f32(u16::from_le_bytes([
4848 bytes[sc_off + so],
4849 bytes[sc_off + so + 1],
4850 ]));
4851 d1 += (buf[j] as i8) as f32 * s * xv;
4852 }
4853 for &(j, xv) in &a2.outliers {
4854 let so = (r * ng + j / GROUP_SIZE) * 2;
4855 let s = f16_to_f32(u16::from_le_bytes([
4856 bytes[sc_off + so],
4857 bytes[sc_off + so + 1],
4858 ]));
4859 d2 += (buf[j] as i8) as f32 * s * xv;
4860 }
4861 (d1, d2)
4862 })
4863 };
4864 for r in start..end {
4865 let (v1, v2) = row_dots(r);
4866 unsafe {
4868 *p1.at(r) = v1;
4869 *p2.at(r) = v2;
4870 }
4871 }
4872}
4873
4874#[allow(clippy::too_many_arguments)]
4878fn vbit_range2_f32(
4879 bytes: &[u8],
4880 offsets: &[usize],
4881 x1: &[f32],
4882 x2: &[f32],
4883 rows: usize,
4884 cols: usize,
4885 p1: SendMut,
4886 p2: SendMut,
4887 start: usize,
4888 end: usize,
4889) {
4890 let ng = cols / GROUP_SIZE;
4891 let bits = &bytes[..rows];
4892 let sc_off = rows;
4893 #[inline(always)]
4894 #[allow(clippy::too_many_arguments)]
4895 fn dot_row2<const B: usize>(
4896 data: &[u8],
4897 bytes: &[u8],
4898 sc_off: usize,
4899 r: usize,
4900 ng: usize,
4901 x1: &[f32],
4902 x2: &[f32],
4903 ) -> (f32, f32) {
4904 let l = ((1i32 << (B - 1)) - 1) as f32;
4905 let gbytes = GROUP_SIZE * B / 8;
4906 let (mut d1, mut d2) = (0f32, 0f32);
4907 for g in 0..ng {
4908 let so = (r * ng + g) * 2;
4909 let s = f16_to_f32(u16::from_le_bytes([
4910 bytes[sc_off + so],
4911 bytes[sc_off + so + 1],
4912 ]));
4913 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4914 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4915 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4916 let (mut g1, mut g2) = (0f32, 0f32);
4917 for blk in 0..GROUP_SIZE / 8 {
4918 let u = unpack8::<B>(&gd0[blk * B..]);
4919 for k in 0..8 {
4920 let w = u[k] as f32 - l;
4921 g1 += w * x1g[blk * 8 + k];
4922 g2 += w * x2g[blk * 8 + k];
4923 }
4924 }
4925 d1 += g1 * s;
4926 d2 += g2 * s;
4927 }
4928 (d1, d2)
4929 }
4930 for r in start..end {
4931 let data = &bytes[offsets[r]..offsets[r + 1]];
4932 let (v1, v2) = match bits[r] {
4933 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4934 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4935 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4936 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4937 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4938 b => unreachable!("vbit bit-width {b} (validated at load)"),
4939 };
4940 unsafe {
4942 *p1.at(r) = v1;
4943 *p2.at(r) = v2;
4944 }
4945 }
4946}
4947
4948#[inline]
4955#[allow(unreachable_code)]
4956fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4957 #[cfg(target_arch = "aarch64")]
4958 unsafe {
4959 return dot_q4t_row_sdot(bytes, r, gpr, xq);
4960 }
4961 #[cfg(target_arch = "x86_64")]
4962 unsafe {
4963 if vnni_tiles_enabled() {
4964 return dot_q4t_row_vnni(bytes, r, gpr, xq);
4965 }
4966 return dot_q4t_row_avx2(bytes, r, gpr, xq);
4967 }
4968 let mut acc = 0f32;
4969 for gi in 0..gpr {
4970 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4971 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4972 let mut d = 0i32;
4973 for (k, &b) in tile[2..].iter().enumerate() {
4974 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4975 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4976 }
4977 acc += d as f32 * s;
4978 }
4979 acc
4980}
4981
4982#[cfg(target_arch = "aarch64")]
4983#[target_feature(enable = "neon,dotprod")]
4984unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4985 unsafe {
4988 use core::arch::aarch64::*;
4989 use core::arch::asm;
4990 let lomask = vdupq_n_u8(0x0F);
4991 let eight = vdupq_n_s8(8);
4992 let mut acc = 0f32;
4993 for gi in 0..gpr {
4994 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4995 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4996 let b = vld1q_u8(t.add(2));
4997 let lo = vandq_u8(b, lomask);
4998 let hi = vshrq_n_u8::<4>(b);
4999 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5000 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5001 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5002 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5003 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5004 asm!(
5005 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5006 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5007 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5008 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5009 options(pure, nomem, nostack),
5010 );
5011 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5012 }
5013 acc
5014 }
5015}
5016
5017#[cfg(target_arch = "x86_64")]
5018#[target_feature(enable = "avx2")]
5019unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5020 unsafe {
5022 use core::arch::x86_64::*;
5023 let lomask = _mm_set1_epi8(0x0F);
5024 let eight = _mm256_set1_epi8(8);
5025 let ones = _mm256_set1_epi16(1);
5026 let mut acc = 0f32;
5027 for gi in 0..gpr {
5028 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5029 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5030 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5031 let lo = _mm_and_si128(b, lomask);
5032 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5033 let w = _mm256_sub_epi8(
5034 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5035 eight,
5036 );
5037 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5038 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5039 let d = _mm256_madd_epi16(p16, ones);
5040 let hi128 = _mm256_extracti128_si256::<1>(d);
5041 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5042 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5043 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5044 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5045 }
5046 acc
5047 }
5048}
5049
5050#[cfg(target_arch = "x86_64")]
5054#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5055unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5056 unsafe {
5058 use core::arch::x86_64::*;
5059 let lomask = _mm_set1_epi8(0x0F);
5060 let eight = _mm256_set1_epi8(8);
5061 let mut acc = 0f32;
5062 for gi in 0..gpr {
5063 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5064 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5065 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5066 let lo = _mm_and_si128(b, lomask);
5067 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5068 let w = _mm256_sub_epi8(
5069 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5070 eight,
5071 );
5072 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5073 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5074 acc += d as f32 * s;
5075 }
5076 acc
5077 }
5078}
5079
5080#[cfg(target_arch = "x86_64")]
5085#[target_feature(enable = "avx2,fma")]
5090unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5091 unsafe {
5093 use core::arch::x86_64::*;
5094 let lomask = _mm_set1_epi8(0x0F);
5095 let eight = _mm256_set1_epi8(8);
5096 let ones = _mm256_set1_epi16(1);
5097 let mut f0 = _mm256_setzero_ps();
5110 let mut f1 = _mm256_setzero_ps();
5111 let mut f2 = _mm256_setzero_ps();
5112 let mut f3 = _mm256_setzero_ps();
5113 for gi in 0..gpr {
5114 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5115 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5116 let sv = _mm256_set1_ps(s);
5117 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5118 let lo = _mm_and_si128(bb, lomask);
5119 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5120 let w = _mm256_sub_epi8(
5121 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5122 eight,
5123 );
5124 let aw = _mm256_abs_epi8(w);
5125 let off = gi * GROUP_SIZE;
5126 let dot = |xq: &[i8]| {
5127 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5128 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5129 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5130 };
5131 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5132 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5133 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5134 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5135 }
5136 [
5137 hsum256_ps(f0),
5138 hsum256_ps(f1),
5139 hsum256_ps(f2),
5140 hsum256_ps(f3),
5141 ]
5142 }
5143}
5144
5145#[cfg(target_arch = "x86_64")]
5148#[target_feature(enable = "avx2")]
5149#[inline]
5150unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5151 unsafe {
5153 use core::arch::x86_64::*;
5154 let hi = _mm256_extractf128_ps::<1>(v);
5155 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5156 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5157 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5158 _mm_cvtss_f32(s)
5159 }
5160}
5161
5162#[cfg(target_arch = "x86_64")]
5164#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5165unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5166 unsafe {
5168 use core::arch::x86_64::*;
5169 let lomask = _mm_set1_epi8(0x0F);
5170 let eight = _mm256_set1_epi8(8);
5171 let mut f0 = _mm256_setzero_ps();
5174 let mut f1 = _mm256_setzero_ps();
5175 let mut f2 = _mm256_setzero_ps();
5176 let mut f3 = _mm256_setzero_ps();
5177 for gi in 0..gpr {
5178 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5179 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5180 let sv = _mm256_set1_ps(s);
5181 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5182 let lo = _mm_and_si128(bb, lomask);
5183 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5184 let w = _mm256_sub_epi8(
5185 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5186 eight,
5187 );
5188 let aw = _mm256_abs_epi8(w);
5189 let off = gi * GROUP_SIZE;
5190 let dot = |xq: &[i8]| {
5191 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5192 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5193 _mm256_setzero_si256(),
5194 aw,
5195 _mm256_sign_epi8(x, w),
5196 ))
5197 };
5198 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5199 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5200 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5201 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5202 }
5203 let acc = [
5204 hsum256_ps(f0),
5205 hsum256_ps(f1),
5206 hsum256_ps(f2),
5207 hsum256_ps(f3),
5208 ];
5209 acc
5210 }
5211}
5212
5213#[cfg(target_arch = "aarch64")]
5218#[target_feature(enable = "neon,dotprod")]
5219unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5220 unsafe {
5222 use core::arch::aarch64::*;
5223 use core::arch::asm;
5224 let lomask = vdupq_n_u8(0x0F);
5225 let eight = vdupq_n_s8(8);
5226 let mut acc = [0f32; 4];
5227 for gi in 0..gpr {
5228 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5229 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5230 let b = vld1q_u8(t.add(2));
5231 let lo = vandq_u8(b, lomask);
5232 let hi = vshrq_n_u8::<4>(b);
5233 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5234 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5235 for (k, xq) in xs.iter().enumerate() {
5236 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5237 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5238 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5239 asm!(
5240 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5241 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5242 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5243 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5244 options(pure, nomem, nostack),
5245 );
5246 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5247 }
5248 }
5249 acc
5250 }
5251}
5252
5253#[inline]
5255fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5256 let gi = j / GROUP_SIZE;
5257 let k = j % GROUP_SIZE;
5258 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5259 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5260 let byte = tile[2 + k / 2];
5261 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5262 ((nib as i32 - 8) as f32, s)
5263}
5264
5265#[inline]
5268fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5269 let mut acc = 0f32;
5270 for gi in 0..gpr {
5271 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5272 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5273 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5274 let mut ga = 0f32;
5275 for (k, &b) in tile[2..].iter().enumerate() {
5276 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5277 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5278 }
5279 acc += ga * s;
5280 }
5281 acc
5282}
5283
5284struct Q4tpView<'a> {
5288 nib: &'a [u8],
5289 params: &'a [u8],
5290 codes: &'a [u8],
5291 stride: usize,
5292 zero_rung: bool,
5294}
5295
5296impl<'a> Q4tpView<'a> {
5297 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5298 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5299 Self {
5300 nib: &bytes[..params_off],
5301 params: &bytes[params_off..codes_off],
5302 codes: &bytes[codes_off..],
5303 stride,
5304 zero_rung: false,
5305 }
5306 }
5307
5308 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5310 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5311 Self {
5312 nib: &bytes[..params_off],
5313 params: &bytes[params_off..codes_off],
5314 codes: &bytes[codes_off..],
5315 stride,
5316 zero_rung: true,
5317 }
5318 }
5319
5320 #[inline]
5336 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5337 let tab = if self.zero_rung {
5338 q2tp_ladder(self.params, r)
5339 } else {
5340 q4tp_ladder(self.params, r)
5341 };
5342 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5343 let out = &mut out[..gpr];
5344 let mut chunks = out.chunks_exact_mut(8);
5345 let mut ci = 0usize;
5346 for c in &mut chunks {
5347 let w = u64::from(codes[ci])
5348 | u64::from(codes[ci + 1]) << 8
5349 | u64::from(codes[ci + 2]) << 16
5350 | u64::from(codes[ci + 3]) << 24
5351 | u64::from(codes[ci + 4]) << 32;
5352 for (k, o) in c.iter_mut().enumerate() {
5353 *o = tab[((w >> (5 * k)) & 31) as usize];
5354 }
5355 ci += 5;
5356 }
5357 let tail = &codes[ci..];
5360 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5361 *o = tab[q4tp_code(tail, k)];
5362 }
5363 }
5364}
5365
5366#[inline]
5367fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5368 #[cfg(target_arch = "aarch64")]
5369 unsafe {
5370 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5371 }
5372 #[cfg(target_arch = "x86_64")]
5373 unsafe {
5374 if vnni_tiles_enabled() {
5375 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5376 }
5377 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5378 }
5379 #[allow(unreachable_code)]
5380 {
5381 let mut acc = 0f32;
5382 for gi in 0..gpr {
5383 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5384 let s = scales[gi];
5385 let mut d = 0i32;
5386 for (k, &b) in tile.iter().enumerate() {
5387 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5388 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5389 }
5390 acc += d as f32 * s;
5391 }
5392 acc
5393 }
5394}
5395
5396#[cfg(target_arch = "aarch64")]
5399#[target_feature(enable = "neon,dotprod")]
5400unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5401 unsafe {
5404 use core::arch::aarch64::*;
5405 use core::arch::asm;
5406 let lomask = vdupq_n_u8(0x0F);
5407 let eight = vdupq_n_s8(8);
5408 let mut acc = 0f32;
5409 for gi in 0..gpr {
5410 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5411 let s = *scales.get_unchecked(gi);
5412 let b = vld1q_u8(t);
5413 let lo = vandq_u8(b, lomask);
5414 let hi = vshrq_n_u8::<4>(b);
5415 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5416 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5417 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5418 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5419 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5420 asm!(
5421 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5422 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5423 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5424 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5425 options(pure, nomem, nostack),
5426 );
5427 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5428 }
5429 acc
5430 }
5431}
5432
5433#[cfg(target_arch = "x86_64")]
5434#[target_feature(enable = "avx2")]
5435unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5436 unsafe {
5438 use core::arch::x86_64::*;
5439 let lomask = _mm_set1_epi8(0x0F);
5440 let eight = _mm256_set1_epi8(8);
5441 let ones = _mm256_set1_epi16(1);
5442 let mut acc = 0f32;
5443 for gi in 0..gpr {
5444 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5445 let s = *scales.get_unchecked(gi);
5446 let b = _mm_loadu_si128(t as *const __m128i);
5447 let lo = _mm_and_si128(b, lomask);
5448 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5449 let w = _mm256_sub_epi8(
5450 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5451 eight,
5452 );
5453 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5454 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5455 let d = _mm256_madd_epi16(p16, ones);
5456 let hi128 = _mm256_extracti128_si256::<1>(d);
5457 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5458 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5459 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5460 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5461 }
5462 acc
5463 }
5464}
5465
5466#[cfg(target_arch = "x86_64")]
5469#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5470unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5471 unsafe {
5473 use core::arch::x86_64::*;
5474 let lomask = _mm_set1_epi8(0x0F);
5475 let eight = _mm256_set1_epi8(8);
5476 let mut acc = 0f32;
5477 for gi in 0..gpr {
5478 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5479 let s = *scales.get_unchecked(gi);
5480 let b = _mm_loadu_si128(t as *const __m128i);
5481 let lo = _mm_and_si128(b, lomask);
5482 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5483 let w = _mm256_sub_epi8(
5484 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5485 eight,
5486 );
5487 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5488 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5489 }
5490 acc
5491 }
5492}
5493
5494#[inline]
5497fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5498 #[cfg(target_arch = "x86_64")]
5499 if avx2_enabled() {
5500 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5502 }
5503 q4tp_row_float_scalar(nib, r, gpr, x, scales)
5504}
5505
5506#[inline]
5507fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5508 let mut acc = 0f32;
5509 for gi in 0..gpr {
5510 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5511 let s = scales[gi];
5512 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5513 let mut ga = 0f32;
5514 for (k, &b) in tile.iter().enumerate() {
5515 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5516 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5517 }
5518 acc += ga * s;
5519 }
5520 acc
5521}
5522
5523#[cfg(target_arch = "x86_64")]
5527#[target_feature(enable = "avx2")]
5528unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5529 unsafe {
5532 use core::arch::x86_64::*;
5533 let mask = _mm_set1_epi8(15);
5534 let eight = _mm_set1_epi8(8);
5535 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5536 let mut acc = 0.0f32;
5537 for gi in 0..gpr {
5538 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5539 let lo = _mm_and_si128(packed, mask);
5540 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5541 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5542 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5543 let xp = x.as_ptr().add(gi * GROUP_SIZE);
5544 let a = _mm256_mul_ps(
5545 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5546 _mm256_loadu_ps(xp),
5547 );
5548 let b = _mm256_mul_ps(
5549 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5550 _mm256_loadu_ps(xp.add(8)),
5551 );
5552 let c = _mm256_mul_ps(
5553 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5554 _mm256_loadu_ps(xp.add(16)),
5555 );
5556 let d = _mm256_mul_ps(
5557 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5558 _mm256_loadu_ps(xp.add(24)),
5559 );
5560 let mut pairs = [0.0f32; 16];
5561 _mm256_storeu_ps(
5562 pairs.as_mut_ptr(),
5563 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5564 );
5565 _mm256_storeu_ps(
5566 pairs.as_mut_ptr().add(8),
5567 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5568 );
5569 let mut ga = 0.0f32;
5570 for v in pairs {
5571 ga += v;
5572 }
5573 acc += ga * scales[gi];
5574 }
5575 acc
5576 }
5577}
5578
5579#[inline]
5582fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5583 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5584 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5585 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5586 ((n as i32 - 8) as f32, scales[gi])
5587}
5588
5589fn q4tp_matvec(
5591 bytes: &[u8],
5592 x: &[f32],
5593 rows: usize,
5594 cols: usize,
5595 out: &mut [f32],
5596 pool: Option<&Pool>,
5597) {
5598 debug_assert_eq!(out.len(), rows);
5599 let gpr = cols / GROUP_SIZE;
5600 let v = Q4tpView::new(bytes, rows, cols);
5601 let out_addr = SendMut(out.as_mut_ptr());
5602 if a8w8_enabled() {
5603 let act = split_act(x);
5604 let run = |start: usize, end: usize| {
5605 with_krow(gpr, |sc| {
5607 for r in start..end {
5608 v.scales_into(r, gpr, sc);
5609 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5610 for &(j, xv) in &act.outliers {
5611 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5612 acc += w * s * xv;
5613 }
5614 unsafe { *out_addr.at(r) = acc };
5616 }
5617 })
5618 };
5619 dispatch_rows(pool, rows, &run);
5620 return;
5621 }
5622 let run = |start: usize, end: usize| {
5623 with_krow(gpr, |sc| {
5624 for r in start..end {
5625 v.scales_into(r, gpr, sc);
5626 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5628 }
5629 })
5630 };
5631 dispatch_rows(pool, rows, &run);
5632}
5633
5634#[allow(clippy::too_many_arguments)]
5637fn q4tp_matvec2(
5638 bytes: &[u8],
5639 x1: &[f32],
5640 x2: &[f32],
5641 rows: usize,
5642 cols: usize,
5643 o1: &mut [f32],
5644 o2: &mut [f32],
5645 pool: Option<&Pool>,
5646) {
5647 let gpr = cols / GROUP_SIZE;
5648 let v = Q4tpView::new(bytes, rows, cols);
5649 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5650 let run = |start: usize, end: usize| {
5651 let mut sc = vec![0f32; gpr];
5652 for r in start..end {
5653 v.scales_into(r, gpr, &mut sc);
5654 unsafe {
5656 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5657 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5658 }
5659 }
5660 };
5661 dispatch_rows(pool, rows, &run);
5662}
5663
5664#[inline]
5667fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5668 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5669 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5670 let c = (byte >> (2 * (k % 4))) & 3;
5671 (c as f32 - 1.5, scales[gi])
5672}
5673
5674#[cfg(target_arch = "x86_64")]
5675const Q2TP_DECODE_U32: [u32; 256] = {
5676 let mut tab = [0u32; 256];
5677 let mut b = 0usize;
5678 while b < 256 {
5679 tab[b] = ((b as u32) & 3)
5680 | ((((b as u32) >> 2) & 3) << 8)
5681 | ((((b as u32) >> 4) & 3) << 16)
5682 | ((((b as u32) >> 6) & 3) << 24);
5683 b += 1;
5684 }
5685 tab
5686};
5687
5688#[cfg(target_arch = "x86_64")]
5692#[target_feature(enable = "avx2")]
5693unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5694 use core::arch::x86_64::*;
5695 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5696 let codes = _mm256_setr_epi32(
5697 Q2TP_DECODE_U32[ch[0] as usize] as i32,
5698 Q2TP_DECODE_U32[ch[1] as usize] as i32,
5699 Q2TP_DECODE_U32[ch[2] as usize] as i32,
5700 Q2TP_DECODE_U32[ch[3] as usize] as i32,
5701 Q2TP_DECODE_U32[ch[4] as usize] as i32,
5702 Q2TP_DECODE_U32[ch[5] as usize] as i32,
5703 Q2TP_DECODE_U32[ch[6] as usize] as i32,
5704 Q2TP_DECODE_U32[ch[7] as usize] as i32,
5705 );
5706 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5707 let pair = _mm256_maddubs_epi16(codes, xv);
5708 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5709 let sum128 = _mm_add_epi32(
5710 _mm256_castsi256_si128(quad),
5711 _mm256_extracti128_si256(quad, 1),
5712 );
5713 let sum64 = _mm_hadd_epi32(sum128, sum128);
5714 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5715}
5716
5717#[cfg(target_arch = "x86_64")]
5718#[inline]
5719fn q2tp_avx2_enabled() -> bool {
5720 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5721 *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
5722}
5723
5724#[inline]
5731fn dot_q2tp_row_i8(
5732 chunks: &[u8],
5733 r: usize,
5734 gpr: usize,
5735 xq: &[i8],
5736 gsum: &[i32],
5737 scales: &[f32],
5738) -> f32 {
5739 let mut acc = 0f32;
5740 let base = r * gpr * Q2TP_CHUNK;
5741 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5742 let mut codes = [0i8; GROUP_SIZE];
5743 #[cfg(target_arch = "x86_64")]
5744 let avx2 = q2tp_avx2_enabled();
5749 for gi in 0..gpr {
5750 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5751 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5752 #[cfg(target_arch = "aarch64")]
5753 let dot = unsafe {
5759 use core::arch::aarch64::*;
5760 let b = vld1_u8(ch.as_ptr());
5761 let three = vdup_n_u8(3);
5762 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5763 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5764 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5765 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5766 let x4 = vld4_s8(xg.as_ptr());
5767 let mut acc4 = vdupq_n_s32(0);
5768 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5769 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5770 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5771 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5772 vaddvq_s32(acc4)
5773 };
5774 #[cfg(target_arch = "x86_64")]
5775 let dot: i32 = if avx2 {
5776 unsafe { q2tp_code_dot_avx2(ch, xg) }
5779 } else {
5780 ch.iter()
5781 .enumerate()
5782 .map(|(k, &b)| {
5783 ((b & 3) as i32) * xg[k * 4] as i32
5784 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5785 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5786 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5787 })
5788 .sum()
5789 };
5790 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5791 let dot: i32 = {
5792 for (k, &b) in ch.iter().enumerate() {
5793 codes[k * 4] = (b & 3) as i8;
5794 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5795 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5796 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5797 }
5798 codes
5799 .iter()
5800 .zip(xg)
5801 .map(|(&c, &x)| c as i32 * x as i32)
5802 .sum()
5803 };
5804 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5805 }
5806 acc
5807}
5808
5809fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5813 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5814}
5815
5816#[inline]
5820fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5821 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
5822}
5823
5824#[inline]
5825fn q2tp_row_exact_center(
5826 chunks: &[u8],
5827 r: usize,
5828 gpr: usize,
5829 x: &[f32],
5830 scales: &[f32],
5831 center: f32,
5832) -> f32 {
5833 let mut acc = 0f32;
5834 for gi in 0..gpr {
5835 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
5836 let s = scales[gi];
5837 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5838 let mut g = 0f32;
5839 for (k, &b) in ch.iter().enumerate() {
5840 g += ((b & 3) as f32 - center) * xb[k * 4]
5841 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
5842 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
5843 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
5844 }
5845 acc += s * g;
5846 }
5847 acc
5848}
5849
5850fn q2tp_matvec(
5851 bytes: &[u8],
5852 x: &[f32],
5853 rows: usize,
5854 cols: usize,
5855 out: &mut [f32],
5856 pool: Option<&Pool>,
5857) {
5858 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
5859}
5860
5861fn q2tp_affine_matvec(
5862 bytes: &[u8],
5863 x: &[f32],
5864 rows: usize,
5865 cols: usize,
5866 out: &mut [f32],
5867 pool: Option<&Pool>,
5868) {
5869 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
5870}
5871
5872fn q2tp_matvec_mode(
5873 bytes: &[u8],
5874 x: &[f32],
5875 rows: usize,
5876 cols: usize,
5877 out: &mut [f32],
5878 pool: Option<&Pool>,
5879 affine: bool,
5880) {
5881 debug_assert_eq!(out.len(), rows);
5882 let gpr = cols / GROUP_SIZE;
5883 let v = Q4tpView::new_q2(bytes, rows, cols);
5884 let out_addr = SendMut(out.as_mut_ptr());
5885 if !affine && a8w8_enabled() {
5890 let act = split_act(x);
5891 let gsum = q1_group_sums(&act.xq, gpr);
5892 let (act, gsum) = (&act, &gsum);
5893 let run = move |start: usize, end: usize| {
5894 with_krow(gpr, |sc| {
5895 for r in start..end {
5896 v.scales_into(r, gpr, sc);
5897 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
5898 for &(j, xv) in &act.outliers {
5899 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
5900 acc += w * s * xv;
5901 }
5902 unsafe { *out_addr.at(r) = acc };
5904 }
5905 })
5906 };
5907 dispatch_rows(pool, rows, &run);
5908 return;
5909 }
5910 let run = |start: usize, end: usize| {
5911 with_krow(gpr, |sc| {
5912 for r in start..end {
5913 v.scales_into(r, gpr, sc);
5914 unsafe {
5916 *out_addr.at(r) = if affine {
5917 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
5918 } else {
5919 q2tp_row_exact(v.nib, r, gpr, x, sc)
5920 }
5921 };
5922 }
5923 })
5924 };
5925 dispatch_rows(pool, rows, &run);
5926}
5927
5928#[allow(clippy::too_many_arguments)]
5930fn q2tp_matvec2(
5931 bytes: &[u8],
5932 x1: &[f32],
5933 x2: &[f32],
5934 rows: usize,
5935 cols: usize,
5936 o1: &mut [f32],
5937 o2: &mut [f32],
5938 pool: Option<&Pool>,
5939) {
5940 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
5941}
5942
5943#[allow(clippy::too_many_arguments)]
5944fn q2tp_affine_matvec2(
5945 bytes: &[u8],
5946 x1: &[f32],
5947 x2: &[f32],
5948 rows: usize,
5949 cols: usize,
5950 o1: &mut [f32],
5951 o2: &mut [f32],
5952 pool: Option<&Pool>,
5953) {
5954 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
5955}
5956
5957#[allow(clippy::too_many_arguments)]
5958fn q2tp_matvec2_mode(
5959 bytes: &[u8],
5960 x1: &[f32],
5961 x2: &[f32],
5962 rows: usize,
5963 cols: usize,
5964 o1: &mut [f32],
5965 o2: &mut [f32],
5966 pool: Option<&Pool>,
5967 affine: bool,
5968) {
5969 let gpr = cols / GROUP_SIZE;
5970 let v = Q4tpView::new_q2(bytes, rows, cols);
5971 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5972 let run = |start: usize, end: usize| {
5973 let mut sc = vec![0f32; gpr];
5974 for r in start..end {
5975 v.scales_into(r, gpr, &mut sc);
5976 unsafe {
5978 *p1.at(r) = if affine {
5979 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
5980 } else {
5981 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
5982 };
5983 *p2.at(r) = if affine {
5984 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
5985 } else {
5986 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
5987 };
5988 }
5989 }
5990 };
5991 dispatch_rows(pool, rows, &run);
5992}
5993
5994pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6001 let gpr = cols / GROUP_SIZE;
6006 let v = Q4tpView::new_q2(bytes, rows, cols);
6007 with_krow(gpr, |sc| {
6008 for r in 0..rows {
6009 v.scales_into(r, gpr, sc);
6010 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6011 }
6012 });
6013}
6014
6015pub fn q2tp_affine_matvec_for_test(
6018 bytes: &[u8],
6019 x: &[f32],
6020 rows: usize,
6021 cols: usize,
6022 out: &mut [f32],
6023) {
6024 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6025}
6026
6027pub fn q2tp_matmat_for_test(
6028 bytes: &[u8],
6029 xs_all: &[f32],
6030 b: usize,
6031 rows: usize,
6032 cols: usize,
6033 out: &mut [f32],
6034) {
6035 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6036}
6037
6038fn q2tp_matmat(
6039 bytes: &[u8],
6040 xs_all: &[f32],
6041 b: usize,
6042 rows: usize,
6043 cols: usize,
6044 out: &mut [f32],
6045 pool: Option<&Pool>,
6046) {
6047 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6048}
6049
6050fn q2tp_affine_matmat(
6051 bytes: &[u8],
6052 xs_all: &[f32],
6053 b: usize,
6054 rows: usize,
6055 cols: usize,
6056 out: &mut [f32],
6057 pool: Option<&Pool>,
6058) {
6059 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6060}
6061
6062fn q2tp_matmat_mode(
6063 bytes: &[u8],
6064 xs_all: &[f32],
6065 b: usize,
6066 rows: usize,
6067 cols: usize,
6068 out: &mut [f32],
6069 pool: Option<&Pool>,
6070 affine: bool,
6071) {
6072 debug_assert_eq!(out.len(), b * rows);
6073 let gpr = cols / GROUP_SIZE;
6074 let v = Q4tpView::new_q2(bytes, rows, cols);
6075 let out_addr = SendMut(out.as_mut_ptr());
6076 let run = |start: usize, end: usize| {
6077 let mut sc = vec![0f32; gpr];
6078 for r in start..end {
6079 v.scales_into(r, gpr, &mut sc);
6080 for bi in 0..b {
6081 let x = &xs_all[bi * cols..(bi + 1) * cols];
6082 unsafe {
6084 *out_addr.at(bi * rows + r) = if affine {
6085 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6086 } else {
6087 q2tp_row_exact(v.nib, r, gpr, x, &sc)
6088 }
6089 };
6090 }
6091 }
6092 };
6093 dispatch_rows(pool, rows, &run);
6094}
6095
6096#[cfg(target_arch = "aarch64")]
6104#[target_feature(enable = "neon,dotprod")]
6105unsafe fn dot_q4tp_row_1x4_sdot_v1(
6106 nib: &[u8],
6107 r: usize,
6108 gpr: usize,
6109 xs: [&[i8]; 4],
6110 scales: &[f32],
6111) -> [f32; 4] {
6112 unsafe {
6113 use core::arch::aarch64::*;
6114 use core::arch::asm;
6115 let lomask = vdupq_n_u8(0x0F);
6116 let eight = vdupq_n_s8(8);
6117 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6118 for gi in 0..gpr {
6119 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6120 let s = *scales.get_unchecked(gi);
6121 let bb = vld1q_u8(t);
6122 let lo = vandq_u8(bb, lomask);
6123 let hi = vshrq_n_u8::<4>(bb);
6124 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6125 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6126 let mut d = [0f32; 4];
6127 for (k, dk) in d.iter_mut().enumerate() {
6128 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6129 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6130 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6131 asm!(
6132 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6133 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6134 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6135 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6136 options(pure, nomem, nostack),
6137 );
6138 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6139 }
6140 f0 += d[0];
6141 f1 += d[1];
6142 f2 += d[2];
6143 f3 += d[3];
6144 }
6145 [f0, f1, f2, f3]
6146 }
6147}
6148
6149#[allow(dead_code)]
6156static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6157
6158#[cfg(test)]
6161static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6162
6163#[cfg(target_arch = "x86_64")]
6169fn q4tp_blocked_x86() -> bool {
6170 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6171 1 => false,
6172 2 => avx512vnni_enabled(),
6177 _ => blocked_enabled() && avx512vnni_enabled(),
6181 }
6182}
6183
6184#[cfg(target_arch = "aarch64")]
6186#[allow(dead_code)]
6187fn q4tp_v1() -> bool {
6188 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6189 1 => true,
6190 2 => false,
6191 _ => {
6192 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6193 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6194 }
6195 }
6196}
6197
6198#[cfg(target_arch = "x86_64")]
6210#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6211unsafe fn dot_q4tp_2x8_avx512(
6212 nib: &[u8],
6213 r0: usize,
6214 gpr: usize,
6215 xs: [&[i8]; 8],
6216 sc0: &[f32],
6217 sc1: &[f32],
6218) -> [[f32; 8]; 2] {
6219 unsafe {
6222 use core::arch::x86_64::*;
6223 let lomask = _mm256_set1_epi8(0x0F);
6224 let eight = _mm256_set1_epi8(8);
6225 let zero = _mm512_setzero_si512();
6226 let mut v0 = [_mm512_setzero_ps(); 8];
6227 let mut v1 = [_mm512_setzero_ps(); 8];
6228 let pairs = gpr / 2;
6229 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6230 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6231 let bb = _mm256_loadu_si256(t as *const __m256i);
6232 let lo = _mm256_and_si256(bb, lomask);
6233 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6234 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6235 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6236 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6237 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6238 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6239 };
6240 for gp in 0..pairs {
6241 let gi = gp * 2;
6242 let (wa0, neg0) = unpack(r0, gi);
6243 let (wa1, neg1) = unpack(r0 + 1, gi);
6244 let off = gi * GROUP_SIZE;
6245 let sv = |sc: &[f32]| {
6246 _mm512_insertf32x8::<1>(
6247 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6248 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6249 )
6250 };
6251 let s0 = sv(sc0);
6252 let s1 = sv(sc1);
6253 for k in 0..8 {
6254 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6255 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6256 zero,
6257 wa0,
6258 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6259 ));
6260 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6261 zero,
6262 wa1,
6263 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6264 ));
6265 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6266 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6267 }
6268 }
6269 let mut acc = [[0f32; 8]; 2];
6270 for k in 0..8 {
6271 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6272 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6273 }
6274 if gpr % 2 == 1 {
6275 let off = (gpr - 1) * GROUP_SIZE;
6276 for j in off..off + GROUP_SIZE {
6277 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6278 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6279 for k in 0..8 {
6280 let x = *xs[k].get_unchecked(j) as f32;
6281 acc[0][k] += w0 * sa * x;
6282 acc[1][k] += w1 * sb * x;
6283 }
6284 }
6285 }
6286 acc
6287 }
6288}
6289
6290#[cfg(target_arch = "x86_64")]
6295#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6296unsafe fn dot_q4tp_row_1x8_avx512(
6297 nib: &[u8],
6298 r: usize,
6299 gpr: usize,
6300 xs: [&[i8]; 8],
6301 scales: &[f32],
6302) -> [f32; 8] {
6303 unsafe {
6305 use core::arch::x86_64::*;
6306 let lomask = _mm256_set1_epi8(0x0F);
6307 let eight = _mm256_set1_epi8(8);
6308 let zero = _mm512_setzero_si512();
6309 let (mut v0, mut v1, mut v2, mut v3) = (
6310 _mm512_setzero_ps(),
6311 _mm512_setzero_ps(),
6312 _mm512_setzero_ps(),
6313 _mm512_setzero_ps(),
6314 );
6315 let (mut v4, mut v5, mut v6, mut v7) = (
6316 _mm512_setzero_ps(),
6317 _mm512_setzero_ps(),
6318 _mm512_setzero_ps(),
6319 _mm512_setzero_ps(),
6320 );
6321 let pairs = gpr / 2;
6322 for gp in 0..pairs {
6323 let gi = gp * 2;
6324 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6325 let bb = _mm256_loadu_si256(t as *const __m256i);
6326 let lo = _mm256_and_si256(bb, lomask);
6327 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6328 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6333 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6334 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6335 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6336 let wabs = _mm512_abs_epi8(w);
6337 let neg = _mm512_movepi8_mask(w);
6338 let off = gi * GROUP_SIZE;
6339 let sv = _mm512_insertf32x8::<1>(
6340 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6341 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6342 );
6343 let dot = |x: &[i8]| -> __m512 {
6344 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6345 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6346 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6347 };
6348 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6349 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6350 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6351 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6352 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6353 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6354 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6355 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6356 }
6357 let mut acc = [
6358 _mm512_reduce_add_ps(v0),
6359 _mm512_reduce_add_ps(v1),
6360 _mm512_reduce_add_ps(v2),
6361 _mm512_reduce_add_ps(v3),
6362 _mm512_reduce_add_ps(v4),
6363 _mm512_reduce_add_ps(v5),
6364 _mm512_reduce_add_ps(v6),
6365 _mm512_reduce_add_ps(v7),
6366 ];
6367 if gpr % 2 == 1 {
6370 let off = (gpr - 1) * GROUP_SIZE;
6371 for j in off..off + GROUP_SIZE {
6372 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6373 let ws = w * s;
6374 for k in 0..8 {
6375 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6376 }
6377 }
6378 }
6379 acc
6380 }
6381}
6382
6383#[cfg(target_arch = "x86_64")]
6396#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6397unsafe fn dot_q4tp_row_1x4_avx512(
6398 nib: &[u8],
6399 r: usize,
6400 gpr: usize,
6401 xs: [&[i8]; 4],
6402 scales: &[f32],
6403) -> [f32; 4] {
6404 unsafe {
6406 use core::arch::x86_64::*;
6407 let lomask = _mm256_set1_epi8(0x0F);
6408 let eight = _mm256_set1_epi8(8);
6409 let zero = _mm512_setzero_si512();
6410 let (mut v0, mut v1, mut v2, mut v3) = (
6411 _mm512_setzero_ps(),
6412 _mm512_setzero_ps(),
6413 _mm512_setzero_ps(),
6414 _mm512_setzero_ps(),
6415 );
6416 let pairs = gpr / 2;
6417 for gp in 0..pairs {
6418 let gi = gp * 2;
6419 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6420 let bb = _mm256_loadu_si256(t as *const __m256i);
6421 let lo = _mm256_and_si256(bb, lomask);
6422 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6423 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6428 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6429 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6430 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6431 let wabs = _mm512_abs_epi8(w);
6432 let neg = _mm512_movepi8_mask(w);
6433 let off = gi * GROUP_SIZE;
6434 let sv = _mm512_insertf32x8::<1>(
6435 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6436 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6437 );
6438 let dot = |x: &[i8]| -> __m512 {
6439 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6440 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6441 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6442 };
6443 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6444 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6445 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6446 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6447 }
6448 let mut acc = [
6449 _mm512_reduce_add_ps(v0),
6450 _mm512_reduce_add_ps(v1),
6451 _mm512_reduce_add_ps(v2),
6452 _mm512_reduce_add_ps(v3),
6453 ];
6454 if gpr % 2 == 1 {
6457 let off = (gpr - 1) * GROUP_SIZE;
6458 for j in off..off + GROUP_SIZE {
6459 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6460 let ws = w * s;
6461 for k in 0..4 {
6462 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6463 }
6464 }
6465 }
6466 acc
6467 }
6468}
6469
6470#[cfg(target_arch = "aarch64")]
6474#[target_feature(enable = "neon,dotprod")]
6475unsafe fn dot_q4tp_row_1x4_sdot(
6476 nib: &[u8],
6477 r: usize,
6478 gpr: usize,
6479 xs: [&[i8]; 4],
6480 scales: &[f32],
6481) -> [f32; 4] {
6482 unsafe {
6484 use core::arch::aarch64::*;
6485 use core::arch::asm;
6486 let lomask = vdupq_n_u8(0x0F);
6487 let eight = vdupq_n_s8(8);
6488 let (mut v0, mut v1, mut v2, mut v3) = (
6504 vdupq_n_f32(0.0),
6505 vdupq_n_f32(0.0),
6506 vdupq_n_f32(0.0),
6507 vdupq_n_f32(0.0),
6508 );
6509 for gi in 0..gpr {
6510 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6511 let s = *scales.get_unchecked(gi);
6512 let bb = vld1q_u8(t);
6513 let lo = vandq_u8(bb, lomask);
6514 let hi = vshrq_n_u8::<4>(bb);
6515 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6516 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6517 let off = gi * GROUP_SIZE;
6518 let dot4 = |x: &[i8]| -> int32x4_t {
6519 let x0 = vld1q_s8(x.as_ptr().add(off));
6520 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6521 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6522 asm!(
6523 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6524 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6525 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6526 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6527 options(pure, nomem, nostack),
6528 );
6529 vaddq_s32(a0, a1)
6530 };
6531 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6532 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6533 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6534 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6535 }
6536 [
6537 vaddvq_f32(v0),
6538 vaddvq_f32(v1),
6539 vaddvq_f32(v2),
6540 vaddvq_f32(v3),
6541 ]
6542 }
6543}
6544
6545fn q4tp_matmat(
6553 bytes: &[u8],
6554 xs_all: &[f32],
6555 b: usize,
6556 rows: usize,
6557 cols: usize,
6558 out: &mut [f32],
6559 pool: Option<&Pool>,
6560) {
6561 q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6562}
6563
6564#[allow(clippy::too_many_arguments)]
6568fn q4tp_matmat_with(
6569 bytes: &[u8],
6570 xs_all: &[f32],
6571 b: usize,
6572 rows: usize,
6573 cols: usize,
6574 out: &mut [f32],
6575 pool: Option<&Pool>,
6576 exact: bool,
6577) {
6578 debug_assert_eq!(out.len(), b * rows);
6579 let gpr = cols / GROUP_SIZE;
6580 let v = Q4tpView::new(bytes, rows, cols);
6581
6582 #[cfg(target_os = "macos")]
6586 if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6587 dequant_matmat_accel(
6588 &|r, dst| {
6589 let mut sc = [0f32; 32];
6590 let mut scv;
6591 let s: &[f32] = if gpr <= 32 {
6592 v.scales_into(r, gpr, &mut sc);
6593 &sc[..gpr]
6594 } else {
6595 scv = vec![0f32; gpr];
6596 v.scales_into(r, gpr, &mut scv);
6597 &scv
6598 };
6599 for gi in 0..gpr {
6600 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6601 for (k, &bb) in tile.iter().enumerate() {
6602 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6603 dst[gi * GROUP_SIZE + k * 2 + 1] =
6604 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6605 }
6606 }
6607 },
6608 xs_all,
6609 b,
6610 rows,
6611 cols,
6612 out,
6613 pool,
6614 );
6615 return;
6616 }
6617
6618 let out_addr = SendMut(out.as_mut_ptr());
6619 if a8w8_enabled() {
6620 let acts: Vec<SplitAct> = (0..b)
6621 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6622 .collect();
6623 let acts = &acts;
6624 #[cfg(target_arch = "aarch64")]
6631 let blocked_ok = sdot_enabled() && blocked_enabled();
6632 #[cfg(target_arch = "x86_64")]
6639 let blocked_ok = q4tp_blocked_x86() && !exact;
6640 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6641 let blocked_ok = {
6642 let _ = exact;
6643 false
6644 };
6645 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6655 .ok()
6656 .and_then(|v| v.parse().ok())
6657 .filter(|v| *v > 0)
6658 .unwrap_or(256);
6659 let run = |start: usize, end: usize| {
6660 for abase in (0..acts.len()).step_by(panel_cols) {
6661 let alen = (acts.len() - abase).min(panel_cols);
6662 let mut sc = vec![0f32; gpr];
6663 #[cfg(target_arch = "x86_64")]
6664 let mut r_lo = start;
6665 #[cfg(target_arch = "x86_64")]
6666 if blocked_ok && alen >= 8 {
6667 let mut sc1 = vec![0f32; gpr];
6668 while r_lo + 2 <= end {
6669 v.scales_into(r_lo, gpr, &mut sc);
6670 v.scales_into(r_lo + 1, gpr, &mut sc1);
6671 let mut bi = 0usize;
6672 while bi + 8 <= alen {
6673 let xs = [
6674 acts[abase + bi].xq.as_slice(),
6675 acts[abase + bi + 1].xq.as_slice(),
6676 acts[abase + bi + 2].xq.as_slice(),
6677 acts[abase + bi + 3].xq.as_slice(),
6678 acts[abase + bi + 4].xq.as_slice(),
6679 acts[abase + bi + 5].xq.as_slice(),
6680 acts[abase + bi + 6].xq.as_slice(),
6681 acts[abase + bi + 7].xq.as_slice(),
6682 ];
6683 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6684 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6685 for k in 0..8 {
6686 let act = &acts[abase + bi + k];
6687 let mut acc = dr[k] * act.sx;
6688 for &(j, xv) in &act.outliers {
6689 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6690 acc += w * s * xv;
6691 }
6692 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6694 }
6695 }
6696 bi += 8;
6697 }
6698 for row in [r_lo, r_lo + 1] {
6701 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6702 for b2 in bi..alen {
6703 let act = &acts[abase + b2];
6704 let xs4 = [
6705 act.xq.as_slice(),
6706 act.xq.as_slice(),
6707 act.xq.as_slice(),
6708 act.xq.as_slice(),
6709 ];
6710 let d =
6711 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6712 let mut acc = d[0] * act.sx;
6713 for &(j, xv) in &act.outliers {
6714 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6715 acc += w * s * xv;
6716 }
6717 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6719 }
6720 }
6721 r_lo += 2;
6722 }
6723 }
6724 #[cfg(target_arch = "x86_64")]
6725 let row_start = r_lo;
6726 #[cfg(not(target_arch = "x86_64"))]
6727 let row_start = start;
6728 for r in row_start..end {
6729 v.scales_into(r, gpr, &mut sc);
6730 let mut bi = 0usize;
6731 #[cfg(target_arch = "x86_64")]
6732 if blocked_ok {
6733 while bi + 8 <= alen {
6734 let xs = [
6735 acts[abase + bi].xq.as_slice(),
6736 acts[abase + bi + 1].xq.as_slice(),
6737 acts[abase + bi + 2].xq.as_slice(),
6738 acts[abase + bi + 3].xq.as_slice(),
6739 acts[abase + bi + 4].xq.as_slice(),
6740 acts[abase + bi + 5].xq.as_slice(),
6741 acts[abase + bi + 6].xq.as_slice(),
6742 acts[abase + bi + 7].xq.as_slice(),
6743 ];
6744 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6745 for k in 0..8 {
6746 let act = &acts[abase + bi + k];
6747 let mut acc = d[k] * act.sx;
6748 for &(j, xv) in &act.outliers {
6749 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6750 acc += w * s * xv;
6751 }
6752 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6754 }
6755 bi += 8;
6756 }
6757 while bi + 4 <= alen {
6758 let xs = [
6759 acts[abase + bi].xq.as_slice(),
6760 acts[abase + bi + 1].xq.as_slice(),
6761 acts[abase + bi + 2].xq.as_slice(),
6762 acts[abase + bi + 3].xq.as_slice(),
6763 ];
6764 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6765 for k in 0..4 {
6766 let act = &acts[abase + bi + k];
6767 let mut acc = d[k] * act.sx;
6768 for &(j, xv) in &act.outliers {
6769 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6770 acc += w * s * xv;
6771 }
6772 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6774 }
6775 bi += 4;
6776 }
6777 }
6778 #[cfg(target_arch = "aarch64")]
6779 if blocked_ok {
6780 while bi + 4 <= alen {
6781 let xs = [
6782 acts[abase + bi].xq.as_slice(),
6783 acts[abase + bi + 1].xq.as_slice(),
6784 acts[abase + bi + 2].xq.as_slice(),
6785 acts[abase + bi + 3].xq.as_slice(),
6786 ];
6787 let d = unsafe {
6788 if exact || q4tp_v1() {
6789 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6790 } else {
6791 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6792 }
6793 };
6794 for k in 0..4 {
6795 let act = &acts[abase + bi + k];
6796 let mut acc = d[k] * act.sx;
6797 for &(j, xv) in &act.outliers {
6798 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6799 acc += w * s * xv;
6800 }
6801 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6803 }
6804 bi += 4;
6805 }
6806 }
6807 let _ = blocked_ok;
6808 while bi < alen {
6809 let act = &acts[abase + bi];
6810 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6811 for &(j, xv) in &act.outliers {
6812 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6813 acc += w * s * xv;
6814 }
6815 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
6817 bi += 1;
6818 }
6819 }
6820 }
6821 };
6822 dispatch_rows(pool, rows, &run);
6823 return;
6824 }
6825
6826 let run = |start: usize, end: usize| {
6827 let mut sc = vec![0f32; gpr];
6828 for r in start..end {
6829 v.scales_into(r, gpr, &mut sc);
6830 for bi in 0..b {
6831 let x = &xs_all[bi * cols..(bi + 1) * cols];
6832 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
6834 }
6835 }
6836 };
6837 dispatch_rows(pool, rows, &run);
6838}
6839
6840fn q4t_matvec(
6842 bytes: &[u8],
6843 x: &[f32],
6844 rows: usize,
6845 cols: usize,
6846 out: &mut [f32],
6847 pool: Option<&Pool>,
6848) {
6849 debug_assert_eq!(out.len(), rows);
6850 let gpr = cols / GROUP_SIZE;
6851 let out_addr = SendMut(out.as_mut_ptr());
6852 if a8w8_enabled() {
6853 let act = split_act(x);
6854 let run = move |start: usize, end: usize| {
6855 for r in start..end {
6856 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6857 for &(j, xv) in &act.outliers {
6858 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6859 acc += w * s * xv;
6860 }
6861 unsafe { *out_addr.at(r) = acc };
6863 }
6864 };
6865 dispatch_rows(pool, rows, &run);
6866 return;
6867 }
6868 let run = move |start: usize, end: usize| {
6869 for r in start..end {
6870 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
6872 }
6873 };
6874 dispatch_rows(pool, rows, &run);
6875}
6876
6877#[allow(clippy::too_many_arguments)]
6879fn q4t_matvec2(
6880 bytes: &[u8],
6881 x1: &[f32],
6882 x2: &[f32],
6883 rows: usize,
6884 cols: usize,
6885 o1: &mut [f32],
6886 o2: &mut [f32],
6887 pool: Option<&Pool>,
6888) {
6889 let gpr = cols / GROUP_SIZE;
6890 let p1 = SendMut(o1.as_mut_ptr());
6891 let p2 = SendMut(o2.as_mut_ptr());
6892 if a8w8_enabled() {
6893 let a1 = split_act(x1);
6894 let a2 = split_act(x2);
6895 let run = move |start: usize, end: usize| {
6896 for r in start..end {
6897 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
6898 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
6899 for &(j, xv) in &a1.outliers {
6900 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6901 v1 += w * s * xv;
6902 }
6903 for &(j, xv) in &a2.outliers {
6904 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6905 v2 += w * s * xv;
6906 }
6907 unsafe {
6909 *p1.at(r) = v1;
6910 *p2.at(r) = v2;
6911 }
6912 }
6913 };
6914 dispatch_rows(pool, rows, &run);
6915 return;
6916 }
6917 let run = move |start: usize, end: usize| {
6918 for r in start..end {
6919 unsafe {
6921 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
6922 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
6923 }
6924 }
6925 };
6926 dispatch_rows(pool, rows, &run);
6927}
6928
6929#[allow(clippy::too_many_arguments)]
6931#[cfg(target_os = "macos")]
6937fn dequant_matmat_accel(
6938 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
6939 xs_all: &[f32],
6940 b: usize,
6941 rows: usize,
6942 cols: usize,
6943 out: &mut [f32],
6944 pool: Option<&Pool>,
6945) {
6946 const TR: usize = 2048;
6947 thread_local! {
6948 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6949 }
6950 WTILE.with(|wt| {
6951 let mut wtile = wt.borrow_mut();
6952 wtile.resize(TR * cols, 0.0);
6953 let mut r0 = 0usize;
6954 while r0 < rows {
6955 let tr = TR.min(rows - r0);
6956 let wt_addr = SendMut(wtile.as_mut_ptr());
6957 let run = |start: usize, end: usize| {
6958 for r in start..end {
6959 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
6961 dequant_row(r0 + r, dst);
6962 }
6963 };
6964 dispatch_rows(pool, tr, &run);
6965 unsafe {
6966 accel_blas::cblas_sgemm(
6967 101, 111, 112, b as i32,
6971 tr as i32,
6972 cols as i32,
6973 1.0,
6974 xs_all.as_ptr(),
6975 cols as i32,
6976 wtile.as_ptr(),
6977 cols as i32,
6978 0.0,
6979 out.as_mut_ptr().add(r0),
6980 rows as i32,
6981 );
6982 }
6983 r0 += tr;
6984 }
6985 });
6986}
6987
6988fn q4t_matmat(
6989 bytes: &[u8],
6990 xs_all: &[f32],
6991 b: usize,
6992 rows: usize,
6993 cols: usize,
6994 out: &mut [f32],
6995 pool: Option<&Pool>,
6996) {
6997 debug_assert_eq!(out.len(), b * rows);
6998 let gpr = cols / GROUP_SIZE;
6999 #[cfg(target_os = "macos")]
7003 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7004 dequant_matmat_accel(
7005 &|r, dst| {
7006 for gi in 0..gpr {
7007 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7008 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7009 for (k, &bb) in tile[2..].iter().enumerate() {
7010 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7011 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7012 }
7013 }
7014 },
7015 xs_all,
7016 b,
7017 rows,
7018 cols,
7019 out,
7020 pool,
7021 );
7022 return;
7023 }
7024 let out_addr = SendMut(out.as_mut_ptr());
7025 if a8w8_enabled() {
7026 let acts: Vec<SplitAct> = (0..b)
7027 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7028 .collect();
7029 let acts = &acts;
7030 #[cfg(target_arch = "x86_64")]
7031 let blocked_ok = avx2_enabled() && blocked_enabled();
7032 #[cfg(target_arch = "aarch64")]
7033 let blocked_ok = sdot_enabled() && blocked_enabled();
7034 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7035 let blocked_ok = false;
7036 let run = move |start: usize, end: usize| {
7037 for r in start..end {
7038 let mut bi = 0usize;
7039 #[cfg(target_arch = "aarch64")]
7040 if blocked_ok {
7041 while bi + 4 <= acts.len() {
7042 let xs = [
7043 acts[bi].xq.as_slice(),
7044 acts[bi + 1].xq.as_slice(),
7045 acts[bi + 2].xq.as_slice(),
7046 acts[bi + 3].xq.as_slice(),
7047 ];
7048 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7049 for k in 0..4 {
7050 let act = &acts[bi + k];
7051 let mut acc = d[k] * act.sx;
7052 for &(j, xv) in &act.outliers {
7053 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7054 acc += w * sc * xv;
7055 }
7056 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7058 }
7059 bi += 4;
7060 }
7061 }
7062 #[cfg(target_arch = "x86_64")]
7063 if blocked_ok {
7064 while bi + 4 <= acts.len() {
7065 let xs = [
7066 acts[bi].xq.as_slice(),
7067 acts[bi + 1].xq.as_slice(),
7068 acts[bi + 2].xq.as_slice(),
7069 acts[bi + 3].xq.as_slice(),
7070 ];
7071 let d = unsafe {
7072 if vnni_tiles_enabled() {
7073 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7074 } else {
7075 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7076 }
7077 };
7078 for k in 0..4 {
7079 let act = &acts[bi + k];
7080 let mut acc = d[k] * act.sx;
7081 for &(j, xv) in &act.outliers {
7082 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7083 acc += w * sc * xv;
7084 }
7085 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7087 }
7088 bi += 4;
7089 }
7090 }
7091 let _ = blocked_ok;
7092 while bi < acts.len() {
7093 let act = &acts[bi];
7094 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7095 for &(j, xv) in &act.outliers {
7096 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7097 acc += w * s * xv;
7098 }
7099 unsafe { *out_addr.at(bi * rows + r) = acc };
7101 bi += 1;
7102 }
7103 }
7104 };
7105 dispatch_rows(pool, rows, &run);
7106 return;
7107 }
7108 let run = move |start: usize, end: usize| {
7109 for r in start..end {
7110 for bi in 0..b {
7111 let x = &xs_all[bi * cols..(bi + 1) * cols];
7112 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7114 }
7115 }
7116 };
7117 dispatch_rows(pool, rows, &run);
7118}
7119
7120fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7129 (0..gpr)
7130 .map(|gi| {
7131 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7132 .iter()
7133 .map(|&v| v as i32)
7134 .sum()
7135 })
7136 .collect()
7137}
7138
7139#[inline]
7143#[allow(unreachable_code)]
7144#[cfg(target_arch = "x86_64")]
7149#[target_feature(enable = "avx2")]
7150unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7151 unsafe {
7153 use core::arch::x86_64::*;
7154 let expand = _mm256_setr_epi8(
7156 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,
7157 3, 3, 3,
7158 );
7159 let bitsel = _mm256_setr_epi8(
7160 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7161 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7162 );
7163 let ones8 = _mm256_set1_epi8(1);
7164 let ones16 = _mm256_set1_epi16(1);
7165 let mut acc = 0f32;
7166 for gi in 0..gpr {
7167 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7168 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7169 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7170 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7171 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7172 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7173 let sel = _mm256_and_si256(x, mask);
7174 let p16 = _mm256_maddubs_epi16(ones8, sel);
7176 let d32 = _mm256_madd_epi16(p16, ones16);
7177 let hi128 = _mm256_extracti128_si256::<1>(d32);
7178 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7179 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7180 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7181 let msum = _mm_cvtsi128_si32(s32);
7182 let d = 2 * msum - gsum[gi];
7185 acc += d as f32 * s;
7186 }
7187 acc
7188 }
7189}
7190
7191#[cfg(target_arch = "x86_64")]
7194#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7195unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7196 unsafe {
7198 use core::arch::x86_64::*;
7199 let expand = _mm256_setr_epi8(
7200 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,
7201 3, 3, 3,
7202 );
7203 let bitsel = _mm256_setr_epi8(
7204 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7205 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7206 );
7207 let ones8 = _mm256_set1_epi8(1);
7208 let mut acc = 0f32;
7209 for gi in 0..gpr {
7210 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7211 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7212 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7213 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7214 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7215 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7216 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7217 let d = 2 * msum - gsum[gi];
7218 acc += d as f32 * s;
7219 }
7220 acc
7221 }
7222}
7223
7224#[cfg(target_arch = "x86_64")]
7226#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7227unsafe fn dot_q1_row_1x4_vnni(
7228 bytes: &[u8],
7229 r: usize,
7230 gpr: usize,
7231 xs: [&[i8]; 4],
7232 gsums: [&[i32]; 4],
7233) -> [f32; 4] {
7234 unsafe {
7236 use core::arch::x86_64::*;
7237 let expand = _mm256_setr_epi8(
7238 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,
7239 3, 3, 3,
7240 );
7241 let bitsel = _mm256_setr_epi8(
7242 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7243 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7244 );
7245 let ones8 = _mm256_set1_epi8(1);
7246 let mut acc = [0f32; 4];
7247 for gi in 0..gpr {
7248 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7249 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7250 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7251 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7252 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7253 for (k, xq) in xs.iter().enumerate() {
7254 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7255 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7256 let d = 2 * msum - gsums[k][gi];
7257 acc[k] += d as f32 * s;
7258 }
7259 }
7260 acc
7261 }
7262}
7263
7264#[cfg(target_arch = "x86_64")]
7267#[target_feature(enable = "avx2")]
7268unsafe fn dot_q1_row_1x4_avx2(
7269 bytes: &[u8],
7270 r: usize,
7271 gpr: usize,
7272 xs: [&[i8]; 4],
7273 gsums: [&[i32]; 4],
7274) -> [f32; 4] {
7275 unsafe {
7277 use core::arch::x86_64::*;
7278 let expand = _mm256_setr_epi8(
7279 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,
7280 3, 3, 3,
7281 );
7282 let bitsel = _mm256_setr_epi8(
7283 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7284 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7285 );
7286 let ones8 = _mm256_set1_epi8(1);
7287 let ones16 = _mm256_set1_epi16(1);
7288 let mut acc = [0f32; 4];
7289 for gi in 0..gpr {
7290 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7291 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7292 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7293 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7294 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7295 for (k, xq) in xs.iter().enumerate() {
7296 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7297 let sel = _mm256_and_si256(x, mask);
7298 let p16 = _mm256_maddubs_epi16(ones8, sel);
7299 let d32 = _mm256_madd_epi16(p16, ones16);
7300 let hi128 = _mm256_extracti128_si256::<1>(d32);
7301 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7302 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7303 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7304 let msum = _mm_cvtsi128_si32(s32);
7305 let d = 2 * msum - gsums[k][gi];
7306 acc[k] += d as f32 * s;
7307 }
7308 }
7309 acc
7310 }
7311}
7312
7313#[allow(unreachable_code)]
7314fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7315 #[cfg(target_arch = "aarch64")]
7316 unsafe {
7317 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7318 }
7319 #[cfg(target_arch = "x86_64")]
7320 if avx2_enabled() {
7321 unsafe {
7322 if vnni_tiles_enabled() {
7323 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7324 }
7325 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7326 }
7327 }
7328 let _ = gsum;
7329 let mut acc = 0f32;
7330 for gi in 0..gpr {
7331 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7332 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7333 let mut d = 0i32;
7334 for (j, &b) in tile[2..].iter().enumerate() {
7335 for k in 0..8 {
7336 let w = ((b >> k) & 1) as i32 * 2 - 1;
7337 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7338 }
7339 }
7340 acc += d as f32 * s;
7341 }
7342 acc
7343}
7344
7345#[cfg(target_arch = "aarch64")]
7354#[target_feature(enable = "neon,dotprod")]
7355unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7356 unsafe {
7359 use core::arch::aarch64::*;
7360 use core::arch::asm;
7361 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7362 let m = vld1q_u8(MASKS.as_ptr());
7363 macro_rules! tile_dot {
7365 ($t:expr, $x:expr) => {{
7366 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7367 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7368 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7369 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7370 let x0 = vld1q_s8($x);
7371 let x1 = vld1q_s8($x.add(16));
7372 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7373 asm!(
7374 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7375 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7376 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7377 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7378 options(pure, nomem, nostack),
7379 );
7380 vaddq_s32(a0, a1)
7381 }};
7382 }
7383 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7392 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7393 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7394 const IW11: [u8; 16] = [
7395 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7396 ];
7397 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7398 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7399 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7400 let isc = vld1_u8(ISC.as_ptr());
7401 macro_rules! tile_dot_tbl {
7403 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7404 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7405 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7406 let x0 = vld1q_s8($x);
7407 let x1 = vld1q_s8($x.add(16));
7408 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7409 asm!(
7410 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7411 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7412 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7413 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7414 options(pure, nomem, nostack),
7415 );
7416 vaddq_s32(a0, a1)
7417 }};
7418 }
7419 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7420 let row_base = r * gpr * Q1_TILE;
7421 let abs_end = bytes.len();
7422 let xp = xq.as_ptr();
7423 let gp = gsum.as_ptr();
7424 let mut accv = vdupq_n_f32(0.0);
7425 let mut gi = 0;
7426 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7429 let t0 = base.add(gi * Q1_TILE);
7430 let ld_a = vld1q_u8(t0);
7431 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7432 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7433 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7434 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7435 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7436 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7438 let g = vld1q_s32(gp.add(gi));
7439 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7440 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7441 let scf: float32x4_t;
7442 asm!(
7443 "fcvtl {o:v}.4s, {i:v}.4h",
7444 o = out(vreg) scf, i = in(vreg) sc16,
7445 options(pure, nomem, nostack),
7446 );
7447 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7448 gi += 4;
7449 }
7450 let mut acc = vaddvq_f32(accv);
7451 while gi < gpr {
7452 let t = base.add(gi * Q1_TILE);
7453 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7454 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7455 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7456 gi += 1;
7457 }
7458 acc
7459 }
7460}
7461
7462#[cfg(target_arch = "aarch64")]
7467#[target_feature(enable = "neon,dotprod")]
7468unsafe fn dot_q1_row_1x4_sdot(
7469 bytes: &[u8],
7470 r: usize,
7471 gpr: usize,
7472 xs: [&[i8]; 4],
7473 gs: [&[i32]; 4],
7474) -> [f32; 4] {
7475 unsafe {
7477 use core::arch::aarch64::*;
7478 use core::arch::asm;
7479 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7480 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7481 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7482 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7483 const IW11: [u8; 16] = [
7484 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7485 ];
7486 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7487 let m = vld1q_u8(MASKS.as_ptr());
7488 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7489 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7490 let isc = vld1_u8(ISC.as_ptr());
7491 macro_rules! sdot2 {
7492 ($w0:expr, $w1:expr, $x:expr) => {{
7493 let x0 = vld1q_s8($x);
7494 let x1 = vld1q_s8($x.add(16));
7495 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7496 asm!(
7497 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7498 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7499 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7500 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7501 options(pure, nomem, nostack),
7502 );
7503 vaddq_s32(a0, a1)
7504 }};
7505 }
7506 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7507 let row_base = r * gpr * Q1_TILE;
7508 let abs_end = bytes.len();
7509 let mut accv = [vdupq_n_f32(0.0); 4];
7510 let mut gi = 0;
7511 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7512 let t0 = base.add(gi * Q1_TILE);
7513 let ld_a = vld1q_u8(t0);
7514 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7515 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7517 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7518 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7519 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7520 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7521 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7522 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7523 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7524 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7525 let scf: float32x4_t;
7526 asm!(
7527 "fcvtl {o:v}.4s, {i:v}.4h",
7528 o = out(vreg) scf, i = in(vreg) sc16,
7529 options(pure, nomem, nostack),
7530 );
7531 for k in 0..4 {
7532 let xp = xs[k].as_ptr();
7533 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7534 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7535 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7536 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7537 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7538 let g = vld1q_s32(gs[k].as_ptr().add(gi));
7539 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7540 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7541 }
7542 gi += 4;
7543 }
7544 let mut acc = [
7545 vaddvq_f32(accv[0]),
7546 vaddvq_f32(accv[1]),
7547 vaddvq_f32(accv[2]),
7548 vaddvq_f32(accv[3]),
7549 ];
7550 while gi < gpr {
7551 let t = base.add(gi * Q1_TILE);
7552 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7553 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7554 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7555 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7556 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7557 for k in 0..4 {
7558 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7559 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7560 }
7561 gi += 1;
7562 }
7563 acc
7564 }
7565}
7566
7567#[inline]
7569fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7570 let gi = j / GROUP_SIZE;
7571 let k = j % GROUP_SIZE;
7572 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7573 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7574 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7575 ((bit as i32 * 2 - 1) as f32, s)
7576}
7577
7578#[inline]
7580fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7581 let mut acc = 0f32;
7582 for gi in 0..gpr {
7583 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7584 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7585 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7586 let mut ga = 0f32;
7587 for (j, &b) in tile[2..].iter().enumerate() {
7588 for k in 0..8 {
7589 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7590 }
7591 }
7592 acc += ga * s;
7593 }
7594 acc
7595}
7596
7597#[allow(clippy::too_many_arguments)]
7600fn q1_range_a8w8(
7601 bytes: &[u8],
7602 gpr: usize,
7603 act: &SplitAct,
7604 gsum: &[i32],
7605 out: SendMut,
7606 start: usize,
7607 end: usize,
7608) {
7609 for r in start..end {
7610 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7611 for &(j, xv) in &act.outliers {
7612 let (w, s) = q1_outlier(bytes, r, gpr, j);
7613 acc += w * s * xv;
7614 }
7615 unsafe { *out.at(r) = acc };
7617 }
7618}
7619
7620fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7622 for r in start..end {
7623 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
7625 }
7626}
7627
7628fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
7633 let entries = base_len + (rows + 1) * 4;
7634 (base_len, entries, entries <= bytes.len())
7635}
7636
7637#[inline]
7639fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
7640 let o = rp_off + r * 4;
7641 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
7642}
7643
7644const SIGN5: [[f32; 5]; 256] = {
7648 let mut lut = [[0.0f32; 5]; 256];
7649 let pow3 = [1u16, 3, 9, 27, 81];
7650 let mut byte = 0usize;
7651 while byte < 256 {
7652 let mut i = 0usize;
7653 while i < 5 {
7654 let code = (byte as u16 / pow3[i]) % 3;
7655 lut[byte][i] = if code == 1 {
7656 1.0
7657 } else if code == 2 {
7658 -1.0
7659 } else {
7660 0.0
7661 };
7662 i += 1;
7663 }
7664 byte += 1;
7665 }
7666 lut
7667};
7668
7669const SIGN5_I8: [[i8; 5]; 256] = {
7671 let mut lut = [[0i8; 5]; 256];
7672 let pow3 = [1u16, 3, 9, 27, 81];
7673 let mut byte = 0usize;
7674 while byte < 256 {
7675 let mut i = 0usize;
7676 while i < 5 {
7677 let code = (byte as u16 / pow3[i]) % 3;
7678 lut[byte][i] = if code == 1 {
7679 1
7680 } else if code == 2 {
7681 -1
7682 } else {
7683 0
7684 };
7685 i += 1;
7686 }
7687 byte += 1;
7688 }
7689 lut
7690};
7691
7692const SIGN5_U64: [u64; 256] = {
7698 let mut lut = [0u64; 256];
7699 let pow3 = [1u16, 3, 9, 27, 81];
7700 let mut byte = 0usize;
7701 while byte < 256 {
7702 let mut v = 0u64;
7703 let mut i = 0usize;
7704 while i < 5 {
7705 let code = (byte as u16 / pow3[i]) % 3;
7706 let s: u8 = if code == 1 {
7707 1
7708 } else if code == 2 {
7709 0xFF
7710 } else {
7711 0
7712 };
7713 v |= (s as u64) << (i * 8);
7714 i += 1;
7715 }
7716 lut[byte] = v;
7717 byte += 1;
7718 }
7719 lut
7720};
7721
7722#[inline]
7727fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
7728 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7729 let off = (r * gpr + j / GROUP_SIZE) * TILE;
7730 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7731 let within = j % GROUP_SIZE;
7732 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
7733}
7734
7735#[cfg(target_arch = "aarch64")]
7738#[target_feature(enable = "neon,dotprod")]
7739#[inline]
7740unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
7741 unsafe {
7743 use core::arch::aarch64::*;
7744 use core::arch::asm;
7745 let w0 = vld1q_s8(w);
7746 let w1 = vld1q_s8(w.add(16));
7747 let x0 = vld1q_s8(x);
7748 let x1 = vld1q_s8(x.add(16));
7749 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7750 asm!(
7751 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7752 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7753 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7754 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7755 options(pure, nomem, nostack),
7756 );
7757 vaddvq_s32(vaddq_s32(a0, a1))
7758 }
7759}
7760
7761#[cfg(target_arch = "x86_64")]
7764#[target_feature(enable = "avx2")]
7765#[inline]
7766unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
7767 unsafe {
7769 use core::arch::x86_64::*;
7770 let wv = _mm256_loadu_si256(w as *const __m256i);
7771 let xv = _mm256_loadu_si256(x as *const __m256i);
7772 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7773 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
7774 let hi128 = _mm256_extracti128_si256::<1>(d);
7775 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7776 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7777 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7778 _mm_cvtsi128_si32(s32)
7779 }
7780}
7781
7782#[inline]
7787fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
7788 debug_assert!(dst.len() >= 40);
7789 unsafe {
7792 let p = dst.as_mut_ptr();
7793 for bi in 0..7 {
7794 core::ptr::write_unaligned(
7795 p.add(bi * 5) as *mut u64,
7796 SIGN5_U64[*codes.add(bi) as usize],
7797 );
7798 }
7799 }
7800}
7801
7802#[inline]
7807fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
7808 #[cfg(target_arch = "aarch64")]
7809 unsafe {
7810 return sdot32_i8(w, x);
7811 }
7812 #[cfg(target_arch = "x86_64")]
7813 unsafe {
7814 return i8dot32_avx2(w, x);
7815 }
7816 #[allow(unreachable_code)]
7817 unsafe {
7818 let mut s = 0i32;
7819 for k in 0..GROUP_SIZE {
7820 s += *w.add(k) as i32 * *x.add(k) as i32;
7821 }
7822 s
7823 }
7824}
7825
7826#[inline]
7827unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
7828 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
7829 (
7830 SIGN5_U64[*codes as usize],
7831 SIGN5_U64[*codes.add(1) as usize],
7832 SIGN5_U64[*codes.add(2) as usize],
7833 SIGN5_U64[*codes.add(3) as usize],
7834 SIGN5_U64[*codes.add(4) as usize],
7835 SIGN5_U64[*codes.add(5) as usize],
7836 SIGN5_U64[*codes.add(6) as usize],
7837 )
7838 };
7839
7840 let u0 = s0 | (s1 << 40);
7841 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
7842 let u2 = (s3 >> 8) | (s4 << 32);
7843 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
7844
7845 (u0, u1, u2, u3)
7846}
7847
7848#[cfg(target_arch = "aarch64")]
7852#[target_feature(enable = "neon,dotprod")]
7853unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7854 use core::arch::aarch64::*;
7855 use core::arch::asm;
7856 unsafe {
7857 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7858 let mut acc = 0f32;
7859 let bytes_ptr = bytes.as_ptr();
7860 let xq_ptr = xq.as_ptr();
7861 let row_off = r * gpr * TILE;
7862
7863 let gpr2 = gpr & !1;
7864 let mut gi = 0;
7865 while gi < gpr2 {
7866 let off0 = row_off + gi * TILE;
7867 let off1 = off0 + TILE;
7868 let s0 = f16_to_f32(u16::from_le_bytes([
7869 *bytes_ptr.add(off0),
7870 *bytes_ptr.add(off0 + 1),
7871 ]));
7872 let s1 = f16_to_f32(u16::from_le_bytes([
7873 *bytes_ptr.add(off1),
7874 *bytes_ptr.add(off1 + 1),
7875 ]));
7876
7877 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7878 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7879
7880 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7881 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7882 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7883 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7884
7885 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7886 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7887 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
7888 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
7889
7890 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
7891 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7892 asm!(
7893 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
7894 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
7895 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
7896 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
7897 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
7898 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
7899 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
7900 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
7901 options(pure, nomem, nostack),
7902 );
7903 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
7904 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
7905 acc += d0 as f32 * s0 + d1 as f32 * s1;
7906 gi += 2;
7907 }
7908
7909 if gi < gpr {
7910 let off = row_off + gi * TILE;
7911 let s = f16_to_f32(u16::from_le_bytes([
7912 *bytes_ptr.add(off),
7913 *bytes_ptr.add(off + 1),
7914 ]));
7915 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7916 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7917 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7918 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7919 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7920 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7921 asm!(
7922 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7923 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7924 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7925 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7926 options(pure, nomem, nostack),
7927 );
7928 let d = vaddvq_s32(vaddq_s32(a0, a1));
7929 acc += d as f32 * s;
7930 }
7931 acc
7932 }
7933}
7934
7935#[cfg(target_arch = "x86_64")]
7937#[target_feature(enable = "avx2")]
7938unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7939 use core::arch::x86_64::*;
7940 unsafe {
7941 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7942 let mut acc = 0f32;
7943 let bytes_ptr = bytes.as_ptr();
7944 let xq_ptr = xq.as_ptr();
7945 let row_off = r * gpr * TILE;
7946
7947 let ones = _mm256_set1_epi16(1);
7948 for gi in 0..gpr {
7949 let off = row_off + gi * TILE;
7950 let s = f16_to_f32(u16::from_le_bytes([
7951 *bytes_ptr.add(off),
7952 *bytes_ptr.add(off + 1),
7953 ]));
7954 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7955 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7956 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7957 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7958 let d256 = _mm256_madd_epi16(p16, ones);
7959 let d128 = _mm_add_epi32(
7960 _mm256_castsi256_si128(d256),
7961 _mm256_extracti128_si256(d256, 1),
7962 );
7963 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
7964 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
7965 acc += d32 as f32 * s;
7966 }
7967 acc
7968 }
7969}
7970
7971#[cfg(target_arch = "x86_64")]
7973#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7974unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7975 use core::arch::x86_64::*;
7976 unsafe {
7978 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7979 let mut acc = 0f32;
7980 let bytes_ptr = bytes.as_ptr();
7981 let xq_ptr = xq.as_ptr();
7982 let row_off = r * gpr * TILE;
7983 for gi in 0..gpr {
7984 let off = row_off + gi * TILE;
7985 let s = f16_to_f32(u16::from_le_bytes([
7986 *bytes_ptr.add(off),
7987 *bytes_ptr.add(off + 1),
7988 ]));
7989 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7990 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7991 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7992 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7993 acc += d as f32 * s;
7994 }
7995 acc
7996 }
7997}
7998
7999#[inline]
8003fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8004 #[cfg(target_arch = "aarch64")]
8005 unsafe {
8006 return q1t_dot_row_sdot(bytes, r, gpr, xq);
8007 }
8008 #[cfg(target_arch = "x86_64")]
8009 unsafe {
8010 if vnni_tiles_enabled() {
8011 return q1t_dot_row_vnni(bytes, r, gpr, xq);
8012 }
8013 return q1t_dot_row_avx2(bytes, r, gpr, xq);
8014 }
8015 #[allow(unreachable_code)]
8016 {
8017 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8018 let mut acc = 0f32;
8019 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
8021 let off = (r * gpr + gi) * TILE;
8022 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8023 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
8024 let mut d = 0i32;
8025 for k in 0..GROUP_SIZE {
8026 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
8027 }
8028 acc += d as f32 * s;
8029 }
8030 acc
8031 }
8032}
8033
8034fn q1t_row_outlier_correction(
8041 bytes: &[u8],
8042 r: usize,
8043 rp_off: usize,
8044 entries_off: usize,
8045 has_ov: bool,
8046 x: &[f32],
8047) -> f32 {
8048 if !has_ov {
8049 return 0.0;
8050 }
8051 let (c0, c1) = (
8052 q1t_rowptr(bytes, rp_off, r),
8053 q1t_rowptr(bytes, rp_off, r + 1),
8054 );
8055 let mut corr = 0f32;
8056 for p in c0..c1 {
8057 let e = entries_off + p * 4;
8058 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8059 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8060 corr += val * x[col];
8061 }
8062 corr
8063}
8064
8065fn q1t_dequant_row(
8069 bytes: &[u8],
8070 r: usize,
8071 gpr: usize,
8072 rp_off: usize,
8073 entries_off: usize,
8074 has_ov: bool,
8075 buf: &mut [f32],
8076) {
8077 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8078 for g in 0..gpr {
8079 let off = (r * gpr + g) * TILE;
8080 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8081 let codes = &bytes[off + 2..off + TILE];
8082 let bc = g * GROUP_SIZE;
8083 for bi in 0..6 {
8085 let lut = &SIGN5[codes[bi] as usize];
8086 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
8087 for i in 0..5 {
8088 d[i] = lut[i] * s;
8089 }
8090 }
8091 let lut = &SIGN5[codes[6] as usize];
8092 buf[bc + 30] = lut[0] * s;
8093 buf[bc + 31] = lut[1] * s;
8094 }
8095 if !has_ov {
8096 return;
8097 }
8098 let (c0, c1) = (
8099 q1t_rowptr(bytes, rp_off, r),
8100 q1t_rowptr(bytes, rp_off, r + 1),
8101 );
8102 for p in c0..c1 {
8103 let e = entries_off + p * 4;
8104 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8105 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8106 }
8107}
8108
8109fn q1t_add_overlay(
8113 bytes: &[u8],
8114 x: &[f32],
8115 rows: usize,
8116 cols: usize,
8117 out: &mut [f32],
8118 pool: Option<&Pool>,
8119) {
8120 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8121 let gpr = cols / GROUP_SIZE;
8122 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8123 if !has_ov {
8124 return;
8125 }
8126 let out_addr = SendMut(out.as_mut_ptr());
8127 let run = move |start: usize, end: usize| {
8128 for r in start..end {
8129 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8130 unsafe { *out_addr.at(r) += corr };
8132 }
8133 };
8134 dispatch_rows(pool, rows, &run);
8135}
8136
8137#[allow(clippy::too_many_arguments)]
8140fn q1t_range_a8w8(
8141 bytes: &[u8],
8142 gpr: usize,
8143 rp_off: usize,
8144 ent_off: usize,
8145 has_ov: bool,
8146 act: &SplitAct,
8147 x: &[f32],
8148 out: SendMut,
8149 start: usize,
8150 end: usize,
8151) {
8152 for r in start..end {
8153 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8154 for &(j, xv) in &act.outliers {
8155 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8156 }
8157 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8158 unsafe { *out.at(r) = acc };
8160 }
8161}
8162
8163#[allow(clippy::too_many_arguments)]
8166fn q1t_range_f32_batch(
8167 bytes: &[u8],
8168 gpr: usize,
8169 rp_off: usize,
8170 ent_off: usize,
8171 has_ov: bool,
8172 x: &[f32],
8173 out: SendMut,
8174 start: usize,
8175 end: usize,
8176) {
8177 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8178 let mut sg = [0f32; GROUP_SIZE];
8179 for r in start..end {
8180 let mut acc = 0f32;
8181 for g in 0..gpr {
8182 let off = (r * gpr + g) * TILE;
8183 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8184 let codes = &bytes[off + 2..off + TILE];
8185 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8186 for bi in 0..6 {
8187 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8188 }
8189 let lut = &SIGN5[codes[6] as usize];
8190 sg[30] = lut[0];
8191 sg[31] = lut[1];
8192 let mut gsum = 0f32;
8193 for k in 0..GROUP_SIZE {
8194 gsum += sg[k] * xg[k];
8195 }
8196 acc += s * gsum;
8197 }
8198 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8199 unsafe { *out.at(r) = acc };
8201 }
8202}
8203
8204fn q1t_matvec(
8208 bytes: &[u8],
8209 x: &[f32],
8210 rows: usize,
8211 cols: usize,
8212 out: &mut [f32],
8213 pool: Option<&Pool>,
8214) {
8215 debug_assert_eq!(out.len(), rows);
8216 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8217 let gpr = cols / GROUP_SIZE;
8218 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8219 let out_addr = SendMut(out.as_mut_ptr());
8220 if a8w8_enabled() {
8224 let act = split_act(x);
8225 let act = &act;
8226 let run = move |start: usize, end: usize| {
8227 for r in start..end {
8228 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8229 for &(j, xv) in &act.outliers {
8230 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8231 }
8232 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8233 unsafe { *out_addr.at(r) = acc };
8235 }
8236 };
8237 dispatch_rows(pool, rows, &run);
8238 return;
8239 }
8240 let run = move |start: usize, end: usize| {
8241 let mut sg = [0f32; GROUP_SIZE];
8245 for r in start..end {
8246 let mut acc = 0f32;
8247 for g in 0..gpr {
8248 let off = (r * gpr + g) * TILE;
8249 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8250 let codes = &bytes[off + 2..off + TILE];
8251 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8252 for bi in 0..6 {
8253 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8254 }
8255 let lut = &SIGN5[codes[6] as usize];
8256 sg[30] = lut[0];
8257 sg[31] = lut[1];
8258 let mut gsum = 0f32;
8259 for k in 0..GROUP_SIZE {
8260 gsum += sg[k] * xg[k];
8261 }
8262 acc += s * gsum;
8263 }
8264 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8265 unsafe { *out_addr.at(r) = acc };
8266 }
8267 };
8268 dispatch_rows(pool, rows, &run);
8269}
8270
8271#[cfg(target_arch = "aarch64")]
8277#[target_feature(enable = "neon,dotprod")]
8278unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
8279 use core::arch::aarch64::*;
8280 use core::arch::asm;
8281 unsafe {
8283 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8284 let bytes_ptr = bytes.as_ptr();
8285 let row_off = r * gpr * TILE;
8286 let xp = [xa.as_ptr(), xb.as_ptr()];
8287 let mut acc = [0f32; 2];
8288 macro_rules! sdot2 {
8289 ($w0:expr, $w1:expr, $x:expr) => {{
8290 let x0 = vld1q_s8($x);
8291 let x1 = vld1q_s8($x.add(16));
8292 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8293 asm!(
8294 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8295 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8296 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8297 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8298 options(pure, nomem, nostack),
8299 );
8300 vaddvq_s32(vaddq_s32(a0, a1))
8301 }};
8302 }
8303 let gpr2 = gpr & !1;
8304 let mut gi = 0;
8305 while gi < gpr2 {
8306 let off0 = row_off + gi * TILE;
8307 let off1 = off0 + TILE;
8308 let s0 = f16_to_f32(u16::from_le_bytes([
8309 *bytes_ptr.add(off0),
8310 *bytes_ptr.add(off0 + 1),
8311 ]));
8312 let s1 = f16_to_f32(u16::from_le_bytes([
8313 *bytes_ptr.add(off1),
8314 *bytes_ptr.add(off1 + 1),
8315 ]));
8316 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8317 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8318 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8319 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8320 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8321 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8322 for k in 0..2 {
8323 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
8324 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
8325 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
8326 }
8327 gi += 2;
8328 }
8329 if gi < gpr {
8330 let off = row_off + gi * TILE;
8331 let s = f16_to_f32(u16::from_le_bytes([
8332 *bytes_ptr.add(off),
8333 *bytes_ptr.add(off + 1),
8334 ]));
8335 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8336 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8337 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8338 for k in 0..2 {
8339 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
8340 acc[k] += d as f32 * s;
8341 }
8342 }
8343 acc
8344 }
8345}
8346
8347fn q1t_matvec2(
8353 bytes: &[u8],
8354 x1: &[f32],
8355 x2: &[f32],
8356 rows: usize,
8357 cols: usize,
8358 o1: &mut [f32],
8359 o2: &mut [f32],
8360 pool: Option<&Pool>,
8361) {
8362 debug_assert_eq!(o1.len(), rows);
8363 debug_assert_eq!(o2.len(), rows);
8364 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8365 let gpr = cols / GROUP_SIZE;
8366 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8367 let out1 = SendMut(o1.as_mut_ptr());
8368 let out2 = SendMut(o2.as_mut_ptr());
8369 if a8w8_enabled() {
8370 let a1 = split_act(x1);
8371 let a2 = split_act(x2);
8372 let (a1, a2) = (&a1, &a2);
8373 let run = move |start: usize, end: usize| {
8374 for r in start..end {
8375 #[cfg(target_arch = "aarch64")]
8376 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
8379 #[cfg(not(target_arch = "aarch64"))]
8380 let ds = [
8381 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
8382 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
8383 ];
8384 let mut acc1 = ds[0] * a1.sx;
8385 for &(j, xv) in &a1.outliers {
8386 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
8387 }
8388 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
8389 let mut acc2 = ds[1] * a2.sx;
8390 for &(j, xv) in &a2.outliers {
8391 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
8392 }
8393 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
8394 unsafe {
8396 *out1.at(r) = acc1;
8397 *out2.at(r) = acc2;
8398 }
8399 }
8400 };
8401 dispatch_rows(pool, rows, &run);
8402 return;
8403 }
8404 let run = move |start: usize, end: usize| {
8405 let mut sg = [0f32; GROUP_SIZE];
8408 for r in start..end {
8409 let mut acc1 = 0f32;
8410 let mut acc2 = 0f32;
8411 for g in 0..gpr {
8412 let off = (r * gpr + g) * TILE;
8413 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8414 let codes = &bytes[off + 2..off + TILE];
8415 for bi in 0..6 {
8416 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8417 }
8418 let lut = &SIGN5[codes[6] as usize];
8419 sg[30] = lut[0];
8420 sg[31] = lut[1];
8421 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8422 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8423 let mut gsum1 = 0f32;
8424 for k in 0..GROUP_SIZE {
8425 gsum1 += sg[k] * xg1[k];
8426 }
8427 acc1 += s * gsum1;
8428 let mut gsum2 = 0f32;
8429 for k in 0..GROUP_SIZE {
8430 gsum2 += sg[k] * xg2[k];
8431 }
8432 acc2 += s * gsum2;
8433 }
8434 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
8435 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
8436 unsafe {
8438 *out1.at(r) = acc1;
8439 *out2.at(r) = acc2;
8440 }
8441 }
8442 };
8443 dispatch_rows(pool, rows, &run);
8444}
8445
8446fn q1t_matmat(
8449 bytes: &[u8],
8450 xs: &[f32],
8451 b: usize,
8452 rows: usize,
8453 cols: usize,
8454 out: &mut [f32],
8455 pool: Option<&Pool>,
8456) {
8457 debug_assert_eq!(out.len(), b * rows);
8458 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8459 let gpr = cols / GROUP_SIZE;
8460 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8461 let out_addr = SendMut(out.as_mut_ptr());
8462 if a8w8_enabled() {
8466 let acts: Vec<SplitAct> = (0..b)
8467 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
8468 .collect();
8469 let acts = &acts;
8470 let run = move |start: usize, end: usize| {
8471 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
8475 for g in 0..gpr {
8476 let off = (r * gpr + g) * TILE;
8477 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8478 q1t_unpack_group_i8(
8479 bytes.as_ptr().wrapping_add(off + 2),
8480 &mut sg[g * GROUP_SIZE..],
8481 );
8482 }
8483 for bi in 0..b {
8484 let act = &acts[bi];
8485 let mut isum = 0f32;
8486 for g in 0..gpr {
8487 let d = q1t_i8dot32(
8488 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
8489 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
8490 );
8491 isum += d as f32 * sc[g];
8492 }
8493 let mut acc = isum * act.sx;
8494 for &(j, xv) in &act.outliers {
8495 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8496 }
8497 accs[bi] = acc;
8498 }
8499 if has_ov {
8503 let (c0, c1) = (
8504 q1t_rowptr(bytes, rp_off, r),
8505 q1t_rowptr(bytes, rp_off, r + 1),
8506 );
8507 for p in c0..c1 {
8508 let e = ent_off + p * 4;
8509 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8510 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8511 for bi in 0..b {
8512 accs[bi] += val * xs[bi * cols + col];
8513 }
8514 }
8515 }
8516 for bi in 0..b {
8517 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
8518 }
8519 }
8520 };
8521 dispatch_rows(pool, rows, &run);
8522 return;
8523 }
8524 let run = move |start: usize, end: usize| {
8525 let mut buf = vec![0f32; cols];
8526 for r in start..end {
8527 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
8528 for bi in 0..b {
8529 let xr = &xs[bi * cols..(bi + 1) * cols];
8530 let mut acc = 0f32;
8531 for j in 0..cols {
8532 acc += buf[j] * xr[j];
8533 }
8534 unsafe { *out_addr.at(bi * rows + r) = acc };
8535 }
8536 }
8537 };
8538 dispatch_rows(pool, rows, &run);
8539}
8540
8541fn q1_matvec(
8542 bytes: &[u8],
8543 x: &[f32],
8544 rows: usize,
8545 cols: usize,
8546 out: &mut [f32],
8547 pool: Option<&Pool>,
8548) {
8549 debug_assert_eq!(out.len(), rows);
8550 let gpr = cols / GROUP_SIZE;
8551 let out_addr = SendMut(out.as_mut_ptr());
8552 if a8w8_enabled() {
8553 let act = split_act(x);
8554 let gsum = q1_group_sums(&act.xq, gpr);
8555 let (act, gsum) = (&act, &gsum);
8556 let run = move |start: usize, end: usize| {
8557 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
8558 };
8559 dispatch_rows(pool, rows, &run);
8560 return;
8561 }
8562 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
8563 dispatch_rows(pool, rows, &run);
8564}
8565
8566#[allow(clippy::too_many_arguments)]
8568fn q1_matvec2(
8569 bytes: &[u8],
8570 x1: &[f32],
8571 x2: &[f32],
8572 rows: usize,
8573 cols: usize,
8574 o1: &mut [f32],
8575 o2: &mut [f32],
8576 pool: Option<&Pool>,
8577) {
8578 let gpr = cols / GROUP_SIZE;
8579 let p1 = SendMut(o1.as_mut_ptr());
8580 let p2 = SendMut(o2.as_mut_ptr());
8581 if a8w8_enabled() {
8582 let a1 = split_act(x1);
8583 let a2 = split_act(x2);
8584 let g1 = q1_group_sums(&a1.xq, gpr);
8585 let g2 = q1_group_sums(&a2.xq, gpr);
8586 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
8587 let run = move |start: usize, end: usize| {
8588 for r in start..end {
8589 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
8590 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
8591 for &(j, xv) in &a1.outliers {
8592 let (w, s) = q1_outlier(bytes, r, gpr, j);
8593 v1 += w * s * xv;
8594 }
8595 for &(j, xv) in &a2.outliers {
8596 let (w, s) = q1_outlier(bytes, r, gpr, j);
8597 v2 += w * s * xv;
8598 }
8599 unsafe {
8601 *p1.at(r) = v1;
8602 *p2.at(r) = v2;
8603 }
8604 }
8605 };
8606 dispatch_rows(pool, rows, &run);
8607 return;
8608 }
8609 let run = move |start: usize, end: usize| {
8610 for r in start..end {
8611 unsafe {
8613 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
8614 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
8615 }
8616 }
8617 };
8618 dispatch_rows(pool, rows, &run);
8619}
8620
8621#[allow(clippy::too_many_arguments)]
8623fn q1_matmat(
8624 bytes: &[u8],
8625 xs_all: &[f32],
8626 b: usize,
8627 rows: usize,
8628 cols: usize,
8629 out: &mut [f32],
8630 pool: Option<&Pool>,
8631) {
8632 debug_assert_eq!(out.len(), b * rows);
8633 let gpr = cols / GROUP_SIZE;
8634 let out_addr = SendMut(out.as_mut_ptr());
8635 if a8w8_enabled() {
8636 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
8637 .map(|bi| {
8638 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
8639 let gsum = q1_group_sums(&act.xq, gpr);
8640 (act, gsum)
8641 })
8642 .collect();
8643 let acts = &acts;
8644 #[cfg(target_arch = "x86_64")]
8645 let blocked_ok = avx2_enabled() && blocked_enabled();
8646 #[cfg(target_arch = "aarch64")]
8647 let blocked_ok = sdot_enabled() && blocked_enabled();
8648 let run = move |start: usize, end: usize| {
8649 for r in start..end {
8650 let mut bi = 0usize;
8651 #[cfg(target_arch = "aarch64")]
8654 if blocked_ok {
8655 while bi + 4 <= acts.len() {
8656 let xs = [
8657 acts[bi].0.xq.as_slice(),
8658 acts[bi + 1].0.xq.as_slice(),
8659 acts[bi + 2].0.xq.as_slice(),
8660 acts[bi + 3].0.xq.as_slice(),
8661 ];
8662 let gs = [
8663 acts[bi].1.as_slice(),
8664 acts[bi + 1].1.as_slice(),
8665 acts[bi + 2].1.as_slice(),
8666 acts[bi + 3].1.as_slice(),
8667 ];
8668 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
8669 for k in 0..4 {
8670 let (act, _) = &acts[bi + k];
8671 let mut acc = d[k] * act.sx;
8672 for &(j, xv) in &act.outliers {
8673 let (w, sc) = q1_outlier(bytes, r, gpr, j);
8674 acc += w * sc * xv;
8675 }
8676 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8678 }
8679 bi += 4;
8680 }
8681 }
8682 #[cfg(target_arch = "x86_64")]
8683 if blocked_ok {
8684 while bi + 4 <= acts.len() {
8685 let xs = [
8686 acts[bi].0.xq.as_slice(),
8687 acts[bi + 1].0.xq.as_slice(),
8688 acts[bi + 2].0.xq.as_slice(),
8689 acts[bi + 3].0.xq.as_slice(),
8690 ];
8691 let gs = [
8692 acts[bi].1.as_slice(),
8693 acts[bi + 1].1.as_slice(),
8694 acts[bi + 2].1.as_slice(),
8695 acts[bi + 3].1.as_slice(),
8696 ];
8697 let d = unsafe {
8698 if vnni_tiles_enabled() {
8699 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
8700 } else {
8701 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
8702 }
8703 };
8704 for k in 0..4 {
8705 let (act, _) = &acts[bi + k];
8706 let mut acc = d[k] * act.sx;
8707 for &(j, xv) in &act.outliers {
8708 let (w, sc) = q1_outlier(bytes, r, gpr, j);
8709 acc += w * sc * xv;
8710 }
8711 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8713 }
8714 bi += 4;
8715 }
8716 }
8717 while bi < acts.len() {
8718 let (act, gsum) = &acts[bi];
8719 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
8720 for &(j, xv) in &act.outliers {
8721 let (w, s) = q1_outlier(bytes, r, gpr, j);
8722 acc += w * s * xv;
8723 }
8724 unsafe { *out_addr.at(bi * rows + r) = acc };
8726 bi += 1;
8727 }
8728 }
8729 };
8730 dispatch_rows(pool, rows, &run);
8731 return;
8732 }
8733 let run = move |start: usize, end: usize| {
8734 for r in start..end {
8735 for bi in 0..b {
8736 let x = &xs_all[bi * cols..(bi + 1) * cols];
8737 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
8739 }
8740 }
8741 };
8742 dispatch_rows(pool, rows, &run);
8743}
8744
8745fn q4matvec(
8751 bytes: &[u8],
8752 x: &[f32],
8753 rows: usize,
8754 cols: usize,
8755 out: &mut [f32],
8756 pool: Option<&Pool>,
8757) {
8758 debug_assert_eq!(out.len(), rows);
8759 let (packed, scales) = q4_split(bytes, rows, cols);
8760 let gpr = cols / GROUP_SIZE;
8761 let out_addr = SendMut(out.as_mut_ptr());
8762
8763 if a8w8_enabled() {
8764 let act = split_act(x);
8765 let run = move |start: usize, end: usize| {
8766 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
8767 };
8768 dispatch_rows(pool, rows, &run);
8769 return;
8770 }
8771
8772 let run =
8773 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
8774 dispatch_rows(pool, rows, &run);
8775}
8776
8777#[inline]
8780#[allow(unreachable_code)]
8781#[cfg(target_arch = "x86_64")]
8786#[target_feature(enable = "avx2")]
8787unsafe fn dot_q4b_row_1x4_avx2(
8788 buf: &[u8],
8789 scales: &[u8],
8790 g0: usize,
8791 gpr: usize,
8792 xs: [&[i8]; 4],
8793) -> [f32; 4] {
8794 unsafe {
8796 use core::arch::x86_64::*;
8797 let ones = _mm256_set1_epi16(1);
8798 let mut acc = [0f32; 4];
8799 for gi in 0..gpr {
8800 let s = f16_to_f32(u16::from_le_bytes([
8801 scales[(g0 + gi) * 2],
8802 scales[(g0 + gi) * 2 + 1],
8803 ]));
8804 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8805 let aw = _mm256_abs_epi8(w);
8806 for (k, xq) in xs.iter().enumerate() {
8807 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8808 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8809 let d = _mm256_madd_epi16(p16, ones);
8810 let hi128 = _mm256_extracti128_si256::<1>(d);
8811 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8812 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8813 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8814 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
8815 }
8816 }
8817 acc
8818 }
8819}
8820
8821#[cfg(target_arch = "x86_64")]
8823#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8824unsafe fn dot_q4b_row_1x4_vnni(
8825 buf: &[u8],
8826 scales: &[u8],
8827 g0: usize,
8828 gpr: usize,
8829 xs: [&[i8]; 4],
8830) -> [f32; 4] {
8831 unsafe {
8833 use core::arch::x86_64::*;
8834 let mut acc = [0f32; 4];
8835 for gi in 0..gpr {
8836 let s = f16_to_f32(u16::from_le_bytes([
8837 scales[(g0 + gi) * 2],
8838 scales[(g0 + gi) * 2 + 1],
8839 ]));
8840 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8841 let aw = _mm256_abs_epi8(w);
8842 for (k, xq) in xs.iter().enumerate() {
8843 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8844 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8845 acc[k] += d as f32 * s;
8846 }
8847 }
8848 acc
8849 }
8850}
8851
8852#[cfg(target_arch = "x86_64")]
8858#[target_feature(enable = "avx2")]
8859unsafe fn dot_q4b_row_1x4_sx_avx2(
8860 buf: &[u8],
8861 scales: &[u8],
8862 g0: usize,
8863 gpr: usize,
8864 xs: [&[i8]; 4],
8865 sxs: [f32; 4],
8866) -> [f32; 4] {
8867 unsafe {
8869 use core::arch::x86_64::*;
8870 let ones = _mm256_set1_epi16(1);
8871 let mut acc = [0f32; 4];
8872 for gi in 0..gpr {
8873 let s = f16_to_f32(u16::from_le_bytes([
8874 scales[(g0 + gi) * 2],
8875 scales[(g0 + gi) * 2 + 1],
8876 ]));
8877 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8878 let aw = _mm256_abs_epi8(w);
8879 for (k, xq) in xs.iter().enumerate() {
8880 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8881 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8882 let d = _mm256_madd_epi16(p16, ones);
8883 let hi128 = _mm256_extracti128_si256::<1>(d);
8884 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8885 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8886 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8887 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
8888 }
8889 }
8890 acc
8891 }
8892}
8893
8894#[cfg(target_arch = "x86_64")]
8897#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8898unsafe fn dot_q4b_row_1x4_sx_vnni(
8899 buf: &[u8],
8900 scales: &[u8],
8901 g0: usize,
8902 gpr: usize,
8903 xs: [&[i8]; 4],
8904 sxs: [f32; 4],
8905) -> [f32; 4] {
8906 unsafe {
8908 use core::arch::x86_64::*;
8909 let mut acc = [0f32; 4];
8910 for gi in 0..gpr {
8911 let s = f16_to_f32(u16::from_le_bytes([
8912 scales[(g0 + gi) * 2],
8913 scales[(g0 + gi) * 2 + 1],
8914 ]));
8915 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8916 let aw = _mm256_abs_epi8(w);
8917 for (k, xq) in xs.iter().enumerate() {
8918 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8919 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8920 acc[k] += (d as f32 * sxs[k]) * s;
8921 }
8922 }
8923 acc
8924 }
8925}
8926
8927#[allow(unreachable_code)]
8928fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8929 #[cfg(target_arch = "aarch64")]
8930 unsafe {
8931 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
8932 }
8933 #[cfg(target_arch = "x86_64")]
8934 unsafe {
8935 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
8936 }
8937 let mut acc = 0f32;
8938 for gi in 0..gpr {
8939 let g = g0 + gi;
8940 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8941 let mut d = 0i32;
8942 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8943 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
8944 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
8945 }
8946 acc += d as f32 * s;
8947 }
8948 acc
8949}
8950
8951#[inline]
8953#[allow(unreachable_code)]
8954fn dot_q4_row_i8_2(
8955 packed: &[u8],
8956 scales: &[u8],
8957 g0: usize,
8958 gpr: usize,
8959 xq1: &[i8],
8960 xq2: &[i8],
8961) -> (f32, f32) {
8962 #[cfg(target_arch = "aarch64")]
8963 unsafe {
8964 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
8965 }
8966 #[cfg(target_arch = "x86_64")]
8967 unsafe {
8968 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
8969 }
8970 (
8971 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
8972 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
8973 )
8974}
8975
8976#[allow(clippy::too_many_arguments)]
8979fn q4_range_a8w8(
8980 packed: &[u8],
8981 scales: &[u8],
8982 gpr: usize,
8983 cols: usize,
8984 act: &SplitAct,
8985 out: SendMut,
8986 start: usize,
8987 end: usize,
8988) {
8989 for r in start..end {
8990 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
8991 for &(j, xv) in &act.outliers {
8993 let flat = r * cols + j;
8994 let byte = packed[flat / 2];
8995 let nib = if flat & 1 == 0 {
8996 byte & 0x0F
8997 } else {
8998 byte >> 4
8999 };
9000 let s = f16_to_f32(u16::from_le_bytes([
9001 scales[(flat / GROUP_SIZE) * 2],
9002 scales[(flat / GROUP_SIZE) * 2 + 1],
9003 ]));
9004 acc += ((nib as i32 - 8) as f32) * s * xv;
9005 }
9006 unsafe { *out.at(r) = acc };
9008 }
9009}
9010
9011#[allow(clippy::too_many_arguments)]
9014fn q4_range2_a8w8(
9015 packed: &[u8],
9016 scales: &[u8],
9017 gpr: usize,
9018 cols: usize,
9019 a1: &SplitAct,
9020 a2: &SplitAct,
9021 p1: SendMut,
9022 p2: SendMut,
9023 start: usize,
9024 end: usize,
9025) {
9026 for r in start..end {
9027 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
9028 let mut acc1 = s1 * a1.sx;
9029 let mut acc2 = s2 * a2.sx;
9030 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
9032 for &(j, xv) in outliers {
9033 let flat = r * cols + j;
9034 let byte = packed[flat / 2];
9035 let nib = if flat & 1 == 0 {
9036 byte & 0x0F
9037 } else {
9038 byte >> 4
9039 };
9040 let s = f16_to_f32(u16::from_le_bytes([
9041 scales[(flat / GROUP_SIZE) * 2],
9042 scales[(flat / GROUP_SIZE) * 2 + 1],
9043 ]));
9044 *acc += ((nib as i32 - 8) as f32) * s * xv;
9045 }
9046 };
9047 fix(&a1.outliers, &mut acc1);
9048 fix(&a2.outliers, &mut acc2);
9049 unsafe {
9051 *p1.at(r) = acc1;
9052 *p2.at(r) = acc2;
9053 }
9054 }
9055}
9056
9057fn q4_range_f32(
9059 packed: &[u8],
9060 scales: &[u8],
9061 gpr: usize,
9062 x: &[f32],
9063 out: SendMut,
9064 start: usize,
9065 end: usize,
9066) {
9067 for r in start..end {
9068 let mut acc = 0f32;
9069 for gi in 0..gpr {
9070 let g = r * gpr + gi;
9071 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9072 let pk = &packed[g * 16..(g + 1) * 16];
9073 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9074 let mut ga = 0f32;
9075 for (k, &b) in pk.iter().enumerate() {
9076 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
9077 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
9078 }
9079 acc += ga * s;
9080 }
9081 unsafe { *out.at(r) = acc };
9083 }
9084}
9085
9086#[allow(clippy::too_many_arguments)]
9090fn q4matvec2(
9091 bytes: &[u8],
9092 x1: &[f32],
9093 x2: &[f32],
9094 rows: usize,
9095 cols: usize,
9096 o1: &mut [f32],
9097 o2: &mut [f32],
9098 pool: Option<&Pool>,
9099) {
9100 debug_assert_eq!(o1.len(), rows);
9101 debug_assert_eq!(o2.len(), rows);
9102 let (packed, scales) = q4_split(bytes, rows, cols);
9103 let gpr = cols / GROUP_SIZE;
9104
9105 if a8w8_enabled() {
9106 let a1 = split_act(x1);
9107 let a2 = split_act(x2);
9108 let p1 = SendMut(o1.as_mut_ptr());
9109 let p2 = SendMut(o2.as_mut_ptr());
9110 let run = move |start: usize, end: usize| {
9111 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
9112 };
9113 dispatch_rows(pool, rows, &run);
9114 return;
9115 }
9116
9117 let p1 = SendMut(o1.as_mut_ptr());
9118 let p2 = SendMut(o2.as_mut_ptr());
9119 let run = move |start: usize, end: usize| {
9120 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
9121 };
9122 dispatch_rows(pool, rows, &run);
9123}
9124
9125#[allow(clippy::too_many_arguments)]
9127fn q4_range2_f32(
9128 packed: &[u8],
9129 scales: &[u8],
9130 gpr: usize,
9131 x1: &[f32],
9132 x2: &[f32],
9133 p1: SendMut,
9134 p2: SendMut,
9135 start: usize,
9136 end: usize,
9137) {
9138 for r in start..end {
9139 let (mut acc1, mut acc2) = (0f32, 0f32);
9140 for gi in 0..gpr {
9141 let g = r * gpr + gi;
9142 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9143 let pk = &packed[g * 16..(g + 1) * 16];
9144 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9145 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9146 let (mut g1, mut g2) = (0f32, 0f32);
9147 for (k, &b) in pk.iter().enumerate() {
9148 let wl = (b & 0x0F) as f32 - 8.0;
9149 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
9150 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
9151 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
9152 }
9153 acc1 += g1 * s;
9154 acc2 += g2 * s;
9155 }
9156 unsafe {
9158 *p1.at(r) = acc1;
9159 *p2.at(r) = acc2;
9160 }
9161 }
9162}
9163
9164thread_local! {
9165 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
9168 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9169}
9170
9171#[allow(clippy::too_many_arguments)]
9177fn q4matmat(
9178 bytes: &[u8],
9179 xs_all: &[f32],
9180 b: usize,
9181 rows: usize,
9182 cols: usize,
9183 out: &mut [f32],
9184 pool: Option<&Pool>,
9185) {
9186 debug_assert_eq!(xs_all.len(), b * cols);
9187 debug_assert_eq!(out.len(), b * rows);
9188 let (packed, scales) = q4_split(bytes, rows, cols);
9189 let gpr = cols / GROUP_SIZE;
9190 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9191
9192 if a8w8_enabled() {
9193 let acts: Vec<SplitAct> = (0..b)
9194 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
9195 .collect();
9196 let acts = &acts;
9197 let out_addr = SendMut(out.as_mut_ptr());
9198 let run = move |start: usize, end: usize| {
9199 ROW_I8.with(|rb| {
9200 let mut buf = rb.borrow_mut();
9201 buf.resize(cols, 0);
9202 for r in start..end {
9203 for gi in 0..gpr {
9207 let g = r * gpr + gi;
9208 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9209 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
9210 buf[gi * GROUP_SIZE + k * 2 + 1] =
9211 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
9212 }
9213 }
9214 let mut bi = 0usize;
9215 #[cfg(target_arch = "x86_64")]
9216 if avx2_enabled() && blocked_enabled() {
9217 while bi + 4 <= acts.len() {
9218 let xs = [
9219 acts[bi].xq.as_slice(),
9220 acts[bi + 1].xq.as_slice(),
9221 acts[bi + 2].xq.as_slice(),
9222 acts[bi + 3].xq.as_slice(),
9223 ];
9224 let d = unsafe {
9225 if vnni_tiles_enabled() {
9226 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
9227 } else {
9228 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
9229 }
9230 };
9231 for k in 0..4 {
9232 let act = &acts[bi + k];
9233 let mut acc = d[k] * act.sx;
9234 for &(j, xv) in &act.outliers {
9235 acc += (buf[j] as i8) as f32
9236 * gscale((r * cols + j) / GROUP_SIZE)
9237 * xv;
9238 }
9239 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9241 }
9242 bi += 4;
9243 }
9244 }
9245 while bi < acts.len() {
9246 let act = &acts[bi];
9247 let mut acc = 0f32;
9248 for gi in 0..gpr {
9249 let d = dot_i8_i8(
9250 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
9251 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
9252 );
9253 acc += d as f32 * gscale(r * gpr + gi);
9254 }
9255 acc *= act.sx;
9256 for &(j, xv) in &act.outliers {
9258 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
9259 }
9260 unsafe { *out_addr.at(bi * rows + r) = acc };
9262 bi += 1;
9263 }
9264 }
9265 })
9266 };
9267 dispatch_rows(pool, rows, &run);
9268 return;
9269 }
9270
9271 let out_addr = SendMut(out.as_mut_ptr());
9272 let run = move |start: usize, end: usize| {
9273 ROW_F32.with(|rb| {
9274 let mut buf = rb.borrow_mut();
9275 buf.resize(cols, 0.0);
9276 for r in start..end {
9277 for gi in 0..gpr {
9280 let g = r * gpr + gi;
9281 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9282 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
9283 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
9284 }
9285 }
9286 for bi in 0..b {
9287 let x = &xs_all[bi * cols..(bi + 1) * cols];
9288 let mut acc = 0f32;
9289 for gi in 0..gpr {
9290 let mut ga = 0f32;
9291 for k in 0..GROUP_SIZE / 2 {
9296 let e = gi * GROUP_SIZE + k * 2;
9297 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
9298 }
9299 acc += ga * gscale(r * gpr + gi);
9300 }
9301 unsafe { *out_addr.at(bi * rows + r) = acc };
9303 }
9304 }
9305 })
9306 };
9307 dispatch_rows(pool, rows, &run);
9308}
9309
9310#[allow(clippy::too_many_arguments)]
9315fn vbitmatmat(
9316 bytes: &[u8],
9317 offsets: &[usize],
9318 xs_all: &[f32],
9319 b: usize,
9320 rows: usize,
9321 cols: usize,
9322 out: &mut [f32],
9323 pool: Option<&Pool>,
9324) {
9325 debug_assert_eq!(xs_all.len(), b * cols);
9326 debug_assert_eq!(out.len(), b * rows);
9327 debug_assert_eq!(offsets.len(), rows + 1);
9328 let ng = cols / GROUP_SIZE;
9329 let bits = &bytes[..rows];
9330 let sc_off = rows;
9331 let gscale = |r: usize, g: usize| {
9332 let so = (r * ng + g) * 2;
9333 f16_to_f32(u16::from_le_bytes([
9334 bytes[sc_off + so],
9335 bytes[sc_off + so + 1],
9336 ]))
9337 };
9338
9339 let decode_f32 = |r: usize, dst: &mut [f32]| {
9341 let bw = bits[r] as usize;
9342 let l = ((1i32 << (bw - 1)) - 1) as f32;
9343 let data = &bytes[offsets[r]..offsets[r + 1]];
9344 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
9345 for d in dst.iter_mut() {
9346 while nbits < bw {
9347 acc = (acc << 8) | data[idx] as u64;
9348 idx += 1;
9349 nbits += 8;
9350 }
9351 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
9352 nbits -= bw;
9353 *d = u - l;
9354 }
9355 };
9356
9357 if a8w8_enabled() {
9358 let acts: Vec<SplitAct> = (0..b)
9359 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
9360 .collect();
9361 let acts = &acts;
9362 let out_addr = SendMut(out.as_mut_ptr());
9363 let run = move |start: usize, end: usize| {
9364 for r in start..end {
9365 let bw = bits[r] as usize;
9366 if bw == 8 {
9367 ROW_F32.with(|rb| {
9370 let mut buf = rb.borrow_mut();
9371 buf.resize(cols, 0.0);
9372 decode_f32(r, &mut buf);
9373 for bi in 0..b {
9374 let x = &xs_all[bi * cols..(bi + 1) * cols];
9375 let mut dot = 0f32;
9376 for g in 0..ng {
9377 let mut gd = 0f32;
9378 for k in 0..GROUP_SIZE {
9379 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9380 }
9381 dot += gd * gscale(r, g);
9382 }
9383 unsafe { *out_addr.at(bi * rows + r) = dot };
9385 }
9386 });
9387 continue;
9388 }
9389 let l = (1i32 << (bw - 1)) - 1;
9390 let data = &bytes[offsets[r]..offsets[r + 1]];
9391 ROW_I8.with(|rb| {
9392 let mut buf = rb.borrow_mut();
9393 buf.resize(cols, 0);
9394 #[inline(always)]
9395 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
9396 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
9397 let u = unpack8::<B>(&data[blk * B..]);
9398 for k in 0..8 {
9399 chunk[k] = (u[k] - l) as i8 as u8;
9400 }
9401 }
9402 }
9403 match bw {
9404 3 => fill::<3>(data, l, &mut buf),
9405 4 => vbit_fill4(data, &mut buf),
9406 5 => fill::<5>(data, l, &mut buf),
9407 6 => fill::<6>(data, l, &mut buf),
9408 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
9409 }
9410 let mut bi = 0usize;
9411 #[cfg(target_arch = "x86_64")]
9415 if avx2_enabled() && blocked_enabled() {
9416 while bi + 4 <= acts.len() {
9417 let xs = [
9418 acts[bi].xq.as_slice(),
9419 acts[bi + 1].xq.as_slice(),
9420 acts[bi + 2].xq.as_slice(),
9421 acts[bi + 3].xq.as_slice(),
9422 ];
9423 let sxs = [
9424 acts[bi].sx,
9425 acts[bi + 1].sx,
9426 acts[bi + 2].sx,
9427 acts[bi + 3].sx,
9428 ];
9429 let d = unsafe {
9430 if vnni_tiles_enabled() {
9431 dot_q4b_row_1x4_sx_vnni(
9432 &buf,
9433 &bytes[sc_off..],
9434 r * ng,
9435 ng,
9436 xs,
9437 sxs,
9438 )
9439 } else {
9440 dot_q4b_row_1x4_sx_avx2(
9441 &buf,
9442 &bytes[sc_off..],
9443 r * ng,
9444 ng,
9445 xs,
9446 sxs,
9447 )
9448 }
9449 };
9450 for k in 0..4 {
9451 let act = &acts[bi + k];
9452 let mut dot = d[k];
9453 for &(j, xv) in &act.outliers {
9454 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9455 }
9456 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
9458 }
9459 bi += 4;
9460 }
9461 }
9462 while bi < acts.len() {
9463 let act = &acts[bi];
9464 let mut dot = 0f32;
9465 for g in 0..ng {
9466 let d = dot_i8_i8(
9467 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9468 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9469 ) as f32
9470 * act.sx;
9471 dot += d * gscale(r, g);
9472 }
9473 for &(j, xv) in &act.outliers {
9474 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9475 }
9476 unsafe { *out_addr.at(bi * rows + r) = dot };
9478 bi += 1;
9479 }
9480 });
9481 }
9482 };
9483 dispatch_rows(pool, rows, &run);
9484 return;
9485 }
9486
9487 let out_addr = SendMut(out.as_mut_ptr());
9488 let run = move |start: usize, end: usize| {
9489 ROW_F32.with(|rb| {
9490 let mut buf = rb.borrow_mut();
9491 buf.resize(cols, 0.0);
9492 for r in start..end {
9493 decode_f32(r, &mut buf);
9494 for bi in 0..b {
9495 let x = &xs_all[bi * cols..(bi + 1) * cols];
9496 let mut dot = 0f32;
9497 for g in 0..ng {
9498 let mut gd = 0f32;
9499 for k in 0..GROUP_SIZE {
9500 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9501 }
9502 dot += gd * gscale(r, g);
9503 }
9504 unsafe { *out_addr.at(bi * rows + r) = dot };
9506 }
9507 }
9508 })
9509 };
9510 dispatch_rows(pool, rows, &run);
9511}
9512
9513pub(crate) fn gpu_batch_job<'a>(
9517 t: &'a QTensor,
9518 x: &[f32],
9519) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
9520 match t {
9521 QTensor::Mapped {
9522 model,
9523 idx,
9524 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
9525 rows,
9526 cols,
9527 row_scale,
9528 col_field,
9529 ..
9530 } => Some((
9531 model.clone(),
9532 crate::gpu::BatchJob {
9533 idx: *idx,
9534 rows: *rows,
9535 cols: *cols,
9536 row_scale,
9537 xs: prescale(x, col_field, *dt).into_owned(),
9538 layout: crate::gpu::BatchLayout::Q8,
9539 },
9540 )),
9541 QTensor::Mapped {
9543 model,
9544 idx,
9545 dtype: TensorDtype::Q1,
9546 rows,
9547 cols,
9548 ..
9549 } => Some((
9550 model.clone(),
9551 crate::gpu::BatchJob {
9552 idx: *idx,
9553 rows: *rows,
9554 cols: *cols,
9555 row_scale: &[],
9556 xs: x.to_vec(),
9557 layout: crate::gpu::BatchLayout::Q1,
9558 },
9559 )),
9560 QTensor::Mapped {
9565 model,
9566 idx,
9567 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
9568 rows,
9569 cols,
9570 ..
9571 } => Some((
9572 model.clone(),
9573 crate::gpu::BatchJob {
9574 idx: *idx,
9575 rows: *rows,
9576 cols: *cols,
9577 row_scale: &[],
9578 xs: x.to_vec(),
9579 layout: if *dt == TensorDtype::Q4Tiled {
9580 crate::gpu::BatchLayout::Q4t
9581 } else {
9582 crate::gpu::BatchLayout::Q4tp
9583 },
9584 },
9585 )),
9586 _ => None,
9587 }
9588}
9589
9590thread_local! {
9591 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9592 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9593}
9594
9595pub(crate) fn prescale<'a>(
9596 x: &'a [f32],
9597 col_field: &[f32],
9598 dtype: TensorDtype,
9599) -> std::borrow::Cow<'a, [f32]> {
9600 if dtype == TensorDtype::Q8_2f {
9601 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
9602 } else {
9603 std::borrow::Cow::Borrowed(x)
9604 }
9605}
9606
9607pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
9610 x: &[f32],
9611 col_field: &[f32],
9612 dtype: TensorDtype,
9613 buf_id: u8,
9614 f: F,
9615) -> R {
9616 if dtype == TensorDtype::Q8_2f {
9617 if buf_id == 1 {
9618 PRESCALE_BUF1.with(|b| {
9619 let mut buf = b.borrow_mut();
9620 buf.clear();
9621 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9622 f(&buf)
9623 })
9624 } else {
9625 PRESCALE_BUF2.with(|b| {
9626 let mut buf = b.borrow_mut();
9627 buf.clear();
9628 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9629 f(&buf)
9630 })
9631 }
9632 } else {
9633 f(x)
9634 }
9635}
9636
9637#[cfg(target_arch = "x86_64")]
9642pub(crate) fn avx2_enabled() -> bool {
9643 use std::sync::OnceLock;
9644 static ON: OnceLock<bool> = OnceLock::new();
9645 *ON.get_or_init(|| {
9646 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
9647 && std::arch::is_x86_feature_detected!("avx2")
9648 && std::arch::is_x86_feature_detected!("fma")
9649 })
9650}
9651
9652#[cfg(target_arch = "x86_64")]
9657fn avx2_a8w8_enabled() -> bool {
9658 if FLOAT_ACTIVATIONS.get() {
9659 return false;
9660 }
9661 use std::sync::OnceLock;
9662 static ON: OnceLock<bool> = OnceLock::new();
9663 *ON.get_or_init(|| {
9664 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
9665 })
9666}
9667
9668thread_local! {
9669 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
9670}
9671
9672pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
9675 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
9676 impl Drop for Restore {
9677 fn drop(&mut self) {
9678 FULL_GPU_Q8.set(self.0);
9679 }
9680 }
9681 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
9682}
9683
9684thread_local! {
9688 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
9689}
9690
9691pub(crate) fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
9692 struct Restore(bool);
9693 impl Drop for Restore {
9694 fn drop(&mut self) {
9695 FLOAT_ACTIVATIONS.set(self.0);
9696 }
9697 }
9698 let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
9699 f()
9700}
9701
9702static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
9714
9715pub(crate) fn row_exact() -> bool {
9716 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
9717}
9718
9719fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
9720 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
9721 impl Drop for Restore<'_> {
9722 fn drop(&mut self) {
9723 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
9724 }
9725 }
9726 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
9727 let _restore = Restore(active);
9728 f()
9729}
9730
9731pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
9733 counted_row_exact_scope(&ROW_EXACT, f)
9734}
9735
9736#[inline]
9740pub(crate) fn a8w8_enabled() -> bool {
9741 #[cfg(target_arch = "aarch64")]
9742 {
9743 sdot_enabled()
9744 }
9745 #[cfg(target_arch = "x86_64")]
9746 {
9747 avx2_a8w8_enabled()
9748 }
9749 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
9750 {
9751 false
9752 }
9753}
9754
9755#[inline]
9758#[allow(unreachable_code)]
9759fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
9760 #[cfg(target_arch = "aarch64")]
9761 unsafe {
9762 return dot_i8_sdot(w, xq);
9763 }
9764 #[cfg(target_arch = "x86_64")]
9765 unsafe {
9766 if avx512vnni_enabled() {
9767 return dot_i8_i8_vnni(w, xq);
9768 }
9769 return dot_i8_i8_avx2(w, xq);
9770 }
9771 w.iter()
9772 .zip(xq)
9773 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
9774 .sum()
9775}
9776
9777#[cfg(target_arch = "x86_64")]
9781fn avx512vnni_enabled() -> bool {
9782 use std::sync::OnceLock;
9783 static ON: OnceLock<bool> = OnceLock::new();
9784 *ON.get_or_init(|| {
9785 std::env::var("CMF_AVX512")
9786 .map(|v| v != "0")
9787 .unwrap_or(true)
9788 && std::arch::is_x86_feature_detected!("avx512f")
9789 && std::arch::is_x86_feature_detected!("avx512bw")
9790 && std::arch::is_x86_feature_detected!("avx512vl")
9791 && std::arch::is_x86_feature_detected!("avx512vnni")
9792 })
9793}
9794
9795#[cfg(target_arch = "x86_64")]
9803fn vnni_tiles_enabled() -> bool {
9804 use std::sync::OnceLock;
9805 static ON: OnceLock<bool> = OnceLock::new();
9806 *ON.get_or_init(|| {
9807 std::env::var("CMF_VNNI_TILES")
9808 .map(|v| v != "0")
9809 .unwrap_or(true)
9810 && avx512vnni_enabled()
9811 })
9812}
9813
9814#[cfg(target_arch = "x86_64")]
9819#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9820#[inline]
9821unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
9822 unsafe {
9824 use core::arch::x86_64::*;
9825 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
9826 let hi128 = _mm256_extracti128_si256::<1>(d);
9827 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9828 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9829 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9830 _mm_cvtsi128_si32(s32)
9831 }
9832}
9833
9834#[cfg(target_arch = "x86_64")]
9839#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9840unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9841 unsafe {
9843 use core::arch::x86_64::*;
9844 let n = w.len();
9845 let mut j = 0usize;
9846 let mut total: i32;
9847 {
9852 #[inline(always)]
9853 unsafe fn step(
9854 w: *const u8,
9855 x: *const i8,
9856 acc: core::arch::x86_64::__m512i,
9857 ) -> core::arch::x86_64::__m512i {
9858 unsafe {
9859 use core::arch::x86_64::*;
9860 let wv = _mm512_loadu_si512(w as *const _);
9861 let xv = _mm512_loadu_si512(x as *const _);
9862 let aw = _mm512_abs_epi8(wv);
9863 let neg = _mm512_movepi8_mask(wv);
9864 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
9865 _mm512_dpbusd_epi32(acc, aw, sx)
9866 }
9867 }
9868 let (mut a0, mut a1, mut a2, mut a3) = (
9869 _mm512_setzero_si512(),
9870 _mm512_setzero_si512(),
9871 _mm512_setzero_si512(),
9872 _mm512_setzero_si512(),
9873 );
9874 while j + 256 <= n {
9875 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9876 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
9877 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
9878 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
9879 j += 256;
9880 }
9881 while j + 64 <= n {
9882 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9883 j += 64;
9884 }
9885 let s01 = _mm512_add_epi32(a0, a1);
9886 let s23 = _mm512_add_epi32(a2, a3);
9887 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
9888 }
9889 if j + 32 <= n {
9891 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9892 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9893 let d = _mm256_dpbusd_epi32(
9894 _mm256_setzero_si256(),
9895 _mm256_abs_epi8(wv),
9896 _mm256_sign_epi8(xv, wv),
9897 );
9898 let hi128 = _mm256_extracti128_si256::<1>(d);
9899 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9900 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9901 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9902 total += _mm_cvtsi128_si32(s32);
9903 j += 32;
9904 }
9905 while j < n {
9906 total += (w[j] as i8) as i32 * xq[j] as i32;
9907 j += 1;
9908 }
9909 total
9910 }
9911}
9912
9913#[cfg(target_arch = "x86_64")]
9915#[target_feature(enable = "avx2,fma")]
9916unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
9917 unsafe {
9919 use core::arch::x86_64::*;
9920 let n = x.len();
9921 let wp = w.as_ptr();
9922 let xp = x.as_ptr();
9923 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
9924 let mut j = 0usize;
9925 while j + 16 <= n {
9926 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
9927 let lo = _mm256_cvtepi8_epi32(wb);
9928 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
9929 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
9930 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
9931 j += 16;
9932 }
9933 let acc = _mm256_add_ps(a0, a1);
9934 let hi128 = _mm256_extractf128_ps::<1>(acc);
9935 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
9936 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
9937 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
9938 let mut sum = _mm_cvtss_f32(s32);
9939 while j < n {
9940 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
9941 j += 1;
9942 }
9943 sum
9944 }
9945}
9946
9947#[cfg(target_arch = "x86_64")]
9952#[target_feature(enable = "avx2")]
9953unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
9954 unsafe {
9956 use core::arch::x86_64::*;
9957 let n = w.len();
9958 let ones = _mm256_set1_epi16(1);
9959 let mut acc = _mm256_setzero_si256();
9960 let mut j = 0usize;
9961 while j + 32 <= n {
9962 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9963 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9964 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9965 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
9966 j += 32;
9967 }
9968 let hi128 = _mm256_extracti128_si256::<1>(acc);
9969 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
9970 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9971 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9972 let mut s = _mm_cvtsi128_si32(s32);
9973 while j < n {
9974 s += (w[j] as i8) as i32 * xq[j] as i32;
9975 j += 1;
9976 }
9977 s
9978 }
9979}
9980
9981#[cfg(target_arch = "aarch64")]
9985#[target_feature(enable = "neon,i8mm")]
9986unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9987 unsafe {
9989 use core::arch::aarch64::*;
9990 use core::arch::asm;
9991 let n = w0.len();
9992 let w0p = w0.as_ptr() as *const i8;
9993 let w1p = w1.as_ptr() as *const i8;
9994 let mut acc01 = vdupq_n_s32(0);
9997 let mut acc23 = vdupq_n_s32(0);
9998 let mut i = 0usize;
9999 while i + 8 <= n {
10000 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
10001 let xb01 = vcombine_s8(
10002 vld1_s8(xs[0].as_ptr().add(i)),
10003 vld1_s8(xs[1].as_ptr().add(i)),
10004 );
10005 let xb23 = vcombine_s8(
10006 vld1_s8(xs[2].as_ptr().add(i)),
10007 vld1_s8(xs[3].as_ptr().add(i)),
10008 );
10009 asm!(
10010 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
10011 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
10012 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
10013 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
10014 options(pure, nomem, nostack),
10015 );
10016 i += 8;
10017 }
10018 let mut out = [[0i32; 4]; 2];
10019 let a01: [i32; 4] = core::mem::transmute(acc01);
10020 let a23: [i32; 4] = core::mem::transmute(acc23);
10021 out[0][0] = a01[0];
10022 out[0][1] = a01[1];
10023 out[1][0] = a01[2];
10024 out[1][1] = a01[3];
10025 out[0][2] = a23[0];
10026 out[0][3] = a23[1];
10027 out[1][2] = a23[2];
10028 out[1][3] = a23[3];
10029 if i < n {
10030 for (k, x) in xs.iter().enumerate() {
10031 for j in i..n {
10032 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
10033 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
10034 }
10035 }
10036 }
10037 out
10038 }
10039}
10040
10041#[cfg(target_arch = "aarch64")]
10045#[target_feature(enable = "neon,dotprod")]
10046unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10047 unsafe {
10049 use core::arch::aarch64::*;
10050 use core::arch::asm;
10051 let n = w0.len();
10052 let w0p = w0.as_ptr() as *const i8;
10053 let w1p = w1.as_ptr() as *const i8;
10054 let mut acc = [[vdupq_n_s32(0); 4]; 2];
10055 let mut i = 0usize;
10056 while i + 16 <= n {
10057 let wv0 = vld1q_s8(w0p.add(i));
10058 let wv1 = vld1q_s8(w1p.add(i));
10059 for (k, x) in xs.iter().enumerate() {
10060 let xv = vld1q_s8(x.as_ptr().add(i));
10061 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
10062 asm!(
10063 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
10064 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
10065 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10066 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
10067 options(pure, nomem, nostack),
10068 );
10069 acc[0][k] = a0;
10070 acc[1][k] = a1;
10071 }
10072 i += 16;
10073 }
10074 let mut out = [[0i32; 4]; 2];
10075 for r in 0..2 {
10076 for k in 0..4 {
10077 out[r][k] = vaddvq_s32(acc[r][k]);
10078 }
10079 }
10080 if i < n {
10081 for (k, x) in xs.iter().enumerate() {
10082 for j in i..n {
10083 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
10084 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
10085 }
10086 }
10087 }
10088 out
10089 }
10090}
10091
10092#[cfg(target_arch = "x86_64")]
10098#[target_feature(enable = "avx2")]
10099unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10100 unsafe {
10102 use core::arch::x86_64::*;
10103 let n = w0.len();
10104 let ones = _mm256_set1_epi16(1);
10105 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
10106 let mut j = 0usize;
10107 while j + 32 <= n {
10108 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
10109 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
10110 let aw0 = _mm256_abs_epi8(wv0);
10111 let aw1 = _mm256_abs_epi8(wv1);
10112 for (k, x) in xs.iter().enumerate() {
10113 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
10114 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
10115 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
10116 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
10117 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
10118 }
10119 j += 32;
10120 }
10121 let mut out = [[0i32; 4]; 2];
10122 for r in 0..2 {
10123 for k in 0..4 {
10124 let a = acc[r][k];
10125 let hi128 = _mm256_extracti128_si256::<1>(a);
10126 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
10127 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10128 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10129 out[r][k] = _mm_cvtsi128_si32(s32);
10130 }
10131 }
10132 if j < n {
10133 for (k, x) in xs.iter().enumerate() {
10134 for i in j..n {
10135 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
10136 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
10137 }
10138 }
10139 }
10140 out
10141 }
10142}
10143
10144#[cfg(target_arch = "x86_64")]
10149#[inline]
10150fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
10151 let dot = if avx512vnni_enabled() && row.len() >= 64 {
10152 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
10153 } else {
10154 unsafe { dot_i8_i8_avx2(row, &act.xq) }
10155 };
10156 let mut acc = dot as f32 * act.sx;
10157 for &(j, xv) in &act.outliers {
10158 acc += (row[j] as i8) as f32 * xv;
10159 }
10160 acc
10161}
10162
10163#[cfg(target_arch = "x86_64")]
10167#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10168unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
10169 unsafe {
10171 use core::arch::x86_64::*;
10172 let n = w.len();
10173 let flip = _mm512_set1_epi8(-128); #[inline(always)]
10175 unsafe fn step(
10176 w: *const u8,
10177 x: *const i8,
10178 flip: core::arch::x86_64::__m512i,
10179 acc: core::arch::x86_64::__m512i,
10180 ) -> core::arch::x86_64::__m512i {
10181 unsafe {
10182 use core::arch::x86_64::*;
10183 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
10184 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
10185 }
10186 }
10187 let (mut a0, mut a1, mut a2, mut a3) = (
10188 _mm512_setzero_si512(),
10189 _mm512_setzero_si512(),
10190 _mm512_setzero_si512(),
10191 _mm512_setzero_si512(),
10192 );
10193 let mut j = 0usize;
10194 while j + 256 <= n {
10195 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
10196 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
10197 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
10198 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
10199 j += 256;
10200 }
10201 while j + 64 <= n {
10202 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
10203 j += 64;
10204 }
10205 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
10206 _mm512_add_epi32(a0, a1),
10207 _mm512_add_epi32(a2, a3),
10208 ));
10209 while j < n {
10211 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
10212 j += 1;
10213 }
10214 total
10215 }
10216}
10217
10218#[cfg(target_arch = "x86_64")]
10224#[target_feature(enable = "avx2")]
10225unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10226 unsafe {
10229 use core::arch::x86_64::*;
10230 let lomask = _mm_set1_epi8(0x0F);
10231 let eight = _mm256_set1_epi8(8);
10232 let ones = _mm256_set1_epi16(1);
10233 let mut acc = 0f32;
10234 for gi in 0..gpr {
10235 let g = g0 + gi;
10236 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10237 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
10238 let lo = _mm_and_si128(b, lomask);
10239 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
10240 let w = _mm256_sub_epi8(
10241 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
10242 eight,
10243 );
10244 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10245 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
10246 let d = _mm256_madd_epi16(p16, ones);
10247 let hi128 = _mm256_extracti128_si256::<1>(d);
10248 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10249 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10250 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10251 acc += _mm_cvtsi128_si32(s32) as f32 * s;
10252 }
10253 acc
10254 }
10255}
10256
10257#[cfg(target_arch = "x86_64")]
10260#[target_feature(enable = "avx2")]
10261unsafe fn dot_q4_row_avx2_2(
10262 packed: &[u8],
10263 scales: &[u8],
10264 g0: usize,
10265 gpr: usize,
10266 xq1: &[i8],
10267 xq2: &[i8],
10268) -> (f32, f32) {
10269 unsafe {
10271 use core::arch::x86_64::*;
10272 let lomask = _mm_set1_epi8(0x0F);
10273 let eight = _mm256_set1_epi8(8);
10274 let ones = _mm256_set1_epi16(1);
10275 let (mut acc1, mut acc2) = (0f32, 0f32);
10276 #[inline(always)]
10277 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
10278 unsafe {
10279 use core::arch::x86_64::*;
10280 let hi128 = _mm256_extracti128_si256::<1>(d);
10281 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10282 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10283 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10284 _mm_cvtsi128_si32(s32)
10285 }
10286 }
10287 for gi in 0..gpr {
10288 let g = g0 + gi;
10289 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10290 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
10291 let lo = _mm_and_si128(b, lomask);
10292 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
10293 let w = _mm256_sub_epi8(
10294 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
10295 eight,
10296 );
10297 let aw = _mm256_abs_epi8(w);
10298 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10299 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10300 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
10301 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
10302 acc1 += hsum(d1) as f32 * s;
10303 acc2 += hsum(d2) as f32 * s;
10304 }
10305 (acc1, acc2)
10306 }
10307}
10308
10309#[cfg(target_arch = "x86_64")]
10311fn q8_range_avx2(
10312 q: &[u8],
10313 row_scale: &[f32],
10314 act: &SplitAct,
10315 cols: usize,
10316 out_addr: SendMut,
10317 start: usize,
10318 end: usize,
10319) {
10320 for o in start..end {
10321 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10322 unsafe { *out_addr.at(o) = v };
10324 }
10325}
10326
10327#[cfg(target_arch = "x86_64")]
10329#[allow(clippy::too_many_arguments)]
10330fn q8_range2_avx2(
10331 q: &[u8],
10332 row_scale: &[f32],
10333 a1: &SplitAct,
10334 a2: &SplitAct,
10335 cols: usize,
10336 p1: SendMut,
10337 p2: SendMut,
10338 start: usize,
10339 end: usize,
10340) {
10341 for o in start..end {
10342 let row = &q[o * cols..(o + 1) * cols];
10343 unsafe {
10345 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
10346 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
10347 }
10348 }
10349}
10350
10351#[cfg(target_arch = "aarch64")]
10362fn i8mm_enabled() -> bool {
10363 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10364 *ON.get_or_init(|| {
10365 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
10366 && std::arch::is_aarch64_feature_detected!("i8mm")
10367 })
10368}
10369
10370#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
10374fn sdot_enabled() -> bool {
10375 if FLOAT_ACTIVATIONS.get() {
10376 return false;
10377 }
10378 use std::sync::OnceLock;
10379 static ON: OnceLock<bool> = OnceLock::new();
10380 *ON.get_or_init(|| {
10381 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
10382 if !want {
10383 return false;
10384 }
10385
10386 #[cfg(target_arch = "aarch64")]
10387 {
10388 if std::arch::is_aarch64_feature_detected!("dotprod") {
10389 return true;
10390 }
10391 #[cfg(target_os = "android")]
10392 {
10393 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
10394 if cpuinfo.lines().any(|l| {
10395 (l.starts_with("Features") || l.starts_with("features"))
10396 && l.contains("asimddp")
10397 }) {
10398 return true;
10399 }
10400 }
10401 }
10402 false
10403 }
10404 #[cfg(not(target_arch = "aarch64"))]
10405 {
10406 false
10407 }
10408 })
10409}
10410
10411struct SplitAct {
10416 xq: Vec<i8>,
10417 sx: f32,
10418 outliers: Vec<(usize, f32)>,
10419 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
10422 xsum: i32,
10423}
10424
10425thread_local! {
10426 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
10429 const { std::cell::RefCell::new(Vec::new()) };
10430}
10431
10432impl Drop for SplitAct {
10433 fn drop(&mut self) {
10434 let buf = std::mem::take(&mut self.xq);
10435 if buf.capacity() > 0 {
10436 XQ_FREE.with(|f| {
10437 let mut f = f.borrow_mut();
10438 if f.len() < 16 {
10439 f.push(buf);
10440 }
10441 });
10442 }
10443 }
10444}
10445
10446thread_local! {
10447 static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
10454 static KROWS: UnsafeCell<[Vec<f32>; 2]> =
10459 const { UnsafeCell::new([Vec::new(), Vec::new()]) };
10460}
10461
10462#[inline]
10465fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
10466 KROW.with(|s| {
10467 let b = unsafe { &mut *s.get() };
10470 if b.len() < n {
10471 b.resize(n, 0.0);
10472 }
10473 f(&mut b[..n])
10474 })
10475}
10476
10477#[inline(always)]
10488fn q8_round(t: f32) -> i8 {
10489 let t = t.clamp(-127.0, 127.0);
10490 let i = t as i32;
10491 let f = t - i as f32;
10492 let r = if f >= 0.5 {
10493 i + 1
10494 } else if f <= -0.5 {
10495 i - 1
10496 } else {
10497 i
10498 };
10499 r as i8
10500}
10501
10502#[inline]
10503fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
10504 KROWS.with(|s| {
10505 let b = unsafe { &mut *s.get() };
10508 for row in b.iter_mut() {
10509 if row.len() < n {
10510 row.resize(n, 0.0);
10511 }
10512 }
10513 let (a, b) = b.split_at_mut(1);
10514 f(&mut a[0][..n], &mut b[0][..n])
10515 })
10516}
10517
10518#[inline]
10519fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
10520 if limit > 0.0 {
10521 up = up.clamp(-limit, limit);
10522 gate = gate.min(limit);
10523 }
10524 gate / (1.0 + (-gate).exp()) * up
10525}
10526
10527fn split_act(x: &[f32]) -> SplitAct {
10528 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10529 let n = x.len();
10530 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
10531 let thr = 8.0 * rms;
10532 let mut outliers: Vec<(usize, f32)> = Vec::new();
10536 let mut amax = 0f32;
10537 for (j, &v) in x.iter().enumerate() {
10538 let a = v.abs();
10539 if a > thr {
10540 outliers.push((j, v));
10541 } else if a > amax {
10542 amax = a;
10543 }
10544 }
10545 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10546 let inv = 1.0 / sx;
10547 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10548 xq.clear();
10549 xq.reserve(n);
10550 if outliers.is_empty() {
10551 xq.extend(
10552 x.iter()
10553 .map(|&v| q8_round(v * inv)),
10554 );
10555 } else {
10556 xq.extend(x.iter().map(|&v| {
10558 if v.abs() > thr {
10559 0
10560 } else {
10561 q8_round(v * inv)
10562 }
10563 }));
10564 }
10565 let xsum = xq.iter().map(|&v| v as i32).sum();
10566 SplitAct {
10567 xq,
10568 sx,
10569 outliers,
10570 xsum,
10571 }
10572}
10573
10574fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
10575 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10576 let n = x.len();
10577 let rms = (x
10578 .iter()
10579 .zip(col)
10580 .map(|(&a, &c)| {
10581 let v = a * c;
10582 (v * v) as f64
10583 })
10584 .sum::<f64>()
10585 / n.max(1) as f64)
10586 .sqrt() as f32;
10587 let thr = 8.0 * rms;
10588
10589 let mut outliers = Vec::new();
10590 let mut amax = 0f32;
10591 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
10592 let v = a * c;
10593 let s = v.abs();
10594 if s > thr {
10595 outliers.push((j, v));
10596 } else if s > amax {
10597 amax = s;
10598 }
10599 }
10600
10601 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10602 let inv = 1.0 / sx;
10603 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10604 xq.clear();
10605 xq.reserve(n);
10606 if outliers.is_empty() {
10607 xq.extend(
10608 x.iter()
10609 .zip(col)
10610 .map(|(&a, &c)| q8_round((a * c) * inv)),
10611 );
10612 } else {
10613 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
10614 let v = a * c;
10615 if v.abs() > thr {
10616 0
10617 } else {
10618 q8_round(v * inv)
10619 }
10620 }));
10621 }
10622 let xsum = xq.iter().map(|&v| v as i32).sum();
10623 SplitAct {
10624 xq,
10625 sx,
10626 outliers,
10627 xsum,
10628 }
10629}
10630
10631#[cfg(target_arch = "aarch64")]
10634#[target_feature(enable = "neon,dotprod")]
10635unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
10636 unsafe {
10638 use core::arch::aarch64::*;
10639 use core::arch::asm;
10640 let wp = w.as_ptr() as *const i8;
10641 let n = w.len();
10642 let (mut a0, mut a1, mut a2, mut a3) = (
10643 vdupq_n_s32(0),
10644 vdupq_n_s32(0),
10645 vdupq_n_s32(0),
10646 vdupq_n_s32(0),
10647 );
10648 let mut i = 0;
10649 while i + 64 <= n {
10650 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10651 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
10652 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
10653 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
10654 asm!(
10655 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10656 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10657 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
10658 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
10659 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10660 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10661 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
10662 options(pure, nomem, nostack),
10663 );
10664 i += 64;
10665 }
10666 while i + 16 <= n {
10667 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10668 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
10669 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
10670 i += 16;
10671 }
10672 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
10673 while i < n {
10674 s += (*wp.add(i)) as i32 * xq[i] as i32;
10675 i += 1;
10676 }
10677 s
10678 }
10679}
10680
10681#[cfg(target_arch = "aarch64")]
10685#[target_feature(enable = "neon,dotprod")]
10686unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
10687 unsafe {
10689 use core::arch::aarch64::*;
10690 use core::arch::asm;
10691 let n = xq.len();
10692 let px = xq.as_ptr();
10693 let (p0, p1, p2, p3) = (
10694 w0.as_ptr() as *const i8,
10695 w1.as_ptr() as *const i8,
10696 w2.as_ptr() as *const i8,
10697 w3.as_ptr() as *const i8,
10698 );
10699 let (mut a0, mut a1, mut a2, mut a3) = (
10700 vdupq_n_s32(0),
10701 vdupq_n_s32(0),
10702 vdupq_n_s32(0),
10703 vdupq_n_s32(0),
10704 );
10705 let mut i = 0;
10706 while i + 16 <= n {
10707 let x = vld1q_s8(px.add(i));
10708 let v0 = vld1q_s8(p0.add(i));
10709 let v1 = vld1q_s8(p1.add(i));
10710 let v2 = vld1q_s8(p2.add(i));
10711 let v3 = vld1q_s8(p3.add(i));
10712 asm!(
10713 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10714 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10715 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10716 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10717 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10718 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10719 options(pure, nomem, nostack),
10720 );
10721 i += 16;
10722 }
10723 let mut r = [
10724 vaddvq_s32(a0),
10725 vaddvq_s32(a1),
10726 vaddvq_s32(a2),
10727 vaddvq_s32(a3),
10728 ];
10729 while i < n {
10730 let xi = *px.add(i) as i32;
10731 r[0] += (*p0.add(i)) as i32 * xi;
10732 r[1] += (*p1.add(i)) as i32 * xi;
10733 r[2] += (*p2.add(i)) as i32 * xi;
10734 r[3] += (*p3.add(i)) as i32 * xi;
10735 i += 1;
10736 }
10737 r
10738 }
10739}
10740
10741#[cfg(target_arch = "aarch64")]
10748#[target_feature(enable = "neon,dotprod")]
10749unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
10750 unsafe {
10753 use core::arch::aarch64::*;
10754 use core::arch::asm;
10755 let n = xq.len();
10756 let px = xq.as_ptr();
10757 let pg = g.as_ptr() as *const i8;
10758 let (mut a0, mut a1, mut a2, mut a3) = (
10759 vdupq_n_s32(0),
10760 vdupq_n_s32(0),
10761 vdupq_n_s32(0),
10762 vdupq_n_s32(0),
10763 );
10764 let mut i = 0;
10765 while i + 16 <= n {
10766 let x = vld1q_s8(px.add(i));
10767 let base = pg.add(4 * i);
10768 let v0 = vld1q_s8(base);
10769 let v1 = vld1q_s8(base.add(16));
10770 let v2 = vld1q_s8(base.add(32));
10771 let v3 = vld1q_s8(base.add(48));
10772 asm!(
10773 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10774 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10775 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10776 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10777 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10778 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10779 options(pure, nomem, nostack),
10780 );
10781 i += 16;
10782 }
10783 [
10784 vaddvq_s32(a0),
10785 vaddvq_s32(a1),
10786 vaddvq_s32(a2),
10787 vaddvq_s32(a3),
10788 ]
10789 }
10790}
10791
10792#[cfg(target_arch = "aarch64")]
10798fn q8_range_sdot(
10799 q: &[u8],
10800 rep: &[u8],
10801 row_scale: &[f32],
10802 act: &SplitAct,
10803 cols: usize,
10804 out_addr: SendMut,
10805 start: usize,
10806 end: usize,
10807) {
10808 let mut o = start;
10809 if !rep.is_empty() {
10812 while o < end && o % 4 != 0 {
10813 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10814 unsafe { *out_addr.at(o) = v };
10815 o += 1;
10816 }
10817 }
10818 while o + 4 <= end {
10819 let r = if rep.is_empty() {
10820 unsafe {
10821 dot_i8_sdot_4rows(
10822 &q[o * cols..(o + 1) * cols],
10823 &q[(o + 1) * cols..(o + 2) * cols],
10824 &q[(o + 2) * cols..(o + 3) * cols],
10825 &q[(o + 3) * cols..(o + 4) * cols],
10826 &act.xq,
10827 )
10828 }
10829 } else {
10830 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
10831 };
10832 for k in 0..4 {
10833 let mut acc = r[k] as f32 * act.sx;
10834 for &(j, xv) in &act.outliers {
10835 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
10836 }
10837 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
10839 }
10840 o += 4;
10841 }
10842 while o < end {
10843 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10844 unsafe { *out_addr.at(o) = v };
10845 o += 1;
10846 }
10847}
10848
10849#[cfg(target_arch = "aarch64")]
10852#[allow(clippy::too_many_arguments)]
10853fn q8_range2_sdot(
10854 q: &[u8],
10855 row_scale: &[f32],
10856 a1: &SplitAct,
10857 a2: &SplitAct,
10858 cols: usize,
10859 p1: SendMut,
10860 p2: SendMut,
10861 start: usize,
10862 end: usize,
10863) {
10864 for o in start..end {
10865 let row = &q[o * cols..(o + 1) * cols];
10866 unsafe {
10868 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
10869 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
10870 }
10871 }
10872}
10873
10874#[allow(clippy::too_many_arguments)]
10876fn q8_range2_f32(
10877 q: &[u8],
10878 row_scale: &[f32],
10879 x1: &[f32],
10880 x2: &[f32],
10881 cols: usize,
10882 p1: SendMut,
10883 p2: SendMut,
10884 start: usize,
10885 end: usize,
10886) {
10887 for o in start..end {
10888 let row = &q[o * cols..(o + 1) * cols];
10889 unsafe {
10891 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
10892 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
10893 }
10894 }
10895}
10896
10897fn q8_range_f32(
10899 q: &[u8],
10900 row_scale: &[f32],
10901 xs: &[f32],
10902 cols: usize,
10903 out_addr: SendMut,
10904 start: usize,
10905 end: usize,
10906) {
10907 for o in start..end {
10908 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10909 unsafe { *out_addr.at(o) = v };
10911 }
10912}
10913
10914#[inline]
10918fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
10919 #[cfg(target_arch = "aarch64")]
10920 return row_dot_sdot(row, act);
10921 #[cfg(target_arch = "x86_64")]
10922 return row_dot_avx2(row, act);
10923 #[allow(unreachable_code)]
10924 q8_row_dot_scalar(row, act)
10925}
10926
10927#[allow(dead_code)]
10928fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
10929 let mut acc = 0i32;
10930 for (k, &b) in row.iter().enumerate() {
10931 acc += (b as i8) as i32 * act.xq[k] as i32;
10932 }
10933 let mut acc = acc as f32 * act.sx;
10934 for &(j, xv) in &act.outliers {
10935 acc += (row[j] as i8) as f32 * xv;
10936 }
10937 acc
10938}
10939
10940#[cfg(target_arch = "aarch64")]
10943#[inline]
10944fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
10945 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
10946 for &(j, xv) in &act.outliers {
10947 acc += (row[j] as i8) as f32 * xv;
10948 }
10949 acc
10950}
10951
10952#[cfg(target_arch = "aarch64")]
10960#[target_feature(enable = "neon,dotprod")]
10961unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10962 unsafe {
10965 use core::arch::aarch64::*;
10966 use core::arch::asm;
10967 let lomask = vdupq_n_u8(0x0F);
10968 let eight = vdupq_n_s8(8);
10969 let mut acc = 0f32;
10970 for gi in 0..gpr {
10971 let g = g0 + gi;
10972 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10973 let b = vld1q_u8(packed.as_ptr().add(g * 16));
10974 let lo = vandq_u8(b, lomask);
10975 let hi = vshrq_n_u8::<4>(b);
10976 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10977 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10978 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
10979 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
10980 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10981 asm!(
10982 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
10983 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
10984 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10985 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
10986 options(pure, nomem, nostack),
10987 );
10988 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10989 }
10990 acc
10991 }
10992}
10993
10994#[cfg(target_arch = "aarch64")]
10999#[target_feature(enable = "neon,dotprod")]
11000unsafe fn dot_q4_row_sdot2(
11001 packed: &[u8],
11002 scales: &[u8],
11003 g0: usize,
11004 gpr: usize,
11005 xq1: &[i8],
11006 xq2: &[i8],
11007) -> (f32, f32) {
11008 unsafe {
11011 use core::arch::aarch64::*;
11012 use core::arch::asm;
11013 let lomask = vdupq_n_u8(0x0F);
11014 let eight = vdupq_n_s8(8);
11015 let (mut acc1, mut acc2) = (0f32, 0f32);
11016 for gi in 0..gpr {
11017 let g = g0 + gi;
11018 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11019 let b = vld1q_u8(packed.as_ptr().add(g * 16));
11020 let lo = vandq_u8(b, lomask);
11021 let hi = vshrq_n_u8::<4>(b);
11022 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11023 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11024 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
11025 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
11026 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
11027 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
11028 let (mut a0, mut a1, mut b0, mut b1) = (
11029 vdupq_n_s32(0),
11030 vdupq_n_s32(0),
11031 vdupq_n_s32(0),
11032 vdupq_n_s32(0),
11033 );
11034 asm!(
11035 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
11036 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
11037 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
11038 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
11039 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11040 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
11041 e0 = in(vreg) e0, e1 = in(vreg) e1,
11042 x10 = in(vreg) x10, x11 = in(vreg) x11,
11043 x20 = in(vreg) x20, x21 = in(vreg) x21,
11044 options(pure, nomem, nostack),
11045 );
11046 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11047 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
11048 }
11049 (acc1, acc2)
11050 }
11051}
11052
11053#[inline]
11058pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
11059 #[cfg(target_arch = "aarch64")]
11060 unsafe {
11061 return axpy_i8_f32_neon(acc, row, w);
11062 }
11063 #[cfg(target_arch = "x86_64")]
11064 if avx2_enabled() {
11065 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
11066 }
11067 #[allow(unreachable_code)]
11068 {
11069 for (a, &b) in acc.iter_mut().zip(row) {
11070 *a += w * b as f32;
11071 }
11072 }
11073}
11074
11075#[cfg(target_arch = "x86_64")]
11077#[target_feature(enable = "avx2,fma")]
11078unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
11079 unsafe {
11081 use core::arch::x86_64::*;
11082 let n = acc.len().min(row.len());
11083 let ap = acc.as_mut_ptr();
11084 let rp = row.as_ptr();
11085 let wv = _mm256_set1_ps(w);
11086 let mut j = 0usize;
11087 while j + 16 <= n {
11088 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
11089 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
11090 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
11091 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
11092 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
11093 _mm256_storeu_ps(ap.add(j), v0);
11094 _mm256_storeu_ps(ap.add(j + 8), v1);
11095 j += 16;
11096 }
11097 while j < n {
11098 *ap.add(j) += w * (*rp.add(j)) as f32;
11099 j += 1;
11100 }
11101 }
11102}
11103
11104#[cfg(target_arch = "aarch64")]
11105#[target_feature(enable = "neon")]
11106unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
11107 unsafe {
11109 use core::arch::aarch64::*;
11110 let n = acc.len().min(row.len());
11111 let ap = acc.as_mut_ptr();
11112 let rp = row.as_ptr();
11113 let wv = vdupq_n_f32(w);
11114 let mut j = 0usize;
11115 while j + 16 <= n {
11116 let rb = vld1q_s8(rp.add(j));
11117 let lo = vmovl_s8(vget_low_s8(rb));
11118 let hi = vmovl_s8(vget_high_s8(rb));
11119 for (off, half) in [(0, lo), (8, hi)] {
11120 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
11121 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
11122 let o = j + off;
11123 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
11124 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
11125 }
11126 j += 16;
11127 }
11128 while j < n {
11129 *ap.add(j) += w * (*rp.add(j)) as f32;
11130 j += 1;
11131 }
11132 }
11133}
11134
11135#[inline]
11138pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
11139 #[cfg(target_arch = "aarch64")]
11140 unsafe {
11141 return dot_i8_f32_neon(w, x);
11142 }
11143 #[cfg(target_arch = "x86_64")]
11144 if avx2_enabled() {
11145 return unsafe { dot_i8_f32_avx2(w, x) };
11146 }
11147 #[allow(unreachable_code)]
11148 {
11149 let mut sum = 0.0f32;
11150 for (j, &b) in w.iter().enumerate() {
11151 sum += (b as i8) as f32 * x[j];
11152 }
11153 sum
11154 }
11155}
11156
11157#[inline]
11161fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
11162 #[cfg(target_arch = "aarch64")]
11163 unsafe {
11164 return dot_i8_col_f32_neon(w, x, col);
11165 }
11166 #[allow(unreachable_code)]
11167 {
11168 let mut sum = 0.0f32;
11169 for (j, &b) in w.iter().enumerate() {
11170 sum += (b as i8) as f32 * x[j] * col[j];
11171 }
11172 sum
11173 }
11174}
11175
11176#[cfg(target_arch = "aarch64")]
11177#[target_feature(enable = "neon")]
11178unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
11179 unsafe {
11181 use core::arch::aarch64::*;
11182 let n = x.len();
11183 let wp = w.as_ptr() as *const i8;
11184 let xp = x.as_ptr();
11185 let cp = col.as_ptr();
11186 let (mut a0, mut a1, mut a2, mut a3) = (
11187 vdupq_n_f32(0.0),
11188 vdupq_n_f32(0.0),
11189 vdupq_n_f32(0.0),
11190 vdupq_n_f32(0.0),
11191 );
11192 let mut j = 0usize;
11193 while j + 16 <= n {
11194 let wb = vld1q_s8(wp.add(j));
11195 let lo = vmovl_s8(vget_low_s8(wb));
11196 let hi = vmovl_s8(vget_high_s8(wb));
11197 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
11198 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
11199 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
11200 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
11201 a0 = vfmaq_f32(
11202 a0,
11203 w0,
11204 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
11205 );
11206 a1 = vfmaq_f32(
11207 a1,
11208 w1,
11209 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
11210 );
11211 a2 = vfmaq_f32(
11212 a2,
11213 w2,
11214 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
11215 );
11216 a3 = vfmaq_f32(
11217 a3,
11218 w3,
11219 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
11220 );
11221 j += 16;
11222 }
11223 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
11224 while j < n {
11225 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
11226 j += 1;
11227 }
11228 sum
11229 }
11230}
11231
11232#[cfg(target_arch = "aarch64")]
11233#[target_feature(enable = "neon")]
11234unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
11235 unsafe {
11237 use core::arch::aarch64::*;
11238 let n = x.len();
11239 let wp = w.as_ptr() as *const i8;
11240 let xp = x.as_ptr();
11241 let (mut a0, mut a1, mut a2, mut a3) = (
11242 vdupq_n_f32(0.0),
11243 vdupq_n_f32(0.0),
11244 vdupq_n_f32(0.0),
11245 vdupq_n_f32(0.0),
11246 );
11247 let mut j = 0usize;
11248 while j + 16 <= n {
11249 let wb = vld1q_s8(wp.add(j));
11250 let lo = vmovl_s8(vget_low_s8(wb));
11251 let hi = vmovl_s8(vget_high_s8(wb));
11252 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
11253 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
11254 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
11255 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
11256 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
11257 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
11258 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
11259 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
11260 j += 16;
11261 }
11262 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
11263 while j < n {
11264 sum += (*wp.add(j)) as f32 * *xp.add(j);
11265 j += 1;
11266 }
11267 sum
11268 }
11269}
11270
11271#[allow(clippy::too_many_arguments)]
11272fn qmatvec(
11273 q: &[u8],
11274 rep: &[u8],
11275 row_scale: &[f32],
11276 x: &[f32],
11277 col_field: &[f32],
11278 dtype: TensorDtype,
11279 rows: usize,
11280 cols: usize,
11281 out: &mut [f32],
11282 pool: Option<&Pool>,
11283) {
11284 debug_assert_eq!(out.len(), rows);
11285 #[cfg(not(target_arch = "aarch64"))]
11286 let _ = rep;
11287
11288 #[cfg(target_arch = "aarch64")]
11289 if sdot_enabled() {
11290 let act = if dtype == TensorDtype::Q8_2f {
11291 split_act_q8_2f(x, col_field)
11292 } else {
11293 split_act(x)
11294 };
11295 let out_addr = SendMut(out.as_mut_ptr());
11296 let run_range = |start: usize, end: usize| {
11297 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
11298 };
11299 match pool {
11300 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11301 _ => run_range(0, rows),
11302 }
11303 return;
11304 }
11305 #[cfg(target_arch = "x86_64")]
11308 if avx2_a8w8_enabled() {
11309 let act = if dtype == TensorDtype::Q8_2f {
11310 split_act_q8_2f(x, col_field)
11311 } else {
11312 split_act(x)
11313 };
11314 let out_addr = SendMut(out.as_mut_ptr());
11315 let run_range = |start: usize, end: usize| {
11316 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
11317 };
11318 match pool {
11319 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11320 _ => run_range(0, rows),
11321 }
11322 return;
11323 }
11324
11325 prescale_with(x, col_field, dtype, 1, |xs| {
11326 let out_addr = SendMut(out.as_mut_ptr());
11327 let run_range = move |start: usize, end: usize| {
11328 for o in start..end {
11329 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
11330 unsafe { *out_addr.at(o) = v };
11332 }
11333 };
11334 match pool {
11335 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11336 _ => run_range(0, rows),
11337 }
11338 });
11339}
11340
11341#[allow(clippy::too_many_arguments)]
11342fn qmatvec2(
11343 q: &[u8],
11344 row_scale: &[f32],
11345 x1: &[f32],
11346 x2: &[f32],
11347 col_field: &[f32],
11348 dtype: TensorDtype,
11349 rows: usize,
11350 cols: usize,
11351 o1: &mut [f32],
11352 o2: &mut [f32],
11353 pool: Option<&Pool>,
11354) {
11355 #[cfg(target_arch = "aarch64")]
11356 if sdot_enabled() {
11357 let a1s = if dtype == TensorDtype::Q8_2f {
11358 split_act_q8_2f(x1, col_field)
11359 } else {
11360 split_act(x1)
11361 };
11362 let a2s = if dtype == TensorDtype::Q8_2f {
11363 split_act_q8_2f(x2, col_field)
11364 } else {
11365 split_act(x2)
11366 };
11367 let p1 = SendMut(o1.as_mut_ptr());
11368 let p2 = SendMut(o2.as_mut_ptr());
11369 let run_range = |start: usize, end: usize| {
11370 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
11371 };
11372 match pool {
11373 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11374 _ => run_range(0, rows),
11375 }
11376 return;
11377 }
11378 #[cfg(target_arch = "x86_64")]
11379 if avx2_a8w8_enabled() {
11380 let a1s = if dtype == TensorDtype::Q8_2f {
11381 split_act_q8_2f(x1, col_field)
11382 } else {
11383 split_act(x1)
11384 };
11385 let a2s = if dtype == TensorDtype::Q8_2f {
11386 split_act_q8_2f(x2, col_field)
11387 } else {
11388 split_act(x2)
11389 };
11390 let p1 = SendMut(o1.as_mut_ptr());
11391 let p2 = SendMut(o2.as_mut_ptr());
11392 let run_range = |start: usize, end: usize| {
11393 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
11394 };
11395 match pool {
11396 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11397 _ => run_range(0, rows),
11398 }
11399 return;
11400 }
11401
11402 prescale_with(x1, col_field, dtype, 1, |x1s| {
11403 prescale_with(x2, col_field, dtype, 2, |x2s| {
11404 let p1 = SendMut(o1.as_mut_ptr());
11405 let p2 = SendMut(o2.as_mut_ptr());
11406 let run_range = move |start: usize, end: usize| {
11407 for o in start..end {
11408 let row = &q[o * cols..(o + 1) * cols];
11409 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
11410 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
11411 unsafe {
11413 *p1.at(o) = s1;
11414 *p2.at(o) = s2;
11415 }
11416 }
11417 };
11418 match pool {
11419 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11420 _ => run_range(0, rows),
11421 }
11422 });
11423 });
11424}
11425
11426#[derive(Clone, Copy)]
11427struct SendMut(*mut f32);
11428unsafe impl Send for SendMut {}
11429unsafe impl Sync for SendMut {}
11430
11431impl SendMut {
11432 #[inline]
11433 fn at(self, i: usize) -> *mut f32 {
11434 unsafe { self.0.add(i) }
11435 }
11436}
11437
11438#[cfg(test)]
11439mod tests {
11440 #[test]
11444 fn q8_round_is_round_clamp() {
11445 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
11446 let mut probe = vec![
11447 0.0f32,
11448 -0.0,
11449 f32::NAN,
11450 f32::INFINITY,
11451 f32::NEG_INFINITY,
11452 f32::MAX,
11453 f32::MIN,
11454 1e30,
11455 -1e30,
11456 f32::MIN_POSITIVE,
11457 -f32::MIN_POSITIVE,
11458 ];
11459 for k in -300i32..=300 {
11460 let h = k as f32 * 0.5;
11461 let up = f32::from_bits(h.to_bits() + 1);
11462 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
11463 for t in [h, up, down] {
11464 probe.push(t);
11465 probe.push(-t);
11466 }
11467 }
11468 let mut t = -140.0f32;
11469 while t < 140.0 {
11470 probe.push(t);
11471 t += 0.000_731;
11472 }
11473 for t in probe {
11474 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
11475 }
11476 }
11477
11478 use super::*;
11479
11480 #[test]
11484 fn f32_matmat_equals_matvec_bitwise() {
11485 let (rows, cols, b) = (5usize, 37usize, 7usize);
11486 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
11487 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
11488 let t = QTensor::from_f32(w, rows, cols);
11489 let mut out = vec![0f32; b * rows];
11490 t.matmat(&xs, b, &mut out, None);
11491 for bi in 0..b {
11492 let mut one = vec![0f32; rows];
11493 t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
11494 for o in 0..rows {
11495 assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
11496 }
11497 }
11498 }
11499
11500 #[test]
11501 fn q2tp_i8_dot_matches_exact_on_grid() {
11502 let (rows, cols) = (5, 64);
11506 let gpr = cols / GROUP_SIZE;
11507 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
11511 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
11512 .collect();
11513 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
11514 let x: Vec<f32> = (0..cols)
11515 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
11516 .collect();
11517 let act = split_act(&x);
11518 assert!(
11519 act.outliers.is_empty(),
11520 "on-grid input must have no outliers"
11521 );
11522 let gsum = q1_group_sums(&act.xq, gpr);
11523 for r in 0..rows {
11524 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
11525 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
11526 assert!(
11527 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
11528 "row {r}: exact {exact} vs i8 {fast}"
11529 );
11530 }
11531 }
11532
11533 #[cfg(target_arch = "x86_64")]
11534 #[test]
11535 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
11536 if !std::arch::is_x86_feature_detected!("avx2") {
11541 return;
11542 }
11543 let mut seed = 0x9e3779b9u32;
11544 let mut next = || {
11545 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
11546 seed
11547 };
11548 for _ in 0..20_000 {
11549 let mut ch = [0u8; Q2TP_CHUNK];
11550 let mut x = [0i8; GROUP_SIZE];
11551 for b in &mut ch {
11552 *b = next() as u8;
11553 }
11554 for v in &mut x {
11555 *v = (next() >> 24) as i8;
11556 }
11557 let mut reference = 0i32;
11558 for (k, &b) in ch.iter().enumerate() {
11559 reference += (b & 3) as i32 * x[k * 4] as i32;
11560 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
11561 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
11562 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
11563 }
11564 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
11567 assert_eq!(got, reference, "packed q2 lane mismatch");
11568 }
11569 }
11570
11571 #[test]
11572 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
11573 let (rows, cols) = (1usize, GROUP_SIZE);
11574 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
11575 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
11579 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
11580 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
11582 let mut raw = vec![0.0f32; rows];
11583 let mut affine = vec![0.0f32; rows];
11584 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
11585 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
11586 assert_eq!(raw, vec![-24.0]);
11587 assert_eq!(affine, vec![-8.0]);
11588 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
11589 }
11590
11591 #[test]
11592 fn q8_row_dot_fast_matches_scalar() {
11593 let cols = 96;
11596 let row: Vec<u8> = (0..cols)
11597 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
11598 .collect();
11599 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
11600 let act = split_act(&x);
11601 let fast = q8_row_dot(&row, &act);
11602 let scalar = q8_row_dot_scalar(&row, &act);
11603 assert!(
11604 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
11605 "fast {fast} vs scalar {scalar}"
11606 );
11607 }
11608
11609 #[test]
11610 fn f32_matvec_matches_matvec_rows_bitexact() {
11611 let (rows, cols) = (300, 40);
11612 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
11613 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
11614 let qt = QTensor::from_f32(w.clone(), rows, cols);
11615
11616 let mut a = vec![0.0f32; rows];
11617 matvec_rows(None, &w, &x, &mut a);
11618 let mut b = vec![0.0f32; rows];
11619 qt.matvec(&x, &mut b, None);
11620 assert_eq!(a, b);
11621 }
11622
11623 #[test]
11624 fn sdot_kernel_exact_on_grid() {
11625 eprintln!("sdot_enabled = {}", sdot_enabled());
11630 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
11632 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11633 .collect();
11634 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
11635 let x: Vec<f32> = (0..cols)
11636 .map(|i| match i % 3 {
11637 0 => 1.0,
11638 1 => -1.0,
11639 _ => 0.0,
11640 })
11641 .collect();
11642 let mut a = vec![0.0f32; rows];
11643 qmatvec(
11644 &w,
11645 &[],
11646 &scales,
11647 &x,
11648 &[],
11649 TensorDtype::Q8Row,
11650 rows,
11651 cols,
11652 &mut a,
11653 None,
11654 );
11655 for o in 0..rows {
11656 let mut acc = 0.0f32;
11657 for j in 0..cols {
11658 acc += (w[o * cols + j] as i8) as f32 * x[j];
11659 }
11660 let expect = acc * scales[o];
11661 assert!(
11662 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11663 "row {o}: {} vs {expect}",
11664 a[o]
11665 );
11666 }
11667 }
11668
11669 #[test]
11670 fn q1_tbl_fast_path_matches_reference() {
11671 let (rows, cols) = (5, 256);
11676 let gpr = cols / GROUP_SIZE;
11677 let mut bytes = Vec::new();
11678 for t in 0..rows * gpr {
11679 let s = 0.007 + (t % 11) as f32 * 0.004;
11680 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11681 for j in 0..4 {
11682 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
11683 }
11684 }
11685 let x: Vec<f32> = (0..cols)
11686 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
11687 .collect();
11688 let mut w = vec![0.0f32; rows * cols];
11689 cortiq_core::quant::dequant_q1(&bytes, &mut w);
11690 let mut got = vec![0.0f32; rows];
11691 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11692 for o in 0..rows {
11693 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11694 assert!(
11695 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11696 "row {o}: {} vs {expect}",
11697 got[o]
11698 );
11699 }
11700 let b = 5usize;
11703 let mut xs_all = Vec::new();
11704 for bi in 0..b {
11705 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
11706 }
11707 let mut mm = vec![0.0f32; b * rows];
11708 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
11709 for bi in 0..b {
11710 let mut single = vec![0.0f32; rows];
11711 q1_matvec(
11712 &bytes,
11713 &xs_all[bi * cols..(bi + 1) * cols],
11714 rows,
11715 cols,
11716 &mut single,
11717 None,
11718 );
11719 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
11720 }
11721 }
11722
11723 #[test]
11724 fn q1_kernels_match_exact_reference() {
11725 let (rows, cols) = (7, 96);
11727 let gpr = cols / GROUP_SIZE;
11728 let mut bytes = Vec::new();
11729 for t in 0..rows * gpr {
11730 let s = 0.01 + (t % 13) as f32 * 0.003;
11731 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11732 for j in 0..4 {
11733 bytes.push(((t * 31 + j * 97) % 251) as u8);
11734 }
11735 }
11736 let x: Vec<f32> = (0..cols)
11738 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
11739 .collect();
11740 let mut w = vec![0.0f32; rows * cols];
11742 cortiq_core::quant::dequant_q1(&bytes, &mut w);
11743 let mut expect = vec![0.0f32; rows];
11744 for o in 0..rows {
11745 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11746 }
11747 let mut got = vec![0.0f32; rows];
11748 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11749 for o in 0..rows {
11750 assert!(
11751 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
11752 "row {o}: {} vs {}",
11753 got[o],
11754 expect[o]
11755 );
11756 }
11757 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
11759 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
11760 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
11761 assert_eq!(a1, got);
11762 let mut xs = x.clone();
11763 xs.extend_from_slice(&x2);
11764 let mut mm = vec![0.0f32; 2 * rows];
11765 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
11766 assert_eq!(&mm[..rows], got.as_slice());
11767 assert_eq!(&mm[rows..], a2.as_slice());
11768 }
11769
11770 #[test]
11771 fn repack_is_bit_identical() {
11772 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
11778 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
11779 .collect();
11780 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
11781 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
11782 let rep = q8_repack_layout(&w, rows, cols);
11783 for g in 0..rows / 4 {
11785 for c in 0..cols / 16 {
11786 for lane in 0..4 {
11787 assert_eq!(
11788 &rep[g * 4 * cols + c * 64 + lane * 16
11789 ..g * 4 * cols + c * 64 + lane * 16 + 16],
11790 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
11791 );
11792 }
11793 }
11794 }
11795 let mut a = vec![0.0f32; rows];
11796 qmatvec(
11797 &w,
11798 &[],
11799 &scales,
11800 &x,
11801 &[],
11802 TensorDtype::Q8Row,
11803 rows,
11804 cols,
11805 &mut a,
11806 None,
11807 );
11808 let mut b = vec![0.0f32; rows];
11809 qmatvec(
11810 &w,
11811 &rep,
11812 &scales,
11813 &x,
11814 &[],
11815 TensorDtype::Q8Row,
11816 rows,
11817 cols,
11818 &mut b,
11819 None,
11820 );
11821 assert_eq!(a, b, "full-range repack output diverged");
11822
11823 #[cfg(target_arch = "aarch64")]
11824 if sdot_enabled() {
11825 let act = split_act(&x);
11827 let mut c1 = vec![0.0f32; rows];
11828 let mut c2 = vec![0.0f32; rows];
11829 q8_range_sdot(
11830 &w,
11831 &[],
11832 &scales,
11833 &act,
11834 cols,
11835 SendMut(c1.as_mut_ptr()),
11836 3,
11837 rows - 2,
11838 );
11839 q8_range_sdot(
11840 &w,
11841 &rep,
11842 &scales,
11843 &act,
11844 cols,
11845 SendMut(c2.as_mut_ptr()),
11846 3,
11847 rows - 2,
11848 );
11849 assert_eq!(c1, c2, "unaligned-range repack output diverged");
11850 }
11851 }
11852
11853 #[test]
11854 fn sdot_a8w8_noise_is_bounded() {
11855 let (rows, cols) = (16, 512);
11859 let w: Vec<u8> = (0..rows * cols)
11860 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11861 .collect();
11862 let scales = vec![0.01f32; rows];
11863 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
11864 let mut a = vec![0.0f32; rows];
11865 qmatvec(
11866 &w,
11867 &[],
11868 &scales,
11869 &x,
11870 &[],
11871 TensorDtype::Q8Row,
11872 rows,
11873 cols,
11874 &mut a,
11875 None,
11876 );
11877 let (mut num, mut den) = (0f64, 0f64);
11878 for o in 0..rows {
11879 let mut acc = 0.0f32;
11880 for j in 0..cols {
11881 acc += (w[o * cols + j] as i8) as f32 * x[j];
11882 }
11883 let expect = acc * scales[o];
11884 num += ((a[o] - expect) as f64).powi(2);
11885 den += (expect as f64).powi(2);
11886 }
11887 let rel = (num / den.max(1e-12)).sqrt();
11888 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
11889 }
11890
11891 #[test]
11892 fn i8_dot_neon_matches_scalar() {
11893 let n = 100;
11894 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
11895 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
11896 let mut scalar = 0.0f32;
11897 for j in 0..n {
11898 scalar += (w[j] as i8) as f32 * x[j];
11899 }
11900 let fast = dot_i8_f32(&w, &x);
11901 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
11902 }
11903
11904 #[test]
11906 fn vbitmatvec_matches_full_dequant() {
11907 let (rows, cols) = (6, 64);
11908 let ng = cols / GROUP_SIZE;
11909 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11911 let mut bytes = bits.clone();
11912 for g in 0..rows * ng {
11913 let s = 0.02 + 0.001 * g as f32;
11914 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11915 }
11916 for r in 0..rows {
11917 let b = bits[r] as usize;
11918 let (mut acc, mut nb) = (0u64, 0usize);
11919 let mut rowbytes = Vec::new();
11920 for i in 0..cols {
11921 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11922 acc = (acc << b) | v;
11923 nb += b;
11924 while nb >= 8 {
11925 nb -= 8;
11926 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11927 }
11928 }
11929 if nb > 0 {
11930 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11931 }
11932 bytes.extend_from_slice(&rowbytes);
11933 }
11934 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11935
11936 let mut reference = vec![0f32; rows * cols];
11937 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
11938 let mut expect = vec![0f32; rows];
11939 for r in 0..rows {
11940 expect[r] = reference[r * cols..(r + 1) * cols]
11941 .iter()
11942 .zip(&x)
11943 .map(|(w, xv)| w * xv)
11944 .sum();
11945 }
11946 let mut got = vec![0f32; rows];
11947 let offsets = vbit_row_offsets(&bytes, rows, cols);
11948 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
11949 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11953 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11954 for r in 0..rows {
11955 assert!(
11956 (got[r] - expect[r]).abs() < tol * scale,
11957 "row {r}: {} vs {}",
11958 got[r],
11959 expect[r]
11960 );
11961 }
11962 }
11963
11964 #[test]
11969 #[cfg(target_arch = "x86_64")]
11970 fn vbit_matmat_blocked_matches_per_row() {
11971 let (rows, cols, b) = (64usize, 128usize, 9usize);
11972 let ng = cols / GROUP_SIZE;
11973 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
11974 let mut bytes = bits.clone();
11975 for g in 0..rows * ng {
11976 let sc = 0.02 + 0.0005 * g as f32;
11977 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11978 }
11979 for r in 0..rows {
11980 let bw = bits[r] as usize;
11981 let (mut acc, mut nb) = (0u64, 0usize);
11982 let mut rowbytes = Vec::new();
11983 for i in 0..cols {
11984 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11985 acc = (acc << bw) | v;
11986 nb += bw;
11987 while nb >= 8 {
11988 nb -= 8;
11989 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11990 }
11991 }
11992 if nb > 0 {
11993 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11994 }
11995 bytes.extend_from_slice(&rowbytes);
11996 }
11997 let x: Vec<f32> = (0..b * cols)
11998 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11999 .collect();
12000 let offsets = vbit_row_offsets(&bytes, rows, cols);
12001 let mut y_a = vec![0f32; b * rows];
12002 let mut y_b = vec![0f32; b * rows];
12003 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
12004 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
12005 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
12006 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
12007 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
12008 let max_d = y_a
12009 .iter()
12010 .zip(&y_b)
12011 .map(|(p, q)| (p - q).abs())
12012 .fold(0.0f32, f32::max);
12013 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
12014 }
12015
12016 #[test]
12024 fn q4t_matmat_blocked_matches_per_row() {
12025 let (rows, cols, b) = (16usize, 64usize, 9usize);
12026 let gpr = cols / GROUP_SIZE;
12027 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
12028 for r in 0..rows {
12029 for g in 0..gpr {
12030 let t = (r * gpr + g) * Q4_TILE;
12031 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
12032 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
12033 for k in 0..16 {
12034 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
12035 }
12036 }
12037 }
12038 let x: Vec<f32> = (0..b * cols)
12039 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12040 .collect();
12041 let mut y_blk = vec![0f32; b * rows];
12042 let mut y_row = vec![0f32; b * rows];
12043 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
12044 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
12045 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
12046 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
12047 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
12048 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
12049 }
12050
12051 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
12058 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
12059 let gpr = cols / GROUP_SIZE;
12060 let stride = q4tp_code_stride(gpr);
12061 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
12062 let mut b = vec![0u8; codes_off + rows * stride];
12063 for r in 0..rows {
12064 for g in 0..gpr {
12065 let t = (r * gpr + g) * Q4TP_NIB;
12066 for k in 0..16 {
12067 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
12068 }
12069 }
12070 let lo = -6.0 - 0.03 * (r % 17) as f32;
12071 let step = 0.01 + 0.004 * (r % 11) as f32;
12072 let p = params_off + r * 4;
12073 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
12074 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
12075 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
12076 for g in 0..gpr {
12077 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
12078 }
12079 }
12080 b
12081 }
12082
12083 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
12087 let gpr = cols / GROUP_SIZE;
12088 let v = Q4tpView::new(bytes, rows, cols);
12089 let mut out = vec![0u8; rows * gpr * Q4_TILE];
12090 let mut sc = vec![0f32; gpr];
12091 for r in 0..rows {
12092 v.scales_into(r, gpr, &mut sc);
12093 for g in 0..gpr {
12094 let t = (r * gpr + g) * Q4_TILE;
12095 let s = sc[g];
12096 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12097 let src = (r * gpr + g) * Q4TP_NIB;
12098 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
12099 }
12100 }
12101 out
12102 }
12103
12104 #[test]
12110 fn q4tp_exact_path_matches_dequant_reference() {
12111 let (rows, cols) = (256usize, 512usize);
12112 let gpr = cols / GROUP_SIZE;
12113 let bytes = synth_q4tp(rows, cols);
12114 let mut w = vec![0f32; rows * cols];
12115 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
12116
12117 let x: Vec<f32> = (0..cols)
12118 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12119 .collect();
12120 let v = Q4tpView::new(&bytes, rows, cols);
12121 let mut sc = vec![0f32; gpr];
12122 for r in 0..rows {
12123 v.scales_into(r, gpr, &mut sc);
12124 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
12125 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
12126 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
12130 assert!(
12131 (got - want).abs() <= 1e-5 * mag,
12132 "row {r}: kernel {got} vs dequant {want}"
12133 );
12134 }
12135 }
12136
12137 #[test]
12143 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
12144 let (rows, cols) = (256usize, 512usize);
12145 let bytes = synth_q4tp(rows, cols);
12146 let twin = q4tp_as_q4t(&bytes, rows, cols);
12147 let x: Vec<f32> = (0..cols)
12148 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12149 .collect();
12150
12151 let mut got = vec![0f32; rows];
12152 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
12153 let mut want = vec![0f32; rows];
12154 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
12155
12156 let mut w = vec![0f32; rows * cols];
12159 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
12160 for r in 0..rows {
12161 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
12162 assert!(
12163 (got[r] - want[r]).abs() <= 1e-3 * mag,
12164 "row {r}: q4tp {} vs q4t {}",
12165 got[r],
12166 want[r]
12167 );
12168 }
12169 }
12170
12171 #[test]
12176 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
12177 let (rows, cols, b) = (256usize, 512usize, 5usize);
12178 let bytes = synth_q4tp(rows, cols);
12179 let twin = q4tp_as_q4t(&bytes, rows, cols);
12180 let xs: Vec<f32> = (0..b * cols)
12181 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
12182 .collect();
12183
12184 let mut got = vec![0f32; b * rows];
12185 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
12186 let mut want = vec![0f32; b * rows];
12187 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
12188
12189 let mut w = vec![0f32; rows * cols];
12190 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
12191 for t in 0..b {
12192 for r in 0..rows {
12193 let mag: f32 = (0..cols)
12194 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
12195 .sum();
12196 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
12197 assert!(
12198 (g - wa).abs() <= 1e-3 * mag,
12199 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
12200 );
12201 }
12202 }
12203 }
12204
12205 #[test]
12206 fn q4tp_matvec2_matches_the_single_stream_kernel() {
12207 let (rows, cols) = (128usize, 256usize);
12208 let gpr = cols / GROUP_SIZE;
12209 let bytes = synth_q4tp(rows, cols);
12210 let xs: Vec<f32> = (0..2 * cols)
12211 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
12212 .collect();
12213
12214 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
12215 q4tp_matvec2(
12216 &bytes,
12217 &xs[..cols],
12218 &xs[cols..],
12219 rows,
12220 cols,
12221 &mut o1,
12222 &mut o2,
12223 None,
12224 );
12225
12226 let v = Q4tpView::new(&bytes, rows, cols);
12229 let mut sc = vec![0f32; gpr];
12230 for r in 0..rows {
12231 v.scales_into(r, gpr, &mut sc);
12232 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
12233 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
12234 }
12235 }
12236
12237 #[test]
12244 fn q4tp_matvec_keeps_pace_with_q4t() {
12245 let (rows, cols) = (4096usize, 3072usize);
12246 let bytes = synth_q4tp(rows, cols);
12247 let twin = q4tp_as_q4t(&bytes, rows, cols);
12248 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
12249 let mut o = vec![0f32; rows];
12250 let n = 12;
12251 let mut best = (f64::MAX, f64::MAX);
12252 for _ in 0..3 {
12255 let t0 = std::time::Instant::now();
12256 for _ in 0..n {
12257 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
12258 }
12259 best.0 = best.0.min(t0.elapsed().as_secs_f64());
12260 let t0 = std::time::Instant::now();
12261 for _ in 0..n {
12262 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
12263 }
12264 best.1 = best.1.min(t0.elapsed().as_secs_f64());
12265 }
12266 let ratio = best.1 / best.0;
12267 println!(
12268 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
12269 best.0 * 1e3 / n as f64,
12270 best.1 * 1e3 / n as f64
12271 );
12272 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
12273 }
12274
12275 #[cfg(target_os = "macos")]
12276 #[test]
12277 fn q4t_matmat_accel_matches_dequant_reference() {
12278 if !accel_gemm_enabled() {
12279 return; }
12281 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
12283 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
12284 for r in 0..rows {
12285 for g in 0..gpr {
12286 let t = (r * gpr + g) * Q4_TILE;
12287 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
12288 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
12289 for k in 0..16 {
12290 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
12291 }
12292 }
12293 }
12294 let x: Vec<f32> = (0..b * cols)
12295 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12296 .collect();
12297 let mut got = vec![0f32; b * rows];
12298 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
12299 let mut w = vec![0f32; rows * cols];
12301 for r in 0..rows {
12302 for g in 0..gpr {
12303 let t = (r * gpr + g) * Q4_TILE;
12304 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
12305 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
12306 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
12307 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
12308 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
12309 }
12310 }
12311 }
12312 for bi in 0..b {
12313 for r in 0..rows {
12314 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
12315 let d = (got[bi * rows + r] - want).abs();
12316 assert!(
12317 d <= want.abs().max(1.0) * 1e-4,
12318 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
12319 got[bi * rows + r]
12320 );
12321 }
12322 }
12323 }
12324
12325 #[test]
12326 fn q4matvec_matches_full_dequant() {
12327 let (rows, cols) = (8, 64);
12328 let groups = rows * cols / GROUP_SIZE;
12329 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12331 for i in 0..groups * 16 {
12332 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12333 }
12334 for g in 0..groups {
12335 let s = 0.01 + 0.003 * g as f32;
12336 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12337 }
12338 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12339
12340 let mut reference = vec![0.0f32; rows * cols];
12341 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12342 let mut expect = vec![0.0f32; rows];
12343 for r in 0..rows {
12344 expect[r] = reference[r * cols..(r + 1) * cols]
12345 .iter()
12346 .zip(&x)
12347 .map(|(w, xv)| w * xv)
12348 .sum();
12349 }
12350
12351 let mut got = vec![0.0f32; rows];
12352 q4matvec(&bytes, &x, rows, cols, &mut got, None);
12353 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
12357 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12358 for r in 0..rows {
12359 assert!(
12360 (got[r] - expect[r]).abs() < tol * scale,
12361 "row {r}: {} vs {}",
12362 got[r],
12363 expect[r]
12364 );
12365 }
12366 }
12367
12368 #[test]
12371 fn vbitmatvec2_equals_two_singles() {
12372 let (rows, cols) = (6, 64);
12373 let ng = cols / GROUP_SIZE;
12374 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
12375 let mut bytes = bits.clone();
12376 for g in 0..rows * ng {
12377 let s = 0.02 + 0.001 * g as f32;
12378 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12379 }
12380 for r in 0..rows {
12381 let b = bits[r] as usize;
12382 let (mut acc, mut nb) = (0u64, 0usize);
12383 let mut rowbytes = Vec::new();
12384 for i in 0..cols {
12385 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
12386 acc = (acc << b) | v;
12387 nb += b;
12388 while nb >= 8 {
12389 nb -= 8;
12390 rowbytes.push(((acc >> nb) & 0xFF) as u8);
12391 }
12392 }
12393 if nb > 0 {
12394 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
12395 }
12396 bytes.extend_from_slice(&rowbytes);
12397 }
12398 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
12399 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
12400 let offsets = vbit_row_offsets(&bytes, rows, cols);
12401
12402 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12403 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
12404 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
12405 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
12406 vbitmatvec2(
12407 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
12408 );
12409 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
12410 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
12411 }
12412
12413 #[test]
12415 fn q4matvec2_equals_two_singles() {
12416 let (rows, cols) = (8, 128);
12417 let groups = rows * cols / GROUP_SIZE;
12418 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12419 for i in 0..groups * 16 {
12420 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12421 }
12422 for g in 0..groups {
12423 let s = 0.01 + 0.003 * g as f32;
12424 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12425 }
12426 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12429 x1[9] = 250.0;
12430 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12431
12432 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12433 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
12434 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
12435 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
12436 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
12437 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
12438 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
12439 }
12440
12441 #[test]
12444 fn matvec_many_equals_separate_matvecs() {
12445 use crate::pool::Pool;
12446 let (r1, r2, cols) = (300, 200, 64);
12447 let mk = |salt: usize, rows: usize| {
12448 QTensor::from_f32(
12449 (0..rows * cols)
12450 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
12451 .collect(),
12452 rows,
12453 cols,
12454 )
12455 };
12456 let (a, b) = (mk(1, r1), mk(5, r2));
12457 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
12458 let pool = Pool::new(3);
12459
12460 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
12461 a.matvec(&x, &mut ea, Some(&pool));
12462 b.matvec(&x, &mut eb, Some(&pool));
12463 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
12464 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
12465 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
12466 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
12467 }
12468
12469 #[test]
12474 fn q4tp_matvec_many_equals_separate_matvecs() {
12475 use crate::pool::Pool;
12476 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
12477
12478 let (r1, r2, cols) = (300usize, 200usize, 64usize);
12479 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
12480 "arch_name": "tiny-q4tp",
12481 "hidden_size": cols,
12482 "intermediate_size": cols * 2,
12483 "num_layers": 1,
12484 "num_attention_heads": 2,
12485 "num_kv_heads": 1,
12486 "head_dim": 32,
12487 "vocab_size": r1,
12488 "layer_types": ["FullAttention"],
12489 "rms_norm_eps": 1e-6,
12490 "max_position_embeddings": 8,
12491 "linear_conv_kernel_dim": 0,
12492 "linear_num_key_heads": 0,
12493 "linear_num_value_heads": 0
12494 }))
12495 .unwrap();
12496 let header = CmfHeader {
12497 format: "cmf".into(),
12498 version: CMF_VERSION,
12499 arch,
12500 quant_type: QuantType::Q4Block,
12501 provenance: None,
12502 tokenizer_config: None,
12503 section_hashes: None,
12504 skills: Vec::new(),
12505 shard: None,
12506 calibration: None,
12507 routing: None,
12508 genome: None,
12509 lineage: Vec::new(),
12510 router: None,
12511 segments: Vec::new(),
12512 };
12513 let specs = [
12514 TensorSpec {
12515 name: "q".into(),
12516 dtype: TensorDtype::Q4TiledP,
12517 shape: vec![r1, cols],
12518 data: synth_q4tp(r1, cols),
12519 },
12520 TensorSpec {
12521 name: "kv".into(),
12522 dtype: TensorDtype::Q4TiledP,
12523 shape: vec![r2, cols],
12524 data: synth_q4tp(r2, cols),
12525 },
12526 ];
12527 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
12528 std::fs::create_dir_all(&dir).unwrap();
12529 let path = dir.join("m.cmf");
12530 CmfModel::write(&path, &header, &specs, None, None).unwrap();
12531 let model = Arc::new(CmfModel::open(&path).unwrap());
12532 let (a, b) = (
12533 QTensor::from_model(&model, "q").unwrap(),
12534 QTensor::from_model(&model, "kv").unwrap(),
12535 );
12536 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
12537 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
12538 let x: Vec<f32> = (0..cols)
12539 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
12540 .collect();
12541 let pool = Pool::new(3);
12542 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
12543 a.matvec(&x, &mut ea, Some(&pool));
12544 b.matvec(&x, &mut eb, Some(&pool));
12545 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
12546 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
12547 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
12548 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
12549 let _ = std::fs::remove_dir_all(&dir);
12550 }
12551
12552 #[test]
12558 fn multi_token_moe_rows_equal_single_token_decode() {
12559 use crate::pool::Pool;
12560 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
12561
12562 let (h, inter, ne) = (64usize, 128usize, 3usize);
12563 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
12564 "arch_name": "tiny-q4tp-moe",
12565 "hidden_size": h,
12566 "intermediate_size": inter,
12567 "num_layers": 1,
12568 "num_attention_heads": 2,
12569 "num_kv_heads": 1,
12570 "head_dim": 32,
12571 "vocab_size": 8,
12572 "layer_types": ["FullAttention"],
12573 "rms_norm_eps": 1e-6,
12574 "max_position_embeddings": 8,
12575 "linear_conv_kernel_dim": 0,
12576 "linear_num_key_heads": 0,
12577 "linear_num_value_heads": 0
12578 }))
12579 .unwrap();
12580 let header = CmfHeader {
12581 format: "cmf".into(),
12582 version: CMF_VERSION,
12583 arch,
12584 quant_type: QuantType::Q4Block,
12585 provenance: None,
12586 tokenizer_config: None,
12587 section_hashes: None,
12588 skills: Vec::new(),
12589 shard: None,
12590 calibration: None,
12591 routing: None,
12592 genome: None,
12593 lineage: Vec::new(),
12594 router: None,
12595 segments: Vec::new(),
12596 };
12597 let mut specs = Vec::new();
12598 for e in 0..ne {
12599 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
12600 .into_iter()
12601 .enumerate()
12602 {
12603 let mut data = synth_q4tp(r, c);
12606 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
12607 .iter_mut()
12608 .enumerate()
12609 {
12610 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
12611 }
12612 specs.push(TensorSpec {
12613 name: format!("{n}{e}"),
12614 dtype: TensorDtype::Q4TiledP,
12615 shape: vec![r, c],
12616 data,
12617 });
12618 }
12619 }
12620 let dir = std::env::temp_dir().join(format!(
12621 "cmf-moe-rows-{}-{}",
12622 std::process::id(),
12623 FLOAT_ACTIVATIONS.get()
12624 ));
12625 std::fs::create_dir_all(&dir).unwrap();
12626 let path = dir.join("m.cmf");
12627 CmfModel::write(&path, &header, &specs, None, None).unwrap();
12628 let model = Arc::new(CmfModel::open(&path).unwrap());
12629 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
12630 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
12631 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
12632 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
12633 let b = 4usize;
12634 let mut xs: Vec<f32> = (0..b * h)
12635 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
12636 .collect();
12637 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
12640 (vec![2, 0], vec![0.6, 0.4]),
12641 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
12642 (vec![1], vec![1.0]),
12643 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
12644 ];
12645 let pool = Pool::new(3);
12646 let mut want = vec![0f32; b * h];
12648 for (tk, (idx, w)) in routes.iter().enumerate() {
12649 let x = &xs[tk * h..(tk + 1) * h];
12650 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
12651 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
12652 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
12653 if FLOAT_ACTIVATIONS.get() {
12654 for (slot, &e) in idx.iter().enumerate() {
12655 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
12656 g[e].matvec(x, &mut gate, Some(&pool));
12657 u[e].matvec(x, &mut up, Some(&pool));
12658 for (v, u) in gate.iter_mut().zip(up) {
12659 *v = (*v / (1.0 + (-*v).exp())) * u;
12660 }
12661 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
12662 }
12663 }
12664 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
12665 assert!(QTensor::moe_down_many(
12666 &downs,
12667 &gs,
12668 w,
12669 &mut want[tk * h..(tk + 1) * h],
12670 Some(&pool)
12671 ));
12672 }
12673 if FLOAT_ACTIVATIONS.get() {
12674 for (tk, (idx, w)) in routes.iter().enumerate() {
12675 let mut scalar = vec![0.0; h];
12676 for (&e, &weight) in idx.iter().zip(w) {
12677 let (mut gate, mut up, mut down) =
12678 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
12679 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
12680 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
12681 for (v, u) in gate.iter_mut().zip(up) {
12682 *v = (*v / (1.0 + (-*v).exp())) * u;
12683 }
12684 d[e].matvec(&gate, &mut down, Some(&pool));
12685 for (v, d) in scalar.iter_mut().zip(down) {
12686 *v += weight * d;
12687 }
12688 }
12689 assert_eq!(
12690 &want[tk * h..(tk + 1) * h],
12691 scalar,
12692 "float many equals scalar experts"
12693 );
12694 }
12695 }
12696 let mut experts: Vec<usize> = Vec::new();
12698 let mut groups: Vec<Vec<usize>> = Vec::new();
12699 for (tk, (idx, _)) in routes.iter().enumerate() {
12700 for &e in idx {
12701 match experts.iter().position(|&x| x == e) {
12702 Some(k) => groups[k].push(tk),
12703 None => {
12704 experts.push(e);
12705 groups.push(vec![tk]);
12706 }
12707 }
12708 }
12709 }
12710 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
12711 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
12712 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
12713 assert!(QTensor::moe_gate_up_rows(
12714 &pairs,
12715 &groups,
12716 &xs,
12717 &mut gs,
12718 Some(&pool)
12719 ));
12720 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
12721 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
12722 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
12723 assert!(QTensor::moe_down_rows(
12724 &downs,
12725 &lens,
12726 &gs,
12727 &mut ds,
12728 Some(&pool)
12729 ));
12730 let slot = |tk: usize, e: usize| {
12731 let k = experts.iter().position(|&x| x == e).unwrap();
12732 groups[..k].iter().map(|g| g.len()).sum::<usize>()
12733 + groups[k].iter().position(|&x| x == tk).unwrap()
12734 };
12735 let mut got = vec![0f32; b * h];
12736 for (tk, (idx, w)) in routes.iter().enumerate() {
12737 for i in 0..h {
12738 let mut acc = 0f32;
12739 for (&e, &we) in idx.iter().zip(w) {
12740 acc += we * ds[slot(tk, e)][i];
12741 }
12742 got[tk * h + i] = acc;
12743 }
12744 }
12745 assert!(want.iter().any(|v| *v != 0.0));
12746 assert_eq!(
12747 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12748 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12749 "multi-token MoE must equal decode bit for bit"
12750 );
12751
12752 let b5 = 5usize;
12755 let x5: Vec<f32> = (0..b5 * h)
12756 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
12757 .collect();
12758 let mut mm = vec![0f32; b5 * inter];
12759 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
12760 for tk in 0..b5 {
12761 let mut mv = vec![0f32; inter];
12762 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
12763 assert_eq!(
12764 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12765 mm[tk * inter..(tk + 1) * inter]
12766 .iter()
12767 .map(|v| v.to_bits())
12768 .collect::<Vec<_>>(),
12769 "row-exact matmat token {tk}"
12770 );
12771 }
12772 let _ = std::fs::remove_dir_all(&dir);
12775 }
12776
12777 #[test]
12785 fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
12786 use crate::pool::Pool;
12787 use std::sync::atomic::Ordering::Relaxed;
12788 let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
12789 Q4TP_ALT.store(2, Relaxed);
12792 let pool = Pool::new(3);
12793 for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
12797 let bytes = synth_q4tp(rows, cols);
12798 let mut xs: Vec<f32> = (0..b * cols)
12799 .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
12800 .collect();
12801 xs[3] = 7.5; let run = |exact: bool| {
12803 let mut out = vec![0f32; b * rows];
12804 q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
12805 out
12806 };
12807 let (fast, exact) = (run(false), run(true));
12808 #[cfg(not(target_arch = "aarch64"))]
12809 let _ = fast;
12810 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
12811 let mut matvecs = vec![0f32; b * rows];
12812 for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
12813 q4tp_matvec(
12814 &bytes,
12815 &xs[bi * cols..(bi + 1) * cols],
12816 rows,
12817 cols,
12818 o,
12819 Some(&pool),
12820 );
12821 }
12822 assert!(matvecs.iter().any(|v| *v != 0.0));
12823 assert_eq!(
12824 bits(&exact),
12825 bits(&matvecs),
12826 "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
12827 );
12828 #[cfg(target_arch = "aarch64")]
12829 {
12830 let gpr = cols / GROUP_SIZE;
12832 let v = Q4tpView::new(&bytes, rows, cols);
12833 let mut old = vec![0f32; b * rows];
12834 let mut sc = vec![0f32; gpr];
12835 let a8w8 = a8w8_enabled();
12836 let blocked = sdot_enabled() && blocked_enabled();
12837 let acts: Vec<SplitAct> = (0..b)
12838 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
12839 .collect();
12840 for r in 0..rows {
12841 v.scales_into(r, gpr, &mut sc);
12842 if !a8w8 {
12843 for bi in 0..b {
12844 let x = &xs[bi * cols..(bi + 1) * cols];
12845 old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
12846 }
12847 continue;
12848 }
12849 let finish = |d: f32, act: &SplitAct| {
12850 let mut acc = d * act.sx;
12851 for &(j, xv) in &act.outliers {
12852 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
12853 acc += w * s * xv;
12854 }
12855 acc
12856 };
12857 let mut bi = 0usize;
12858 while blocked && bi + 4 <= b {
12859 let xs4 = [
12860 acts[bi].xq.as_slice(),
12861 acts[bi + 1].xq.as_slice(),
12862 acts[bi + 2].xq.as_slice(),
12863 acts[bi + 3].xq.as_slice(),
12864 ];
12865 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
12866 for k in 0..4 {
12867 old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
12868 }
12869 bi += 4;
12870 }
12871 for (bi, act) in acts.iter().enumerate().skip(bi) {
12872 let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
12873 old[bi * rows + r] = finish(d, act);
12874 }
12875 }
12876 assert_eq!(
12877 bits(&fast),
12878 bits(&old),
12879 "{rows}x{cols} b={b}: the fast path changed outside row_exact"
12880 );
12881 if blocked {
12884 assert_ne!(
12885 bits(&fast),
12886 bits(&matvecs),
12887 "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
12888 );
12889 }
12890 }
12891 }
12892 Q4TP_ALT.store(0, Relaxed);
12893 }
12894
12895 #[test]
12896 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
12897 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
12898 let active = AtomicUsize::new(0);
12901 counted_row_exact_scope(&active, || {
12902 assert_eq!(active.load(Ordering::Acquire), 1);
12903 counted_row_exact_scope(&active, || {
12904 assert_eq!(active.load(Ordering::Acquire), 2);
12905 });
12906 assert_eq!(active.load(Ordering::Acquire), 1);
12907 });
12908 assert_eq!(active.load(Ordering::Acquire), 0);
12909
12910 let both_entered = Barrier::new(2);
12911 let release_last = Barrier::new(2);
12912 std::thread::scope(|s| {
12913 let first = s.spawn(|| counted_row_exact_scope(&active, || {
12914 both_entered.wait();
12915 }));
12916 let last = s.spawn(|| counted_row_exact_scope(&active, || {
12917 both_entered.wait();
12918 release_last.wait();
12919 }));
12920 first.join().unwrap();
12921 let after_first = active.load(Ordering::Acquire);
12922 release_last.wait();
12923 last.join().unwrap();
12924 assert_eq!(after_first, 1, "second request must remain exact");
12925 });
12926 assert_eq!(active.load(Ordering::Acquire), 0);
12927 let panic = std::panic::catch_unwind(|| {
12928 counted_row_exact_scope(&active, || panic!("scope unwind"));
12929 });
12930 assert!(panic.is_err());
12931 assert_eq!(active.load(Ordering::Acquire), 0);
12932 }
12933
12934 #[test]
12935 #[cfg(target_arch = "x86_64")]
12936 fn q4tp_float_avx2_is_bitwise_scalar() {
12937 if !avx2_enabled() {
12938 return;
12939 }
12940 for cols in [32, 64, 96, 2048, 4096] {
12941 let rows = 9;
12942 let bytes = synth_q4tp(rows, cols);
12943 let v = Q4tpView::new(&bytes, rows, cols);
12944 let gpr = cols / GROUP_SIZE;
12945 let mut sc = vec![0.0; gpr];
12946 for seed in 1..=5 {
12947 let xs: Vec<f32> = (0..cols)
12948 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
12949 .collect();
12950 for r in 0..rows {
12951 v.scales_into(r, gpr, &mut sc);
12952 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
12953 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
12954 assert_eq!(
12955 scalar.to_bits(),
12956 vector.to_bits(),
12957 "cols={cols} row={r} seed={seed}"
12958 );
12959 }
12960 }
12961 }
12962 }
12963
12964 #[test]
12965 fn multi_token_moe_rows_float_equal_single_token_decode() {
12966 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
12967 }
12968
12969 #[test]
12970 fn full_gpu_q8_scope_is_nested_and_thread_local() {
12971 assert!(!FULL_GPU_Q8.get());
12972 let before = gpu_split_frac();
12973 {
12974 let _guard = enter_full_gpu_q8_scope();
12975 assert_eq!(gpu_split_frac(), 1.0);
12976 {
12977 let _nested = enter_full_gpu_q8_scope();
12978 }
12979 assert_eq!(gpu_split_frac(), 1.0);
12980 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
12981 }
12982 assert!(!FULL_GPU_Q8.get());
12983 assert_eq!(gpu_split_frac(), before);
12984 }
12985
12986 #[test]
12987 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
12988 assert!(!FLOAT_ACTIVATIONS.get());
12989 let before = a8w8_enabled();
12990 float_activations_scope(|| {
12991 assert!(!a8w8_enabled());
12992 float_activations_scope(|| assert!(!a8w8_enabled()));
12993 assert!(FLOAT_ACTIVATIONS.get());
12994 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
12995 .join()
12996 .unwrap();
12997 });
12998 assert!(!FLOAT_ACTIVATIONS.get());
12999 assert_eq!(a8w8_enabled(), before);
13000 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
13001 assert!(!FLOAT_ACTIVATIONS.get());
13002 }
13003
13004 #[test]
13007 fn batched_matmat_equals_per_position_matvec() {
13008 let (rows, cols, b) = (8, 64, 5);
13009 let groups = rows * cols / GROUP_SIZE;
13011 let mut q4 = Vec::new();
13012 for i in 0..groups * 16 {
13013 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13014 }
13015 for g in 0..groups {
13016 q4.extend_from_slice(
13017 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
13018 );
13019 }
13020 let ng = cols / GROUP_SIZE;
13022 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
13023 let mut vb = bits.clone();
13024 for g in 0..rows * ng {
13025 vb.extend_from_slice(
13026 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
13027 );
13028 }
13029 for r in 0..rows {
13030 let bw = bits[r] as usize;
13031 let (mut acc, mut nb) = (0u64, 0usize);
13032 let mut rowbytes = Vec::new();
13033 for i in 0..cols {
13034 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13035 acc = (acc << bw) | v;
13036 nb += bw;
13037 while nb >= 8 {
13038 nb -= 8;
13039 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13040 }
13041 }
13042 if nb > 0 {
13043 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13044 }
13045 vb.extend_from_slice(&rowbytes);
13046 }
13047 let offsets = vbit_row_offsets(&vb, rows, cols);
13048
13049 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
13050
13051 let mut got = vec![0f32; b * rows];
13053 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
13054 for bi in 0..b {
13055 let mut expect = vec![0f32; rows];
13056 q4matvec(
13057 &q4,
13058 &xs[bi * cols..(bi + 1) * cols],
13059 rows,
13060 cols,
13061 &mut expect,
13062 None,
13063 );
13064 assert_eq!(
13065 &got[bi * rows..(bi + 1) * rows],
13066 &expect[..],
13067 "q4 batch pos {bi}"
13068 );
13069 }
13070
13071 let mut got = vec![0f32; b * rows];
13073 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
13074 for bi in 0..b {
13075 let mut expect = vec![0f32; rows];
13076 vbitmatvec(
13077 &vb,
13078 &offsets,
13079 &xs[bi * cols..(bi + 1) * cols],
13080 rows,
13081 cols,
13082 &mut expect,
13083 None,
13084 );
13085 assert_eq!(
13086 &got[bi * rows..(bi + 1) * rows],
13087 &expect[..],
13088 "vbit batch pos {bi}"
13089 );
13090 }
13091 }
13092
13093 #[test]
13097 fn q4_tiled_matches_q4_block_bitexact() {
13098 let (rows, cols, b) = (8usize, 128usize, 3usize);
13099 let groups = rows * cols / GROUP_SIZE;
13100 let mut split = Vec::with_capacity(groups * 18);
13101 for i in 0..groups * 16 {
13102 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13103 }
13104 for g in 0..groups {
13105 split.extend_from_slice(
13106 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
13107 );
13108 }
13109 let (packed, scales) = split.split_at(groups * 16);
13111 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
13112 for g in 0..groups {
13113 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
13114 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
13115 }
13116
13117 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13118 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
13120
13121 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
13122 q4matvec(&split, &x1, rows, cols, &mut a, None);
13123 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
13124 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
13125
13126 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13127 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
13128 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
13129 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
13130 assert_eq!(a1, t1);
13131 assert_eq!(a2, t2);
13132
13133 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
13134 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
13135 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
13136 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
13137 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
13138 }
13139
13140 #[test]
13147 fn q4matvec_sdot_outlier_exact() {
13148 let (rows, cols) = (4, 128);
13149 let groups = rows * cols / GROUP_SIZE;
13150 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13151 for i in 0..groups * 16 {
13152 bytes.push(((i * 11 + 5) % 256) as u8);
13153 }
13154 for g in 0..groups {
13155 let s = 0.02 + 0.002 * g as f32;
13156 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13157 }
13158 let mut x: Vec<f32> = (0..cols)
13159 .map(|i| match i % 3 {
13160 0 => 1.0,
13161 1 => -1.0,
13162 _ => 0.0,
13163 })
13164 .collect();
13165 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
13168 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13169 let mut expect = vec![0.0f32; rows];
13170 for r in 0..rows {
13171 expect[r] = reference[r * cols..(r + 1) * cols]
13172 .iter()
13173 .zip(&x)
13174 .map(|(w, xv)| w * xv)
13175 .sum();
13176 }
13177 let mut got = vec![0.0f32; rows];
13178 q4matvec(&bytes, &x, rows, cols, &mut got, None);
13179 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13180 for r in 0..rows {
13181 assert!(
13182 (got[r] - expect[r]).abs() < 2e-3 * scale,
13183 "row {r}: {} vs {} (outlier term must be exact)",
13184 got[r],
13185 expect[r]
13186 );
13187 }
13188 }
13189
13190 #[test]
13194 fn q1t_matvec_matches_reference() {
13195 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
13196 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
13198 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
13199 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
13201 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
13202 let mut bytes = Vec::new();
13203 for r in 0..rows {
13204 for g in 0..gpr {
13205 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
13206 let mut c = [0u8; 7];
13207 for k in 0..GROUP_SIZE {
13208 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
13210 0
13211 } else {
13212 ((k + r * 3 + g) % 3) as u8 };
13214 cortiq_core::quant::q1t_pack(&mut c, k, code);
13215 }
13216 bytes.extend_from_slice(&c);
13217 }
13218 }
13219 let mut row_ptr = vec![0u32; rows + 1];
13222 for &(idx, _) in &outliers {
13223 row_ptr[idx as usize / cols + 1] += 1;
13224 }
13225 for r in 0..rows {
13226 row_ptr[r + 1] += row_ptr[r];
13227 }
13228 for &p in &row_ptr {
13229 bytes.extend_from_slice(&p.to_le_bytes());
13230 }
13231 for &(idx, v) in &outliers {
13232 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
13233 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
13234 }
13235
13236 let mut refw = vec![0f32; rows * cols];
13237 dequant_q1t(&bytes, rows, cols, &mut refw);
13238 let x: Vec<f32> = (0..cols)
13241 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
13242 .collect();
13243 let mut expect = vec![0f32; rows];
13244 for r in 0..rows {
13245 let mut a = 0.0f32;
13246 for j in 0..cols {
13247 a += refw[r * cols + j] * x[j];
13248 }
13249 expect[r] = a;
13250 }
13251 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
13252 let mut got = vec![0f32; rows];
13253 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
13254 for r in 0..rows {
13255 assert!(
13256 (got[r] - expect[r]).abs() < tol(expect[r]),
13257 "row {r}: {} vs {}",
13258 got[r],
13259 expect[r]
13260 );
13261 }
13262 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
13264 let mut gm = vec![0f32; 2 * rows];
13265 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
13266 for r in 0..rows {
13267 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
13268 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
13269 }
13270 let xb: Vec<f32> = (0..cols)
13274 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
13275 .collect();
13276 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
13277 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
13278 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
13279 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
13280 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
13281 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
13282 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
13283 }
13284
13285 #[test]
13288 fn q1t_matvec2_odd_gpr_matches_singles() {
13289 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
13290 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
13292 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
13293 for r in 0..rows {
13294 for g in 0..gpr {
13295 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
13296 let mut c = [0u8; 7];
13297 for k in 0..GROUP_SIZE {
13298 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
13299 }
13300 bytes.extend_from_slice(&c);
13301 }
13302 }
13303 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
13304 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
13305 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
13306 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
13307 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
13308 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
13309 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
13310 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
13311 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
13312 }
13313
13314 #[test]
13318 #[ignore]
13319 fn q1t_matvec2_speed() {
13320 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
13321 use std::time::Instant;
13322 let (rows, cols) = (8192usize, 4096usize);
13323 let gpr = cols / GROUP_SIZE;
13324 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
13325 for r in 0..rows {
13326 for g in 0..gpr {
13327 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
13328 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
13329 let mut c = [0u8; 7];
13330 for k in 0..GROUP_SIZE {
13331 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
13332 }
13333 bytes.extend_from_slice(&c);
13334 }
13335 }
13336 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
13337 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
13338 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
13339 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
13340 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
13342 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
13343 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
13344 for _ in 0..8 {
13345 let t0 = Instant::now();
13346 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
13347 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
13348 let t1 = Instant::now();
13349 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
13350 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
13351 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
13352 }
13353 assert_eq!(p1, s1);
13354 assert_eq!(p2, s2);
13355 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
13356 }
13357
13358 #[test]
13362 #[ignore]
13363 fn q1t_matvec_speed() {
13364 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
13365 use std::time::Instant;
13366 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
13368 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
13369 for r in 0..rows {
13370 for g in 0..gpr {
13371 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
13372 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
13373 let mut c = [0u8; 7];
13374 for k in 0..GROUP_SIZE {
13375 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
13376 }
13377 bytes.extend_from_slice(&c);
13378 }
13379 }
13380 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
13382 let mut idx = 0usize;
13383 while idx < n {
13384 row_ptr[idx / cols + 1] += 1;
13385 idx += stride;
13386 }
13387 for r in 0..rows {
13388 row_ptr[r + 1] += row_ptr[r];
13389 }
13390 for &p in &row_ptr {
13391 bytes.extend_from_slice(&p.to_le_bytes());
13392 }
13393 let mut idx = 0usize;
13394 while idx < n {
13395 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
13396 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
13397 idx += stride;
13398 }
13399 let x: Vec<f32> = (0..cols)
13402 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
13403 .collect();
13404 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
13405
13406 let slow = |out: &mut [f32]| {
13408 let mut buf = vec![0f32; cols];
13409 for r in 0..rows {
13410 for g in 0..gpr {
13411 let off = (r * gpr + g) * Q1T_TILE;
13412 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
13413 let codes = &bytes[off + 2..off + Q1T_TILE];
13414 for k in 0..GROUP_SIZE {
13415 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
13416 1 => s,
13417 2 => -s,
13418 _ => 0.0,
13419 };
13420 }
13421 }
13422 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
13423 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
13424 }
13425 };
13426 let iters = 5;
13427 let mut a = vec![0f32; rows];
13428 slow(&mut a); let t = Instant::now();
13430 for _ in 0..iters {
13431 slow(&mut a);
13432 }
13433 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
13434
13435 let mut b = vec![0f32; rows];
13436 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
13438 for _ in 0..iters {
13439 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
13440 }
13441 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
13442
13443 for r in 0..rows {
13444 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
13445 }
13446 println!(
13447 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
13448 slow_ms / fast_ms
13449 );
13450 }
13451}
13452
13453#[cfg(test)]
13454mod gemm_bench {
13455 #[test]
13466 #[ignore]
13467 fn q4tp_matmat_throughput() {
13468 let _alt = super::Q4TP_ALT_TEST_LOCK
13469 .lock()
13470 .unwrap_or_else(|e| e.into_inner());
13471 let b: usize = std::env::var("CMF_BENCH_B")
13475 .ok()
13476 .and_then(|v| v.parse().ok())
13477 .unwrap_or(296);
13478 let (rows, cols) = (9216usize, 2304usize);
13479 let (_, _, _) = (rows, cols, b);
13480 let total =
13481 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
13482 .unwrap();
13483 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13488 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13489 let lo = cortiq_core::quant::f32_to_f16(-4.0);
13490 let step = cortiq_core::quant::f32_to_f16(0.1);
13491 for r in 0..rows {
13492 let o = params_off + r * 4;
13493 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13494 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13495 }
13496 let _ = codes_off;
13497 let xs: Vec<f32> = (0..b * cols)
13498 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13499 .collect();
13500 let mut out = vec![0f32; b * rows];
13501 let pool = crate::pool::Pool::from_env();
13502 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13508 let reps: usize = std::env::var("CMF_BENCH_REPS")
13509 .ok()
13510 .and_then(|v| v.parse().ok())
13511 .unwrap_or(10);
13512 let mut best = [f64::MAX; 2];
13513 let mut sums = [0f32; 2];
13514 for _ in 0..reps {
13515 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
13516 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
13517 let t = std::time::Instant::now();
13518 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13519 best[k] = best[k].min(t.elapsed().as_secs_f64());
13520 sums[k] = out.iter().take(64).sum::<f32>();
13521 }
13522 }
13523 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
13524 for (k, name) in ["previous", "tuned "].iter().enumerate() {
13525 println!(
13526 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
13527 best[k] * 1e3,
13528 flops / best[k] / 1e9,
13529 sums[k]
13530 );
13531 }
13532 assert!(
13533 (sums[0] - sums[1]).abs() < 1e-2,
13534 "the tuned kernel changed the result: {} vs {}",
13535 sums[0],
13536 sums[1]
13537 );
13538 }
13539
13540 #[test]
13546 fn q4tp_matmat_blocked_matches_scalar() {
13547 use std::sync::atomic::Ordering::Relaxed;
13548 let _alt = super::Q4TP_ALT_TEST_LOCK
13549 .lock()
13550 .unwrap_or_else(|e| e.into_inner());
13551 for &(rows, cols, b) in &[
13559 (64usize, 128usize, 7usize),
13560 (33, 96, 4),
13561 (16, 256, 9),
13562 (192, 2304, 37),
13563 ] {
13564 let total = cortiq_core::quant::expected_nbytes(
13565 cortiq_core::TensorDtype::Q4TiledP,
13566 &[rows, cols],
13567 )
13568 .unwrap();
13569 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13570 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
13571 let lo = cortiq_core::quant::f32_to_f16(-4.0);
13572 let step = cortiq_core::quant::f32_to_f16(0.1);
13573 for r in 0..rows {
13574 let o = params_off + r * 4;
13575 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13576 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13577 }
13578 let xs: Vec<f32> = (0..b * cols)
13579 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
13580 .collect();
13581 let mut got = vec![0f32; b * rows];
13582 let mut want = vec![0f32; b * rows];
13583 let gpr = cols / 32;
13584 let view = super::Q4tpView::new(&bytes, rows, cols);
13585 let pool = crate::pool::Pool::from_env();
13586 super::Q4TP_ALT.store(2, Relaxed);
13587 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
13588 super::Q4TP_ALT.store(1, Relaxed);
13589 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
13590 super::Q4TP_ALT.store(0, Relaxed);
13591 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13598 let (mut worst, mut at) = (0f32, 0usize);
13599 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
13600 if (g - w).abs() > worst {
13601 worst = (g - w).abs();
13602 at = i;
13603 }
13604 }
13605 assert!(
13606 worst <= 1e-4 * scale,
13607 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
13608 (scale {scale:.3e}) at cell {at}: {} vs {}",
13609 got[at],
13610 want[at]
13611 );
13612
13613 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
13620 for bi in 0..b {
13621 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
13622 for r in 0..rows {
13623 let mut sc = vec![0f32; gpr];
13624 view.scales_into(r, gpr, &mut sc);
13625 let mut exact = 0f64;
13626 for j in 0..cols {
13627 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
13628 exact += w as f64 * sq as f64 * act.xq[j] as f64;
13629 }
13630 exact *= act.sx as f64;
13631 for &(j, xv) in &act.outliers {
13632 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
13633 exact += w as f64 * sq as f64 * xv as f64;
13634 }
13635 let i = bi * rows + r;
13636 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
13637 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
13638 }
13639 }
13640 println!(
13641 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
13642 per-column {e_scalar:.3e}"
13643 );
13644 assert!(
13649 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
13650 "{rows}x{cols} b={b}: error against f64 too large — blocked \
13651 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
13652 );
13653 }
13654 }
13655
13656 #[test]
13662 #[ignore]
13663 fn q4tp_matmat_row_exact_cost() {
13664 let _alt = super::Q4TP_ALT_TEST_LOCK
13665 .lock()
13666 .unwrap_or_else(|e| e.into_inner());
13667 let pool = crate::pool::Pool::from_env();
13668 let reps: usize = std::env::var("CMF_BENCH_REPS")
13669 .ok()
13670 .and_then(|v| v.parse().ok())
13671 .unwrap_or(30);
13672 for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
13673 let total = cortiq_core::quant::expected_nbytes(
13674 cortiq_core::TensorDtype::Q4TiledP,
13675 &[rows, cols],
13676 )
13677 .unwrap();
13678 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13679 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13680 let lo = cortiq_core::quant::f32_to_f16(-4.0);
13681 let step = cortiq_core::quant::f32_to_f16(0.1);
13682 for r in 0..rows {
13683 let o = params_off + r * 4;
13684 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13685 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13686 }
13687 for &b in &[2usize, 4, 5, 8] {
13688 let xs: Vec<f32> = (0..b * cols)
13689 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13690 .collect();
13691 let mut out = vec![0f32; b * rows];
13692 let mut best = [f64::MAX; 3];
13693 for _ in 0..reps {
13694 for (k, best_k) in best.iter_mut().enumerate() {
13695 let t = std::time::Instant::now();
13696 match k {
13697 0 | 1 => super::q4tp_matmat_with(
13698 &bytes,
13699 &xs,
13700 b,
13701 rows,
13702 cols,
13703 &mut out,
13704 pool.as_deref(),
13705 k == 1,
13706 ),
13707 _ => {
13708 for (bi, o) in out.chunks_mut(rows).enumerate() {
13709 super::q4tp_matvec(
13710 &bytes,
13711 &xs[bi * cols..(bi + 1) * cols],
13712 rows,
13713 cols,
13714 o,
13715 pool.as_deref(),
13716 );
13717 }
13718 }
13719 }
13720 *best_k = best_k.min(t.elapsed().as_secs_f64());
13721 }
13722 }
13723 println!(
13724 "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
13725 {b} matvecs {:.3} ms",
13726 best[0] * 1e3,
13727 best[1] * 1e3,
13728 best[2] * 1e3
13729 );
13730 }
13731 }
13732 }
13733
13734 #[test]
13738 #[ignore]
13739 fn q4t_matmat_throughput() {
13740 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
13741 let total =
13742 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
13743 .unwrap();
13744 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13747 let sc = cortiq_core::quant::f32_to_f16(0.02);
13748 for t in bytes.chunks_mut(super::Q4_TILE) {
13749 t[..2].copy_from_slice(&sc.to_le_bytes());
13750 }
13751 let xs: Vec<f32> = (0..b * cols)
13752 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13753 .collect();
13754 let mut out = vec![0f32; b * rows];
13755 let pool = crate::pool::Pool::from_env();
13756 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13757 let reps: usize = std::env::var("CMF_BENCH_REPS")
13758 .ok()
13759 .and_then(|v| v.parse().ok())
13760 .unwrap_or(10);
13761 let mut best = f64::MAX;
13762 for _ in 0..reps {
13763 let t = std::time::Instant::now();
13764 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13765 best = best.min(t.elapsed().as_secs_f64());
13766 }
13767 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
13768 println!(
13769 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
13770 best * 1e3,
13771 flops / best / 1e9,
13772 out.iter().take(64).sum::<f32>()
13773 );
13774 }
13775}