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;
22use std::sync::atomic::{AtomicU64, Ordering};
23
24mod mxfp4;
26
27pub static MATMAT_GATE: AtomicU64 = AtomicU64::new(0);
30pub static MATMAT_GPU_HIT: AtomicU64 = AtomicU64::new(0);
31
32pub enum QTensor {
33 F32 {
34 data: Vec<f32>,
35 rows: usize,
36 cols: usize,
37 },
38 Mapped {
39 model: Arc<CmfModel>,
40 idx: usize,
42 dtype: TensorDtype,
43 rows: usize,
44 cols: usize,
45 row_scale: Vec<f32>,
47 col_field: Vec<f32>,
49 vbit_offsets: Vec<usize>,
53 repack: Vec<u8>,
61 },
62 Blr(Box<crate::blr::BlrTensor>),
68}
69
70fn repack_enabled() -> bool {
77 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
78 *ON.get_or_init(|| {
79 std::env::var("CMF_REPACK")
80 .map(|v| v == "1")
81 .unwrap_or(cfg!(target_os = "android"))
82 })
83}
84
85fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
89 #[cfg(target_arch = "aarch64")]
90 let arch_ok = sdot_enabled();
91 #[cfg(not(target_arch = "aarch64"))]
92 let arch_ok = false;
93 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
94 return Vec::new();
95 }
96 q8_repack_layout(bytes, rows, cols)
97}
98
99fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
102 let groups = rows / 4;
103 let mut rep = vec![0u8; groups * 4 * cols];
104 for g in 0..groups {
105 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
106 for c in 0..cols / 16 {
107 for lane in 0..4 {
108 let src = (g * 4 + lane) * cols + c * 16;
109 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
110 .copy_from_slice(&bytes[src..src + 16]);
111 }
112 }
113 }
114 rep
115}
116
117fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
120 let ng = cols / GROUP_SIZE;
121 let bits = &bytes[..rows];
122 let mut offsets = Vec::with_capacity(rows + 1);
123 let mut off = rows + rows * ng * 2;
124 for r in 0..rows {
125 offsets.push(off);
126 off += (cols * bits[r] as usize).div_ceil(8);
127 }
128 offsets.push(off);
129 offsets
130}
131
132fn blocked_enabled() -> bool {
138 use std::sync::atomic::Ordering::Relaxed;
139 match BLOCKED_OVERRIDE.load(Relaxed) {
140 1 => false,
141 2 => true,
142 _ => {
143 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
144 *ON.get_or_init(|| {
145 std::env::var("CMF_X86_BLOCKED")
146 .map(|v| v != "0")
147 .unwrap_or(true)
148 })
149 }
150 }
151}
152
153static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
154
155pub fn set_blocked_override(on: Option<bool>) {
162 let v = match on {
163 None => 0,
164 Some(false) => 1,
165 Some(true) => 2,
166 };
167 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
168}
169
170fn gpu_lmhead_enabled() -> bool {
171 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
172 *ON.get_or_init(|| {
173 std::env::var("CMF_GPU_LMHEAD")
174 .map(|v| v != "0")
175 .unwrap_or(true)
176 })
177}
178
179fn gpu_split_frac() -> f32 {
180 if FULL_GPU_Q8.get() {
183 return 1.0;
184 }
185 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
186 *F.get_or_init(|| {
187 std::env::var("CMF_GPU_SPLIT")
188 .ok()
189 .and_then(|v| v.parse::<f32>().ok())
190 .unwrap_or(0.5)
191 .clamp(0.0, 1.0)
192 })
193}
194
195impl QTensor {
196 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
197 debug_assert_eq!(data.len(), rows * cols);
198 Self::F32 { data, rows, cols }
199 }
200
201 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
204 let idx = model
207 .tensor_index(name)
208 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
209 let entry = &model.tensors[idx];
210 if entry.shape.len() != 2 {
211 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
212 }
213 let (rows, cols) = (entry.shape[0], entry.shape[1]);
214 let bytes = model.entry_bytes(entry);
215
216 match entry.dtype {
217 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
218 let n = rows * cols;
219 let scales_off = n;
220 let row_scale: Vec<f32> = (0..rows)
221 .map(|o| {
222 f16_to_f32(u16::from_le_bytes([
223 bytes[scales_off + o * 2],
224 bytes[scales_off + o * 2 + 1],
225 ]))
226 })
227 .collect();
228 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
229 let col_off = n + rows * 2;
230 (0..cols)
231 .map(|i| {
232 f16_to_f32(u16::from_le_bytes([
233 bytes[col_off + i * 2],
234 bytes[col_off + i * 2 + 1],
235 ]))
236 })
237 .collect()
238 } else {
239 Vec::new()
240 };
241 Ok(Self::Mapped {
242 model: model.clone(),
243 idx,
244 dtype: entry.dtype,
245 rows,
246 cols,
247 row_scale,
248 col_field,
249 vbit_offsets: Vec::new(),
250 repack: q8_repack(bytes, rows, cols),
251 })
252 }
253 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
255 model: model.clone(),
256 idx,
257 dtype: entry.dtype,
258 rows,
259 cols,
260 row_scale: Vec::new(),
261 col_field: Vec::new(),
262 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
263 repack: Vec::new(),
264 }),
265 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
269 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
270 let offsets: Vec<usize> = (0..=rows)
271 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
272 .collect();
273 Ok(Self::Mapped {
274 model: model.clone(),
275 idx,
276 dtype: entry.dtype,
277 rows,
278 cols,
279 row_scale: Vec::new(),
280 col_field: Vec::new(),
281 vbit_offsets: offsets,
282 repack: Vec::new(),
283 })
284 }
285 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
291 model: model.clone(),
292 idx,
293 dtype: entry.dtype,
294 rows,
295 cols,
296 row_scale: Vec::new(),
297 col_field: Vec::new(),
298 vbit_offsets: Vec::new(),
299 repack: Vec::new(),
300 }),
301 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
304 model: model.clone(),
305 idx,
306 dtype: entry.dtype,
307 rows,
308 cols,
309 row_scale: Vec::new(),
310 col_field: Vec::new(),
311 vbit_offsets: Vec::new(),
312 repack: Vec::new(),
313 }),
314 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
316 model: model.clone(),
317 idx,
318 dtype: entry.dtype,
319 rows,
320 cols,
321 row_scale: Vec::new(),
322 col_field: Vec::new(),
323 vbit_offsets: Vec::new(),
324 repack: Vec::new(),
325 }),
326 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
327 model: model.clone(),
328 idx,
329 dtype: entry.dtype,
330 rows,
331 cols,
332 row_scale: Vec::new(),
333 col_field: Vec::new(),
334 vbit_offsets: Vec::new(),
335 repack: Vec::new(),
336 }),
337 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
340 model: model.clone(),
341 idx,
342 dtype: entry.dtype,
343 rows,
344 cols,
345 row_scale: Vec::new(),
346 col_field: Vec::new(),
347 vbit_offsets: Vec::new(),
348 repack: Vec::new(),
349 }),
350 TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
351 model: model.clone(),
352 idx,
353 dtype: entry.dtype,
354 rows,
355 cols,
356 row_scale: Vec::new(),
357 col_field: Vec::new(),
358 vbit_offsets: Vec::new(),
359 repack: Vec::new(),
360 }),
361 TensorDtype::Mxfp4 if cols % cortiq_core::mxfp4::GROUP == 0 => Ok(Self::Mapped {
364 model: model.clone(),
365 idx,
366 dtype: entry.dtype,
367 rows,
368 cols,
369 row_scale: Vec::new(),
370 col_field: Vec::new(),
371 vbit_offsets: Vec::new(),
372 repack: Vec::new(),
373 }),
374 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
378 model: model.clone(),
379 idx,
380 dtype: entry.dtype,
381 rows,
382 cols,
383 row_scale: Vec::new(),
384 col_field: Vec::new(),
385 vbit_offsets: Vec::new(),
386 repack: Vec::new(),
387 }),
388 TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
391 Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
392 model, idx,
393 )?)))
394 }
395 _ => {
397 let mut data = vec![0.0f32; rows * cols];
398 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
399 Ok(Self::from_f32(data, rows, cols))
400 }
401 }
402 }
403
404 pub fn is_blr(&self) -> bool {
407 matches!(self, Self::Blr(_))
408 }
409
410 pub fn is_mxfp4(&self) -> bool {
413 matches!(
414 self,
415 Self::Mapped {
416 dtype: TensorDtype::Mxfp4,
417 ..
418 }
419 )
420 }
421
422 fn per_tensor_only(&self) -> bool {
425 self.is_blr() || self.is_mxfp4()
426 }
427
428 pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
434 match self {
435 Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
436 _ => None,
437 }
438 }
439
440 pub(crate) fn graph_weight_or_blr(
444 &self,
445 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
446 self.graph_weight()
447 .or_else(|| self.graph_weight_descriptor())
448 .or_else(|| {
449 self.graph_blr()
450 .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
451 })
452 }
453
454 pub(crate) fn is_q1(&self) -> bool {
458 matches!(
459 self,
460 Self::Mapped {
461 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
462 ..
463 }
464 )
465 }
466
467 fn bin_w(&self) -> Option<BinW<'_>> {
470 match self {
471 Self::Mapped {
472 dtype: TensorDtype::Q1,
473 ..
474 } => Some(BinW::Q1(self.quant_bytes())),
475 Self::Mapped {
476 dtype: TensorDtype::Q1G,
477 rows,
478 cols,
479 ..
480 } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
481 _ => None,
482 }
483 }
484
485 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
488 match self {
489 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
490 _ => None,
491 }
492 }
493
494 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
501 if self.has_prism_contract() {
502 return None;
503 }
504 match self {
505 #[cfg(target_os = "macos")]
506 Self::Mapped {
507 dtype: TensorDtype::Q1T,
508 ..
509 } if !crate::gpu::metal_q1t_enabled() => None,
510 Self::Mapped {
511 idx,
512 dtype:
513 TensorDtype::Q1
514 | TensorDtype::Q1G
515 | TensorDtype::Q1T
516 | TensorDtype::Q4Block
517 | TensorDtype::Q4Tiled
518 | TensorDtype::Q4TiledP
522 | TensorDtype::Q8Row
523 | TensorDtype::Q8_2f,
524 rows,
525 cols,
526 ..
527 } => Some((*idx, *rows, *cols)),
528 _ => None,
529 }
530 }
531
532 #[cfg(target_os = "macos")]
539 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
540 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
541 let name = &model.tensors[idx].name;
542 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
543 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
544 if forward && affine {
545 let e = model.tensors.get(idx)?;
546 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
547 }
548 }
549 self.q1_parts()
550 }
551
552 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
558 if self.has_prism_contract() {
559 return None;
560 }
561 match self {
562 Self::Mapped {
563 idx,
564 dtype: TensorDtype::Q4Tiled,
565 rows,
566 cols,
567 ..
568 } => Some((*idx, *rows, *cols)),
569 _ => None,
570 }
571 }
572
573 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
577 if self.has_prism_contract() {
578 return None;
579 }
580 match self {
581 Self::Mapped {
582 idx,
583 dtype: TensorDtype::Q4TiledP,
584 rows,
585 cols,
586 ..
587 } => Some((*idx, *rows, *cols)),
588 _ => None,
589 }
590 }
591
592 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
597 if self.has_prism_contract() {
598 return None;
599 }
600 match self {
601 Self::Mapped {
602 idx,
603 dtype: TensorDtype::Q8Row,
604 rows,
605 cols,
606 row_scale,
607 col_field,
608 ..
609 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
610 _ => None,
611 }
612 }
613
614 pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
617 if self.has_prism_contract() {
618 return None;
619 }
620 match self {
621 Self::Mapped {
622 idx,
623 dtype: TensorDtype::Q8_2f,
624 rows,
625 cols,
626 row_scale,
627 col_field,
628 ..
629 } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
630 _ => None,
631 }
632 }
633
634 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
638 match self {
639 Self::Mapped { dtype, .. } => Some(*dtype),
640 _ => None,
641 }
642 }
643
644 pub fn model_idx(&self) -> Option<usize> {
649 match self {
650 Self::Mapped { idx, .. } => Some(*idx),
651 _ => None,
652 }
653 }
654
655 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
660 match self {
661 Self::Mapped { model, .. } => Some(model.clone()),
662 _ => None,
663 }
664 }
665
666 pub(crate) fn has_prism_contract(&self) -> bool {
671 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
672 }
673
674 pub fn rows(&self) -> usize {
675 match self {
676 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
677 Self::Blr(t) => t.rows,
678 }
679 }
680
681 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
684 if self.has_prism_contract() {
685 return None;
686 }
687 match self {
688 Self::Mapped {
689 model,
690 idx,
691 dtype: TensorDtype::Q4Tiled,
692 ..
693 } => Some((model, *idx)),
694 _ => None,
695 }
696 }
697
698 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
701 if self.has_prism_contract() {
702 return None;
703 }
704 match self {
705 Self::Mapped {
706 model,
707 idx,
708 dtype: TensorDtype::Q4TiledP,
709 ..
710 } => Some((model, *idx)),
711 _ => None,
712 }
713 }
714
715 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
723 if self.has_prism_contract() {
724 return None;
725 }
726 match self {
727 Self::Mapped {
728 model,
729 idx,
730 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
731 ..
732 } => Some((model, *idx)),
733 _ => None,
734 }
735 }
736
737 pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
740 if self.has_prism_contract() {
741 return None;
742 }
743 match self {
744 Self::Mapped {
745 model,
746 idx,
747 dtype: TensorDtype::Q8_2f,
748 ..
749 } => Some((model, *idx)),
750 _ => None,
751 }
752 }
753
754 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
757 if self.has_prism_contract() {
758 return None;
759 }
760 match self {
761 Self::Mapped {
762 model,
763 idx,
764 dtype: TensorDtype::Q2TiledP,
765 ..
766 } => Some((model, *idx)),
767 _ => None,
768 }
769 }
770
771 pub fn cols(&self) -> usize {
772 match self {
773 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
774 Self::Blr(t) => t.cols,
775 }
776 }
777
778 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
781 if self.has_prism_contract() {
782 return None;
783 }
784 match self {
785 Self::Mapped {
786 model,
787 idx,
788 dtype: TensorDtype::Q1,
789 ..
790 } => Some((model, *idx)),
791 _ => None,
792 }
793 }
794
795 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
806 if self.has_prism_contract() {
807 return None;
808 }
809 self.graph_weight_descriptor()
810 }
811
812 pub(crate) fn graph_weight_descriptor(
817 &self,
818 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
819 match self {
820 Self::Mapped {
821 model,
822 idx,
823 dtype: TensorDtype::Q8Row,
824 row_scale,
825 ..
826 } => Some((model, *idx, 0, row_scale.as_slice())),
827 Self::Mapped {
828 model,
829 idx,
830 dtype: TensorDtype::Q1,
831 ..
832 } => Some((model, *idx, 1, &[])),
833 Self::Mapped {
836 model,
837 idx,
838 dtype: TensorDtype::Q1G,
839 ..
840 } => Some((model, *idx, 10, &[])),
841 Self::Mapped {
846 model,
847 idx,
848 dtype: TensorDtype::Q4Tiled,
849 ..
850 } => Some((model, *idx, 5, &[])),
851 Self::Mapped {
855 model,
856 idx,
857 dtype: TensorDtype::Q4TiledP,
858 ..
859 } => Some((model, *idx, 6, &[])),
860 Self::Mapped {
861 model,
862 idx,
863 dtype: TensorDtype::Q4Block,
864 ..
865 } => Some((model, *idx, 2, &[])),
866 Self::Mapped {
871 model,
872 idx,
873 dtype: TensorDtype::Q8_2f,
874 ..
875 } => Some((model, *idx, 7, &[])),
876 Self::Mapped {
877 model,
878 idx,
879 dtype: TensorDtype::Q1T,
880 ..
881 } => Some((model, *idx, 3, &[])),
882 Self::Mapped {
886 model,
887 idx,
888 dtype: TensorDtype::Q2TiledP,
889 ..
890 } => Some((model, *idx, 9, &[])),
891 _ => None,
892 }
893 }
894
895 pub fn as_f32(&self) -> Option<&[f32]> {
898 match self {
899 Self::F32 { data, .. } => Some(data),
900 Self::Mapped { .. } | Self::Blr(_) => None,
901 }
902 }
903
904 fn quant_bytes(&self) -> &[u8] {
905 match self {
906 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
907 Self::Blr(t) => t.bytes(),
908 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
909 }
910 }
911
912 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
914 let cols = self.cols();
915 debug_assert_eq!(dst.len(), cols);
916 match self {
917 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
918 Self::Blr(t) => t.row_f32(r, dst),
919 Self::Mapped {
920 model,
921 idx,
922 dtype,
923 row_scale,
924 col_field,
925 vbit_offsets,
926 ..
927 } => {
928 if *dtype == TensorDtype::Q4Tiled {
929 let bytes = self.quant_bytes();
930 let gpr = cols / GROUP_SIZE;
931 for gi in 0..gpr {
932 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
933 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
934 for (k, &b) in tile[2..].iter().enumerate() {
935 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
936 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
937 }
938 }
939 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
940 crate::prism::inverse_embedding(model, dst);
941 }
942 return;
943 }
944 if *dtype == TensorDtype::Q4TiledP {
945 let bytes = self.quant_bytes();
946 let gpr = cols / GROUP_SIZE;
947 let v = Q4tpView::new(bytes, self.rows(), cols);
948 let mut sc = vec![0f32; gpr];
949 v.scales_into(r, gpr, &mut sc);
950 for gi in 0..gpr {
951 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
952 let s = sc[gi];
953 for (k, &b) in tile.iter().enumerate() {
954 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
955 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
956 }
957 }
958 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
959 crate::prism::inverse_embedding(model, dst);
960 }
961 return;
962 }
963 if *dtype == TensorDtype::Q2TiledP {
964 let bytes = self.quant_bytes();
965 let gpr = cols / GROUP_SIZE;
966 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
967 let mut sc = vec![0f32; gpr];
968 v.scales_into(r, gpr, &mut sc);
969 for gi in 0..gpr {
970 let ch =
971 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
972 let s = sc[gi];
973 for (k, &b) in ch.iter().enumerate() {
974 for j in 0..4 {
975 let center = if crate::prism::is_affine_target(
976 model,
977 &model.tensors[*idx].name,
978 ) {
979 1.0
980 } else {
981 1.5
982 };
983 dst[gi * GROUP_SIZE + k * 4 + j] =
984 (((b >> (2 * j)) & 3) as f32 - center) * s;
985 }
986 }
987 }
988 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
989 crate::prism::inverse_embedding(model, dst);
990 }
991 return;
992 }
993 if *dtype == TensorDtype::Q4Block {
994 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
995 let gpr = cols / GROUP_SIZE;
996 for gi in 0..gpr {
997 let g = r * gpr + gi;
998 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
999 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1000 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
1001 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
1002 }
1003 }
1004 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1005 crate::prism::inverse_embedding(model, dst);
1006 }
1007 return;
1008 }
1009 if *dtype == TensorDtype::Mxfp4 {
1010 mxfp4::Mx::new(self.quant_bytes(), self.rows(), cols).row_f32(r, dst);
1011 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1012 crate::prism::inverse_embedding(model, dst);
1013 }
1014 return;
1015 }
1016 if *dtype == TensorDtype::Q1 {
1017 let bytes = self.quant_bytes();
1018 let gpr = cols / GROUP_SIZE;
1019 for gi in 0..gpr {
1020 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
1021 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1022 for (j, &b) in tile[2..].iter().enumerate() {
1023 for k in 0..8 {
1024 dst[gi * GROUP_SIZE + j * 8 + k] =
1025 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1026 }
1027 }
1028 }
1029 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1030 crate::prism::inverse_embedding(model, dst);
1031 }
1032 return;
1033 }
1034 if *dtype == TensorDtype::Q1G {
1035 let bytes = self.quant_bytes();
1038 let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
1039 let gpr = cols / GROUP_SIZE;
1040 for gi in 0..gpr {
1041 let t = r * gpr + gi;
1042 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
1043 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
1044 for k in 0..8 {
1045 dst[gi * GROUP_SIZE + j * 8 + k] =
1046 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1047 }
1048 }
1049 }
1050 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1051 crate::prism::inverse_embedding(model, dst);
1052 }
1053 return;
1054 }
1055 if *dtype == TensorDtype::Q1T {
1056 let bytes = self.quant_bytes();
1057 let gpr = cols / GROUP_SIZE;
1058 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1059 for gi in 0..gpr {
1060 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1061 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1062 bytes[off],
1063 bytes[off + 1],
1064 ]));
1065 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1066 for k in 0..GROUP_SIZE {
1067 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1068 {
1069 1 => s,
1070 2 => -s,
1071 _ => 0.0,
1072 };
1073 }
1074 }
1075 let rows = self.rows();
1077 let entries = base_len + (rows + 1) * 4;
1078 if entries <= bytes.len() {
1079 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1080 let r0 = u32::from_le_bytes([
1081 ptrs[r * 4],
1082 ptrs[r * 4 + 1],
1083 ptrs[r * 4 + 2],
1084 ptrs[r * 4 + 3],
1085 ]) as usize;
1086 let r1 = u32::from_le_bytes([
1087 ptrs[(r + 1) * 4],
1088 ptrs[(r + 1) * 4 + 1],
1089 ptrs[(r + 1) * 4 + 2],
1090 ptrs[(r + 1) * 4 + 3],
1091 ]) as usize;
1092 let off = entries + r0 * 4;
1093 for i in 0..r1 - r0 {
1094 let item = &bytes[off + i * 4..off + i * 4 + 4];
1095 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1096 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1097 item[2], item[3],
1098 ]));
1099 if c < cols {
1100 dst[c] = v;
1101 }
1102 }
1103 }
1104 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1105 crate::prism::inverse_embedding(model, dst);
1106 }
1107 return;
1108 }
1109 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1110 let bytes = self.quant_bytes();
1111 let rows = self.rows();
1112 let ng = cols / GROUP_SIZE;
1113 let bits = &bytes[..rows];
1114 let sc_off = rows;
1115 let off = vbit_offsets[r];
1118 let b = bits[r] as usize;
1119 let l = ((1usize << (b - 1)) - 1) as f32;
1120 let data = &bytes[off..];
1121 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1122 for (i, d) in dst.iter_mut().enumerate() {
1123 while nbits < b {
1124 acc = (acc << 8) | data[byte_idx] as u64;
1125 byte_idx += 1;
1126 nbits += 8;
1127 }
1128 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1129 nbits -= b;
1130 let so = (r * ng + i / GROUP_SIZE) * 2;
1131 let sv = f16_to_f32(u16::from_le_bytes([
1132 bytes[sc_off + so],
1133 bytes[sc_off + so + 1],
1134 ]));
1135 *d = (u - l) * sv;
1136 }
1137 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1138 crate::prism::inverse_embedding(model, dst);
1139 }
1140 return;
1141 }
1142 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1143 let s = row_scale[r];
1144 match dtype {
1145 TensorDtype::Q8Row => {
1146 for (d, &b) in dst.iter_mut().zip(q) {
1147 *d = (b as i8) as f32 * s;
1148 }
1149 }
1150 TensorDtype::Q8_2f => {
1151 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1152 *d = (b as i8) as f32 * s * col_field[i];
1153 }
1154 }
1155 _ => unreachable!(),
1156 }
1157 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1158 crate::prism::inverse_embedding(model, dst);
1159 }
1160 }
1161 }
1162 }
1163
1164 pub fn sparse_col_ok(&self) -> bool {
1169 match self {
1170 Self::F32 { .. } => true,
1171 Self::Blr(_) => false,
1172 Self::Mapped { dtype, .. } => {
1173 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1174 }
1175 }
1176 }
1177
1178 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1182 let inter = self.cols();
1183 let hidden = self.rows();
1184 debug_assert_eq!(out.len(), hidden);
1185 match self {
1186 Self::F32 { data, .. } => {
1187 for (k, o) in out.iter_mut().enumerate() {
1188 *o += w * data[k * inter + c];
1189 }
1190 }
1191 Self::Blr(t) => t.add_col_scaled(c, w, out),
1192 Self::Mapped {
1193 dtype,
1194 row_scale,
1195 col_field,
1196 ..
1197 } => {
1198 let q = self.quant_bytes();
1199 let colf = if *dtype == TensorDtype::Q8_2f {
1200 col_field[c]
1201 } else {
1202 1.0
1203 };
1204 let wc = w * colf;
1205 for (k, o) in out.iter_mut().enumerate() {
1206 let b = q[k * inter + c] as i8 as f32;
1207 *o += wc * b * row_scale[k];
1208 }
1209 }
1210 }
1211 }
1212
1213 #[inline]
1222 pub fn prefetch_row(&self, r: usize) {
1223 let Self::Mapped { dtype, .. } = self else {
1224 return;
1225 };
1226 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1227 return;
1228 }
1229 let cols = self.cols();
1230 let q = self.quant_bytes();
1231 let (a, b) = (r * cols, (r + 1) * cols);
1232 if b > q.len() {
1233 return;
1234 }
1235 let mut j = a;
1236 while j < b {
1237 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1238 j += 512;
1239 }
1240 }
1241
1242 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1251 let cols = self.cols();
1252 debug_assert_eq!(out.len(), cols);
1253 match self {
1254 Self::F32 { data, .. } => {
1255 let row = &data[r * cols..(r + 1) * cols];
1256 for (o, v) in out.iter_mut().zip(row) {
1257 *o += w * v;
1258 }
1259 }
1260 Self::Blr(t) => {
1261 let tmp = &mut scratch[..cols];
1262 t.row_f32(r, tmp);
1263 for (o, v) in out.iter_mut().zip(tmp.iter()) {
1264 *o += w * v;
1265 }
1266 }
1267 Self::Mapped {
1268 dtype,
1269 row_scale,
1270 col_field,
1271 ..
1272 } => match dtype {
1273 TensorDtype::Q8Row => {
1274 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1275 let ws = w * row_scale[r];
1276 let row: &[i8] =
1277 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1278 axpy_i8_f32(out, row, ws);
1279 }
1280 TensorDtype::Q8_2f => {
1281 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1282 let ws = w * row_scale[r];
1283 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1284 *o += ws * c * (*b as i8 as f32);
1285 }
1286 }
1287 _ => {
1288 self.row_f32(r, scratch);
1289 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1290 *o += w * v;
1291 }
1292 }
1293 },
1294 }
1295 }
1296
1297 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1301 let cols = self.cols();
1302 match self {
1303 Self::F32 { data, .. } => {
1304 let row = &data[r * cols..(r + 1) * cols];
1305 row.iter().zip(x).map(|(w, v)| w * v).sum()
1306 }
1307 Self::Blr(t) => {
1308 let tmp = &mut scratch[..cols];
1309 t.row_f32(r, tmp);
1310 tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1311 }
1312 Self::Mapped {
1313 model,
1314 idx,
1315 dtype,
1316 row_scale,
1317 col_field,
1318 ..
1319 } => {
1320 let prism_forward =
1321 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1322 if prism_forward {
1323 let transformed = crate::prism::forward(model, &x[..cols]);
1324 let gpr = cols / GROUP_SIZE;
1325 match dtype {
1326 TensorDtype::Q2TiledP => {
1327 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1328 let mut sc = vec![0f32; gpr];
1329 v.scales_into(r, gpr, &mut sc);
1330 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1331 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1332 }
1333 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1334 }
1335 _ => {
1336 self.row_f32(r, scratch);
1337 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1338 }
1339 }
1340 }
1341 match dtype {
1342 TensorDtype::Q8Row => {
1343 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1344 dot_i8_f32(q, x) * row_scale[r]
1345 }
1346 TensorDtype::Q8_2f => {
1347 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1348 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1349 }
1350 _ => {
1351 self.row_f32(r, scratch);
1352 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1353 }
1354 }
1355 }
1356 }
1357 }
1358
1359 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1362 if self.flint_exact_cpu_default() {
1363 return float_activations_scope(|| self.matvec(x, out, pool));
1364 }
1365 match self {
1366 Self::F32 { data, .. } => {
1370 if !crate::f32_backend::matvec(data, x, out) {
1371 matvec_rows(pool, data, x, out);
1372 }
1373 }
1374 Self::Blr(t) => {
1375 if crate::gpu::enabled_here()
1376 && crate::blr::gpu_wanted()
1377 && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1378 {
1379 return;
1380 }
1381 t.matvec(x, out, pool)
1382 }
1383 Self::Mapped {
1384 model,
1385 idx,
1386 dtype,
1387 rows,
1388 cols,
1389 row_scale,
1390 col_field,
1391 vbit_offsets,
1392 repack,
1393 } => {
1394 let _ = (model, idx);
1395 assert!(
1405 out.len() >= *rows && x.len() >= *cols,
1406 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1407 out.len(),
1408 x.len(),
1409 );
1410 let prism_forward =
1411 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1412 if *dtype == TensorDtype::Q2TiledP
1413 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1414 {
1415 use std::sync::atomic::{AtomicUsize, Ordering};
1416 static N: AtomicUsize = AtomicUsize::new(0);
1417 let n = N.fetch_add(1, Ordering::Relaxed);
1418 if n < 128 {
1419 eprintln!(
1420 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1421 model.tensors[*idx].name,
1422 prism_forward,
1423 rows,
1424 cols,
1425 crate::gpu::enabled_here(),
1426 crate::gpu::q2tp_gpu_opt_in(),
1427 crate::gpu::cur_layer(),
1428 );
1429 }
1430 }
1431 if prism_forward {
1436 let transformed = crate::prism::forward(model, &x[..*cols]);
1437 match dtype {
1438 TensorDtype::Q4Block => {
1439 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1440 }
1441 TensorDtype::Q4Tiled => {
1442 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1443 }
1444 TensorDtype::Q4TiledP => {
1445 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1446 }
1447 TensorDtype::Q2TiledP => {
1448 let affine =
1449 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1450 if *rows * *cols >= 8_388_608
1451 && crate::gpu::enabled_here()
1452 && crate::gpu::q2tp_gpu_opt_in()
1453 {
1454 let gpu_ok = if affine {
1455 crate::gpu::q2tp_affine_matvec(
1456 model,
1457 *idx,
1458 &transformed,
1459 *rows,
1460 *cols,
1461 out,
1462 )
1463 } else {
1464 crate::gpu::q2tp_matvec(
1465 model,
1466 *idx,
1467 &transformed,
1468 *rows,
1469 *cols,
1470 out,
1471 )
1472 };
1473 if gpu_ok {
1474 return;
1475 }
1476 }
1477 if affine {
1478 q2tp_affine_matvec(
1479 self.quant_bytes(),
1480 &transformed,
1481 *rows,
1482 *cols,
1483 out,
1484 pool,
1485 )
1486 } else {
1487 q2tp_matvec(
1488 self.quant_bytes(),
1489 &transformed,
1490 *rows,
1491 *cols,
1492 out,
1493 pool,
1494 )
1495 }
1496 }
1497 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1498 self.bin_w().expect("binary dtype"),
1499 &transformed,
1500 *rows,
1501 *cols,
1502 out,
1503 pool,
1504 ),
1505 TensorDtype::Q1T => {
1506 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1507 }
1508 TensorDtype::Mxfp4 => {
1509 mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1510 }
1511 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1512 self.quant_bytes(),
1513 vbit_offsets,
1514 &transformed,
1515 *rows,
1516 *cols,
1517 out,
1518 pool,
1519 ),
1520 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1521 self.quant_bytes(),
1522 repack,
1523 row_scale,
1524 &transformed,
1525 col_field,
1526 *dtype,
1527 *rows,
1528 *cols,
1529 out,
1530 pool,
1531 ),
1532 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1533 }
1534 return;
1535 }
1536 if *dtype == TensorDtype::Q4Block {
1537 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1541 let t0 = std::time::Instant::now();
1542 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1543 crate::gpu::ProbeArm::Gpu => {
1544 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1545 crate::gpu::probe_record(
1546 crate::gpu::OpClass::Matvec,
1547 true,
1548 t0.elapsed(),
1549 );
1550 return;
1551 }
1552 }
1553 crate::gpu::ProbeArm::CpuTimed => {
1554 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1555 crate::gpu::probe_record(
1556 crate::gpu::OpClass::Matvec,
1557 false,
1558 t0.elapsed(),
1559 );
1560 return;
1561 }
1562 crate::gpu::ProbeArm::Cpu => {}
1563 }
1564 }
1565 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1566 return;
1567 }
1568 if *dtype == TensorDtype::Q4Tiled {
1569 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1574 let t0 = std::time::Instant::now();
1575 let cls = crate::gpu::matvec_class(*rows, *cols);
1576 match crate::gpu::probe_arm(cls) {
1577 crate::gpu::ProbeArm::Gpu => {
1578 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1579 crate::gpu::probe_record(cls, true, t0.elapsed());
1580 return;
1581 }
1582 }
1583 crate::gpu::ProbeArm::CpuTimed => {
1584 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1585 crate::gpu::probe_record(cls, false, t0.elapsed());
1586 return;
1587 }
1588 crate::gpu::ProbeArm::Cpu => {}
1589 }
1590 }
1591 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1592 return;
1593 }
1594 if *dtype == TensorDtype::Q4TiledP {
1595 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1601 let t0 = std::time::Instant::now();
1602 let cls = crate::gpu::matvec_class(*rows, *cols);
1603 match crate::gpu::probe_arm(cls) {
1604 crate::gpu::ProbeArm::Gpu => {
1605 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1606 crate::gpu::probe_record(cls, true, t0.elapsed());
1607 return;
1608 }
1609 }
1610 crate::gpu::ProbeArm::CpuTimed => {
1611 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1612 crate::gpu::probe_record(cls, false, t0.elapsed());
1613 return;
1614 }
1615 crate::gpu::ProbeArm::Cpu => {}
1616 }
1617 }
1618 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1619 return;
1620 }
1621 if *dtype == TensorDtype::Q2TiledP {
1622 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1623 return;
1624 }
1625 if *dtype == TensorDtype::Mxfp4 {
1626 self.mxfp4_matvec(x, out, pool);
1627 return;
1628 }
1629 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1630 let w = self.bin_w().expect("binary dtype");
1635 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1636 let t0 = std::time::Instant::now();
1637 let arm = if crate::gpu::q1_force() {
1638 crate::gpu::ProbeArm::Gpu
1639 } else {
1640 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1641 };
1642 match arm {
1643 crate::gpu::ProbeArm::Gpu => {
1644 let ok = if *dtype == TensorDtype::Q1G {
1645 crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1646 } else {
1647 crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1648 };
1649 if ok {
1650 crate::gpu::probe_record(
1651 crate::gpu::OpClass::Matvec,
1652 true,
1653 t0.elapsed(),
1654 );
1655 return;
1656 }
1657 }
1658 crate::gpu::ProbeArm::CpuTimed => {
1659 bin_matvec(w, x, *rows, *cols, out, pool);
1660 crate::gpu::probe_record(
1661 crate::gpu::OpClass::Matvec,
1662 false,
1663 t0.elapsed(),
1664 );
1665 return;
1666 }
1667 crate::gpu::ProbeArm::Cpu => {}
1668 }
1669 }
1670 bin_matvec(w, x, *rows, *cols, out, pool);
1671 return;
1672 }
1673 if *dtype == TensorDtype::Q1T {
1674 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1678 let t0 = std::time::Instant::now();
1679 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1680 crate::gpu::ProbeArm::Gpu => {
1681 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1682 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1683 crate::gpu::probe_record(
1684 crate::gpu::OpClass::Matvec,
1685 true,
1686 t0.elapsed(),
1687 );
1688 return;
1689 }
1690 }
1691 crate::gpu::ProbeArm::CpuTimed => {
1692 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1693 crate::gpu::probe_record(
1694 crate::gpu::OpClass::Matvec,
1695 false,
1696 t0.elapsed(),
1697 );
1698 return;
1699 }
1700 crate::gpu::ProbeArm::Cpu => {}
1701 }
1702 }
1703 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1704 return;
1705 }
1706 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1707 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1708 return;
1709 }
1710 let xs = prescale(x, col_field, *dtype);
1711 if *rows >= crate::gpu::min_rows()
1716 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1717 && gpu_lmhead_enabled()
1718 && crate::gpu::enabled_here()
1719 {
1720 let t0 = std::time::Instant::now();
1723 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1724 crate::gpu::ProbeArm::Gpu => {}
1725 crate::gpu::ProbeArm::CpuTimed => {
1726 qmatvec(
1727 self.quant_bytes(),
1728 repack,
1729 row_scale,
1730 x,
1731 col_field,
1732 *dtype,
1733 *rows,
1734 *cols,
1735 out,
1736 pool,
1737 );
1738 crate::gpu::probe_record(
1739 crate::gpu::OpClass::Matvec,
1740 false,
1741 t0.elapsed(),
1742 );
1743 return;
1744 }
1745 crate::gpu::ProbeArm::Cpu => {
1746 qmatvec(
1747 self.quant_bytes(),
1748 repack,
1749 row_scale,
1750 x,
1751 col_field,
1752 *dtype,
1753 *rows,
1754 *cols,
1755 out,
1756 pool,
1757 );
1758 return;
1759 }
1760 }
1761 let frac = gpu_split_frac();
1762 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1763 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1764 let bytes = self.quant_bytes();
1765 let ok = std::thread::scope(|sc| {
1766 let g = sc.spawn(|| {
1767 crate::gpu::q8_matvec_range(
1768 model,
1769 *idx,
1770 cpu_rows,
1771 &row_scale[cpu_rows..],
1772 &xs,
1773 *rows - cpu_rows,
1774 *cols,
1775 out_gpu,
1776 )
1777 });
1778 if cpu_rows > 0 {
1779 let rep_cpu = if repack.is_empty() {
1782 &[][..]
1783 } else {
1784 &repack[..(cpu_rows / 4) * 4 * *cols]
1785 };
1786 qmatvec(
1787 &bytes[..cpu_rows * *cols],
1788 rep_cpu,
1789 &row_scale[..cpu_rows],
1790 x,
1791 col_field,
1792 *dtype,
1793 cpu_rows,
1794 *cols,
1795 out_cpu,
1796 pool,
1797 );
1798 }
1799 g.join().unwrap_or(false)
1800 });
1801 if ok {
1802 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1803 return;
1804 }
1805 qmatvec(
1808 &bytes[cpu_rows * *cols..(*rows) * *cols],
1809 &[],
1810 &row_scale[cpu_rows..],
1811 x,
1812 col_field,
1813 *dtype,
1814 *rows - cpu_rows,
1815 *cols,
1816 out_gpu,
1817 pool,
1818 );
1819 return;
1820 }
1821 qmatvec(
1822 self.quant_bytes(),
1823 repack,
1824 row_scale,
1825 x,
1826 col_field,
1827 *dtype,
1828 *rows,
1829 *cols,
1830 out,
1831 pool,
1832 );
1833 }
1834 }
1835 }
1836
1837 pub fn matvec2(
1839 &self,
1840 x1: &[f32],
1841 x2: &[f32],
1842 o1: &mut [f32],
1843 o2: &mut [f32],
1844 pool: Option<&Pool>,
1845 ) {
1846 if self.flint_exact_cpu_default() {
1847 return float_activations_scope(|| self.matvec2(x1, x2, o1, o2, pool));
1848 }
1849 match self {
1850 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1851 Self::Blr(_) => {
1852 self.matvec(x1, o1, pool);
1853 self.matvec(x2, o2, pool);
1854 }
1855 Self::Mapped {
1856 model,
1857 idx,
1858 dtype,
1859 rows,
1860 cols,
1861 row_scale,
1862 col_field,
1863 vbit_offsets,
1864 ..
1865 } => {
1866 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1867 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1868 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1869 match dtype {
1870 TensorDtype::Q4Block => {
1871 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1872 }
1873 TensorDtype::Q4Tiled => {
1874 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1875 }
1876 TensorDtype::Q4TiledP => {
1877 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1878 }
1879 TensorDtype::Q2TiledP => {
1880 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1881 q2tp_affine_matvec2(
1882 self.quant_bytes(),
1883 &tx1,
1884 &tx2,
1885 *rows,
1886 *cols,
1887 o1,
1888 o2,
1889 pool,
1890 )
1891 } else {
1892 q2tp_matvec2(
1893 self.quant_bytes(),
1894 &tx1,
1895 &tx2,
1896 *rows,
1897 *cols,
1898 o1,
1899 o2,
1900 pool,
1901 )
1902 }
1903 }
1904 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1905 self.bin_w().expect("binary dtype"),
1906 &tx1,
1907 &tx2,
1908 *rows,
1909 *cols,
1910 o1,
1911 o2,
1912 pool,
1913 ),
1914 TensorDtype::Q1T => {
1915 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1916 }
1917 TensorDtype::Mxfp4 => mxfp4::matvec2(
1918 self.quant_bytes(),
1919 &tx1,
1920 &tx2,
1921 *rows,
1922 *cols,
1923 o1,
1924 o2,
1925 pool,
1926 ),
1927 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1928 self.quant_bytes(),
1929 vbit_offsets,
1930 &tx1,
1931 &tx2,
1932 *rows,
1933 *cols,
1934 o1,
1935 o2,
1936 pool,
1937 ),
1938 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1939 self.quant_bytes(),
1940 row_scale,
1941 &tx1,
1942 &tx2,
1943 col_field,
1944 *dtype,
1945 *rows,
1946 *cols,
1947 o1,
1948 o2,
1949 pool,
1950 ),
1951 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1952 }
1953 return;
1954 }
1955 if *dtype == TensorDtype::Q4Block {
1956 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1957 return;
1958 }
1959 if *dtype == TensorDtype::Q4Tiled {
1960 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1961 return;
1962 }
1963 if *dtype == TensorDtype::Q4TiledP {
1964 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1965 return;
1966 }
1967 if *dtype == TensorDtype::Q2TiledP {
1968 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1969 return;
1970 }
1971 if *dtype == TensorDtype::Mxfp4 {
1972 mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1973 return;
1974 }
1975 if let Some(w) = self.bin_w() {
1976 bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1977 return;
1978 }
1979 if *dtype == TensorDtype::Q1T {
1980 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1986 return;
1987 }
1988 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1989 vbitmatvec2(
1990 self.quant_bytes(),
1991 vbit_offsets,
1992 x1,
1993 x2,
1994 *rows,
1995 *cols,
1996 o1,
1997 o2,
1998 pool,
1999 );
2000 return;
2001 }
2002 qmatvec2(
2003 self.quant_bytes(),
2004 row_scale,
2005 x1,
2006 x2,
2007 col_field,
2008 *dtype,
2009 *rows,
2010 *cols,
2011 o1,
2012 o2,
2013 pool,
2014 );
2015 }
2016 }
2017 }
2018}
2019
2020impl QTensor {
2021 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2029 if self.has_prism_contract() {
2030 return None;
2031 }
2032 match self {
2033 Self::Mapped {
2034 model, idx, dtype, ..
2035 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2036 _ => None,
2037 }
2038 }
2039
2040 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2041 if self.flint_exact_cpu_default() {
2042 return float_activations_scope(|| self.matmat(xs_all, b, out, pool));
2043 }
2044 let cols = self.cols();
2045 let rows = self.rows();
2046 debug_assert_eq!(xs_all.len(), b * cols);
2047 debug_assert_eq!(out.len(), b * rows);
2048 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2049 if crate::gptq_capture::capturing() {
2053 if let Self::Mapped { model, idx, .. } = self {
2054 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2055 }
2056 if let Self::Blr(t) = self {
2057 crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2058 }
2059 }
2060 match self {
2061 Self::Blr(t) => {
2062 if crate::gpu::enabled_here()
2063 && crate::blr::gpu_wanted()
2064 && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2065 {
2066 return;
2067 }
2068 t.matmat(xs_all, b, out, pool)
2069 }
2070 Self::F32 { data, .. } => {
2071 if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2072 return;
2073 }
2074 let out_addr = SendMut(out.as_mut_ptr());
2075 let run = |start: usize, end: usize| {
2082 let mut bi = start;
2083 while bi < end {
2084 let n = (end - bi).min(4);
2085 for o in 0..rows {
2086 let row = &data[o * cols..(o + 1) * cols];
2087 if n == 4 {
2088 let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2089 let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2090 let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2091 for j in 0..cols {
2092 let w = row[j];
2093 a0 += w * x0[j];
2094 a1 += w * x1[j];
2095 a2 += w * x2[j];
2096 a3 += w * x3[j];
2097 }
2098 for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2099 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2100 }
2101 } else {
2102 for k in 0..n {
2103 let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2104 let mut acc = 0f32;
2105 for j in 0..cols {
2106 acc += row[j] * x[j];
2107 }
2108 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2109 }
2110 }
2111 }
2112 bi += n;
2113 }
2114 };
2115 dispatch_rows(pool, b, &run);
2116 }
2117 Self::Mapped {
2118 model,
2119 idx,
2120 dtype,
2121 row_scale,
2122 col_field,
2123 vbit_offsets,
2124 ..
2125 } => {
2126 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2127 let mut transformed = Vec::with_capacity(xs_all.len());
2128 for bi in 0..b {
2129 transformed.extend_from_slice(&crate::prism::forward(
2130 model,
2131 &xs_all[bi * cols..(bi + 1) * cols],
2132 ));
2133 }
2134 match dtype {
2135 TensorDtype::Q4Block => {
2136 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2137 }
2138 TensorDtype::Q4Tiled => {
2139 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2140 }
2141 TensorDtype::Q4TiledP => {
2142 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2143 }
2144 TensorDtype::Q2TiledP => {
2145 let affine =
2146 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2147 let gpu_batch_ok = if affine {
2153 b >= 2
2154 } else {
2155 b >= 32 && b * rows * cols >= 128_000_000
2156 };
2157 if gpu_batch_ok
2158 && cols % 32 == 0
2159 && crate::gpu::enabled_here()
2160 && crate::gpu::q2tp_gpu_opt_in()
2161 {
2162 let gpu_ok = if affine {
2163 crate::gpu::q2tp_affine_matmat(
2164 model,
2165 *idx,
2166 &transformed,
2167 b,
2168 rows,
2169 cols,
2170 out,
2171 )
2172 } else {
2173 crate::gpu::q2tp_matmat(
2174 model,
2175 *idx,
2176 &transformed,
2177 b,
2178 rows,
2179 cols,
2180 out,
2181 )
2182 };
2183 if gpu_ok {
2184 return;
2185 }
2186 }
2187 if affine
2191 && b == 1
2192 && cols % 32 == 0
2193 && crate::gpu::enabled_here()
2194 && crate::gpu::q2tp_gpu_opt_in()
2195 && crate::gpu::q2tp_affine_matvec(
2196 model,
2197 *idx,
2198 &transformed[..cols],
2199 rows,
2200 cols,
2201 &mut out[..rows],
2202 )
2203 {
2204 return;
2205 }
2206 if affine {
2207 q2tp_affine_matmat(
2208 self.quant_bytes(),
2209 &transformed,
2210 b,
2211 rows,
2212 cols,
2213 out,
2214 pool,
2215 )
2216 } else {
2217 q2tp_matmat(
2218 self.quant_bytes(),
2219 &transformed,
2220 b,
2221 rows,
2222 cols,
2223 out,
2224 pool,
2225 )
2226 }
2227 }
2228 TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2229 self.bin_w().expect("binary dtype"),
2230 &transformed,
2231 b,
2232 rows,
2233 cols,
2234 out,
2235 pool,
2236 ),
2237 TensorDtype::Q1T => {
2238 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2239 }
2240 TensorDtype::Mxfp4 => {
2241 mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2242 }
2243 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2244 self.quant_bytes(),
2245 vbit_offsets,
2246 &transformed,
2247 b,
2248 rows,
2249 cols,
2250 out,
2251 pool,
2252 ),
2253 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2254 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2255 .map(|bi| {
2256 prescale(
2257 &transformed[bi * cols..(bi + 1) * cols],
2258 col_field,
2259 *dtype,
2260 )
2261 })
2262 .collect();
2263 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2264 }
2265 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2266 }
2267 return;
2268 }
2269 if *dtype == TensorDtype::Q4Block {
2270 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2271 return;
2272 }
2273 if *dtype == TensorDtype::Mxfp4 {
2274 self.mxfp4_matmat(xs_all, b, out, pool);
2275 return;
2276 }
2277 if *dtype == TensorDtype::Q4TiledP {
2278 if b >= 32
2292 && b * rows * cols >= 128_000_000
2293 && cols % 32 == 0
2294 && !row_exact()
2295 && !crate::gpu::mm_killed()
2296 && crate::gpu::enabled_here()
2297 {
2298 let class = if b >= 128 {
2299 crate::gpu::OpClass::MatmatWide
2300 } else {
2301 crate::gpu::OpClass::Matmat
2302 };
2303 if let Self::Mapped { model, idx, .. } = self {
2304 if crate::mm_ab::on() {
2316 let mut g = vec![0f32; b * rows];
2317 let t = std::time::Instant::now();
2318 let took = crate::gpu::q4tp_matmat(
2319 model, *idx, xs_all, b, rows, cols, &mut g,
2320 );
2321 let dg = t.elapsed();
2322 let t = std::time::Instant::now();
2323 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2324 let dc = t.elapsed();
2325 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2326 return;
2327 }
2328 let t0 = std::time::Instant::now();
2329 let resident = crate::gpu::weight_is_resident(model, *idx);
2333 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2334 crate::gpu::ProbeArm::Gpu => {
2335 if crate::gpu::q4tp_matmat(
2336 model, *idx, xs_all, b, rows, cols, out,
2337 ) {
2338 let el = t0.elapsed();
2339 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2349 let budget = std::time::Duration::from_secs_f64(
2350 flops / 1.5e12 * 8.0 + 0.020,
2351 );
2352 crate::gpu::mm_budget_check(
2353 "q4tp matmat",
2354 el,
2355 budget,
2356 crate::gpu::probe_was_cold() || !resident,
2357 );
2358 crate::gpu::probe_record(class, true, el);
2359 return;
2360 }
2361 }
2362 crate::gpu::ProbeArm::CpuTimed => {
2363 q4tp_matmat(
2364 self.quant_bytes(),
2365 xs_all,
2366 b,
2367 rows,
2368 cols,
2369 out,
2370 pool,
2371 );
2372 crate::gpu::probe_record(class, false, t0.elapsed());
2373 return;
2374 }
2375 crate::gpu::ProbeArm::Cpu => {}
2376 }
2377 }
2378 }
2379 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2380 return;
2381 }
2382 if *dtype == TensorDtype::Q2TiledP {
2383 if b >= 32
2389 && b * rows * cols >= 128_000_000
2390 && cols % 32 == 0
2391 && !crate::gpu::mm_killed()
2392 && crate::gpu::enabled_here()
2393 {
2394 let class = if b >= 128 {
2395 crate::gpu::OpClass::MatmatWide
2396 } else {
2397 crate::gpu::OpClass::Matmat
2398 };
2399 if let Self::Mapped { model, idx, .. } = self {
2400 let t0 = std::time::Instant::now();
2401 match crate::gpu::probe_arm(class) {
2402 crate::gpu::ProbeArm::Gpu => {
2403 if crate::gpu::q2tp_matmat(
2404 model, *idx, xs_all, b, rows, cols, out,
2405 ) {
2406 crate::gpu::probe_record(class, true, t0.elapsed());
2407 return;
2408 }
2409 }
2410 crate::gpu::ProbeArm::CpuTimed => {
2411 q2tp_matmat(
2412 self.quant_bytes(),
2413 xs_all,
2414 b,
2415 rows,
2416 cols,
2417 out,
2418 pool,
2419 );
2420 crate::gpu::probe_record(class, false, t0.elapsed());
2421 return;
2422 }
2423 crate::gpu::ProbeArm::Cpu => {}
2424 }
2425 }
2426 }
2427 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2432 return;
2433 }
2434 if *dtype == TensorDtype::Q4Tiled {
2435 if b >= 32
2447 && b * rows * cols >= 128_000_000
2448 && cols % 32 == 0
2449 && !crate::gpu::mm_killed()
2450 && crate::gpu::enabled_here()
2451 {
2452 let class = if b >= 128 {
2453 crate::gpu::OpClass::MatmatWide
2454 } else {
2455 crate::gpu::OpClass::Matmat
2456 };
2457 if let Self::Mapped { model, idx, .. } = self {
2458 let t0 = std::time::Instant::now();
2459 match crate::gpu::probe_arm(class) {
2460 crate::gpu::ProbeArm::Gpu => {
2461 if crate::gpu::q4t_matmat(
2462 model, *idx, xs_all, b, rows, cols, out,
2463 ) {
2464 let el = t0.elapsed();
2465 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2475 let budget = std::time::Duration::from_secs_f64(
2476 flops / 1.5e12 * 8.0 + 0.020,
2477 );
2478 crate::gpu::mm_budget_check(
2479 "q4t matmat",
2480 el,
2481 budget,
2482 crate::gpu::probe_was_cold(),
2483 );
2484 crate::gpu::probe_record(class, true, el);
2485 return;
2486 }
2487 }
2488 crate::gpu::ProbeArm::CpuTimed => {
2489 q4t_matmat(
2490 self.quant_bytes(),
2491 xs_all,
2492 b,
2493 rows,
2494 cols,
2495 out,
2496 pool,
2497 );
2498 crate::gpu::probe_record(class, false, t0.elapsed());
2499 return;
2500 }
2501 crate::gpu::ProbeArm::Cpu => {}
2502 }
2503 }
2504 }
2505 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2506 return;
2507 }
2508 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2509 let w = self.bin_w().expect("binary dtype");
2513 if b >= 32
2514 && b * rows * cols >= 128_000_000
2515 && cols % 64 == 0
2516 && crate::gpu::enabled_here()
2517 {
2518 if let Self::Mapped { model, idx, .. } = self {
2519 let t0 = std::time::Instant::now();
2520 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2521 crate::gpu::ProbeArm::Gpu => {
2522 let ok = if *dtype == TensorDtype::Q1G {
2523 crate::gpu::q1g_matmat(
2524 model, *idx, xs_all, b, rows, cols, out,
2525 )
2526 } else {
2527 crate::gpu::q1_matmat(
2528 model, *idx, xs_all, b, rows, cols, out,
2529 )
2530 };
2531 if ok {
2532 crate::gpu::probe_record(
2533 crate::gpu::OpClass::Matmat,
2534 true,
2535 t0.elapsed(),
2536 );
2537 return;
2538 }
2539 }
2540 crate::gpu::ProbeArm::CpuTimed => {
2541 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2542 crate::gpu::probe_record(
2543 crate::gpu::OpClass::Matmat,
2544 false,
2545 t0.elapsed(),
2546 );
2547 return;
2548 }
2549 crate::gpu::ProbeArm::Cpu => {}
2550 }
2551 }
2552 }
2553 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2554 return;
2555 }
2556 if *dtype == TensorDtype::Q1T {
2557 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2560 if let Self::Mapped { model, idx, .. } = self {
2561 let t0 = std::time::Instant::now();
2562 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2563 crate::gpu::ProbeArm::Gpu => {
2564 if crate::gpu::q1t_matmat(
2565 model, *idx, xs_all, b, rows, cols, out,
2566 ) {
2567 crate::gpu::probe_record(
2568 crate::gpu::OpClass::Matmat,
2569 true,
2570 t0.elapsed(),
2571 );
2572 return;
2573 }
2574 }
2575 crate::gpu::ProbeArm::CpuTimed => {
2576 q1t_matmat(
2577 self.quant_bytes(),
2578 xs_all,
2579 b,
2580 rows,
2581 cols,
2582 out,
2583 pool,
2584 );
2585 crate::gpu::probe_record(
2586 crate::gpu::OpClass::Matmat,
2587 false,
2588 t0.elapsed(),
2589 );
2590 return;
2591 }
2592 crate::gpu::ProbeArm::Cpu => {}
2593 }
2594 }
2595 }
2596 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2597 return;
2598 }
2599 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2600 vbitmatmat(
2601 self.quant_bytes(),
2602 vbit_offsets,
2603 xs_all,
2604 b,
2605 rows,
2606 cols,
2607 out,
2608 pool,
2609 );
2610 return;
2611 }
2612 let _pmp = crate::cpuprof::spans::span("mm.prescale");
2613 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2614 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2615 .collect();
2616 drop(_pmp);
2617 if row_exact()
2623 && (1..=4).contains(&b)
2624 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2625 && crate::gpu::enabled_here()
2626 && crate::gpu::wgpu_active()
2627 {
2628 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2629 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2630 return;
2631 }
2632 }
2633 if b >= 8
2645 && b * rows * cols >= 128_000_000
2646 && crate::gpu::enabled_here()
2647 && std::env::var("CMF_MATMAT_GPU").as_deref() != Ok("0")
2648 {
2649 if let Self::Mapped { model, idx, .. } = self {
2650 MATMAT_GATE.fetch_add(1, Ordering::Relaxed);
2651 let t0 = std::time::Instant::now();
2652 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2653 crate::gpu::ProbeArm::Gpu
2654 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2655 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2656 {
2657 let q = self.quant_bytes();
2661 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2662 return;
2663 }
2664 crate::gpu::ProbeArm::Gpu => {
2665 if *dtype == TensorDtype::Q8_2f
2673 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2674 && crate::gpu::q8_matmat_2f(
2675 model, *idx, row_scale, col_field, xs_all, b, rows,
2676 cols, out,
2677 )
2678 {
2679 MATMAT_GPU_HIT.fetch_add(1, Ordering::Relaxed);
2680 crate::gpu::probe_record(
2681 crate::gpu::OpClass::Matmat,
2682 true,
2683 t0.elapsed(),
2684 );
2685 return;
2686 }
2687 let flat: Vec<f32> =
2688 pre.iter().flat_map(|v| v.iter().copied()).collect();
2689 if crate::gpu::q8_matmat(
2690 model, *idx, row_scale, &flat, b, rows, cols, out,
2691 ) {
2692 crate::gpu::probe_record(
2693 crate::gpu::OpClass::Matmat,
2694 true,
2695 t0.elapsed(),
2696 );
2697 return;
2698 }
2699 }
2700 crate::gpu::ProbeArm::CpuTimed => {
2701 let q = self.quant_bytes();
2702 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2703 crate::gpu::probe_record(
2704 crate::gpu::OpClass::Matmat,
2705 false,
2706 t0.elapsed(),
2707 );
2708 return;
2709 }
2710 crate::gpu::ProbeArm::Cpu => {}
2711 }
2712 }
2713 }
2714 let q = self.quant_bytes();
2715 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2716 }
2717 }
2718 }
2719}
2720
2721impl QTensor {
2722 fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2726 let (rows, cols) = (self.rows(), self.cols());
2727 if let Self::Mapped { model, idx, .. } = self {
2728 if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2729 let t0 = std::time::Instant::now();
2730 let cls = crate::gpu::matvec_class(rows, cols);
2731 match crate::gpu::probe_arm(cls) {
2732 crate::gpu::ProbeArm::Gpu => {
2733 if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2734 crate::gpu::probe_record(cls, true, t0.elapsed());
2735 return;
2736 }
2737 }
2738 crate::gpu::ProbeArm::CpuTimed => {
2739 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2740 crate::gpu::probe_record(cls, false, t0.elapsed());
2741 return;
2742 }
2743 crate::gpu::ProbeArm::Cpu => {}
2744 }
2745 }
2746 }
2747 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2748 }
2749
2750 fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2755 let (rows, cols) = (self.rows(), self.cols());
2756 if let Self::Mapped { model, idx, .. } = self {
2757 if b >= 2
2758 && b * rows * cols >= 33_554_432
2759 && !row_exact()
2760 && !crate::gpu::mm_killed()
2761 && crate::gpu::enabled_here()
2762 {
2763 let class = if b >= 128 {
2764 crate::gpu::OpClass::MatmatWide
2765 } else {
2766 crate::gpu::OpClass::Matmat
2767 };
2768 let t0 = std::time::Instant::now();
2769 match crate::gpu::probe_arm(class) {
2770 crate::gpu::ProbeArm::Gpu => {
2771 if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2772 crate::gpu::probe_record(class, true, t0.elapsed());
2773 return;
2774 }
2775 }
2776 crate::gpu::ProbeArm::CpuTimed => {
2777 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2778 crate::gpu::probe_record(class, false, t0.elapsed());
2779 return;
2780 }
2781 crate::gpu::ProbeArm::Cpu => {}
2782 }
2783 }
2784 }
2785 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2786 }
2787}
2788
2789impl QTensor {
2790 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2800 let (rows, cols) = (self.rows(), self.cols());
2801 let Self::Mapped {
2802 model,
2803 idx,
2804 dtype,
2805 row_scale,
2806 col_field,
2807 ..
2808 } = self
2809 else {
2810 return false;
2811 };
2812 if crate::prism::has_contract(model) {
2813 return false;
2814 }
2815 match *dtype {
2816 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2817 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2821 if *dtype == TensorDtype::Q8_2f
2824 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2825 && crate::gpu::q8_matmat_2f(
2826 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2827 )
2828 {
2829 return true;
2830 }
2831 let flat: Vec<f32> = (0..b)
2832 .flat_map(|bi| {
2833 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2834 })
2835 .collect();
2836 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2837 }
2838 _ => false,
2839 }
2840 }
2841
2842 pub fn matvec_many<const N: usize>(
2849 ts: [&QTensor; N],
2850 x: &[f32],
2851 mut outs: [&mut [f32]; N],
2852 pool: Option<&Pool>,
2853 ) {
2854 if Self::any_flint_exact_cpu_default(&ts) {
2855 return float_activations_scope(|| Self::matvec_many(ts, x, outs, pool));
2856 }
2857 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2858 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2859 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2864 t.matvec(x, o, pool);
2865 }
2866 return;
2867 }
2868 let uniform_q8 = ts.iter().all(|t| {
2869 matches!(
2870 t,
2871 Self::Mapped {
2872 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2873 ..
2874 }
2875 )
2876 });
2877 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2878 if uniform_f32 && crate::f32_backend::active() {
2879 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2880 t.matvec(x, o, pool);
2881 }
2882 return;
2883 }
2884 let uniform_q4 = ts.iter().all(|t| {
2885 matches!(
2886 t,
2887 Self::Mapped {
2888 dtype: TensorDtype::Q4Block,
2889 ..
2890 }
2891 )
2892 });
2893 let uniform_vbit = ts.iter().all(|t| {
2894 matches!(
2895 t,
2896 Self::Mapped {
2897 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2898 ..
2899 }
2900 )
2901 });
2902 let uniform_q1 =
2905 ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2906 let uniform_q1t = ts.iter().all(|t| {
2907 matches!(
2908 t,
2909 Self::Mapped {
2910 dtype: TensorDtype::Q1T,
2911 ..
2912 }
2913 )
2914 });
2915 let uniform_q4tp = ts.iter().all(|t| {
2920 matches!(
2921 t,
2922 Self::Mapped {
2923 dtype: TensorDtype::Q4TiledP,
2924 ..
2925 }
2926 )
2927 }) && ts
2928 .iter()
2929 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2930 let Some(pool) = pool else {
2931 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2932 t.matvec(x, o, None);
2933 }
2934 return;
2935 };
2936 if total_rows < 256
2937 || !(uniform_q8
2938 || uniform_f32
2939 || uniform_q4
2940 || uniform_vbit
2941 || uniform_q1
2942 || uniform_q1t
2943 || uniform_q4tp)
2944 {
2945 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2946 t.matvec(x, o, Some(pool));
2947 }
2948 return;
2949 }
2950
2951 if uniform_q4tp {
2952 let cols = ts[0].cols();
2958 let gpr = cols / GROUP_SIZE;
2959 let views: [Q4tpView; N] =
2960 std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2961 let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2962 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2963 let locate = |flat: usize| -> (usize, usize) {
2965 let mut acc = 0;
2966 for (i, &r) in rows_of.iter().enumerate() {
2967 if flat < acc + r {
2968 return (i, flat - acc);
2969 }
2970 acc += r;
2971 }
2972 (rows_of.len() - 1, 0)
2973 };
2974 let (views, outs_addr) = (&views, &outs_addr);
2975 if a8w8_enabled() {
2976 let act = split_act(x);
2977 let act = &act;
2978 #[cfg(target_arch = "x86_64")]
2979 let lanes = q4tp_lanes();
2980 let run = |start: usize, end: usize| {
2981 #[cfg(target_arch = "x86_64")]
2982 if lanes > 1 && act.xg8.len() == gpr {
2983 with_krow(gpr * lanes, |sc| {
2986 unsafe {
2989 match lanes {
2990 2 => q4tp_many_lanes::<2>(
2991 &views[..],
2992 &locate,
2993 start,
2994 end,
2995 gpr,
2996 &act.xq,
2997 &act.xg8,
2998 act.sx,
2999 &act.outliers,
3000 sc,
3001 &outs_addr[..],
3002 ),
3003 8 => q4tp_many_lanes::<8>(
3004 &views[..],
3005 &locate,
3006 start,
3007 end,
3008 gpr,
3009 &act.xq,
3010 &act.xg8,
3011 act.sx,
3012 &act.outliers,
3013 sc,
3014 &outs_addr[..],
3015 ),
3016 _ => q4tp_many_lanes::<4>(
3017 &views[..],
3018 &locate,
3019 start,
3020 end,
3021 gpr,
3022 &act.xq,
3023 &act.xg8,
3024 act.sx,
3025 &act.outliers,
3026 sc,
3027 &outs_addr[..],
3028 ),
3029 }
3030 }
3031 });
3032 return;
3033 }
3034 crate::cpuprof::count_plain_rows((end - start) as u64);
3035 with_krow(gpr, |sc| {
3036 for flat in start..end {
3037 let (t, r) = locate(flat);
3038 let v = &views[t];
3039 v.scales_into(r, gpr, sc);
3040 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
3041 for &(j, xv) in &act.outliers {
3042 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3043 acc += w * s * xv;
3044 }
3045 unsafe { *outs_addr[t].at(r) = acc };
3047 }
3048 });
3049 };
3050 pool.run_rows(total_rows, &run);
3051 } else {
3052 let run = |start: usize, end: usize| {
3053 with_krow(gpr, |sc| {
3054 for flat in start..end {
3055 let (t, r) = locate(flat);
3056 let v = &views[t];
3057 v.scales_into(r, gpr, sc);
3058 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
3060 }
3061 });
3062 };
3063 pool.run_rows(total_rows, &run);
3064 }
3065 return;
3066 }
3067
3068 if uniform_q1 {
3069 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3072 if a8w8_enabled() {
3073 let act = split_act(x);
3074 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
3075 let (act, gsum) = (&act, &gsum);
3076 let closures: [_; N] = std::array::from_fn(|i| {
3077 let (w, gpr, out) = (
3078 ts[i].bin_w().expect("uniform_q1 checked"),
3079 ts[i].cols() / GROUP_SIZE,
3080 outs_addr[i],
3081 );
3082 move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
3083 });
3084 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3085 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3086 pool.run_many(&parts);
3087 } else {
3088 let closures: [_; N] = std::array::from_fn(|i| {
3089 let (w, gpr, out) = (
3090 ts[i].bin_w().expect("uniform_q1 checked"),
3091 ts[i].cols() / GROUP_SIZE,
3092 outs_addr[i],
3093 );
3094 move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
3095 });
3096 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3097 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3098 pool.run_many(&parts);
3099 }
3100 return;
3101 }
3102
3103 if uniform_q1t {
3104 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3108 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3109 if a8w8_enabled() {
3110 let act = split_act(x);
3111 let act = &act;
3112 let x_ref = x;
3113 let closures: [_; N] = std::array::from_fn(|i| {
3114 let bytes = ts[i].quant_bytes();
3115 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3116 let gpr = cols / GROUP_SIZE;
3117 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3118 let out = outs_addr[i];
3119 move |s: usize, e: usize| {
3120 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
3121 }
3122 });
3123 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3124 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3125 pool.run_many(&parts);
3126 } else {
3127 let x_ref = x;
3128 let closures: [_; N] = std::array::from_fn(|i| {
3129 let bytes = ts[i].quant_bytes();
3130 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3131 let gpr = cols / GROUP_SIZE;
3132 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3133 let out = outs_addr[i];
3134 move |s: usize, e: usize| {
3135 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3136 }
3137 });
3138 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3139 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3140 pool.run_many(&parts);
3141 }
3142 return;
3143 }
3144
3145 if uniform_q4 || uniform_vbit {
3146 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3147 if a8w8_enabled() {
3149 let act = split_act(x);
3150 let act = &act;
3151 if uniform_q4 {
3152 let closures: [_; N] = std::array::from_fn(|i| {
3153 let (packed, scales) =
3154 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3155 let (gpr, cols, out) =
3156 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3157 move |s: usize, e: usize| {
3158 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3159 }
3160 });
3161 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3162 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3163 pool.run_many(&parts);
3164 } else {
3165 let closures: [_; N] = std::array::from_fn(|i| {
3166 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3167 unreachable!()
3168 };
3169 let (bytes, rows, cols, out) = (
3170 ts[i].quant_bytes(),
3171 ts[i].rows(),
3172 ts[i].cols(),
3173 outs_addr[i],
3174 );
3175 move |s: usize, e: usize| {
3176 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3177 }
3178 });
3179 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3180 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3181 pool.run_many(&parts);
3182 }
3183 return;
3184 }
3185 if uniform_q4 {
3186 let closures: [_; N] = std::array::from_fn(|i| {
3187 let (packed, scales) =
3188 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3189 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3190 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3191 });
3192 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3193 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3194 pool.run_many(&parts);
3195 } else {
3196 let closures: [_; N] = std::array::from_fn(|i| {
3197 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3198 unreachable!()
3199 };
3200 let (bytes, rows, cols, out) = (
3201 ts[i].quant_bytes(),
3202 ts[i].rows(),
3203 ts[i].cols(),
3204 outs_addr[i],
3205 );
3206 move |s: usize, e: usize| {
3207 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3208 }
3209 });
3210 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3211 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3212 pool.run_many(&parts);
3213 }
3214 return;
3215 }
3216
3217 if uniform_f32 {
3218 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3219 let closures: [_; N] = std::array::from_fn(|i| {
3220 let Self::F32 { data, cols, .. } = ts[i] else {
3221 unreachable!()
3222 };
3223 let out = outs_addr[i];
3224 move |start: usize, end: usize| {
3225 for o in start..end {
3226 let row = &data[o * cols..(o + 1) * cols];
3227 let mut sum = 0.0f32;
3228 for j in 0..*cols {
3229 sum += row[j] * x[j];
3230 }
3231 unsafe { *out.at(o) = sum };
3233 }
3234 }
3235 });
3236 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3237 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3238 pool.run_many(&parts);
3239 return;
3240 }
3241
3242 struct Ctx<'a> {
3245 bytes: &'a [u8],
3246 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3247 rep: &'a [u8],
3248 row_scale: &'a [f32],
3249 cols: usize,
3250 xs: std::borrow::Cow<'a, [f32]>,
3251 }
3252 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3253 let Self::Mapped {
3254 dtype,
3255 cols,
3256 row_scale,
3257 col_field,
3258 repack,
3259 ..
3260 } = ts[i]
3261 else {
3262 unreachable!()
3263 };
3264 Ctx {
3265 bytes: ts[i].quant_bytes(),
3266 rep: repack,
3267 row_scale,
3268 cols: *cols,
3269 xs: prescale(x, col_field, *dtype),
3270 }
3271 });
3272 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3273 #[cfg(target_arch = "aarch64")]
3274 if sdot_enabled() {
3275 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3276 let closures: [_; N] = std::array::from_fn(|i| {
3277 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3278 move |start: usize, end: usize| {
3279 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3280 }
3281 });
3282 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3283 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3284 pool.run_many(&parts);
3285 return;
3286 }
3287 #[cfg(target_arch = "x86_64")]
3288 if avx2_a8w8_enabled() {
3289 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3290 let closures: [_; N] = std::array::from_fn(|i| {
3291 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3292 move |start: usize, end: usize| {
3293 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3294 }
3295 });
3296 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3297 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3298 pool.run_many(&parts);
3299 return;
3300 }
3301 let closures: [_; N] = std::array::from_fn(|i| {
3302 let (c, out) = (&ctxs[i], outs_addr[i]);
3303 move |start: usize, end: usize| {
3304 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3305 }
3306 });
3307 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3308 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3309 pool.run_many(&parts);
3310 }
3311}
3312
3313impl QTensor {
3314 #[allow(clippy::needless_range_loop)]
3319 pub fn matvec2_many<const N: usize>(
3320 ts: [&QTensor; N],
3321 x1: &[f32],
3322 x2: &[f32],
3323 mut o1s: [&mut [f32]; N],
3324 mut o2s: [&mut [f32]; N],
3325 pool: Option<&Pool>,
3326 ) {
3327 if Self::any_flint_exact_cpu_default(&ts) {
3328 return float_activations_scope(|| {
3329 Self::matvec2_many(ts, x1, x2, o1s, o2s, pool)
3330 });
3331 }
3332 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3333 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3334 for i in 0..N {
3335 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3336 }
3337 return;
3338 }
3339 let uniform_q8 = ts.iter().all(|t| {
3340 matches!(
3341 t,
3342 Self::Mapped {
3343 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3344 ..
3345 }
3346 )
3347 });
3348 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3349 let uniform_q4 = ts.iter().all(|t| {
3350 matches!(
3351 t,
3352 Self::Mapped {
3353 dtype: TensorDtype::Q4Block,
3354 ..
3355 }
3356 )
3357 });
3358 let uniform_vbit = ts.iter().all(|t| {
3359 matches!(
3360 t,
3361 Self::Mapped {
3362 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3363 ..
3364 }
3365 )
3366 });
3367 let fusable = pool.is_some()
3368 && total_rows >= 256
3369 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3370 if !fusable {
3371 for i in 0..N {
3372 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3373 }
3374 return;
3375 }
3376 let pool = pool.unwrap();
3377
3378 if uniform_q4 || uniform_vbit {
3379 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3380 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3381 if a8w8_enabled() {
3383 let a1 = split_act(x1);
3384 let a2 = split_act(x2);
3385 let (a1, a2) = (&a1, &a2);
3386 if uniform_q4 {
3387 let closures: [_; N] = std::array::from_fn(|i| {
3388 let (packed, scales) =
3389 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3390 let (gpr, cols, o1, o2) =
3391 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3392 move |s: usize, e: usize| {
3393 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3394 }
3395 });
3396 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3397 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3398 pool.run_many(&parts);
3399 } else {
3400 let closures: [_; N] = std::array::from_fn(|i| {
3401 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3402 unreachable!()
3403 };
3404 let (bytes, rows, cols, o1, o2) = (
3405 ts[i].quant_bytes(),
3406 ts[i].rows(),
3407 ts[i].cols(),
3408 p1[i],
3409 p2[i],
3410 );
3411 move |s: usize, e: usize| {
3412 vbit_range2_a8w8(
3413 bytes,
3414 vbit_offsets,
3415 x1,
3416 x2,
3417 a1,
3418 a2,
3419 rows,
3420 cols,
3421 o1,
3422 o2,
3423 s,
3424 e,
3425 )
3426 }
3427 });
3428 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3429 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3430 pool.run_many(&parts);
3431 }
3432 return;
3433 }
3434 if uniform_q4 {
3435 let closures: [_; N] = std::array::from_fn(|i| {
3436 let (packed, scales) =
3437 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3438 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3439 move |s: usize, e: usize| {
3440 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3441 }
3442 });
3443 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3444 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3445 pool.run_many(&parts);
3446 } else {
3447 let closures: [_; N] = std::array::from_fn(|i| {
3448 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3449 unreachable!()
3450 };
3451 let (bytes, rows, cols, o1, o2) = (
3452 ts[i].quant_bytes(),
3453 ts[i].rows(),
3454 ts[i].cols(),
3455 p1[i],
3456 p2[i],
3457 );
3458 move |s: usize, e: usize| {
3459 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3460 }
3461 });
3462 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3463 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3464 pool.run_many(&parts);
3465 }
3466 return;
3467 }
3468
3469 if uniform_f32 {
3470 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3471 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3472 let closures: [_; N] = std::array::from_fn(|i| {
3473 let Self::F32 { data, cols, .. } = ts[i] else {
3474 unreachable!()
3475 };
3476 let (o1, o2) = (p1[i], p2[i]);
3477 move |start: usize, end: usize| {
3478 for o in start..end {
3479 let row = &data[o * cols..(o + 1) * cols];
3480 let (mut s1, mut s2) = (0.0f32, 0.0f32);
3481 for j in 0..*cols {
3482 s1 += row[j] * x1[j];
3483 s2 += row[j] * x2[j];
3484 }
3485 unsafe {
3487 *o1.at(o) = s1;
3488 *o2.at(o) = s2;
3489 }
3490 }
3491 }
3492 });
3493 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3494 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3495 pool.run_many(&parts);
3496 return;
3497 }
3498
3499 struct Ctx<'a> {
3500 bytes: &'a [u8],
3501 row_scale: &'a [f32],
3502 cols: usize,
3503 xs1: std::borrow::Cow<'a, [f32]>,
3504 xs2: std::borrow::Cow<'a, [f32]>,
3505 }
3506 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3507 let Self::Mapped {
3508 dtype,
3509 cols,
3510 row_scale,
3511 col_field,
3512 ..
3513 } = ts[i]
3514 else {
3515 unreachable!()
3516 };
3517 Ctx {
3518 bytes: ts[i].quant_bytes(),
3519 row_scale,
3520 cols: *cols,
3521 xs1: prescale(x1, col_field, *dtype),
3522 xs2: prescale(x2, col_field, *dtype),
3523 }
3524 });
3525 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3526 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3527 #[cfg(target_arch = "aarch64")]
3528 if sdot_enabled() {
3529 let acts: [(SplitAct, SplitAct); N] =
3530 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3531 let closures: [_; N] = std::array::from_fn(|i| {
3532 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3533 move |start: usize, end: usize| {
3534 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3535 }
3536 });
3537 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3538 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3539 pool.run_many(&parts);
3540 return;
3541 }
3542 #[cfg(target_arch = "x86_64")]
3543 if avx2_a8w8_enabled() {
3544 let acts: [(SplitAct, SplitAct); N] =
3545 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3546 let closures: [_; N] = std::array::from_fn(|i| {
3547 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3548 move |start: usize, end: usize| {
3549 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3550 }
3551 });
3552 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3553 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3554 pool.run_many(&parts);
3555 return;
3556 }
3557 let closures: [_; N] = std::array::from_fn(|i| {
3558 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3559 move |start: usize, end: usize| {
3560 q8_range2_f32(
3561 c.bytes,
3562 c.row_scale,
3563 &c.xs1,
3564 &c.xs2,
3565 c.cols,
3566 o1,
3567 o2,
3568 start,
3569 end,
3570 )
3571 }
3572 });
3573 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3574 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3575 pool.run_many(&parts);
3576 }
3577
3578 pub fn matvec_silu_mul(
3583 gate: &QTensor,
3584 up: &QTensor,
3585 x: &[f32],
3586 out: &mut [f32],
3587 pool: Option<&Pool>,
3588 ) -> bool {
3589 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3590 }
3591
3592 pub fn matvec_silu_mul_limited(
3598 gate: &QTensor,
3599 up: &QTensor,
3600 x: &[f32],
3601 out: &mut [f32],
3602 limit: f32,
3603 pool: Option<&Pool>,
3604 ) -> bool {
3605 if gate.flint_exact_cpu_default() || up.flint_exact_cpu_default() {
3606 return float_activations_scope(|| {
3607 Self::matvec_silu_mul_limited(gate, up, x, out, limit, pool)
3608 });
3609 }
3610 if gate.has_prism_contract()
3611 || up.has_prism_contract()
3612 || gate.per_tensor_only()
3613 || up.per_tensor_only()
3614 {
3615 return false;
3619 }
3620 let inter = gate.rows();
3621 debug_assert_eq!(up.rows(), inter);
3622 debug_assert_eq!(out.len(), inter);
3623 debug_assert_eq!(gate.cols(), up.cols());
3624 if !a8w8_enabled() {
3625 return false;
3626 }
3627 let act = split_act(x);
3628 let act = &act;
3629 let x_ref = x;
3630 let out_addr = SendMut(out.as_mut_ptr());
3631
3632 match (gate, up) {
3633 (
3635 Self::Mapped {
3636 dtype: TensorDtype::Q4Block,
3637 ..
3638 },
3639 Self::Mapped {
3640 dtype: TensorDtype::Q4Block,
3641 ..
3642 },
3643 ) => {
3644 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3645 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3646 let gpr = gate.cols() / GROUP_SIZE;
3647 let cols = gate.cols();
3648 let run = move |start: usize, end: usize| {
3649 for r in start..end {
3650 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3651 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3652 for &(j, xv) in &act.outliers {
3653 let flat = r * cols + j;
3654 let gb = gp[flat / 2];
3655 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3656 let gsc = f16_to_f32(u16::from_le_bytes([
3657 gs[(flat / GROUP_SIZE) * 2],
3658 gs[(flat / GROUP_SIZE) * 2 + 1],
3659 ]));
3660 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3661 let ub = up_p[flat / 2];
3662 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3663 let usc = f16_to_f32(u16::from_le_bytes([
3664 up_s[(flat / GROUP_SIZE) * 2],
3665 up_s[(flat / GROUP_SIZE) * 2 + 1],
3666 ]));
3667 uv += ((un as i32 - 8) as f32) * usc * xv;
3668 }
3669 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3671 }
3672 };
3673 dispatch_rows(pool, inter, &run);
3674 true
3675 }
3676 (
3680 Self::Mapped {
3681 dtype: TensorDtype::Q4Tiled,
3682 ..
3683 },
3684 Self::Mapped {
3685 dtype: TensorDtype::Q4Tiled,
3686 ..
3687 },
3688 ) => {
3689 let g_bytes = gate.quant_bytes();
3690 let u_bytes = up.quant_bytes();
3691 let gpr = gate.cols() / GROUP_SIZE;
3692 let run = move |start: usize, end: usize| {
3693 for r in start..end {
3694 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3695 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3696 for &(j, xv) in &act.outliers {
3697 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3698 gv += w * s * xv;
3699 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3700 uv += w * s * xv;
3701 }
3702 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3704 }
3705 };
3706 dispatch_rows(pool, inter, &run);
3707 true
3708 }
3709 (
3712 Self::Mapped {
3713 dtype: TensorDtype::Q4TiledP,
3714 ..
3715 },
3716 Self::Mapped {
3717 dtype: TensorDtype::Q4TiledP,
3718 ..
3719 },
3720 ) => {
3721 let cols = gate.cols();
3722 let gpr = cols / GROUP_SIZE;
3723 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3724 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3725 #[cfg(target_arch = "x86_64")]
3726 let lanes = q4tp_lanes();
3727 let run = |start: usize, end: usize| {
3728 #[cfg(target_arch = "x86_64")]
3729 if lanes > 1 && act.xg8.len() == gpr {
3730 with_krows(gpr * lanes, |gsc, usc| {
3733 unsafe {
3735 match lanes {
3736 2 => q4tp_gateup_lanes::<2>(
3737 &gv_view, &uv_view, start, end, gpr, &act.xq,
3738 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3739 out_addr,
3740 ),
3741 8 => q4tp_gateup_lanes::<8>(
3742 &gv_view, &uv_view, start, end, gpr, &act.xq,
3743 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3744 out_addr,
3745 ),
3746 _ => q4tp_gateup_lanes::<4>(
3747 &gv_view, &uv_view, start, end, gpr, &act.xq,
3748 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3749 out_addr,
3750 ),
3751 }
3752 }
3753 });
3754 return;
3755 }
3756 crate::cpuprof::count_plain_rows((end - start) as u64);
3757 with_krows(gpr, |gsc, usc| {
3758 for r in start..end {
3759 gv_view.scales_into(r, gpr, gsc);
3760 uv_view.scales_into(r, gpr, usc);
3761 let mut gv =
3762 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3763 let mut uv =
3764 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3765 for &(j, xv) in &act.outliers {
3766 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3767 gv += w * s * xv;
3768 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3769 uv += w * s * xv;
3770 }
3771 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3773 }
3774 });
3775 };
3776 dispatch_rows(pool, inter, &run);
3777 true
3778 }
3779 (
3786 Self::Mapped {
3787 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3788 ..
3789 },
3790 Self::Mapped {
3791 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3792 ..
3793 },
3794 ) => {
3795 let gw = gate.bin_w().expect("binary dtype matched");
3796 let uw = up.bin_w().expect("binary dtype matched");
3797 let gpr = gate.cols() / GROUP_SIZE;
3798 let gsum = q1_group_sums(&act.xq, gpr);
3799 let gsum = &gsum;
3800 let run = move |start: usize, end: usize| {
3801 for r in start..end {
3802 let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3803 let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3804 for &(j, xv) in &act.outliers {
3805 let (w, s) = gw.outlier(r, gpr, j);
3806 gv += w * s * xv;
3807 let (w, s) = uw.outlier(r, gpr, j);
3808 uv += w * s * xv;
3809 }
3810 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3812 }
3813 };
3814 dispatch_rows(pool, inter, &run);
3815 true
3816 }
3817 (
3821 Self::Mapped {
3822 dtype: TensorDtype::Q2TiledP,
3823 ..
3824 },
3825 Self::Mapped {
3826 dtype: TensorDtype::Q2TiledP,
3827 ..
3828 },
3829 ) => {
3830 let cols = gate.cols();
3831 let gpr = cols / GROUP_SIZE;
3832 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3833 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3834 let gsum = q1_group_sums(&act.xq, gpr);
3835 let gsum = &gsum;
3836 let run = move |start: usize, end: usize| {
3837 with_krows(gpr, |gsc, usc| {
3838 for r in start..end {
3839 gv_view.scales_into(r, gpr, gsc);
3840 uv_view.scales_into(r, gpr, usc);
3841 let mut gv =
3842 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3843 let mut uv =
3844 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3845 for &(j, xv) in &act.outliers {
3846 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3847 gv += w * s * xv;
3848 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3849 uv += w * s * xv;
3850 }
3851 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3853 }
3854 });
3855 };
3856 dispatch_rows(pool, inter, &run);
3857 true
3858 }
3859 (
3864 Self::Mapped {
3865 dtype: TensorDtype::Q8Row,
3866 row_scale: g_rs,
3867 ..
3868 },
3869 Self::Mapped {
3870 dtype: TensorDtype::Q8Row,
3871 row_scale: u_rs,
3872 ..
3873 },
3874 ) => {
3875 let g_bytes = gate.quant_bytes();
3876 let u_bytes = up.quant_bytes();
3877 let cols = gate.cols();
3878 let run = move |start: usize, end: usize| {
3879 for r in start..end {
3880 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3881 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3882 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3884 }
3885 };
3886 dispatch_rows(pool, inter, &run);
3887 true
3888 }
3889 (
3891 Self::Mapped {
3892 dtype: TensorDtype::Q1T,
3893 ..
3894 },
3895 Self::Mapped {
3896 dtype: TensorDtype::Q1T,
3897 ..
3898 },
3899 ) => {
3900 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3901 let g_bytes = gate.quant_bytes();
3902 let u_bytes = up.quant_bytes();
3903 let gpr = gate.cols() / GROUP_SIZE;
3904 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3905 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3906 let run = move |start: usize, end: usize| {
3907 for r in start..end {
3908 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3909 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3910 for &(j, xv) in &act.outliers {
3911 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3912 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3913 }
3914 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3915 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3916 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3918 }
3919 };
3920 dispatch_rows(pool, inter, &run);
3921 true
3922 }
3923 _ => false,
3924 }
3925 }
3926
3927 pub fn moe_gate_up_many(
3942 pairs: &[(&QTensor, &QTensor)],
3943 x: &[f32],
3944 outs: &mut [Vec<f32>],
3945 pool: Option<&Pool>,
3946 ) -> bool {
3947 Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3948 }
3949
3950 pub fn moe_gate_up_many_limited(
3954 pairs: &[(&QTensor, &QTensor)],
3955 x: &[f32],
3956 outs: &mut [Vec<f32>],
3957 limit: f32,
3958 pool: Option<&Pool>,
3959 ) -> bool {
3960 if Self::any_flint_exact_cpu_default_pairs(pairs) {
3961 return float_activations_scope(|| {
3962 Self::moe_gate_up_many_limited(pairs, x, outs, limit, pool)
3963 });
3964 }
3965 if pairs.is_empty() || pairs.len() != outs.len() {
3966 return false;
3967 }
3968 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
3969 return false;
3970 }
3971 if !a8w8_enabled() {
3972 if limit > 0.0 {
3973 return false;
3976 }
3977 let groups = vec![vec![0]; pairs.len()];
3978 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3979 }
3980 let inter = pairs[0].0.rows();
3981 let cols = pairs[0].0.cols();
3982 if cols % GROUP_SIZE != 0 {
3983 return false;
3984 }
3985 let gpr = cols / GROUP_SIZE;
3986 let q2 = matches!(
3989 pairs[0].0,
3990 Self::Mapped {
3991 dtype: TensorDtype::Q2TiledP,
3992 ..
3993 }
3994 );
3995 let want = if q2 {
3996 TensorDtype::Q2TiledP
3997 } else {
3998 TensorDtype::Q4TiledP
3999 };
4000 let mut views = Vec::with_capacity(pairs.len() * 2);
4001 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
4002 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
4003 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
4004 if !both
4005 || g.rows() != inter
4006 || u.rows() != inter
4007 || g.cols() != cols
4008 || u.cols() != cols
4009 || o.len() != inter
4010 {
4011 return false;
4012 }
4013 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
4014 views.push(mk(g.quant_bytes(), inter, cols));
4015 views.push(mk(u.quant_bytes(), inter, cols));
4016 }
4017 let act = split_act(x);
4018 let gsum = if q2 {
4019 q1_group_sums(&act.xq, gpr)
4020 } else {
4021 Vec::new()
4022 };
4023 let (act, gsum) = (&act, &gsum);
4024 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4025 let (views, ptrs) = (&views, &ptrs);
4026 let run = |start: usize, end: usize| {
4027 with_krows(gpr, |gsc, usc| {
4028 for flat in start..end {
4029 let (e, r) = (flat / inter, flat % inter);
4030 let gv_view = &views[e * 2];
4031 let uv_view = &views[e * 2 + 1];
4032 gv_view.scales_into(r, gpr, gsc);
4033 uv_view.scales_into(r, gpr, usc);
4034 let (mut gv, mut uv) = if q2 {
4035 (
4036 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
4037 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
4038 )
4039 } else {
4040 (
4041 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
4042 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
4043 )
4044 };
4045 for &(j, xv) in &act.outliers {
4046 let (og, ou) = if q2 {
4047 (
4048 q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
4049 q2tp_outlier(uv_view.nib, r, gpr, j, usc),
4050 )
4051 } else {
4052 (
4053 q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
4054 q4tp_outlier(uv_view.nib, r, gpr, j, usc),
4055 )
4056 };
4057 gv += og.0 * og.1 * xv;
4058 uv += ou.0 * ou.1 * xv;
4059 }
4060 unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
4062 }
4063 });
4064 };
4065 dispatch_rows(pool, pairs.len() * inter, &run);
4066 true
4067 }
4068
4069 pub fn moe_down_many(
4078 downs: &[&QTensor],
4079 gs: &[Vec<f32>],
4080 weights: &[f32],
4081 out: &mut [f32],
4082 pool: Option<&Pool>,
4083 ) -> bool {
4084 if Self::any_flint_exact_cpu_default(downs) {
4085 return float_activations_scope(|| {
4086 Self::moe_down_many(downs, gs, weights, out, pool)
4087 });
4088 }
4089 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
4090 return false;
4091 }
4092 if downs.iter().any(|d| d.per_tensor_only()) {
4093 return false;
4094 }
4095 if !a8w8_enabled() {
4096 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
4097 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
4098 return false;
4099 }
4100 out.fill(0.0);
4101 for (row, &w) in terms.iter().zip(weights) {
4102 for (o, &v) in out.iter_mut().zip(row) {
4103 *o += w * v;
4104 }
4105 }
4106 return true;
4107 }
4108 let rows = out.len();
4109 let cols = downs[0].cols();
4110 if cols % GROUP_SIZE != 0 {
4111 return false;
4112 }
4113 let gpr = cols / GROUP_SIZE;
4114 let mut views = Vec::with_capacity(downs.len());
4115 for (d, g) in downs.iter().zip(gs.iter()) {
4116 if !matches!(
4117 d,
4118 Self::Mapped {
4119 dtype: TensorDtype::Q4TiledP,
4120 ..
4121 }
4122 ) || d.rows() != rows
4123 || d.cols() != cols
4124 || g.len() != cols
4125 {
4126 return false;
4127 }
4128 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4129 }
4130 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
4132 let out_addr = SendMut(out.as_mut_ptr());
4140 let (views, acts, weights) = (&views, &acts, &weights);
4141 let run = |start: usize, end: usize| {
4142 with_krow(gpr, |sc| {
4143 for r in start..end {
4144 let mut acc = 0f32;
4145 for (e, v) in views.iter().enumerate() {
4146 v.scales_into(r, gpr, sc);
4147 let a = &acts[e];
4148 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4149 for &(j, xv) in &a.outliers {
4150 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4151 d += w * s * xv;
4152 }
4153 acc += weights[e] * d;
4154 }
4155 unsafe { *out_addr.at(r) = acc };
4157 }
4158 });
4159 };
4160 dispatch_rows(pool, rows, &run);
4161 true
4162 }
4163
4164 pub fn moe_gate_up_rows(
4176 pairs: &[(&QTensor, &QTensor)],
4177 groups: &[Vec<usize>],
4178 xs: &[f32],
4179 outs: &mut [Vec<f32>],
4180 pool: Option<&Pool>,
4181 ) -> bool {
4182 if Self::any_flint_exact_cpu_default_pairs(pairs) {
4183 return float_activations_scope(|| {
4184 Self::moe_gate_up_rows(pairs, groups, xs, outs, pool)
4185 });
4186 }
4187 if pairs.is_empty() || pairs.len() != groups.len() {
4188 return false;
4189 }
4190 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4191 return false;
4192 }
4193 let inter = pairs[0].0.rows();
4194 let cols = pairs[0].0.cols();
4195 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4196 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4197 return false;
4198 }
4199 let b = xs.len() / cols;
4200 let gpr = cols / GROUP_SIZE;
4201 let mut views = Vec::with_capacity(pairs.len() * 2);
4202 for (g, u) in pairs {
4203 let q4tp = |t: &QTensor| {
4204 matches!(
4205 t,
4206 Self::Mapped {
4207 dtype: TensorDtype::Q4TiledP,
4208 ..
4209 }
4210 )
4211 };
4212 if g.has_prism_contract()
4213 || u.has_prism_contract()
4214 || !q4tp(g)
4215 || !q4tp(u)
4216 || g.rows() != inter
4217 || u.rows() != inter
4218 || g.cols() != cols
4219 || u.cols() != cols
4220 {
4221 return false;
4222 }
4223 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4224 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4225 }
4226 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4227 return false;
4228 }
4229 let quantized = a8w8_enabled();
4230 let acts: Vec<SplitAct> = if quantized {
4231 (0..b)
4232 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4233 .collect()
4234 } else {
4235 Vec::new()
4236 };
4237 let mut offs = Vec::with_capacity(groups.len());
4238 let mut o = 0usize;
4239 for g in groups {
4240 offs.push(o);
4241 o += g.len();
4242 }
4243 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4244 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4245 let run = |start: usize, end: usize| {
4246 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4247 for flat in start..end {
4248 let (e, r) = (flat / inter, flat % inter);
4249 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4250 gv_view.scales_into(r, gpr, &mut gsc);
4251 uv_view.scales_into(r, gpr, &mut usc);
4252 for (k, &t) in groups[e].iter().enumerate() {
4253 if !quantized {
4254 let x = &xs[t * cols..(t + 1) * cols];
4255 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4256 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4257 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4258 continue;
4259 }
4260 let act = &acts[t];
4261 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4262 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4263 for &(j, xv) in &act.outliers {
4264 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4265 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4266 gv += og.0 * og.1 * xv;
4267 uv += ou.0 * ou.1 * xv;
4268 }
4269 let silu_g = gv / (1.0 + (-gv).exp());
4270 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4273 }
4274 }
4275 };
4276 dispatch_rows(pool, pairs.len() * inter, &run);
4277 true
4278 }
4279
4280 pub fn moe_down_rows(
4287 downs: &[&QTensor],
4288 group_lens: &[usize],
4289 gs: &[Vec<f32>],
4290 outs: &mut [Vec<f32>],
4291 pool: Option<&Pool>,
4292 ) -> bool {
4293 if Self::any_flint_exact_cpu_default(downs) {
4294 return float_activations_scope(|| {
4295 Self::moe_down_rows(downs, group_lens, gs, outs, pool)
4296 });
4297 }
4298 if downs.is_empty() || downs.len() != group_lens.len() {
4299 return false;
4300 }
4301 if downs.iter().any(|d| d.per_tensor_only()) {
4302 return false;
4303 }
4304 let rows = downs[0].rows();
4305 let cols = downs[0].cols();
4306 let n_pairs: usize = group_lens.iter().sum();
4307 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4308 return false;
4309 }
4310 let gpr = cols / GROUP_SIZE;
4311 let mut views = Vec::with_capacity(downs.len());
4312 for d in downs {
4313 if d.has_prism_contract()
4314 || !matches!(
4315 d,
4316 Self::Mapped {
4317 dtype: TensorDtype::Q4TiledP,
4318 ..
4319 }
4320 )
4321 || d.rows() != rows
4322 || d.cols() != cols
4323 {
4324 return false;
4325 }
4326 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4327 }
4328 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4329 return false;
4330 }
4331 let quantized = a8w8_enabled();
4332 let acts: Vec<SplitAct> = if quantized {
4333 gs.iter().map(|g| split_act(g)).collect()
4334 } else {
4335 Vec::new()
4336 };
4337 let mut offs = Vec::with_capacity(group_lens.len());
4338 let mut o = 0usize;
4339 for &l in group_lens {
4340 offs.push(o);
4341 o += l;
4342 }
4343 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4344 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4345 let run = |start: usize, end: usize| {
4346 let mut sc = vec![0f32; gpr];
4347 for flat in start..end {
4348 let (e, r) = (flat / rows, flat % rows);
4349 let v = &views[e];
4350 v.scales_into(r, gpr, &mut sc);
4351 for k in 0..group_lens[e] {
4352 if !quantized {
4353 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4354 unsafe { *ptrs[offs[e] + k].at(r) = d };
4355 continue;
4356 }
4357 let a = &acts[offs[e] + k];
4358 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4359 for &(j, xv) in &a.outliers {
4360 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4361 d += w * s * xv;
4362 }
4363 unsafe { *ptrs[offs[e] + k].at(r) = d };
4365 }
4366 }
4367 };
4368 dispatch_rows(pool, downs.len() * rows, &run);
4369 true
4370 }
4371}
4372
4373#[cfg(target_os = "macos")]
4378mod accel_blas {
4379 #[link(name = "Accelerate", kind = "framework")]
4380 unsafe extern "C" {
4381 pub fn cblas_sgemm(
4382 order: i32,
4383 trans_a: i32,
4384 trans_b: i32,
4385 m: i32,
4386 n: i32,
4387 k: i32,
4388 alpha: f32,
4389 a: *const f32,
4390 lda: i32,
4391 b: *const f32,
4392 ldb: i32,
4393 beta: f32,
4394 c: *mut f32,
4395 ldc: i32,
4396 );
4397 }
4398}
4399
4400#[cfg(target_os = "macos")]
4401pub(crate) fn accel_gemm_enabled() -> bool {
4402 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4403 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4404}
4405
4406#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4409pub(crate) fn accel_gemm_enabled() -> bool {
4410 true
4411}
4412
4413#[cfg(target_arch = "aarch64")]
4420#[allow(clippy::too_many_arguments)]
4421pub(crate) fn neon_gemm_rm(
4422 m: usize,
4423 n: usize,
4424 k: usize,
4425 alpha: f32,
4426 a: &[f32],
4427 lda: usize,
4428 b_mat: &[f32],
4429 ldb: usize,
4430 b_rows_are_n: bool,
4431 c: &mut [f32],
4432 ldc: usize,
4433) {
4434 debug_assert!(a.len() >= (m - 1) * lda + k);
4435 debug_assert!(c.len() >= (m - 1) * ldc + n);
4436 unsafe {
4438 use core::arch::aarch64::*;
4439 let mut i = 0usize;
4440 while i < m {
4441 let mi = (m - i).min(4);
4442 let mut j = 0usize;
4443 while j < n {
4444 let nj = (n - j).min(8);
4445 if mi == 4 && nj == 8 {
4446 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4447 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4448 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4449 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4450 for p in 0..k {
4451 let (b0, b1) = if b_rows_are_n {
4452 let base = b_mat.as_ptr().add(j * ldb + p);
4455 let g = |o: usize| *base.add(o * ldb);
4456 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4457 } else {
4458 let base = b_mat.as_ptr().add(p * ldb + j);
4459 (
4460 [*base, *base.add(1), *base.add(2), *base.add(3)],
4461 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4462 )
4463 };
4464 let bv0 = vld1q_f32(b0.as_ptr());
4465 let bv1 = vld1q_f32(b1.as_ptr());
4466 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4467 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4468 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4469 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4470 c0a = vfmaq_f32(c0a, a0, bv0);
4471 c0b = vfmaq_f32(c0b, a0, bv1);
4472 c1a = vfmaq_f32(c1a, a1, bv0);
4473 c1b = vfmaq_f32(c1b, a1, bv1);
4474 c2a = vfmaq_f32(c2a, a2, bv0);
4475 c2b = vfmaq_f32(c2b, a2, bv1);
4476 c3a = vfmaq_f32(c3a, a3, bv0);
4477 c3b = vfmaq_f32(c3b, a3, bv1);
4478 }
4479 let al = vdupq_n_f32(alpha);
4480 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4481 .iter()
4482 .enumerate()
4483 {
4484 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4485 vst1q_f32(dst, vmulq_f32(*ca, al));
4486 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4487 }
4488 } else {
4489 for r in 0..mi {
4490 for q in 0..nj {
4491 let mut acc = 0f32;
4492 for p in 0..k {
4493 let bv = if b_rows_are_n {
4494 b_mat[(j + q) * ldb + p]
4495 } else {
4496 b_mat[p * ldb + j + q]
4497 };
4498 acc += a[(i + r) * lda + p] * bv;
4499 }
4500 c[(i + r) * ldc + j + q] = acc * alpha;
4501 }
4502 }
4503 }
4504 j += nj;
4505 }
4506 i += mi;
4507 }
4508 }
4509}
4510
4511#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4513#[allow(clippy::too_many_arguments)]
4514pub(crate) fn sgemm_rm(
4515 m: usize,
4516 n: usize,
4517 k: usize,
4518 alpha: f32,
4519 a: &[f32],
4520 lda: usize,
4521 b_mat: &[f32],
4522 ldb: usize,
4523 b_rows_are_n: bool,
4524 c: &mut [f32],
4525 ldc: usize,
4526) {
4527 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4528}
4529
4530#[allow(clippy::too_many_arguments)]
4534pub fn sgemm_public(
4535 m: usize,
4536 n: usize,
4537 k: usize,
4538 alpha: f32,
4539 a: &[f32],
4540 lda: usize,
4541 b_mat: &[f32],
4542 ldb: usize,
4543 b_rows_are_n: bool,
4544 c: &mut [f32],
4545 ldc: usize,
4546) {
4547 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4548 {
4549 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4550 }
4551 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4556 {
4557 if m == 0 || n == 0 {
4558 return;
4559 }
4560 assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4561 let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4562 assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4563 unsafe {
4565 matrixmultiply::sgemm(
4566 m,
4567 k,
4568 n,
4569 alpha,
4570 a.as_ptr(),
4571 lda as isize,
4572 1,
4573 b_mat.as_ptr(),
4574 rsb,
4575 csb,
4576 0.0,
4577 c.as_mut_ptr(),
4578 ldc as isize,
4579 1,
4580 );
4581 }
4582 }
4583}
4584
4585#[cfg(target_os = "macos")]
4588#[allow(clippy::too_many_arguments)]
4589pub(crate) fn sgemm_rm(
4590 m: usize,
4591 n: usize,
4592 k: usize,
4593 alpha: f32,
4594 a: &[f32],
4595 lda: usize,
4596 b_mat: &[f32],
4597 ldb: usize,
4598 b_rows_are_n: bool,
4599 c: &mut [f32],
4600 ldc: usize,
4601) {
4602 debug_assert!(a.len() >= (m - 1) * lda + k);
4603 debug_assert!(c.len() >= (m - 1) * ldc + n);
4604 #[cfg(target_arch = "aarch64")]
4609 if std::env::var("CMF_FORCE_NEON_GEMM")
4610 .map(|v| v == "1")
4611 .unwrap_or(false)
4612 {
4613 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4614 }
4615 unsafe {
4616 accel_blas::cblas_sgemm(
4617 101, 111, if b_rows_are_n { 112 } else { 111 },
4620 m as i32,
4621 n as i32,
4622 k as i32,
4623 alpha,
4624 a.as_ptr(),
4625 lda as i32,
4626 b_mat.as_ptr(),
4627 ldb as i32,
4628 0.0,
4629 c.as_mut_ptr(),
4630 ldc as i32,
4631 );
4632 }
4633}
4634
4635#[cfg(target_os = "macos")]
4642fn qmatmat_accel(
4643 q: &[u8],
4644 row_scale: &[f32],
4645 pre: &[std::borrow::Cow<'_, [f32]>],
4646 rows: usize,
4647 cols: usize,
4648 out: &mut [f32],
4649 pool: Option<&Pool>,
4650) {
4651 const TR: usize = 2048;
4656 let b = pre.len();
4657 thread_local! {
4658 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4659 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4660 }
4661 XPANEL.with(|xp| {
4662 WTILE.with(|wt| {
4663 let mut xpanel = xp.borrow_mut();
4664 xpanel.clear();
4665 for x in pre {
4666 xpanel.extend_from_slice(x);
4667 }
4668 let mut wtile = wt.borrow_mut();
4669 wtile.resize(TR * cols, 0.0);
4670 let mut r0 = 0usize;
4671 while r0 < rows {
4672 let tr = TR.min(rows - r0);
4673 let wt_addr = SendMut(wtile.as_mut_ptr());
4675 let run = |start: usize, end: usize| {
4676 for r in start..end {
4677 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4678 let s = row_scale[r0 + r];
4679 let dst =
4681 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4682 for (d, &v) in dst.iter_mut().zip(row) {
4683 *d = (v as i8) as f32 * s;
4684 }
4685 }
4686 };
4687 dispatch_rows(pool, tr, &run);
4688 unsafe {
4690 accel_blas::cblas_sgemm(
4691 101, 111, 112, b as i32,
4695 tr as i32,
4696 cols as i32,
4697 1.0,
4698 xpanel.as_ptr(),
4699 cols as i32,
4700 wtile.as_ptr(),
4701 cols as i32,
4702 0.0,
4703 out.as_mut_ptr().add(r0),
4704 rows as i32,
4705 );
4706 }
4707 r0 += tr;
4708 }
4709 })
4710 });
4711}
4712
4713fn qmatmat(
4714 q: &[u8],
4715 row_scale: &[f32],
4716 pre: &[std::borrow::Cow<'_, [f32]>],
4717 rows: usize,
4718 cols: usize,
4719 out: &mut [f32],
4720 pool: Option<&Pool>,
4721) {
4722 let b = pre.len();
4723 debug_assert_eq!(out.len(), b * rows);
4724 let _p = crate::cpuprof::spans::span("mm.qmatmat");
4725 #[cfg(target_os = "macos")]
4730 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4731 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4732 return;
4733 }
4734 #[cfg(target_arch = "aarch64")]
4735 if sdot_enabled() {
4736 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4737 let out_addr = SendMut(out.as_mut_ptr());
4738 let blocked_ok = blocked_enabled();
4741 let use_i8mm = i8mm_enabled();
4742 if blocked_ok {
4743 let run = |start: usize, end: usize| {
4744 let mut o = start;
4745 while o < end {
4746 if o + 2 <= end {
4747 let r0 = &q[o * cols..(o + 1) * cols];
4748 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4749 let mut bi = 0usize;
4750 while bi + 4 <= acts.len() {
4751 let xs = [
4752 acts[bi].xq.as_slice(),
4753 acts[bi + 1].xq.as_slice(),
4754 acts[bi + 2].xq.as_slice(),
4755 acts[bi + 3].xq.as_slice(),
4756 ];
4757 let d = if use_i8mm {
4758 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4759 } else {
4760 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4761 };
4762 for (r, row) in [r0, r1].into_iter().enumerate() {
4763 for k in 0..4 {
4764 let act = &acts[bi + k];
4765 let mut v = d[r][k] as f32 * act.sx;
4766 for &(j, xv) in &act.outliers {
4767 v += (row[j] as i8) as f32 * xv;
4768 }
4769 unsafe {
4770 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4771 };
4772 }
4773 }
4774 bi += 4;
4775 }
4776 while bi < acts.len() {
4777 for (r, row) in [r0, r1].into_iter().enumerate() {
4778 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4779 unsafe { *out_addr.at(bi * rows + o + r) = v };
4780 }
4781 bi += 1;
4782 }
4783 o += 2;
4784 } else {
4785 let row = &q[o * cols..(o + 1) * cols];
4786 for (bi, act) in acts.iter().enumerate() {
4787 let v = row_dot_sdot(row, act) * row_scale[o];
4788 unsafe { *out_addr.at(bi * rows + o) = v };
4789 }
4790 o += 1;
4791 }
4792 }
4793 };
4794 dispatch_rows(pool, rows, &run);
4795 return;
4796 }
4797 let run = |start: usize, end: usize| {
4798 for o in start..end {
4799 let row = &q[o * cols..(o + 1) * cols];
4800 for (bi, act) in acts.iter().enumerate() {
4801 let v = row_dot_sdot(row, act) * row_scale[o];
4802 unsafe { *out_addr.at(bi * rows + o) = v };
4803 }
4804 }
4805 };
4806 dispatch_rows(pool, rows, &run);
4807 return;
4808 }
4809 #[cfg(target_arch = "x86_64")]
4814 if avx2_a8w8_enabled() {
4815 let acts = {
4816 let _psa = crate::cpuprof::spans::span("mm.split_act");
4817 let a: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4818 drop(_psa);
4819 a
4820 };
4821 let out_addr = SendMut(out.as_mut_ptr());
4822 let blocked_ok = blocked_enabled();
4825 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4826 let run = |start: usize, end: usize| {
4827 let mut o = start;
4828 while o < end {
4829 if o + 2 <= end {
4830 let r0 = &q[o * cols..(o + 1) * cols];
4831 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4832 let mut bi = 0usize;
4833 while bi + 4 <= acts.len() {
4834 let xs = [
4835 acts[bi].xq.as_slice(),
4836 acts[bi + 1].xq.as_slice(),
4837 acts[bi + 2].xq.as_slice(),
4838 acts[bi + 3].xq.as_slice(),
4839 ];
4840 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4841 for (r, row) in [r0, r1].into_iter().enumerate() {
4842 for k in 0..4 {
4843 let act = &acts[bi + k];
4844 let mut v = d[r][k] as f32 * act.sx;
4845 for &(j, xv) in &act.outliers {
4846 v += (row[j] as i8) as f32 * xv;
4847 }
4848 unsafe {
4849 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4850 };
4851 }
4852 }
4853 bi += 4;
4854 }
4855 while bi < acts.len() {
4856 for (r, row) in [r0, r1].into_iter().enumerate() {
4857 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4858 unsafe { *out_addr.at(bi * rows + o + r) = v };
4859 }
4860 bi += 1;
4861 }
4862 o += 2;
4863 } else {
4864 let row = &q[o * cols..(o + 1) * cols];
4865 for (bi, act) in acts.iter().enumerate() {
4866 let v = row_dot_avx2(row, act) * row_scale[o];
4867 unsafe { *out_addr.at(bi * rows + o) = v };
4868 }
4869 o += 1;
4870 }
4871 }
4872 };
4873 dispatch_rows(pool, rows, &run);
4874 return;
4875 }
4876 let run = |start: usize, end: usize| {
4877 for o in start..end {
4878 let row = &q[o * cols..(o + 1) * cols];
4879 for (bi, act) in acts.iter().enumerate() {
4880 let v = row_dot_avx2(row, act) * row_scale[o];
4881 unsafe { *out_addr.at(bi * rows + o) = v };
4882 }
4883 }
4884 };
4885 dispatch_rows(pool, rows, &run);
4886 return;
4887 }
4888 let out_addr = SendMut(out.as_mut_ptr());
4889 let run = |start: usize, end: usize| {
4890 for o in start..end {
4891 let row = &q[o * cols..(o + 1) * cols];
4892 for (bi, x) in pre.iter().enumerate() {
4893 let mut acc = 0f32;
4894 for j in 0..cols {
4895 acc += (row[j] as i8) as f32 * x[j];
4896 }
4897 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4898 }
4899 }
4900 };
4901 dispatch_rows(pool, rows, &run);
4902}
4903
4904fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4907 match pool {
4908 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4909 _ => run(0, rows),
4910 }
4911}
4912
4913fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4915 let groups = rows * cols / GROUP_SIZE;
4916 bytes.split_at(groups * 16)
4917}
4918
4919#[inline]
4924fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4925 #[cfg(target_arch = "aarch64")]
4926 unsafe {
4927 return vbit_fill4_neon(data, buf);
4928 }
4929 #[cfg(target_arch = "x86_64")]
4930 if avx2_enabled() {
4931 return unsafe { vbit_fill4_avx2(data, buf) };
4932 }
4933 #[allow(unreachable_code)]
4934 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4935 let u = unpack8::<4>(&data[blk * 4..]);
4936 for k in 0..8 {
4937 chunk[k] = (u[k] - 7) as i8 as u8;
4938 }
4939 }
4940}
4941
4942#[cfg(target_arch = "aarch64")]
4943#[target_feature(enable = "neon")]
4944unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4945 unsafe {
4948 use core::arch::aarch64::*;
4949 let n = buf.len();
4950 let mask = vdupq_n_u8(0x0F);
4951 let seven = vdupq_n_s8(7);
4952 let mut g = 0usize;
4953 while g * 32 + 32 <= n {
4954 let b = vld1q_u8(data.as_ptr().add(g * 16));
4955 let hi = vshrq_n_u8::<4>(b);
4956 let lo = vandq_u8(b, mask);
4957 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4958 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4959 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4960 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4961 g += 1;
4962 }
4963 }
4964}
4965
4966#[cfg(target_arch = "x86_64")]
4967#[target_feature(enable = "avx2")]
4968unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4969 unsafe {
4971 use core::arch::x86_64::*;
4972 let n = buf.len();
4973 let mask = _mm_set1_epi8(0x0F);
4974 let seven = _mm256_set1_epi8(7);
4975 let mut g = 0usize;
4976 while g * 32 + 32 <= n {
4977 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4978 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4979 let lo = _mm_and_si128(b, mask);
4980 let z = _mm256_sub_epi8(
4981 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4982 seven,
4983 );
4984 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4985 g += 1;
4986 }
4987 }
4988}
4989
4990#[inline(always)]
4995fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4996 let mut acc = 0u64;
4997 for i in 0..B {
4998 acc = (acc << 8) | data[i] as u64;
4999 }
5000 let mask = (1u64 << B) - 1;
5001 let mut out = [0i32; 8];
5002 for (k, o) in out.iter_mut().enumerate() {
5003 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
5004 }
5005 out
5006}
5007
5008#[allow(clippy::too_many_arguments)]
5014fn vbitmatvec(
5015 bytes: &[u8],
5016 offsets: &[usize],
5017 x: &[f32],
5018 rows: usize,
5019 cols: usize,
5020 out: &mut [f32],
5021 pool: Option<&Pool>,
5022) {
5023 debug_assert_eq!(out.len(), rows);
5024 debug_assert_eq!(offsets.len(), rows + 1);
5025
5026 if a8w8_enabled() {
5030 let act = split_act(x);
5031 let out_addr = SendMut(out.as_mut_ptr());
5032 let run = move |start: usize, end: usize| {
5033 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
5034 };
5035 dispatch_rows(pool, rows, &run);
5036 return;
5037 }
5038
5039 let out_addr = SendMut(out.as_mut_ptr());
5040 let run = move |start: usize, end: usize| {
5041 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
5042 };
5043 dispatch_rows(pool, rows, &run);
5044}
5045
5046#[allow(clippy::too_many_arguments)]
5050fn vbit_range_a8w8(
5051 bytes: &[u8],
5052 offsets: &[usize],
5053 x: &[f32],
5054 act: &SplitAct,
5055 rows: usize,
5056 cols: usize,
5057 out: SendMut,
5058 start: usize,
5059 end: usize,
5060) {
5061 let ng = cols / GROUP_SIZE;
5062 let bits = &bytes[..rows];
5063 let sc_off = rows;
5064 let row_dot = |r: usize| -> f32 {
5065 let b = bits[r] as usize;
5066 let l = (1i32 << (b - 1)) - 1;
5067 let mask = (1u64 << b) - 1;
5068 let data = &bytes[offsets[r]..offsets[r + 1]];
5069 if b == 8 {
5070 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5072 let mut dot = 0f32;
5073 for g in 0..ng {
5074 let so = (r * ng + g) * 2;
5075 let sgf = f16_to_f32(u16::from_le_bytes([
5076 bytes[sc_off + so],
5077 bytes[sc_off + so + 1],
5078 ]));
5079 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5080 let mut gd = 0f32;
5081 for &xv in xg.iter() {
5082 if nbits < 8 {
5083 acc = (acc << 8) | data[idx] as u64;
5084 idx += 1;
5085 nbits += 8;
5086 }
5087 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5088 nbits -= 8;
5089 gd += (u - l) as f32 * xv;
5090 }
5091 dot += gd * sgf;
5092 }
5093 return dot;
5094 }
5095 thread_local! {
5099 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
5100 const { std::cell::RefCell::new(Vec::new()) };
5101 }
5102 #[inline(always)]
5103 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5104 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5105 let u = unpack8::<B>(&data[blk * B..]);
5106 for k in 0..8 {
5107 chunk[k] = (u[k] - l) as i8 as u8;
5108 }
5109 }
5110 }
5111 let _ = mask;
5112 VBIT_SCRATCH.with(|scratch| {
5113 let mut buf = scratch.borrow_mut();
5114 buf.resize(cols, 0);
5115 match b {
5116 3 => fill::<3>(data, l, &mut buf),
5117 4 => vbit_fill4(data, &mut buf),
5118 5 => fill::<5>(data, l, &mut buf),
5119 6 => fill::<6>(data, l, &mut buf),
5120 _ => unreachable!(),
5121 }
5122 let mut dot = 0f32;
5123 for g in 0..ng {
5124 let so = (r * ng + g) * 2;
5125 let s = f16_to_f32(u16::from_le_bytes([
5126 bytes[sc_off + so],
5127 bytes[sc_off + so + 1],
5128 ]));
5129 let d = dot_i8_i8(
5130 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5131 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5132 ) as f32
5133 * act.sx;
5134 dot += d * s;
5135 }
5136 for &(j, xv) in &act.outliers {
5137 let so = (r * ng + j / GROUP_SIZE) * 2;
5138 let s = f16_to_f32(u16::from_le_bytes([
5139 bytes[sc_off + so],
5140 bytes[sc_off + so + 1],
5141 ]));
5142 dot += (buf[j] as i8) as f32 * s * xv;
5144 }
5145 dot
5146 })
5147 };
5148 for r in start..end {
5149 unsafe { *out.at(r) = row_dot(r) };
5151 }
5152}
5153
5154#[allow(clippy::too_many_arguments)]
5156fn vbit_range_f32(
5157 bytes: &[u8],
5158 offsets: &[usize],
5159 x: &[f32],
5160 rows: usize,
5161 cols: usize,
5162 out: SendMut,
5163 start: usize,
5164 end: usize,
5165) {
5166 let ng = cols / GROUP_SIZE;
5167 let bits = &bytes[..rows];
5168 let sc_off = rows;
5169 #[inline(always)]
5173 fn dot_row<const B: usize>(
5174 data: &[u8],
5175 bytes: &[u8],
5176 sc_off: usize,
5177 r: usize,
5178 ng: usize,
5179 x: &[f32],
5180 ) -> f32 {
5181 let l = ((1i32 << (B - 1)) - 1) as f32;
5182 let gbytes = GROUP_SIZE * B / 8;
5183 let mut dot = 0f32;
5184 for g in 0..ng {
5185 let so = (r * ng + g) * 2;
5186 let s = f16_to_f32(u16::from_le_bytes([
5187 bytes[sc_off + so],
5188 bytes[sc_off + so + 1],
5189 ]));
5190 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5191 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5192 let mut gd = 0f32;
5193 for blk in 0..GROUP_SIZE / 8 {
5194 let u = unpack8::<B>(&gd0[blk * B..]);
5195 let xb = &xg[blk * 8..blk * 8 + 8];
5196 for k in 0..8 {
5197 gd += (u[k] as f32 - l) * xb[k];
5198 }
5199 }
5200 dot += gd * s;
5201 }
5202 dot
5203 }
5204 for r in start..end {
5205 let data = &bytes[offsets[r]..offsets[r + 1]];
5206 let v = match bits[r] {
5207 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5208 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5209 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5210 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5211 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5212 b => unreachable!("vbit bit-width {b} (validated at load)"),
5213 };
5214 unsafe { *out.at(r) = v };
5216 }
5217}
5218
5219#[allow(clippy::too_many_arguments)]
5224fn vbitmatvec2(
5225 bytes: &[u8],
5226 offsets: &[usize],
5227 x1: &[f32],
5228 x2: &[f32],
5229 rows: usize,
5230 cols: usize,
5231 o1: &mut [f32],
5232 o2: &mut [f32],
5233 pool: Option<&Pool>,
5234) {
5235 debug_assert_eq!(o1.len(), rows);
5236 debug_assert_eq!(o2.len(), rows);
5237
5238 if a8w8_enabled() {
5239 let a1 = split_act(x1);
5240 let a2 = split_act(x2);
5241 let p1 = SendMut(o1.as_mut_ptr());
5242 let p2 = SendMut(o2.as_mut_ptr());
5243 let run = move |start: usize, end: usize| {
5244 vbit_range2_a8w8(
5245 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5246 )
5247 };
5248 dispatch_rows(pool, rows, &run);
5249 return;
5250 }
5251
5252 let p1 = SendMut(o1.as_mut_ptr());
5253 let p2 = SendMut(o2.as_mut_ptr());
5254 let run = move |start: usize, end: usize| {
5255 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5256 };
5257 dispatch_rows(pool, rows, &run);
5258}
5259
5260#[allow(clippy::too_many_arguments)]
5264fn vbit_range2_a8w8(
5265 bytes: &[u8],
5266 offsets: &[usize],
5267 x1: &[f32],
5268 x2: &[f32],
5269 a1: &SplitAct,
5270 a2: &SplitAct,
5271 rows: usize,
5272 cols: usize,
5273 p1: SendMut,
5274 p2: SendMut,
5275 start: usize,
5276 end: usize,
5277) {
5278 let ng = cols / GROUP_SIZE;
5279 let bits = &bytes[..rows];
5280 let sc_off = rows;
5281 let row_dots = |r: usize| -> (f32, f32) {
5282 let b = bits[r] as usize;
5283 let l = (1i32 << (b - 1)) - 1;
5284 let data = &bytes[offsets[r]..offsets[r + 1]];
5285 if b == 8 {
5286 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5289 let (mut d1, mut d2) = (0f32, 0f32);
5290 for g in 0..ng {
5291 let so = (r * ng + g) * 2;
5292 let sgf = f16_to_f32(u16::from_le_bytes([
5293 bytes[sc_off + so],
5294 bytes[sc_off + so + 1],
5295 ]));
5296 let (mut g1, mut g2) = (0f32, 0f32);
5297 for k in 0..GROUP_SIZE {
5298 if nbits < 8 {
5299 acc = (acc << 8) | data[idx] as u64;
5300 idx += 1;
5301 nbits += 8;
5302 }
5303 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5304 nbits -= 8;
5305 let w = (u - l) as f32;
5306 g1 += w * x1[g * GROUP_SIZE + k];
5307 g2 += w * x2[g * GROUP_SIZE + k];
5308 }
5309 d1 += g1 * sgf;
5310 d2 += g2 * sgf;
5311 }
5312 return (d1, d2);
5313 }
5314 thread_local! {
5315 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5316 const { std::cell::RefCell::new(Vec::new()) };
5317 }
5318 #[inline(always)]
5319 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5320 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5321 let u = unpack8::<B>(&data[blk * B..]);
5322 for k in 0..8 {
5323 chunk[k] = (u[k] - l) as i8 as u8;
5324 }
5325 }
5326 }
5327 VBIT_SCRATCH2.with(|scratch| {
5328 let mut buf = scratch.borrow_mut();
5329 buf.resize(cols, 0);
5330 match b {
5331 3 => fill::<3>(data, l, &mut buf),
5332 4 => vbit_fill4(data, &mut buf),
5333 5 => fill::<5>(data, l, &mut buf),
5334 6 => fill::<6>(data, l, &mut buf),
5335 _ => unreachable!(),
5336 }
5337 let (mut d1, mut d2) = (0f32, 0f32);
5338 for g in 0..ng {
5339 let so = (r * ng + g) * 2;
5340 let s = f16_to_f32(u16::from_le_bytes([
5341 bytes[sc_off + so],
5342 bytes[sc_off + so + 1],
5343 ]));
5344 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5345 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5346 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5347 d1 += v1 * s;
5348 d2 += v2 * s;
5349 }
5350 for &(j, xv) in &a1.outliers {
5351 let so = (r * ng + j / GROUP_SIZE) * 2;
5352 let s = f16_to_f32(u16::from_le_bytes([
5353 bytes[sc_off + so],
5354 bytes[sc_off + so + 1],
5355 ]));
5356 d1 += (buf[j] as i8) as f32 * s * xv;
5357 }
5358 for &(j, xv) in &a2.outliers {
5359 let so = (r * ng + j / GROUP_SIZE) * 2;
5360 let s = f16_to_f32(u16::from_le_bytes([
5361 bytes[sc_off + so],
5362 bytes[sc_off + so + 1],
5363 ]));
5364 d2 += (buf[j] as i8) as f32 * s * xv;
5365 }
5366 (d1, d2)
5367 })
5368 };
5369 for r in start..end {
5370 let (v1, v2) = row_dots(r);
5371 unsafe {
5373 *p1.at(r) = v1;
5374 *p2.at(r) = v2;
5375 }
5376 }
5377}
5378
5379#[allow(clippy::too_many_arguments)]
5383fn vbit_range2_f32(
5384 bytes: &[u8],
5385 offsets: &[usize],
5386 x1: &[f32],
5387 x2: &[f32],
5388 rows: usize,
5389 cols: usize,
5390 p1: SendMut,
5391 p2: SendMut,
5392 start: usize,
5393 end: usize,
5394) {
5395 let ng = cols / GROUP_SIZE;
5396 let bits = &bytes[..rows];
5397 let sc_off = rows;
5398 #[inline(always)]
5399 #[allow(clippy::too_many_arguments)]
5400 fn dot_row2<const B: usize>(
5401 data: &[u8],
5402 bytes: &[u8],
5403 sc_off: usize,
5404 r: usize,
5405 ng: usize,
5406 x1: &[f32],
5407 x2: &[f32],
5408 ) -> (f32, f32) {
5409 let l = ((1i32 << (B - 1)) - 1) as f32;
5410 let gbytes = GROUP_SIZE * B / 8;
5411 let (mut d1, mut d2) = (0f32, 0f32);
5412 for g in 0..ng {
5413 let so = (r * ng + g) * 2;
5414 let s = f16_to_f32(u16::from_le_bytes([
5415 bytes[sc_off + so],
5416 bytes[sc_off + so + 1],
5417 ]));
5418 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5419 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5420 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5421 let (mut g1, mut g2) = (0f32, 0f32);
5422 for blk in 0..GROUP_SIZE / 8 {
5423 let u = unpack8::<B>(&gd0[blk * B..]);
5424 for k in 0..8 {
5425 let w = u[k] as f32 - l;
5426 g1 += w * x1g[blk * 8 + k];
5427 g2 += w * x2g[blk * 8 + k];
5428 }
5429 }
5430 d1 += g1 * s;
5431 d2 += g2 * s;
5432 }
5433 (d1, d2)
5434 }
5435 for r in start..end {
5436 let data = &bytes[offsets[r]..offsets[r + 1]];
5437 let (v1, v2) = match bits[r] {
5438 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5439 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5440 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5441 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5442 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5443 b => unreachable!("vbit bit-width {b} (validated at load)"),
5444 };
5445 unsafe {
5447 *p1.at(r) = v1;
5448 *p2.at(r) = v2;
5449 }
5450 }
5451}
5452
5453#[inline]
5460#[allow(unreachable_code)]
5461fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5462 #[cfg(target_arch = "aarch64")]
5463 unsafe {
5464 return dot_q4t_row_sdot(bytes, r, gpr, xq);
5465 }
5466 #[cfg(target_arch = "x86_64")]
5467 unsafe {
5468 if vnni_tiles_enabled() {
5469 return dot_q4t_row_vnni(bytes, r, gpr, xq);
5470 }
5471 return dot_q4t_row_avx2(bytes, r, gpr, xq);
5472 }
5473 let mut acc = 0f32;
5474 for gi in 0..gpr {
5475 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5476 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5477 let mut d = 0i32;
5478 for (k, &b) in tile[2..].iter().enumerate() {
5479 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5480 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5481 }
5482 acc += d as f32 * s;
5483 }
5484 acc
5485}
5486
5487#[cfg(target_arch = "aarch64")]
5488#[target_feature(enable = "neon,dotprod")]
5489unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5490 unsafe {
5493 use core::arch::aarch64::*;
5494 use core::arch::asm;
5495 let lomask = vdupq_n_u8(0x0F);
5496 let eight = vdupq_n_s8(8);
5497 let mut acc = 0f32;
5498 for gi in 0..gpr {
5499 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5500 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5501 let b = vld1q_u8(t.add(2));
5502 let lo = vandq_u8(b, lomask);
5503 let hi = vshrq_n_u8::<4>(b);
5504 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5505 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5506 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5507 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5508 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5509 asm!(
5510 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5511 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5512 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5513 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5514 options(pure, nomem, nostack),
5515 );
5516 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5517 }
5518 acc
5519 }
5520}
5521
5522#[cfg(target_arch = "x86_64")]
5523#[target_feature(enable = "avx2")]
5524unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5525 unsafe {
5527 use core::arch::x86_64::*;
5528 let lomask = _mm_set1_epi8(0x0F);
5529 let eight = _mm256_set1_epi8(8);
5530 let ones = _mm256_set1_epi16(1);
5531 let mut acc = 0f32;
5532 for gi in 0..gpr {
5533 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5534 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5535 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5536 let lo = _mm_and_si128(b, lomask);
5537 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5538 let w = _mm256_sub_epi8(
5539 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5540 eight,
5541 );
5542 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5543 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5544 let d = _mm256_madd_epi16(p16, ones);
5545 let hi128 = _mm256_extracti128_si256::<1>(d);
5546 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5547 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5548 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5549 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5550 }
5551 acc
5552 }
5553}
5554
5555#[cfg(target_arch = "x86_64")]
5559#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5560unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5561 unsafe {
5563 use core::arch::x86_64::*;
5564 let lomask = _mm_set1_epi8(0x0F);
5565 let eight = _mm256_set1_epi8(8);
5566 let mut acc = 0f32;
5567 for gi in 0..gpr {
5568 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5569 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5570 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5571 let lo = _mm_and_si128(b, lomask);
5572 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5573 let w = _mm256_sub_epi8(
5574 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5575 eight,
5576 );
5577 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5578 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5579 acc += d as f32 * s;
5580 }
5581 acc
5582 }
5583}
5584
5585#[cfg(target_arch = "x86_64")]
5590#[target_feature(enable = "avx2,fma")]
5595unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5596 unsafe {
5598 use core::arch::x86_64::*;
5599 let lomask = _mm_set1_epi8(0x0F);
5600 let eight = _mm256_set1_epi8(8);
5601 let ones = _mm256_set1_epi16(1);
5602 let mut f0 = _mm256_setzero_ps();
5615 let mut f1 = _mm256_setzero_ps();
5616 let mut f2 = _mm256_setzero_ps();
5617 let mut f3 = _mm256_setzero_ps();
5618 for gi in 0..gpr {
5619 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5620 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5621 let sv = _mm256_set1_ps(s);
5622 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5623 let lo = _mm_and_si128(bb, lomask);
5624 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5625 let w = _mm256_sub_epi8(
5626 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5627 eight,
5628 );
5629 let aw = _mm256_abs_epi8(w);
5630 let off = gi * GROUP_SIZE;
5631 let dot = |xq: &[i8]| {
5632 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5633 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5634 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5635 };
5636 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5637 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5638 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5639 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5640 }
5641 [
5642 hsum256_ps(f0),
5643 hsum256_ps(f1),
5644 hsum256_ps(f2),
5645 hsum256_ps(f3),
5646 ]
5647 }
5648}
5649
5650#[cfg(target_arch = "x86_64")]
5653#[target_feature(enable = "avx2")]
5654#[inline]
5655unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5656 unsafe {
5658 use core::arch::x86_64::*;
5659 let hi = _mm256_extractf128_ps::<1>(v);
5660 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5661 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5662 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5663 _mm_cvtss_f32(s)
5664 }
5665}
5666
5667#[cfg(target_arch = "x86_64")]
5669#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5670unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5671 unsafe {
5673 use core::arch::x86_64::*;
5674 let lomask = _mm_set1_epi8(0x0F);
5675 let eight = _mm256_set1_epi8(8);
5676 let mut f0 = _mm256_setzero_ps();
5679 let mut f1 = _mm256_setzero_ps();
5680 let mut f2 = _mm256_setzero_ps();
5681 let mut f3 = _mm256_setzero_ps();
5682 for gi in 0..gpr {
5683 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5684 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5685 let sv = _mm256_set1_ps(s);
5686 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5687 let lo = _mm_and_si128(bb, lomask);
5688 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5689 let w = _mm256_sub_epi8(
5690 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5691 eight,
5692 );
5693 let aw = _mm256_abs_epi8(w);
5694 let off = gi * GROUP_SIZE;
5695 let dot = |xq: &[i8]| {
5696 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5697 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5698 _mm256_setzero_si256(),
5699 aw,
5700 _mm256_sign_epi8(x, w),
5701 ))
5702 };
5703 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5704 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5705 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5706 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5707 }
5708 let acc = [
5709 hsum256_ps(f0),
5710 hsum256_ps(f1),
5711 hsum256_ps(f2),
5712 hsum256_ps(f3),
5713 ];
5714 acc
5715 }
5716}
5717
5718#[cfg(target_arch = "aarch64")]
5723#[target_feature(enable = "neon,dotprod")]
5724unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5725 unsafe {
5727 use core::arch::aarch64::*;
5728 use core::arch::asm;
5729 let lomask = vdupq_n_u8(0x0F);
5730 let eight = vdupq_n_s8(8);
5731 let mut acc = [0f32; 4];
5732 for gi in 0..gpr {
5733 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5734 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5735 let b = vld1q_u8(t.add(2));
5736 let lo = vandq_u8(b, lomask);
5737 let hi = vshrq_n_u8::<4>(b);
5738 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5739 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5740 for (k, xq) in xs.iter().enumerate() {
5741 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5742 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5743 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5744 asm!(
5745 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5746 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5747 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5748 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5749 options(pure, nomem, nostack),
5750 );
5751 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5752 }
5753 }
5754 acc
5755 }
5756}
5757
5758#[inline]
5760fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5761 let gi = j / GROUP_SIZE;
5762 let k = j % GROUP_SIZE;
5763 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5764 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5765 let byte = tile[2 + k / 2];
5766 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5767 ((nib as i32 - 8) as f32, s)
5768}
5769
5770#[inline]
5773fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5774 let mut acc = 0f32;
5775 for gi in 0..gpr {
5776 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5777 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5778 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5779 let mut ga = 0f32;
5780 for (k, &b) in tile[2..].iter().enumerate() {
5781 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5782 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5783 }
5784 acc += ga * s;
5785 }
5786 acc
5787}
5788
5789struct Q4tpView<'a> {
5793 nib: &'a [u8],
5794 params: &'a [u8],
5795 codes: &'a [u8],
5796 stride: usize,
5797 zero_rung: bool,
5799}
5800
5801impl<'a> Q4tpView<'a> {
5802 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5803 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5804 Self {
5805 nib: &bytes[..params_off],
5806 params: &bytes[params_off..codes_off],
5807 codes: &bytes[codes_off..],
5808 stride,
5809 zero_rung: false,
5810 }
5811 }
5812
5813 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5815 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5816 Self {
5817 nib: &bytes[..params_off],
5818 params: &bytes[params_off..codes_off],
5819 codes: &bytes[codes_off..],
5820 stride,
5821 zero_rung: true,
5822 }
5823 }
5824
5825 #[cfg(target_arch = "x86_64")]
5848 #[inline]
5849 fn ladder_via_avx2(&self) -> bool {
5850 avx2_enabled()
5851 }
5852 #[cfg(not(target_arch = "x86_64"))]
5853 #[inline]
5854 fn ladder_via_avx2(&self) -> bool {
5855 false
5856 }
5857
5858 #[inline]
5859 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5860 let tab = if self.zero_rung {
5861 q2tp_ladder(self.params, r)
5862 } else if self.ladder_via_avx2() {
5863 #[cfg(target_arch = "x86_64")]
5870 let tab = {
5871 let p = &self.params[r * 4..r * 4 + 4];
5872 let lo = f16_to_f32(u16::from_le_bytes([p[0], p[1]]));
5873 let st = f16_to_f32(u16::from_le_bytes([p[2], p[3]]));
5874 let mut tabc = [0f32; 128];
5875 unsafe { ladder_4(&[lo; 4], &[st; 4], &mut tabc) };
5876 let mut t = [0f32; 32];
5877 for (g, o) in t.iter_mut().enumerate() {
5878 *o = tabc[g * 4];
5879 }
5880 t
5881 };
5882 #[cfg(not(target_arch = "x86_64"))]
5886 let tab = q4tp_ladder(self.params, r);
5887 tab
5888 } else {
5889 q4tp_ladder(self.params, r)
5890 };
5891 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5892 let out = &mut out[..gpr];
5893 let mut chunks = out.chunks_exact_mut(8);
5894 let mut ci = 0usize;
5895 for c in &mut chunks {
5896 let w = u64::from(codes[ci])
5897 | u64::from(codes[ci + 1]) << 8
5898 | u64::from(codes[ci + 2]) << 16
5899 | u64::from(codes[ci + 3]) << 24
5900 | u64::from(codes[ci + 4]) << 32;
5901 for (k, o) in c.iter_mut().enumerate() {
5902 *o = tab[((w >> (5 * k)) & 31) as usize];
5903 }
5904 ci += 5;
5905 }
5906 let tail = &codes[ci..];
5909 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5910 *o = tab[q4tp_code(tail, k)];
5911 }
5912 }
5913}
5914
5915#[inline]
5916fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5917 #[cfg(target_arch = "aarch64")]
5918 unsafe {
5919 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5920 }
5921 #[cfg(target_arch = "x86_64")]
5922 unsafe {
5923 if vnni_tiles_enabled() {
5924 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5925 }
5926 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5927 }
5928 #[allow(unreachable_code)]
5929 {
5930 let mut acc = 0f32;
5931 for gi in 0..gpr {
5932 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5933 let s = scales[gi];
5934 let mut d = 0i32;
5935 for (k, &b) in tile.iter().enumerate() {
5936 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5937 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5938 }
5939 acc += d as f32 * s;
5940 }
5941 acc
5942 }
5943}
5944
5945#[cfg(target_arch = "aarch64")]
5948#[target_feature(enable = "neon,dotprod")]
5949unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5950 unsafe {
5953 use core::arch::aarch64::*;
5954 use core::arch::asm;
5955 let lomask = vdupq_n_u8(0x0F);
5956 let eight = vdupq_n_s8(8);
5957 let mut acc = 0f32;
5958 for gi in 0..gpr {
5959 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5960 let s = *scales.get_unchecked(gi);
5961 let b = vld1q_u8(t);
5962 let lo = vandq_u8(b, lomask);
5963 let hi = vshrq_n_u8::<4>(b);
5964 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5965 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5966 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5967 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5968 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5969 asm!(
5970 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5971 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5972 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5973 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5974 options(pure, nomem, nostack),
5975 );
5976 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5977 }
5978 acc
5979 }
5980}
5981
5982#[cfg(target_arch = "x86_64")]
5983#[target_feature(enable = "avx2")]
5984unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5985 unsafe {
5987 use core::arch::x86_64::*;
5988 let lomask = _mm_set1_epi8(0x0F);
5989 let eight = _mm256_set1_epi8(8);
5990 let ones = _mm256_set1_epi16(1);
5991 let mut acc = 0f32;
5992 for gi in 0..gpr {
5993 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5994 let s = *scales.get_unchecked(gi);
5995 let b = _mm_loadu_si128(t as *const __m128i);
5996 let lo = _mm_and_si128(b, lomask);
5997 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5998 let w = _mm256_sub_epi8(
5999 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6000 eight,
6001 );
6002 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6003 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6004 let d = _mm256_madd_epi16(p16, ones);
6005 let hi128 = _mm256_extracti128_si256::<1>(d);
6006 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6007 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6008 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6009 acc += _mm_cvtsi128_si32(s32) as f32 * s;
6010 }
6011 acc
6012 }
6013}
6014
6015#[cfg(target_arch = "x86_64")]
6018#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6019unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6020 unsafe {
6022 use core::arch::x86_64::*;
6023 let lomask = _mm_set1_epi8(0x0F);
6024 let eight = _mm256_set1_epi8(8);
6025 let mut acc = 0f32;
6026 for gi in 0..gpr {
6027 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6028 let s = *scales.get_unchecked(gi);
6029 let b = _mm_loadu_si128(t as *const __m128i);
6030 let lo = _mm_and_si128(b, lomask);
6031 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6032 let w = _mm256_sub_epi8(
6033 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6034 eight,
6035 );
6036 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6037 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
6038 }
6039 acc
6040 }
6041}
6042
6043#[cfg(target_arch = "x86_64")]
6050#[target_feature(enable = "avx2")]
6051#[inline]
6052unsafe fn exp2_128(x: core::arch::x86_64::__m128) -> core::arch::x86_64::__m128 {
6053 unsafe {
6054 use core::arch::x86_64::*;
6055 let tr = _mm_cvttps_epi32(x);
6057 let below = _mm_and_si128(
6060 _mm_castps_si128(_mm_cmplt_ps(x, _mm_cvtepi32_ps(tr))),
6061 _mm_set1_epi32(1),
6062 );
6063 let n = _mm_sub_epi32(tr, below); let r = _mm_sub_ps(x, _mm_cvtepi32_ps(n)); #[allow(clippy::excessive_precision, clippy::approx_constant)]
6074 let a = [
6075 6.931_471_805_599_453e-1,
6076 2.402_265_069_591_007e-1,
6077 5.550_410_866_482_158e-2,
6078 9.618_129_107_628_477e-3,
6079 1.333_355_814_642_844e-3,
6080 1.540_353_039_338_161e-4,
6081 1.525_273_380_405_984e-5,
6082 1.321_548_679_014_431e-6,
6083 1.017_808_600_923_969_6e-7,
6084 ];
6085 let mut p = _mm_set1_ps(a[8]);
6086 for &c in a[..8].iter().rev() {
6087 p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(c));
6088 }
6089 p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(1.0));
6090 let ncl = _mm_min_epi32(
6096 _mm_max_epi32(n, _mm_set1_epi32(-126)),
6097 _mm_set1_epi32(127),
6098 );
6099 let mut bits = _mm_add_epi32(_mm_castps_si128(p), _mm_slli_epi32(ncl, 23));
6100 let over = _mm_cmpgt_epi32(n, _mm_set1_epi32(127));
6101 bits = _mm_or_si128(
6102 _mm_andnot_si128(over, bits),
6103 _mm_and_si128(over, _mm_set1_epi32(0x007F_8000)),
6104 );
6105 let zero = _mm_cmplt_ps(x, _mm_set1_ps(-149.0));
6106 _mm_andnot_ps(zero, _mm_castsi128_ps(bits))
6107 }
6108}
6109
6110#[cfg(target_arch = "x86_64")]
6112#[inline]
6113unsafe fn lanes4(v: core::arch::x86_64::__m128) -> [f32; 4] {
6114 use core::arch::x86_64::*;
6115 unsafe {
6116 [
6117 _mm_cvtss_f32(v),
6118 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0x55)),
6119 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xAA)),
6120 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xFF)),
6121 ]
6122 }
6123}
6124
6125#[cfg(target_arch = "x86_64")]
6133#[target_feature(enable = "avx2")]
6134#[inline]
6135unsafe fn ladder_4(lo4: &[f32; 4], st4: &[f32; 4], tabc: &mut [f32; 128]) {
6136 unsafe {
6137 use core::arch::x86_64::*;
6138 let ratio = exp2_128(_mm_loadu_ps(st4.as_ptr()));
6139 let mut t = exp2_128(_mm_loadu_ps(lo4.as_ptr()));
6140 tabc[0..4].copy_from_slice(&lanes4(t));
6141 for c in 1..32 {
6142 t = _mm_mul_ps(t, ratio);
6143 tabc[c * 4..c * 4 + 4].copy_from_slice(&lanes4(t));
6144 }
6145 }
6146}
6147
6148#[cfg(target_arch = "x86_64")]
6163#[target_feature(enable = "avx2")]
6164unsafe fn q4tp_scales_lanes(
6165 vs: &[&Q4tpView<'_>],
6166 rows: &[usize],
6167 gpr: usize,
6168 out: &mut [f32],
6169) {
6170 use core::arch::x86_64::*;
6171 debug_assert_eq!(vs.len(), rows.len());
6172 debug_assert!(out.len() >= rows.len() * gpr);
6173 let mut done = 0usize;
6174 for (ci, chunk) in rows.chunks(4).enumerate() {
6175 let vchunk = &vs[ci * 4..ci * 4 + chunk.len()];
6176 if chunk.len() < 4 || vchunk.iter().any(|v| v.zero_rung) {
6177 for (l, &r) in chunk.iter().enumerate() {
6178 vchunk[l].scales_into(r, gpr, &mut out[(done + l) * gpr..(done + l + 1) * gpr]);
6179 }
6180 done += chunk.len();
6181 continue;
6182 }
6183 let (r0, r1, r2, r3) = (chunk[0], chunk[1], chunk[2], chunk[3]);
6184 let v0 = vchunk[0];
6185 let v1 = vchunk[1];
6186 let v2 = vchunk[2];
6187 let v3 = vchunk[3];
6188 let mut lo = [0f32; 4];
6190 let mut st = [0f32; 4];
6191 for (l, r) in [r0, r1, r2, r3].iter().enumerate() {
6192 let o = *r * 4;
6193 let vv = vchunk[l];
6194 lo[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6195 vv.params[o],
6196 vv.params[o + 1],
6197 ]));
6198 st[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6199 vv.params[o + 2],
6200 vv.params[o + 3],
6201 ]));
6202 }
6203 let lo_v = _mm_loadu_ps(lo.as_ptr());
6204 let st_v = _mm_loadu_ps(st.as_ptr());
6205 let ratio = exp2_128(st_v);
6206 let mut t = exp2_128(lo_v);
6207 let mut tabc = [0f32; 32 * 4];
6209 tabc[0..4].copy_from_slice(&[
6210 _mm_cvtss_f32(t),
6211 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6212 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6213 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6214 ]);
6215 for c in 1..32 {
6216 t = _mm_mul_ps(t, ratio);
6217 tabc[c * 4..c * 4 + 4].copy_from_slice(&[
6218 _mm_cvtss_f32(t),
6219 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6220 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6221 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6222 ]);
6223 }
6224 let vs4 = [v0, v1, v2, v3];
6226 for (l, &r) in [r0, r1, r2, r3].iter().enumerate() {
6227 let v = vs4[l];
6228 let codes = &v.codes[r * v.stride..(r + 1) * v.stride];
6229 let dst = &mut out[(done + l) * gpr..(done + l + 1) * gpr];
6230 let mut chunks = dst.chunks_exact_mut(8);
6231 let mut ci = 0usize;
6232 for c in &mut chunks {
6233 let w = u64::from(codes[ci])
6234 | u64::from(codes[ci + 1]) << 8
6235 | u64::from(codes[ci + 2]) << 16
6236 | u64::from(codes[ci + 3]) << 24
6237 | u64::from(codes[ci + 4]) << 32;
6238 for (k, o) in c.iter_mut().enumerate() {
6239 *o = tabc[(((w >> (5 * k)) & 31) as usize) * 4 + l];
6240 }
6241 ci += 5;
6242 }
6243 let tail = &codes[ci..];
6244 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6245 *o = tabc[q4tp_code(tail, k) * 4 + l];
6246 }
6247 }
6248 done += 4;
6249 }
6250}
6251
6252fn q4tp_lanes() -> usize {
6260 #[cfg(not(target_arch = "x86_64"))]
6261 {
6262 1
6263 }
6264 #[cfg(target_arch = "x86_64")]
6265 {
6266 static ON: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6267 *ON.get_or_init(|| {
6268 let n: usize = std::env::var("CMF_Q4TP_LANES")
6269 .ok()
6270 .and_then(|v| v.parse().ok())
6271 .unwrap_or(4);
6272 match n {
6273 1 => 1,
6274 2 => 2,
6275 8 => 8,
6276 _ => 4,
6277 }
6278 })
6279 }
6280}
6281
6282#[cfg(target_arch = "x86_64")]
6293#[target_feature(enable = "avx2")]
6294#[inline]
6295unsafe fn q4tp_unpack_w(tile: *const u8) -> core::arch::x86_64::__m256i {
6296 unsafe {
6297 use core::arch::x86_64::*;
6298 let lomask = _mm_set1_epi8(0x0F);
6299 let b = _mm_loadu_si128(tile as *const __m128i);
6300 let lo = _mm_and_si128(b, lomask);
6301 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6302 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi))
6303 }
6304}
6305
6306#[cfg(target_arch = "x86_64")]
6313#[target_feature(enable = "avx2")]
6314#[inline]
6315unsafe fn q4tp_madd_wx(
6316 w: core::arch::x86_64::__m256i,
6317 xp: *const i8,
6318) -> core::arch::x86_64::__m256i {
6319 unsafe {
6320 use core::arch::x86_64::*;
6321 let ones = _mm256_set1_epi16(1);
6322 let x = _mm256_loadu_si256(xp as *const __m256i);
6323 _mm256_madd_epi16(_mm256_maddubs_epi16(w, x), ones)
6324 }
6325}
6326
6327#[cfg(target_arch = "x86_64")]
6332#[target_feature(enable = "avx2")]
6333#[inline]
6334unsafe fn q4tp_hsum256(acc: core::arch::x86_64::__m256) -> f32 {
6335 unsafe {
6336 use core::arch::x86_64::*;
6337 let a = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps::<1>(acc));
6338 let b = _mm_add_ps(a, _mm_movehl_ps(a, a));
6339 let c = _mm_add_ss(b, _mm_shuffle_ps(b, b, 0x55));
6340 _mm_cvtss_f32(c)
6341 }
6342}
6343
6344#[cfg(target_arch = "x86_64")]
6352#[target_feature(enable = "avx2")]
6353#[inline]
6354unsafe fn q4tp_cell_dot_v(
6355 nib: *const u8,
6356 r: usize,
6357 gpr: usize,
6358 xp: *const i8,
6359 sc: *const f32,
6360 xg8: *const f32,
6361) -> f32 {
6362 unsafe {
6363 use core::arch::x86_64::*;
6364 let mut acc = _mm256_setzero_ps();
6365 for gi in 0..gpr {
6366 let w = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
6367 let d = q4tp_madd_wx(w, xp.add(gi * GROUP_SIZE));
6368 let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), _mm256_set1_ps(*xg8.add(gi)));
6369 acc = _mm256_add_ps(acc, _mm256_mul_ps(t, _mm256_set1_ps(*sc.add(gi))));
6370 }
6371 q4tp_hsum256(acc)
6372 }
6373}
6374
6375#[cfg(target_arch = "x86_64")]
6388#[target_feature(enable = "avx2")]
6389unsafe fn q4tp_matmat_row_cols_v(
6390 v: &Q4tpView<'_>,
6391 r: usize,
6392 gpr: usize,
6393 acts: &[SplitAct],
6394 sc: &[f32],
6395 abase: usize,
6396 rows: usize,
6397 out: SendMut,
6398) {
6399 use core::arch::x86_64::*;
6400 debug_assert_eq!(sc.len(), gpr);
6401 let n = acts.len();
6402 let mut bi = 0usize;
6403 while bi + 4 <= n {
6404 let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
6405 let mut accs = [_mm256_setzero_ps(); 4];
6406 for gi in 0..gpr {
6407 let w = q4tp_unpack_w(v.nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB));
6408 let sb = _mm256_set1_ps(*sc.get_unchecked(gi));
6409 for (k, act) in a.iter().enumerate() {
6410 let xp = act.xq.as_ptr().add(gi * GROUP_SIZE);
6411 let d = q4tp_madd_wx(w, xp);
6412 let xgb = _mm256_set1_ps(*act.xg8.get_unchecked(gi));
6413 let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6414 accs[k] = _mm256_add_ps(accs[k], _mm256_mul_ps(t, sb));
6415 }
6416 }
6417 for k in 0..4 {
6418 let act = a[k];
6419 let mut acc = q4tp_hsum256(accs[k]) * act.sx;
6420 for &(j, xv) in &act.outliers {
6421 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6422 acc += w * s * xv;
6423 }
6424 unsafe { *out.at((abase + bi + k) * rows + r) = acc };
6426 }
6427 bi += 4;
6428 }
6429 while bi < n {
6431 let act = &acts[bi];
6432 let mut acc = q4tp_cell_dot_v(
6433 v.nib.as_ptr(),
6434 r,
6435 gpr,
6436 act.xq.as_ptr(),
6437 sc.as_ptr(),
6438 act.xg8.as_ptr(),
6439 ) * act.sx;
6440 for &(j, xv) in &act.outliers {
6441 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6442 acc += w * s * xv;
6443 }
6444 unsafe { *out.at((abase + bi) * rows + r) = acc };
6446 bi += 1;
6447 }
6448}
6449
6450#[cfg(target_arch = "x86_64")]
6459#[target_feature(enable = "avx2")]
6460unsafe fn q4tp_matvec_lanes<const L: usize>(
6461 v: &Q4tpView<'_>,
6462 start: usize,
6463 end: usize,
6464 gpr: usize,
6465 xq: &[i8],
6466 xg8: &[f32],
6467 sx: f32,
6468 outliers: &[(usize, f32)],
6469 sc: &mut [f32],
6470 out: SendMut,
6471) {
6472 use core::arch::x86_64::*;
6473 debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
6474 crate::cpuprof::count_lane_rows((end - start) as u64);
6475 let mut base = start;
6476 while base + L <= end {
6477 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
6478 let mut rows = [0usize; L];
6479 for l in 0..L {
6480 rows[l] = base + l;
6481 }
6482 let v4: [&Q4tpView; L] = [v; L];
6483 q4tp_scales_lanes(&v4, &rows, gpr, sc);
6484 drop(_t0);
6485 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
6486 let mut accs = [_mm256_setzero_ps(); L];
6487 for gi in 0..gpr {
6488 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
6489 let xgb = _mm256_set1_ps(*xg8.get_unchecked(gi));
6490 for l in 0..L {
6491 let w = q4tp_unpack_w(v.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
6492 let d = q4tp_madd_wx(w, xp);
6493 let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6494 let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
6495 accs[l] = _mm256_add_ps(accs[l], _mm256_mul_ps(t, sb));
6496 }
6497 }
6498 drop(_t1);
6499 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
6500 for l in 0..L {
6501 let r = rows[l];
6502 let scl = &sc[l * gpr..(l + 1) * gpr];
6503 let mut acc = q4tp_hsum256(accs[l]) * sx;
6504 for &(j, xv) in outliers {
6505 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
6506 acc += w * s * xv;
6507 }
6508 unsafe { *out.at(r) = acc };
6510 }
6511 drop(_t2);
6512 base += L;
6513 }
6514 for r in base..end {
6518 let scl = &mut sc[..gpr];
6519 v.scales_into(r, gpr, scl);
6520 let scl: &[f32] = scl;
6521 let mut acc = q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr())
6522 * sx;
6523 for &(j, xv) in outliers {
6524 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
6525 acc += w * s * xv;
6526 }
6527 unsafe { *out.at(r) = acc };
6529 }
6530}
6531
6532#[cfg(target_arch = "x86_64")]
6540#[target_feature(enable = "avx2")]
6541unsafe fn q4tp_gateup_lanes<const L: usize>(
6542 gv: &Q4tpView<'_>,
6543 uv: &Q4tpView<'_>,
6544 start: usize,
6545 end: usize,
6546 gpr: usize,
6547 xq: &[i8],
6548 xg8: &[f32],
6549 sx: f32,
6550 limit: f32,
6551 outliers: &[(usize, f32)],
6552 gsc: &mut [f32],
6553 usc: &mut [f32],
6554 out: SendMut,
6555) {
6556 use core::arch::x86_64::*;
6557 debug_assert!(
6558 L >= 2 && gsc.len() >= L * gpr && usc.len() >= L * gpr && xg8.len() == gpr
6559 );
6560 crate::cpuprof::count_lane_rows((end - start) as u64);
6561 let mut base = start;
6562 while base + L <= end {
6563 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
6564 let mut rows = [0usize; L];
6565 for l in 0..L {
6566 rows[l] = base + l;
6567 }
6568 let gv4: [&Q4tpView; L] = [gv; L];
6569 let uv4: [&Q4tpView; L] = [uv; L];
6570 q4tp_scales_lanes(&gv4, &rows, gpr, gsc);
6571 q4tp_scales_lanes(&uv4, &rows, gpr, usc);
6572 drop(_t0);
6573 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
6574 let mut ag = [_mm256_setzero_ps(); L];
6575 let mut au = [_mm256_setzero_ps(); L];
6576 for gi in 0..gpr {
6577 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
6578 let xgb = _mm256_set1_ps(*xg8.get_unchecked(gi));
6579 for l in 0..L {
6580 let w = q4tp_unpack_w(gv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
6581 let d = q4tp_madd_wx(w, xp);
6582 let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6583 let sb = _mm256_set1_ps(*gsc.get_unchecked(l * gpr + gi));
6584 ag[l] = _mm256_add_ps(ag[l], _mm256_mul_ps(t, sb));
6585 }
6586 for l in 0..L {
6587 let w = q4tp_unpack_w(uv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
6588 let d = q4tp_madd_wx(w, xp);
6589 let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6590 let sb = _mm256_set1_ps(*usc.get_unchecked(l * gpr + gi));
6591 au[l] = _mm256_add_ps(au[l], _mm256_mul_ps(t, sb));
6592 }
6593 }
6594 drop(_t1);
6595 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
6596 for l in 0..L {
6597 let r = rows[l];
6598 let gs = &gsc[l * gpr..(l + 1) * gpr];
6599 let us = &usc[l * gpr..(l + 1) * gpr];
6600 let mut g = q4tp_hsum256(ag[l]) * sx;
6601 let mut u = q4tp_hsum256(au[l]) * sx;
6602 for &(j, xv) in outliers {
6603 let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, gs);
6604 g += w * s * xv;
6605 let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, us);
6606 u += w * s * xv;
6607 }
6608 unsafe { *out.at(r) = silu_mul_limited(g, u, limit) };
6610 }
6611 drop(_t2);
6612 base += L;
6613 }
6614 for r in base..end {
6616 let g = &mut gsc[..gpr];
6617 let u = &mut usc[..gpr];
6618 gv.scales_into(r, gpr, g);
6619 uv.scales_into(r, gpr, u);
6620 let (g, u): (&[f32], &[f32]) = (g, u);
6621 let mut gv_ =
6622 q4tp_cell_dot_v(gv.nib.as_ptr(), r, gpr, xq.as_ptr(), g.as_ptr(), xg8.as_ptr()) * sx;
6623 let mut uv_ =
6624 q4tp_cell_dot_v(uv.nib.as_ptr(), r, gpr, xq.as_ptr(), u.as_ptr(), xg8.as_ptr()) * sx;
6625 for &(j, xv) in outliers {
6626 let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, g);
6627 gv_ += w * s * xv;
6628 let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, u);
6629 uv_ += w * s * xv;
6630 }
6631 unsafe { *out.at(r) = silu_mul_limited(gv_, uv_, limit) };
6633 }
6634}
6635
6636#[cfg(target_arch = "x86_64")]
6644#[target_feature(enable = "avx2")]
6645unsafe fn q4tp_many_lanes<const L: usize>(
6646 views: &[Q4tpView<'_>],
6647 locate: &(dyn Fn(usize) -> (usize, usize)),
6648 start: usize,
6649 end: usize,
6650 gpr: usize,
6651 xq: &[i8],
6652 xg8: &[f32],
6653 sx: f32,
6654 outliers: &[(usize, f32)],
6655 sc: &mut [f32],
6656 outs: &[SendMut],
6657) {
6658 use core::arch::x86_64::*;
6659 debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
6660 crate::cpuprof::count_lane_rows((end - start) as u64);
6661 let mut base = start;
6662 while base + L <= end {
6663 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
6664 let mut tk = [0usize; L];
6665 let mut rows = [0usize; L];
6666 for l in 0..L {
6667 let (t, r) = locate(base + l);
6668 tk[l] = t;
6669 rows[l] = r;
6670 }
6671 let v4: [&Q4tpView; L] = std::array::from_fn(|l| &views[tk[l]]);
6672 q4tp_scales_lanes(&v4, &rows, gpr, sc);
6673 drop(_t0);
6674 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
6675 let mut accs = [_mm256_setzero_ps(); L];
6676 for gi in 0..gpr {
6677 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
6678 let xgb = _mm256_set1_ps(*xg8.get_unchecked(gi));
6679 for l in 0..L {
6680 let w = q4tp_unpack_w(
6681 views[tk[l]].nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB),
6682 );
6683 let d = q4tp_madd_wx(w, xp);
6684 let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6685 let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
6686 accs[l] = _mm256_add_ps(accs[l], _mm256_mul_ps(t, sb));
6687 }
6688 }
6689 drop(_t1);
6690 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
6691 for l in 0..L {
6692 let (t, r) = (tk[l], rows[l]);
6693 let scl = &sc[l * gpr..(l + 1) * gpr];
6694 let mut acc = q4tp_hsum256(accs[l]) * sx;
6695 for &(j, xv) in outliers {
6696 let (w, s) = q4tp_outlier(views[t].nib, r, gpr, j, scl);
6697 acc += w * s * xv;
6698 }
6699 unsafe { *outs[t].at(r) = acc };
6701 }
6702 drop(_t2);
6703 base += L;
6704 }
6705 for flat in base..end {
6707 let (t, r) = locate(flat);
6708 let v = &views[t];
6709 let scl = &mut sc[..gpr];
6710 v.scales_into(r, gpr, scl);
6711 let scl: &[f32] = scl;
6712 let mut acc =
6713 q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr()) * sx;
6714 for &(j, xv) in outliers {
6715 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
6716 acc += w * s * xv;
6717 }
6718 unsafe { *outs[t].at(r) = acc };
6720 }
6721}
6722
6723#[inline]
6726fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6727 #[cfg(target_arch = "x86_64")]
6728 if avx2_enabled() {
6729 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
6731 }
6732 q4tp_row_float_scalar(nib, r, gpr, x, scales)
6733}
6734
6735#[inline]
6736fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6737 let mut acc = 0f32;
6738 for gi in 0..gpr {
6739 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6740 let s = scales[gi];
6741 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6742 let mut ga = 0f32;
6743 for (k, &b) in tile.iter().enumerate() {
6744 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6745 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6746 }
6747 acc += ga * s;
6748 }
6749 acc
6750}
6751
6752#[cfg(target_arch = "x86_64")]
6756#[target_feature(enable = "avx2")]
6757unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6758 unsafe {
6761 use core::arch::x86_64::*;
6762 let mask = _mm_set1_epi8(15);
6763 let eight = _mm_set1_epi8(8);
6764 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
6765 let mut acc = 0.0f32;
6766 for gi in 0..gpr {
6767 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
6768 let lo = _mm_and_si128(packed, mask);
6769 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
6770 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
6771 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
6772 let xp = x.as_ptr().add(gi * GROUP_SIZE);
6773 let a = _mm256_mul_ps(
6774 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
6775 _mm256_loadu_ps(xp),
6776 );
6777 let b = _mm256_mul_ps(
6778 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
6779 _mm256_loadu_ps(xp.add(8)),
6780 );
6781 let c = _mm256_mul_ps(
6782 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
6783 _mm256_loadu_ps(xp.add(16)),
6784 );
6785 let d = _mm256_mul_ps(
6786 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
6787 _mm256_loadu_ps(xp.add(24)),
6788 );
6789 let mut pairs = [0.0f32; 16];
6790 _mm256_storeu_ps(
6791 pairs.as_mut_ptr(),
6792 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
6793 );
6794 _mm256_storeu_ps(
6795 pairs.as_mut_ptr().add(8),
6796 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
6797 );
6798 let mut ga = 0.0f32;
6799 for v in pairs {
6800 ga += v;
6801 }
6802 acc += ga * scales[gi];
6803 }
6804 acc
6805 }
6806}
6807
6808#[inline]
6811fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
6812 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
6813 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
6814 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
6815 ((n as i32 - 8) as f32, scales[gi])
6816}
6817
6818fn q4tp_matvec(
6820 bytes: &[u8],
6821 x: &[f32],
6822 rows: usize,
6823 cols: usize,
6824 out: &mut [f32],
6825 pool: Option<&Pool>,
6826) {
6827 debug_assert_eq!(out.len(), rows);
6828 let gpr = cols / GROUP_SIZE;
6829 let v = Q4tpView::new(bytes, rows, cols);
6830 let out_addr = SendMut(out.as_mut_ptr());
6831 if a8w8_enabled() {
6832 let act = split_act(x);
6833 #[cfg(target_arch = "x86_64")]
6834 let lanes = q4tp_lanes();
6835 let run = |start: usize, end: usize| {
6836 #[cfg(target_arch = "x86_64")]
6837 if lanes > 1 && act.xg8.len() == gpr {
6838 with_krow(gpr * lanes, |sc| {
6841 unsafe {
6843 match lanes {
6844 2 => q4tp_matvec_lanes::<2>(
6845 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
6846 &act.outliers, sc, out_addr,
6847 ),
6848 8 => q4tp_matvec_lanes::<8>(
6849 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
6850 &act.outliers, sc, out_addr,
6851 ),
6852 _ => q4tp_matvec_lanes::<4>(
6853 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
6854 &act.outliers, sc, out_addr,
6855 ),
6856 }
6857 }
6858 });
6859 return;
6860 }
6861 crate::cpuprof::count_plain_rows((end - start) as u64);
6862 with_krow(gpr, |sc| {
6864 for r in start..end {
6865 v.scales_into(r, gpr, sc);
6866 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
6867 for &(j, xv) in &act.outliers {
6868 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6869 acc += w * s * xv;
6870 }
6871 unsafe { *out_addr.at(r) = acc }
6873 }
6874 })
6875 };
6876 dispatch_rows(pool, rows, &run);
6877 return;
6878 }
6879 let run = |start: usize, end: usize| {
6880 with_krow(gpr, |sc| {
6881 for r in start..end {
6882 v.scales_into(r, gpr, sc);
6883 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
6885 }
6886 })
6887 };
6888 dispatch_rows(pool, rows, &run);
6889}
6890
6891#[allow(clippy::too_many_arguments)]
6894fn q4tp_matvec2(
6895 bytes: &[u8],
6896 x1: &[f32],
6897 x2: &[f32],
6898 rows: usize,
6899 cols: usize,
6900 o1: &mut [f32],
6901 o2: &mut [f32],
6902 pool: Option<&Pool>,
6903) {
6904 let gpr = cols / GROUP_SIZE;
6905 let v = Q4tpView::new(bytes, rows, cols);
6906 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6907 let run = |start: usize, end: usize| {
6908 let mut sc = vec![0f32; gpr];
6909 for r in start..end {
6910 v.scales_into(r, gpr, &mut sc);
6911 unsafe {
6913 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
6914 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
6915 }
6916 }
6917 };
6918 dispatch_rows(pool, rows, &run);
6919}
6920
6921#[inline]
6924fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
6925 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
6926 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
6927 let c = (byte >> (2 * (k % 4))) & 3;
6928 (c as f32 - 1.5, scales[gi])
6929}
6930
6931#[cfg(target_arch = "x86_64")]
6932const Q2TP_DECODE_U32: [u32; 256] = {
6933 let mut tab = [0u32; 256];
6934 let mut b = 0usize;
6935 while b < 256 {
6936 tab[b] = ((b as u32) & 3)
6937 | ((((b as u32) >> 2) & 3) << 8)
6938 | ((((b as u32) >> 4) & 3) << 16)
6939 | ((((b as u32) >> 6) & 3) << 24);
6940 b += 1;
6941 }
6942 tab
6943};
6944
6945#[cfg(target_arch = "x86_64")]
6949#[target_feature(enable = "avx2")]
6950unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
6951 use core::arch::x86_64::*;
6952 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
6953 let codes = _mm256_setr_epi32(
6954 Q2TP_DECODE_U32[ch[0] as usize] as i32,
6955 Q2TP_DECODE_U32[ch[1] as usize] as i32,
6956 Q2TP_DECODE_U32[ch[2] as usize] as i32,
6957 Q2TP_DECODE_U32[ch[3] as usize] as i32,
6958 Q2TP_DECODE_U32[ch[4] as usize] as i32,
6959 Q2TP_DECODE_U32[ch[5] as usize] as i32,
6960 Q2TP_DECODE_U32[ch[6] as usize] as i32,
6961 Q2TP_DECODE_U32[ch[7] as usize] as i32,
6962 );
6963 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
6964 let pair = _mm256_maddubs_epi16(codes, xv);
6965 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
6966 let sum128 = _mm_add_epi32(
6967 _mm256_castsi256_si128(quad),
6968 _mm256_extracti128_si256(quad, 1),
6969 );
6970 let sum64 = _mm_hadd_epi32(sum128, sum128);
6971 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
6972}
6973
6974#[cfg(target_arch = "x86_64")]
6975#[inline]
6976fn q2tp_avx2_enabled() -> bool {
6977 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6978 *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
6979}
6980
6981#[inline]
6988fn dot_q2tp_row_i8(
6989 chunks: &[u8],
6990 r: usize,
6991 gpr: usize,
6992 xq: &[i8],
6993 gsum: &[i32],
6994 scales: &[f32],
6995) -> f32 {
6996 let mut acc = 0f32;
6997 let base = r * gpr * Q2TP_CHUNK;
6998 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6999 let mut codes = [0i8; GROUP_SIZE];
7000 #[cfg(target_arch = "x86_64")]
7001 let avx2 = q2tp_avx2_enabled();
7006 for gi in 0..gpr {
7007 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
7008 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7009 #[cfg(target_arch = "aarch64")]
7010 let dot = unsafe {
7016 use core::arch::aarch64::*;
7017 let b = vld1_u8(ch.as_ptr());
7018 let three = vdup_n_u8(3);
7019 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
7020 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
7021 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
7022 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
7023 let x4 = vld4_s8(xg.as_ptr());
7024 let mut acc4 = vdupq_n_s32(0);
7025 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
7026 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
7027 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
7028 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
7029 vaddvq_s32(acc4)
7030 };
7031 #[cfg(target_arch = "x86_64")]
7032 let dot: i32 = if avx2 {
7033 unsafe { q2tp_code_dot_avx2(ch, xg) }
7036 } else {
7037 ch.iter()
7038 .enumerate()
7039 .map(|(k, &b)| {
7040 ((b & 3) as i32) * xg[k * 4] as i32
7041 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
7042 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
7043 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
7044 })
7045 .sum()
7046 };
7047 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7048 let dot: i32 = {
7049 for (k, &b) in ch.iter().enumerate() {
7050 codes[k * 4] = (b & 3) as i8;
7051 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
7052 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
7053 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
7054 }
7055 codes
7056 .iter()
7057 .zip(xg)
7058 .map(|(&c, &x)| c as i32 * x as i32)
7059 .sum()
7060 };
7061 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
7062 }
7063 acc
7064}
7065
7066fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7070 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
7071}
7072
7073#[inline]
7077fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7078 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
7079}
7080
7081#[inline]
7082fn q2tp_row_exact_center(
7083 chunks: &[u8],
7084 r: usize,
7085 gpr: usize,
7086 x: &[f32],
7087 scales: &[f32],
7088 center: f32,
7089) -> f32 {
7090 let mut acc = 0f32;
7091 for gi in 0..gpr {
7092 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
7093 let s = scales[gi];
7094 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7095 let mut g = 0f32;
7096 for (k, &b) in ch.iter().enumerate() {
7097 g += ((b & 3) as f32 - center) * xb[k * 4]
7098 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
7099 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
7100 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
7101 }
7102 acc += s * g;
7103 }
7104 acc
7105}
7106
7107fn q2tp_matvec(
7108 bytes: &[u8],
7109 x: &[f32],
7110 rows: usize,
7111 cols: usize,
7112 out: &mut [f32],
7113 pool: Option<&Pool>,
7114) {
7115 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
7116}
7117
7118fn q2tp_affine_matvec(
7119 bytes: &[u8],
7120 x: &[f32],
7121 rows: usize,
7122 cols: usize,
7123 out: &mut [f32],
7124 pool: Option<&Pool>,
7125) {
7126 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
7127}
7128
7129fn q2tp_matvec_mode(
7130 bytes: &[u8],
7131 x: &[f32],
7132 rows: usize,
7133 cols: usize,
7134 out: &mut [f32],
7135 pool: Option<&Pool>,
7136 affine: bool,
7137) {
7138 debug_assert_eq!(out.len(), rows);
7139 let gpr = cols / GROUP_SIZE;
7140 let v = Q4tpView::new_q2(bytes, rows, cols);
7141 let out_addr = SendMut(out.as_mut_ptr());
7142 if !affine && a8w8_enabled() {
7147 let act = split_act(x);
7148 let gsum = q1_group_sums(&act.xq, gpr);
7149 let (act, gsum) = (&act, &gsum);
7150 let run = move |start: usize, end: usize| {
7151 with_krow(gpr, |sc| {
7152 for r in start..end {
7153 v.scales_into(r, gpr, sc);
7154 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
7155 for &(j, xv) in &act.outliers {
7156 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
7157 acc += w * s * xv;
7158 }
7159 unsafe { *out_addr.at(r) = acc };
7161 }
7162 })
7163 };
7164 dispatch_rows(pool, rows, &run);
7165 return;
7166 }
7167 let run = |start: usize, end: usize| {
7168 with_krow(gpr, |sc| {
7169 for r in start..end {
7170 v.scales_into(r, gpr, sc);
7171 unsafe {
7173 *out_addr.at(r) = if affine {
7174 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
7175 } else {
7176 q2tp_row_exact(v.nib, r, gpr, x, sc)
7177 }
7178 };
7179 }
7180 })
7181 };
7182 dispatch_rows(pool, rows, &run);
7183}
7184
7185#[allow(clippy::too_many_arguments)]
7187fn q2tp_matvec2(
7188 bytes: &[u8],
7189 x1: &[f32],
7190 x2: &[f32],
7191 rows: usize,
7192 cols: usize,
7193 o1: &mut [f32],
7194 o2: &mut [f32],
7195 pool: Option<&Pool>,
7196) {
7197 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
7198}
7199
7200#[allow(clippy::too_many_arguments)]
7201fn q2tp_affine_matvec2(
7202 bytes: &[u8],
7203 x1: &[f32],
7204 x2: &[f32],
7205 rows: usize,
7206 cols: usize,
7207 o1: &mut [f32],
7208 o2: &mut [f32],
7209 pool: Option<&Pool>,
7210) {
7211 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
7212}
7213
7214#[allow(clippy::too_many_arguments)]
7215fn q2tp_matvec2_mode(
7216 bytes: &[u8],
7217 x1: &[f32],
7218 x2: &[f32],
7219 rows: usize,
7220 cols: usize,
7221 o1: &mut [f32],
7222 o2: &mut [f32],
7223 pool: Option<&Pool>,
7224 affine: bool,
7225) {
7226 let gpr = cols / GROUP_SIZE;
7227 let v = Q4tpView::new_q2(bytes, rows, cols);
7228 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7229 let run = |start: usize, end: usize| {
7230 let mut sc = vec![0f32; gpr];
7231 for r in start..end {
7232 v.scales_into(r, gpr, &mut sc);
7233 unsafe {
7235 *p1.at(r) = if affine {
7236 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
7237 } else {
7238 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
7239 };
7240 *p2.at(r) = if affine {
7241 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
7242 } else {
7243 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
7244 };
7245 }
7246 }
7247 };
7248 dispatch_rows(pool, rows, &run);
7249}
7250
7251pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
7258 let gpr = cols / GROUP_SIZE;
7263 let v = Q4tpView::new_q2(bytes, rows, cols);
7264 with_krow(gpr, |sc| {
7265 for r in 0..rows {
7266 v.scales_into(r, gpr, sc);
7267 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
7268 }
7269 });
7270}
7271
7272pub fn q2tp_affine_matvec_for_test(
7275 bytes: &[u8],
7276 x: &[f32],
7277 rows: usize,
7278 cols: usize,
7279 out: &mut [f32],
7280) {
7281 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
7282}
7283
7284pub fn q2tp_matmat_for_test(
7285 bytes: &[u8],
7286 xs_all: &[f32],
7287 b: usize,
7288 rows: usize,
7289 cols: usize,
7290 out: &mut [f32],
7291) {
7292 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
7293}
7294
7295fn q2tp_matmat(
7296 bytes: &[u8],
7297 xs_all: &[f32],
7298 b: usize,
7299 rows: usize,
7300 cols: usize,
7301 out: &mut [f32],
7302 pool: Option<&Pool>,
7303) {
7304 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
7305}
7306
7307fn q2tp_affine_matmat(
7308 bytes: &[u8],
7309 xs_all: &[f32],
7310 b: usize,
7311 rows: usize,
7312 cols: usize,
7313 out: &mut [f32],
7314 pool: Option<&Pool>,
7315) {
7316 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
7317}
7318
7319fn q2tp_matmat_mode(
7320 bytes: &[u8],
7321 xs_all: &[f32],
7322 b: usize,
7323 rows: usize,
7324 cols: usize,
7325 out: &mut [f32],
7326 pool: Option<&Pool>,
7327 affine: bool,
7328) {
7329 debug_assert_eq!(out.len(), b * rows);
7330 let gpr = cols / GROUP_SIZE;
7331 let v = Q4tpView::new_q2(bytes, rows, cols);
7332 let out_addr = SendMut(out.as_mut_ptr());
7333 let run = |start: usize, end: usize| {
7334 let mut sc = vec![0f32; gpr];
7335 for r in start..end {
7336 v.scales_into(r, gpr, &mut sc);
7337 for bi in 0..b {
7338 let x = &xs_all[bi * cols..(bi + 1) * cols];
7339 unsafe {
7341 *out_addr.at(bi * rows + r) = if affine {
7342 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
7343 } else {
7344 q2tp_row_exact(v.nib, r, gpr, x, &sc)
7345 }
7346 };
7347 }
7348 }
7349 };
7350 dispatch_rows(pool, rows, &run);
7351}
7352
7353#[cfg(target_arch = "aarch64")]
7361#[target_feature(enable = "neon,dotprod")]
7362unsafe fn dot_q4tp_row_1x4_sdot_v1(
7363 nib: &[u8],
7364 r: usize,
7365 gpr: usize,
7366 xs: [&[i8]; 4],
7367 scales: &[f32],
7368) -> [f32; 4] {
7369 unsafe {
7370 use core::arch::aarch64::*;
7371 use core::arch::asm;
7372 let lomask = vdupq_n_u8(0x0F);
7373 let eight = vdupq_n_s8(8);
7374 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
7375 for gi in 0..gpr {
7376 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7377 let s = *scales.get_unchecked(gi);
7378 let bb = vld1q_u8(t);
7379 let lo = vandq_u8(bb, lomask);
7380 let hi = vshrq_n_u8::<4>(bb);
7381 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7382 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7383 let mut d = [0f32; 4];
7384 for (k, dk) in d.iter_mut().enumerate() {
7385 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
7386 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
7387 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7388 asm!(
7389 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7390 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7391 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7392 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7393 options(pure, nomem, nostack),
7394 );
7395 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7396 }
7397 f0 += d[0];
7398 f1 += d[1];
7399 f2 += d[2];
7400 f3 += d[3];
7401 }
7402 [f0, f1, f2, f3]
7403 }
7404}
7405
7406#[allow(dead_code)]
7413static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
7414
7415#[cfg(test)]
7418static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
7419
7420#[cfg(target_arch = "x86_64")]
7426fn q4tp_blocked_x86() -> bool {
7427 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
7428 1 => false,
7429 2 => avx512vnni_enabled(),
7434 _ => blocked_enabled() && avx512vnni_enabled(),
7438 }
7439}
7440
7441#[cfg(target_arch = "aarch64")]
7443#[allow(dead_code)]
7444fn q4tp_v1() -> bool {
7445 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
7446 1 => true,
7447 2 => false,
7448 _ => {
7449 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7450 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
7451 }
7452 }
7453}
7454
7455#[cfg(target_arch = "x86_64")]
7467#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
7468unsafe fn dot_q4tp_2x8_avx512(
7469 nib: &[u8],
7470 r0: usize,
7471 gpr: usize,
7472 xs: [&[i8]; 8],
7473 sc0: &[f32],
7474 sc1: &[f32],
7475) -> [[f32; 8]; 2] {
7476 unsafe {
7479 use core::arch::x86_64::*;
7480 let lomask = _mm256_set1_epi8(0x0F);
7481 let eight = _mm256_set1_epi8(8);
7482 let zero = _mm512_setzero_si512();
7483 let mut v0 = [_mm512_setzero_ps(); 8];
7484 let mut v1 = [_mm512_setzero_ps(); 8];
7485 let pairs = gpr / 2;
7486 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
7487 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7488 let bb = _mm256_loadu_si256(t as *const __m256i);
7489 let lo = _mm256_and_si256(bb, lomask);
7490 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
7491 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
7492 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
7493 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
7494 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
7495 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
7496 };
7497 for gp in 0..pairs {
7498 let gi = gp * 2;
7499 let (wa0, neg0) = unpack(r0, gi);
7500 let (wa1, neg1) = unpack(r0 + 1, gi);
7501 let off = gi * GROUP_SIZE;
7502 let sv = |sc: &[f32]| {
7503 _mm512_insertf32x8::<1>(
7504 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
7505 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
7506 )
7507 };
7508 let s0 = sv(sc0);
7509 let s1 = sv(sc1);
7510 for k in 0..8 {
7511 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
7512 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
7513 zero,
7514 wa0,
7515 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
7516 ));
7517 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
7518 zero,
7519 wa1,
7520 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
7521 ));
7522 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
7523 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
7524 }
7525 }
7526 let mut acc = [[0f32; 8]; 2];
7527 for k in 0..8 {
7528 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
7529 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
7530 }
7531 if gpr % 2 == 1 {
7532 let off = (gpr - 1) * GROUP_SIZE;
7533 for j in off..off + GROUP_SIZE {
7534 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
7535 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
7536 for k in 0..8 {
7537 let x = *xs[k].get_unchecked(j) as f32;
7538 acc[0][k] += w0 * sa * x;
7539 acc[1][k] += w1 * sb * x;
7540 }
7541 }
7542 }
7543 acc
7544 }
7545}
7546
7547#[cfg(target_arch = "x86_64")]
7552#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
7553unsafe fn dot_q4tp_row_1x8_avx512(
7554 nib: &[u8],
7555 r: usize,
7556 gpr: usize,
7557 xs: [&[i8]; 8],
7558 scales: &[f32],
7559) -> [f32; 8] {
7560 unsafe {
7562 use core::arch::x86_64::*;
7563 let lomask = _mm256_set1_epi8(0x0F);
7564 let eight = _mm256_set1_epi8(8);
7565 let zero = _mm512_setzero_si512();
7566 let (mut v0, mut v1, mut v2, mut v3) = (
7567 _mm512_setzero_ps(),
7568 _mm512_setzero_ps(),
7569 _mm512_setzero_ps(),
7570 _mm512_setzero_ps(),
7571 );
7572 let (mut v4, mut v5, mut v6, mut v7) = (
7573 _mm512_setzero_ps(),
7574 _mm512_setzero_ps(),
7575 _mm512_setzero_ps(),
7576 _mm512_setzero_ps(),
7577 );
7578 let pairs = gpr / 2;
7579 for gp in 0..pairs {
7580 let gi = gp * 2;
7581 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7582 let bb = _mm256_loadu_si256(t as *const __m256i);
7583 let lo = _mm256_and_si256(bb, lomask);
7584 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
7585 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
7590 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
7591 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
7592 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
7593 let wabs = _mm512_abs_epi8(w);
7594 let neg = _mm512_movepi8_mask(w);
7595 let off = gi * GROUP_SIZE;
7596 let sv = _mm512_insertf32x8::<1>(
7597 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
7598 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
7599 );
7600 let dot = |x: &[i8]| -> __m512 {
7601 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
7602 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
7603 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
7604 };
7605 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
7606 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
7607 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
7608 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
7609 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
7610 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
7611 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
7612 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
7613 }
7614 let mut acc = [
7615 _mm512_reduce_add_ps(v0),
7616 _mm512_reduce_add_ps(v1),
7617 _mm512_reduce_add_ps(v2),
7618 _mm512_reduce_add_ps(v3),
7619 _mm512_reduce_add_ps(v4),
7620 _mm512_reduce_add_ps(v5),
7621 _mm512_reduce_add_ps(v6),
7622 _mm512_reduce_add_ps(v7),
7623 ];
7624 if gpr % 2 == 1 {
7627 let off = (gpr - 1) * GROUP_SIZE;
7628 for j in off..off + GROUP_SIZE {
7629 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
7630 let ws = w * s;
7631 for k in 0..8 {
7632 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
7633 }
7634 }
7635 }
7636 acc
7637 }
7638}
7639
7640#[cfg(target_arch = "x86_64")]
7653#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
7654unsafe fn dot_q4tp_row_1x4_avx512(
7655 nib: &[u8],
7656 r: usize,
7657 gpr: usize,
7658 xs: [&[i8]; 4],
7659 scales: &[f32],
7660) -> [f32; 4] {
7661 unsafe {
7663 use core::arch::x86_64::*;
7664 let lomask = _mm256_set1_epi8(0x0F);
7665 let eight = _mm256_set1_epi8(8);
7666 let zero = _mm512_setzero_si512();
7667 let (mut v0, mut v1, mut v2, mut v3) = (
7668 _mm512_setzero_ps(),
7669 _mm512_setzero_ps(),
7670 _mm512_setzero_ps(),
7671 _mm512_setzero_ps(),
7672 );
7673 let pairs = gpr / 2;
7674 for gp in 0..pairs {
7675 let gi = gp * 2;
7676 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7677 let bb = _mm256_loadu_si256(t as *const __m256i);
7678 let lo = _mm256_and_si256(bb, lomask);
7679 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
7680 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
7685 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
7686 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
7687 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
7688 let wabs = _mm512_abs_epi8(w);
7689 let neg = _mm512_movepi8_mask(w);
7690 let off = gi * GROUP_SIZE;
7691 let sv = _mm512_insertf32x8::<1>(
7692 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
7693 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
7694 );
7695 let dot = |x: &[i8]| -> __m512 {
7696 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
7697 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
7698 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
7699 };
7700 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
7701 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
7702 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
7703 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
7704 }
7705 let mut acc = [
7706 _mm512_reduce_add_ps(v0),
7707 _mm512_reduce_add_ps(v1),
7708 _mm512_reduce_add_ps(v2),
7709 _mm512_reduce_add_ps(v3),
7710 ];
7711 if gpr % 2 == 1 {
7714 let off = (gpr - 1) * GROUP_SIZE;
7715 for j in off..off + GROUP_SIZE {
7716 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
7717 let ws = w * s;
7718 for k in 0..4 {
7719 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
7720 }
7721 }
7722 }
7723 acc
7724 }
7725}
7726
7727#[cfg(target_arch = "aarch64")]
7731#[target_feature(enable = "neon,dotprod")]
7732unsafe fn dot_q4tp_row_1x4_sdot(
7733 nib: &[u8],
7734 r: usize,
7735 gpr: usize,
7736 xs: [&[i8]; 4],
7737 scales: &[f32],
7738) -> [f32; 4] {
7739 unsafe {
7741 use core::arch::aarch64::*;
7742 use core::arch::asm;
7743 let lomask = vdupq_n_u8(0x0F);
7744 let eight = vdupq_n_s8(8);
7745 let (mut v0, mut v1, mut v2, mut v3) = (
7761 vdupq_n_f32(0.0),
7762 vdupq_n_f32(0.0),
7763 vdupq_n_f32(0.0),
7764 vdupq_n_f32(0.0),
7765 );
7766 for gi in 0..gpr {
7767 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7768 let s = *scales.get_unchecked(gi);
7769 let bb = vld1q_u8(t);
7770 let lo = vandq_u8(bb, lomask);
7771 let hi = vshrq_n_u8::<4>(bb);
7772 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7773 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7774 let off = gi * GROUP_SIZE;
7775 let dot4 = |x: &[i8]| -> int32x4_t {
7776 let x0 = vld1q_s8(x.as_ptr().add(off));
7777 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
7778 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7779 asm!(
7780 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7781 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7782 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7783 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7784 options(pure, nomem, nostack),
7785 );
7786 vaddq_s32(a0, a1)
7787 };
7788 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
7789 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
7790 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
7791 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
7792 }
7793 [
7794 vaddvq_f32(v0),
7795 vaddvq_f32(v1),
7796 vaddvq_f32(v2),
7797 vaddvq_f32(v3),
7798 ]
7799 }
7800}
7801
7802fn q4tp_matmat(
7810 bytes: &[u8],
7811 xs_all: &[f32],
7812 b: usize,
7813 rows: usize,
7814 cols: usize,
7815 out: &mut [f32],
7816 pool: Option<&Pool>,
7817) {
7818 q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
7819}
7820
7821#[allow(clippy::too_many_arguments)]
7825fn q4tp_matmat_with(
7826 bytes: &[u8],
7827 xs_all: &[f32],
7828 b: usize,
7829 rows: usize,
7830 cols: usize,
7831 out: &mut [f32],
7832 pool: Option<&Pool>,
7833 exact: bool,
7834) {
7835 debug_assert_eq!(out.len(), b * rows);
7836 let gpr = cols / GROUP_SIZE;
7837 let v = Q4tpView::new(bytes, rows, cols);
7838
7839 #[cfg(target_os = "macos")]
7843 if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7844 dequant_matmat_accel(
7845 &|r, dst| {
7846 let mut sc = [0f32; 32];
7847 let mut scv;
7848 let s: &[f32] = if gpr <= 32 {
7849 v.scales_into(r, gpr, &mut sc);
7850 &sc[..gpr]
7851 } else {
7852 scv = vec![0f32; gpr];
7853 v.scales_into(r, gpr, &mut scv);
7854 &scv
7855 };
7856 for gi in 0..gpr {
7857 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
7858 for (k, &bb) in tile.iter().enumerate() {
7859 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
7860 dst[gi * GROUP_SIZE + k * 2 + 1] =
7861 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
7862 }
7863 }
7864 },
7865 xs_all,
7866 b,
7867 rows,
7868 cols,
7869 out,
7870 pool,
7871 );
7872 return;
7873 }
7874
7875 let out_addr = SendMut(out.as_mut_ptr());
7876 if a8w8_enabled() {
7877 let acts: Vec<SplitAct> = (0..b)
7878 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7879 .collect();
7880 let acts = &acts;
7881 #[cfg(target_arch = "aarch64")]
7888 let blocked_ok = sdot_enabled() && blocked_enabled();
7889 #[cfg(target_arch = "x86_64")]
7896 let blocked_ok = q4tp_blocked_x86() && !exact;
7897 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7898 let blocked_ok = {
7899 let _ = exact;
7900 false
7901 };
7902 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
7912 .ok()
7913 .and_then(|v| v.parse().ok())
7914 .filter(|v| *v > 0)
7915 .unwrap_or(256);
7916 let run = |start: usize, end: usize| {
7917 for abase in (0..acts.len()).step_by(panel_cols) {
7918 let alen = (acts.len() - abase).min(panel_cols);
7919 let mut sc = vec![0f32; gpr];
7920 #[cfg(target_arch = "x86_64")]
7921 let mut r_lo = start;
7922 #[cfg(target_arch = "x86_64")]
7923 if blocked_ok && alen >= 8 {
7924 let mut sc1 = vec![0f32; gpr];
7925 while r_lo + 2 <= end {
7926 v.scales_into(r_lo, gpr, &mut sc);
7927 v.scales_into(r_lo + 1, gpr, &mut sc1);
7928 let mut bi = 0usize;
7929 while bi + 8 <= alen {
7930 let xs = [
7931 acts[abase + bi].xq.as_slice(),
7932 acts[abase + bi + 1].xq.as_slice(),
7933 acts[abase + bi + 2].xq.as_slice(),
7934 acts[abase + bi + 3].xq.as_slice(),
7935 acts[abase + bi + 4].xq.as_slice(),
7936 acts[abase + bi + 5].xq.as_slice(),
7937 acts[abase + bi + 6].xq.as_slice(),
7938 acts[abase + bi + 7].xq.as_slice(),
7939 ];
7940 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
7941 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
7942 for k in 0..8 {
7943 let act = &acts[abase + bi + k];
7944 let mut acc = dr[k] * act.sx;
7945 for &(j, xv) in &act.outliers {
7946 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7947 acc += w * s * xv;
7948 }
7949 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
7951 }
7952 }
7953 bi += 8;
7954 }
7955 for row in [r_lo, r_lo + 1] {
7958 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
7959 for b2 in bi..alen {
7960 let act = &acts[abase + b2];
7961 let xs4 = [
7962 act.xq.as_slice(),
7963 act.xq.as_slice(),
7964 act.xq.as_slice(),
7965 act.xq.as_slice(),
7966 ];
7967 let d =
7968 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
7969 let mut acc = d[0] * act.sx;
7970 for &(j, xv) in &act.outliers {
7971 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7972 acc += w * s * xv;
7973 }
7974 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
7976 }
7977 }
7978 r_lo += 2;
7979 }
7980 }
7981 #[cfg(target_arch = "x86_64")]
7982 let row_start = r_lo;
7983 #[cfg(not(target_arch = "x86_64"))]
7984 let row_start = start;
7985 for r in row_start..end {
7986 v.scales_into(r, gpr, &mut sc);
7987 let mut bi = 0usize;
7988 #[cfg(target_arch = "x86_64")]
7989 if blocked_ok {
7990 while bi + 8 <= alen {
7991 let xs = [
7992 acts[abase + bi].xq.as_slice(),
7993 acts[abase + bi + 1].xq.as_slice(),
7994 acts[abase + bi + 2].xq.as_slice(),
7995 acts[abase + bi + 3].xq.as_slice(),
7996 acts[abase + bi + 4].xq.as_slice(),
7997 acts[abase + bi + 5].xq.as_slice(),
7998 acts[abase + bi + 6].xq.as_slice(),
7999 acts[abase + bi + 7].xq.as_slice(),
8000 ];
8001 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
8002 for k in 0..8 {
8003 let act = &acts[abase + bi + k];
8004 let mut acc = d[k] * act.sx;
8005 for &(j, xv) in &act.outliers {
8006 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8007 acc += w * s * xv;
8008 }
8009 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8011 }
8012 bi += 8;
8013 }
8014 while bi + 4 <= alen {
8015 let xs = [
8016 acts[abase + bi].xq.as_slice(),
8017 acts[abase + bi + 1].xq.as_slice(),
8018 acts[abase + bi + 2].xq.as_slice(),
8019 acts[abase + bi + 3].xq.as_slice(),
8020 ];
8021 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
8022 for k in 0..4 {
8023 let act = &acts[abase + bi + k];
8024 let mut acc = d[k] * act.sx;
8025 for &(j, xv) in &act.outliers {
8026 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8027 acc += w * s * xv;
8028 }
8029 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8031 }
8032 bi += 4;
8033 }
8034 }
8035 #[cfg(target_arch = "aarch64")]
8036 if blocked_ok {
8037 while bi + 4 <= alen {
8038 let xs = [
8039 acts[abase + bi].xq.as_slice(),
8040 acts[abase + bi + 1].xq.as_slice(),
8041 acts[abase + bi + 2].xq.as_slice(),
8042 acts[abase + bi + 3].xq.as_slice(),
8043 ];
8044 let d = unsafe {
8045 if exact || q4tp_v1() {
8046 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
8047 } else {
8048 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
8049 }
8050 };
8051 for k in 0..4 {
8052 let act = &acts[abase + bi + k];
8053 let mut acc = d[k] * act.sx;
8054 for &(j, xv) in &act.outliers {
8055 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8056 acc += w * s * xv;
8057 }
8058 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8060 }
8061 bi += 4;
8062 }
8063 }
8064 let _ = blocked_ok;
8065 #[cfg(target_arch = "x86_64")]
8071 if acts[abase..abase + alen]
8072 .iter()
8073 .all(|a| a.xg8.len() == gpr)
8074 {
8075 unsafe {
8078 q4tp_matmat_row_cols_v(
8079 &v,
8080 r,
8081 gpr,
8082 &acts[abase + bi..abase + alen],
8083 &sc,
8084 abase + bi,
8085 rows,
8086 out_addr,
8087 )
8088 };
8089 bi = alen;
8090 }
8091 while bi < alen {
8092 let act = &acts[abase + bi];
8093 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
8094 for &(j, xv) in &act.outliers {
8095 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8096 acc += w * s * xv;
8097 }
8098 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
8100 bi += 1;
8101 }
8102 }
8103 }
8104 };
8105 dispatch_rows(pool, rows, &run);
8106 return;
8107 }
8108
8109 let run = |start: usize, end: usize| {
8110 let mut sc = vec![0f32; gpr];
8111 for r in start..end {
8112 v.scales_into(r, gpr, &mut sc);
8113 for bi in 0..b {
8114 let x = &xs_all[bi * cols..(bi + 1) * cols];
8115 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
8117 }
8118 }
8119 };
8120 dispatch_rows(pool, rows, &run);
8121}
8122
8123fn q4t_matvec(
8125 bytes: &[u8],
8126 x: &[f32],
8127 rows: usize,
8128 cols: usize,
8129 out: &mut [f32],
8130 pool: Option<&Pool>,
8131) {
8132 debug_assert_eq!(out.len(), rows);
8133 let gpr = cols / GROUP_SIZE;
8134 let out_addr = SendMut(out.as_mut_ptr());
8135 if a8w8_enabled() {
8136 let act = split_act(x);
8137 let run = move |start: usize, end: usize| {
8138 for r in start..end {
8139 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8140 for &(j, xv) in &act.outliers {
8141 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8142 acc += w * s * xv;
8143 }
8144 unsafe { *out_addr.at(r) = acc };
8146 }
8147 };
8148 dispatch_rows(pool, rows, &run);
8149 return;
8150 }
8151 let run = move |start: usize, end: usize| {
8152 for r in start..end {
8153 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
8155 }
8156 };
8157 dispatch_rows(pool, rows, &run);
8158}
8159
8160#[allow(clippy::too_many_arguments)]
8162fn q4t_matvec2(
8163 bytes: &[u8],
8164 x1: &[f32],
8165 x2: &[f32],
8166 rows: usize,
8167 cols: usize,
8168 o1: &mut [f32],
8169 o2: &mut [f32],
8170 pool: Option<&Pool>,
8171) {
8172 let gpr = cols / GROUP_SIZE;
8173 let p1 = SendMut(o1.as_mut_ptr());
8174 let p2 = SendMut(o2.as_mut_ptr());
8175 if a8w8_enabled() {
8176 let a1 = split_act(x1);
8177 let a2 = split_act(x2);
8178 let run = move |start: usize, end: usize| {
8179 for r in start..end {
8180 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
8181 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
8182 for &(j, xv) in &a1.outliers {
8183 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8184 v1 += w * s * xv;
8185 }
8186 for &(j, xv) in &a2.outliers {
8187 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8188 v2 += w * s * xv;
8189 }
8190 unsafe {
8192 *p1.at(r) = v1;
8193 *p2.at(r) = v2;
8194 }
8195 }
8196 };
8197 dispatch_rows(pool, rows, &run);
8198 return;
8199 }
8200 let run = move |start: usize, end: usize| {
8201 for r in start..end {
8202 unsafe {
8204 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
8205 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
8206 }
8207 }
8208 };
8209 dispatch_rows(pool, rows, &run);
8210}
8211
8212#[allow(clippy::too_many_arguments)]
8214#[cfg(target_os = "macos")]
8220fn dequant_matmat_accel(
8221 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
8222 xs_all: &[f32],
8223 b: usize,
8224 rows: usize,
8225 cols: usize,
8226 out: &mut [f32],
8227 pool: Option<&Pool>,
8228) {
8229 const TR: usize = 2048;
8230 thread_local! {
8231 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8232 }
8233 WTILE.with(|wt| {
8234 let mut wtile = wt.borrow_mut();
8235 wtile.resize(TR * cols, 0.0);
8236 let mut r0 = 0usize;
8237 while r0 < rows {
8238 let tr = TR.min(rows - r0);
8239 let wt_addr = SendMut(wtile.as_mut_ptr());
8240 let run = |start: usize, end: usize| {
8241 for r in start..end {
8242 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
8244 dequant_row(r0 + r, dst);
8245 }
8246 };
8247 dispatch_rows(pool, tr, &run);
8248 unsafe {
8249 accel_blas::cblas_sgemm(
8250 101, 111, 112, b as i32,
8254 tr as i32,
8255 cols as i32,
8256 1.0,
8257 xs_all.as_ptr(),
8258 cols as i32,
8259 wtile.as_ptr(),
8260 cols as i32,
8261 0.0,
8262 out.as_mut_ptr().add(r0),
8263 rows as i32,
8264 );
8265 }
8266 r0 += tr;
8267 }
8268 });
8269}
8270
8271fn q4t_matmat(
8272 bytes: &[u8],
8273 xs_all: &[f32],
8274 b: usize,
8275 rows: usize,
8276 cols: usize,
8277 out: &mut [f32],
8278 pool: Option<&Pool>,
8279) {
8280 debug_assert_eq!(out.len(), b * rows);
8281 let gpr = cols / GROUP_SIZE;
8282 #[cfg(target_os = "macos")]
8286 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8287 dequant_matmat_accel(
8288 &|r, dst| {
8289 for gi in 0..gpr {
8290 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
8291 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8292 for (k, &bb) in tile[2..].iter().enumerate() {
8293 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
8294 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
8295 }
8296 }
8297 },
8298 xs_all,
8299 b,
8300 rows,
8301 cols,
8302 out,
8303 pool,
8304 );
8305 return;
8306 }
8307 let out_addr = SendMut(out.as_mut_ptr());
8308 if a8w8_enabled() {
8309 let acts: Vec<SplitAct> = (0..b)
8310 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8311 .collect();
8312 let acts = &acts;
8313 #[cfg(target_arch = "x86_64")]
8314 let blocked_ok = avx2_enabled() && blocked_enabled();
8315 #[cfg(target_arch = "aarch64")]
8316 let blocked_ok = sdot_enabled() && blocked_enabled();
8317 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
8318 let blocked_ok = false;
8319 let run = move |start: usize, end: usize| {
8320 for r in start..end {
8321 let mut bi = 0usize;
8322 #[cfg(target_arch = "aarch64")]
8323 if blocked_ok {
8324 while bi + 4 <= acts.len() {
8325 let xs = [
8326 acts[bi].xq.as_slice(),
8327 acts[bi + 1].xq.as_slice(),
8328 acts[bi + 2].xq.as_slice(),
8329 acts[bi + 3].xq.as_slice(),
8330 ];
8331 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
8332 for k in 0..4 {
8333 let act = &acts[bi + k];
8334 let mut acc = d[k] * act.sx;
8335 for &(j, xv) in &act.outliers {
8336 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
8337 acc += w * sc * xv;
8338 }
8339 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8341 }
8342 bi += 4;
8343 }
8344 }
8345 #[cfg(target_arch = "x86_64")]
8346 if blocked_ok {
8347 while bi + 4 <= acts.len() {
8348 let xs = [
8349 acts[bi].xq.as_slice(),
8350 acts[bi + 1].xq.as_slice(),
8351 acts[bi + 2].xq.as_slice(),
8352 acts[bi + 3].xq.as_slice(),
8353 ];
8354 let d = unsafe {
8355 if vnni_tiles_enabled() {
8356 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
8357 } else {
8358 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
8359 }
8360 };
8361 for k in 0..4 {
8362 let act = &acts[bi + k];
8363 let mut acc = d[k] * act.sx;
8364 for &(j, xv) in &act.outliers {
8365 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
8366 acc += w * sc * xv;
8367 }
8368 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8370 }
8371 bi += 4;
8372 }
8373 }
8374 let _ = blocked_ok;
8375 while bi < acts.len() {
8376 let act = &acts[bi];
8377 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8378 for &(j, xv) in &act.outliers {
8379 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8380 acc += w * s * xv;
8381 }
8382 unsafe { *out_addr.at(bi * rows + r) = acc };
8384 bi += 1;
8385 }
8386 }
8387 };
8388 dispatch_rows(pool, rows, &run);
8389 return;
8390 }
8391 let run = move |start: usize, end: usize| {
8392 for r in start..end {
8393 for bi in 0..b {
8394 let x = &xs_all[bi * cols..(bi + 1) * cols];
8395 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
8397 }
8398 }
8399 };
8400 dispatch_rows(pool, rows, &run);
8401}
8402
8403fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
8412 (0..gpr)
8413 .map(|gi| {
8414 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
8415 .iter()
8416 .map(|&v| v as i32)
8417 .sum()
8418 })
8419 .collect()
8420}
8421
8422#[inline]
8426#[allow(unreachable_code)]
8427#[cfg(target_arch = "x86_64")]
8432#[target_feature(enable = "avx2")]
8433unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8434 unsafe {
8436 use core::arch::x86_64::*;
8437 let expand = _mm256_setr_epi8(
8439 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,
8440 3, 3, 3,
8441 );
8442 let bitsel = _mm256_setr_epi8(
8443 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8444 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8445 );
8446 let ones8 = _mm256_set1_epi8(1);
8447 let ones16 = _mm256_set1_epi16(1);
8448 let mut acc = 0f32;
8449 for gi in 0..gpr {
8450 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
8451 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8452 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
8453 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8454 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8455 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8456 let sel = _mm256_and_si256(x, mask);
8457 let p16 = _mm256_maddubs_epi16(ones8, sel);
8459 let d32 = _mm256_madd_epi16(p16, ones16);
8460 let hi128 = _mm256_extracti128_si256::<1>(d32);
8461 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
8462 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8463 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8464 let msum = _mm_cvtsi128_si32(s32);
8465 let d = 2 * msum - gsum[gi];
8468 acc += d as f32 * s;
8469 }
8470 acc
8471 }
8472}
8473
8474#[cfg(target_arch = "x86_64")]
8477#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8478unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8479 unsafe {
8481 use core::arch::x86_64::*;
8482 let expand = _mm256_setr_epi8(
8483 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,
8484 3, 3, 3,
8485 );
8486 let bitsel = _mm256_setr_epi8(
8487 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8488 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8489 );
8490 let ones8 = _mm256_set1_epi8(1);
8491 let mut acc = 0f32;
8492 for gi in 0..gpr {
8493 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
8494 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8495 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
8496 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8497 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8498 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8499 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
8500 let d = 2 * msum - gsum[gi];
8501 acc += d as f32 * s;
8502 }
8503 acc
8504 }
8505}
8506
8507#[cfg(target_arch = "x86_64")]
8509#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8510unsafe fn dot_q1_row_1x4_vnni(
8511 bytes: &[u8],
8512 r: usize,
8513 gpr: usize,
8514 xs: [&[i8]; 4],
8515 gsums: [&[i32]; 4],
8516) -> [f32; 4] {
8517 unsafe {
8519 use core::arch::x86_64::*;
8520 let expand = _mm256_setr_epi8(
8521 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,
8522 3, 3, 3,
8523 );
8524 let bitsel = _mm256_setr_epi8(
8525 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8526 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8527 );
8528 let ones8 = _mm256_set1_epi8(1);
8529 let mut acc = [0f32; 4];
8530 for gi in 0..gpr {
8531 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
8532 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8533 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
8534 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8535 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8536 for (k, xq) in xs.iter().enumerate() {
8537 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8538 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
8539 let d = 2 * msum - gsums[k][gi];
8540 acc[k] += d as f32 * s;
8541 }
8542 }
8543 acc
8544 }
8545}
8546
8547#[cfg(target_arch = "x86_64")]
8550#[target_feature(enable = "avx2")]
8551unsafe fn dot_q1_row_1x4_avx2(
8552 bytes: &[u8],
8553 r: usize,
8554 gpr: usize,
8555 xs: [&[i8]; 4],
8556 gsums: [&[i32]; 4],
8557) -> [f32; 4] {
8558 unsafe {
8560 use core::arch::x86_64::*;
8561 let expand = _mm256_setr_epi8(
8562 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,
8563 3, 3, 3,
8564 );
8565 let bitsel = _mm256_setr_epi8(
8566 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8567 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8568 );
8569 let ones8 = _mm256_set1_epi8(1);
8570 let ones16 = _mm256_set1_epi16(1);
8571 let mut acc = [0f32; 4];
8572 for gi in 0..gpr {
8573 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
8574 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8575 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
8576 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8577 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8578 for (k, xq) in xs.iter().enumerate() {
8579 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8580 let sel = _mm256_and_si256(x, mask);
8581 let p16 = _mm256_maddubs_epi16(ones8, sel);
8582 let d32 = _mm256_madd_epi16(p16, ones16);
8583 let hi128 = _mm256_extracti128_si256::<1>(d32);
8584 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
8585 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8586 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8587 let msum = _mm_cvtsi128_si32(s32);
8588 let d = 2 * msum - gsums[k][gi];
8589 acc[k] += d as f32 * s;
8590 }
8591 }
8592 acc
8593 }
8594}
8595
8596#[allow(unreachable_code)]
8597fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8598 #[cfg(target_arch = "aarch64")]
8599 unsafe {
8600 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
8601 }
8602 #[cfg(target_arch = "x86_64")]
8603 if avx2_enabled() {
8604 unsafe {
8605 if vnni_tiles_enabled() {
8606 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
8607 }
8608 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
8609 }
8610 }
8611 let _ = gsum;
8612 let mut acc = 0f32;
8613 for gi in 0..gpr {
8614 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
8615 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8616 let mut d = 0i32;
8617 for (j, &b) in tile[2..].iter().enumerate() {
8618 for k in 0..8 {
8619 let w = ((b >> k) & 1) as i32 * 2 - 1;
8620 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
8621 }
8622 }
8623 acc += d as f32 * s;
8624 }
8625 acc
8626}
8627
8628#[cfg(target_arch = "aarch64")]
8637#[target_feature(enable = "neon,dotprod")]
8638unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8639 unsafe {
8642 use core::arch::aarch64::*;
8643 use core::arch::asm;
8644 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8645 let m = vld1q_u8(MASKS.as_ptr());
8646 macro_rules! tile_dot {
8648 ($t:expr, $x:expr) => {{
8649 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
8650 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
8651 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8652 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8653 let x0 = vld1q_s8($x);
8654 let x1 = vld1q_s8($x.add(16));
8655 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8656 asm!(
8657 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8658 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8659 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8660 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8661 options(pure, nomem, nostack),
8662 );
8663 vaddq_s32(a0, a1)
8664 }};
8665 }
8666 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
8675 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
8676 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
8677 const IW11: [u8; 16] = [
8678 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8679 ];
8680 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
8681 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
8682 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
8683 let isc = vld1_u8(ISC.as_ptr());
8684 macro_rules! tile_dot_tbl {
8686 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
8687 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
8688 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
8689 let x0 = vld1q_s8($x);
8690 let x1 = vld1q_s8($x.add(16));
8691 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8692 asm!(
8693 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8694 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8695 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8696 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8697 options(pure, nomem, nostack),
8698 );
8699 vaddq_s32(a0, a1)
8700 }};
8701 }
8702 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
8703 let row_base = r * gpr * Q1_TILE;
8704 let abs_end = bytes.len();
8705 let xp = xq.as_ptr();
8706 let gp = gsum.as_ptr();
8707 let mut accv = vdupq_n_f32(0.0);
8708 let mut gi = 0;
8709 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
8712 let t0 = base.add(gi * Q1_TILE);
8713 let ld_a = vld1q_u8(t0);
8714 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
8715 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
8716 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
8717 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
8718 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
8719 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8721 let g = vld1q_s32(gp.add(gi));
8722 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8723 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
8724 let scf: float32x4_t;
8725 asm!(
8726 "fcvtl {o:v}.4s, {i:v}.4h",
8727 o = out(vreg) scf, i = in(vreg) sc16,
8728 options(pure, nomem, nostack),
8729 );
8730 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8731 gi += 4;
8732 }
8733 let mut acc = vaddvq_f32(accv);
8734 while gi < gpr {
8735 let t = base.add(gi * Q1_TILE);
8736 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8737 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
8738 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8739 gi += 1;
8740 }
8741 acc
8742 }
8743}
8744
8745#[cfg(target_arch = "aarch64")]
8750#[target_feature(enable = "neon,dotprod")]
8751unsafe fn dot_q1_row_1x4_sdot(
8752 bytes: &[u8],
8753 r: usize,
8754 gpr: usize,
8755 xs: [&[i8]; 4],
8756 gs: [&[i32]; 4],
8757) -> [f32; 4] {
8758 unsafe {
8760 use core::arch::aarch64::*;
8761 use core::arch::asm;
8762 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8763 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
8764 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
8765 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
8766 const IW11: [u8; 16] = [
8767 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8768 ];
8769 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
8770 let m = vld1q_u8(MASKS.as_ptr());
8771 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
8772 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
8773 let isc = vld1_u8(ISC.as_ptr());
8774 macro_rules! sdot2 {
8775 ($w0:expr, $w1:expr, $x:expr) => {{
8776 let x0 = vld1q_s8($x);
8777 let x1 = vld1q_s8($x.add(16));
8778 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8779 asm!(
8780 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8781 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8782 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8783 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8784 options(pure, nomem, nostack),
8785 );
8786 vaddq_s32(a0, a1)
8787 }};
8788 }
8789 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
8790 let row_base = r * gpr * Q1_TILE;
8791 let abs_end = bytes.len();
8792 let mut accv = [vdupq_n_f32(0.0); 4];
8793 let mut gi = 0;
8794 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
8795 let t0 = base.add(gi * Q1_TILE);
8796 let ld_a = vld1q_u8(t0);
8797 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
8798 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
8800 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
8801 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
8802 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
8803 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
8804 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
8805 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
8806 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
8807 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
8808 let scf: float32x4_t;
8809 asm!(
8810 "fcvtl {o:v}.4s, {i:v}.4h",
8811 o = out(vreg) scf, i = in(vreg) sc16,
8812 options(pure, nomem, nostack),
8813 );
8814 for k in 0..4 {
8815 let xp = xs[k].as_ptr();
8816 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
8817 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
8818 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
8819 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
8820 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8821 let g = vld1q_s32(gs[k].as_ptr().add(gi));
8822 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8823 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8824 }
8825 gi += 4;
8826 }
8827 let mut acc = [
8828 vaddvq_f32(accv[0]),
8829 vaddvq_f32(accv[1]),
8830 vaddvq_f32(accv[2]),
8831 vaddvq_f32(accv[3]),
8832 ];
8833 while gi < gpr {
8834 let t = base.add(gi * Q1_TILE);
8835 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8836 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8837 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
8838 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8839 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8840 for k in 0..4 {
8841 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8842 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8843 }
8844 gi += 1;
8845 }
8846 acc
8847 }
8848}
8849
8850#[inline]
8852fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
8853 let gi = j / GROUP_SIZE;
8854 let k = j % GROUP_SIZE;
8855 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
8856 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8857 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
8858 ((bit as i32 * 2 - 1) as f32, s)
8859}
8860
8861#[inline]
8863fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
8864 let mut acc = 0f32;
8865 for gi in 0..gpr {
8866 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
8867 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8868 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8869 let mut ga = 0f32;
8870 for (j, &b) in tile[2..].iter().enumerate() {
8871 for k in 0..8 {
8872 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8873 }
8874 }
8875 acc += ga * s;
8876 }
8877 acc
8878}
8879
8880#[allow(clippy::too_many_arguments)]
8883fn bin_range_a8w8(
8884 w: BinW,
8885 gpr: usize,
8886 act: &SplitAct,
8887 gsum: &[i32],
8888 out: SendMut,
8889 start: usize,
8890 end: usize,
8891) {
8892 for r in start..end {
8893 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
8894 for &(j, xv) in &act.outliers {
8895 let (wv, s) = w.outlier(r, gpr, j);
8896 acc += wv * s * xv;
8897 }
8898 unsafe { *out.at(r) = acc };
8900 }
8901}
8902
8903fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
8905 for r in start..end {
8906 unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
8908 }
8909}
8910
8911#[derive(Clone, Copy)]
8919pub(crate) enum BinW<'a> {
8920 Q1(&'a [u8]),
8921 Q1g {
8922 bytes: &'a [u8],
8923 sc_off: usize,
8925 rows: usize,
8928 },
8929}
8930
8931impl<'a> BinW<'a> {
8932 fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
8933 Self::Q1g {
8934 bytes,
8935 sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
8936 rows,
8937 }
8938 }
8939
8940 #[inline]
8941 fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8942 match self {
8943 Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
8944 Self::Q1g {
8945 bytes,
8946 sc_off,
8947 rows,
8948 } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
8949 }
8950 }
8951
8952 #[inline]
8953 fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
8954 match self {
8955 Self::Q1(b) => q1_row_exact(b, r, gpr, x),
8956 Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
8957 }
8958 }
8959
8960 #[inline]
8961 fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8962 match self {
8963 Self::Q1(b) => q1_outlier(b, r, gpr, j),
8964 Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
8965 }
8966 }
8967
8968 #[cfg(target_arch = "aarch64")]
8973 #[inline]
8974 unsafe fn dot_1x4_sdot(
8975 self,
8976 r: usize,
8977 gpr: usize,
8978 xs: [&[i8]; 4],
8979 gs: [&[i32]; 4],
8980 ) -> [f32; 4] {
8981 unsafe {
8983 match self {
8984 Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
8985 Self::Q1g {
8986 bytes,
8987 sc_off,
8988 rows,
8989 } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
8990 }
8991 }
8992 }
8993
8994 #[cfg(target_arch = "x86_64")]
8999 #[inline]
9000 unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
9001 unsafe {
9003 match (self, vnni_tiles_enabled()) {
9004 (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
9005 (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
9006 (Self::Q1g { bytes, sc_off, .. }, true) => {
9007 dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
9008 }
9009 (Self::Q1g { bytes, sc_off, .. }, false) => {
9010 dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
9011 }
9012 }
9013 }
9014 }
9015}
9016
9017#[inline]
9027fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
9028 u32::from_le_bytes([
9029 bytes[t * 4],
9030 bytes[t * 4 + 1],
9031 bytes[t * 4 + 2],
9032 bytes[t * 4 + 3],
9033 ])
9034}
9035
9036#[inline]
9038fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
9039 let o = sc_off + (t / 4) * 2;
9040 u16::from_le_bytes([bytes[o], bytes[o + 1]])
9041}
9042
9043#[allow(unreachable_code)]
9045#[allow(clippy::too_many_arguments)]
9046fn dot_q1g_row_i8(
9047 bytes: &[u8],
9048 sc_off: usize,
9049 rows: usize,
9050 r: usize,
9051 gpr: usize,
9052 xq: &[i8],
9053 gsum: &[i32],
9054) -> f32 {
9055 #[cfg(target_arch = "aarch64")]
9056 unsafe {
9057 return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
9058 }
9059 #[cfg(target_arch = "x86_64")]
9060 if avx2_enabled() {
9061 unsafe {
9062 if vnni_tiles_enabled() {
9063 return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
9064 }
9065 return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
9066 }
9067 }
9068 let _ = (gsum, rows);
9069 let mut acc = 0f32;
9070 for gi in 0..gpr {
9071 let t = r * gpr + gi;
9072 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9073 let bits = q1g_bits(bytes, t);
9074 let mut d = 0i32;
9075 for k in 0..GROUP_SIZE {
9076 let w = ((bits >> k) & 1) as i32 * 2 - 1;
9077 d += w * xq[gi * GROUP_SIZE + k] as i32;
9078 }
9079 acc += d as f32 * s;
9080 }
9081 acc
9082}
9083
9084#[cfg(target_arch = "x86_64")]
9089#[target_feature(enable = "avx2")]
9090#[inline]
9091unsafe fn q1g_hsum4(
9092 d0: core::arch::x86_64::__m256i,
9093 d1: core::arch::x86_64::__m256i,
9094 d2: core::arch::x86_64::__m256i,
9095 d3: core::arch::x86_64::__m256i,
9096) -> core::arch::x86_64::__m128i {
9097 unsafe {
9099 use core::arch::x86_64::*;
9100 let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
9101 _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
9102 }
9103}
9104
9105#[cfg(target_arch = "x86_64")]
9109#[target_feature(enable = "avx2")]
9110#[inline]
9111unsafe fn q1g_group_acc(
9112 acc: f32,
9113 msum4: core::arch::x86_64::__m128i,
9114 gsum4: &[i32],
9115 s: f32,
9116) -> f32 {
9117 unsafe {
9119 use core::arch::x86_64::*;
9120 let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
9121 let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
9122 let mut f = [0f32; 4];
9123 _mm_storeu_ps(f.as_mut_ptr(), d);
9124 let mut acc = acc;
9125 for v in f {
9126 acc += v * s;
9127 }
9128 acc
9129 }
9130}
9131
9132#[cfg(target_arch = "x86_64")]
9137#[target_feature(enable = "avx2")]
9138unsafe fn dot_q1g_row_avx2(
9139 bytes: &[u8],
9140 sc_off: usize,
9141 r: usize,
9142 gpr: usize,
9143 xq: &[i8],
9144 gsum: &[i32],
9145) -> f32 {
9146 unsafe {
9149 use core::arch::x86_64::*;
9150 let expand = _mm256_setr_epi8(
9151 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,
9152 3, 3, 3,
9153 );
9154 let bitsel = _mm256_setr_epi8(
9155 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9156 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9157 );
9158 let ones8 = _mm256_set1_epi8(1);
9159 let ones16 = _mm256_set1_epi16(1);
9160 let bp = bytes.as_ptr().add(r * gpr * 4);
9161 let mut acc = 0f32;
9162 let mut gi = 0;
9163 while gi < gpr {
9164 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9165 let mut d = [_mm256_setzero_si256(); 4];
9166 for (k, dk) in d.iter_mut().enumerate() {
9167 let t = gi + k;
9168 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9169 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9170 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9171 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9172 let sel = _mm256_and_si256(x, mask);
9173 *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9174 }
9175 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9176 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9177 gi += 4;
9178 }
9179 acc
9180 }
9181}
9182
9183#[cfg(target_arch = "x86_64")]
9186#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9187unsafe fn dot_q1g_row_vnni(
9188 bytes: &[u8],
9189 sc_off: usize,
9190 r: usize,
9191 gpr: usize,
9192 xq: &[i8],
9193 gsum: &[i32],
9194) -> f32 {
9195 unsafe {
9197 use core::arch::x86_64::*;
9198 let expand = _mm256_setr_epi8(
9199 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,
9200 3, 3, 3,
9201 );
9202 let bitsel = _mm256_setr_epi8(
9203 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9204 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9205 );
9206 let ones8 = _mm256_set1_epi8(1);
9207 let bp = bytes.as_ptr().add(r * gpr * 4);
9208 let mut acc = 0f32;
9209 let mut gi = 0;
9210 while gi < gpr {
9211 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9212 let mut d = [_mm256_setzero_si256(); 4];
9213 for (k, dk) in d.iter_mut().enumerate() {
9214 let t = gi + k;
9215 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9216 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9217 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9218 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9219 *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
9220 }
9221 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9222 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9223 gi += 4;
9224 }
9225 acc
9226 }
9227}
9228
9229#[cfg(target_arch = "x86_64")]
9233#[target_feature(enable = "avx2")]
9234unsafe fn dot_q1g_row_1x4_avx2(
9235 bytes: &[u8],
9236 sc_off: usize,
9237 r: usize,
9238 gpr: usize,
9239 xs: [&[i8]; 4],
9240 gsums: [&[i32]; 4],
9241) -> [f32; 4] {
9242 unsafe {
9244 use core::arch::x86_64::*;
9245 let expand = _mm256_setr_epi8(
9246 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,
9247 3, 3, 3,
9248 );
9249 let bitsel = _mm256_setr_epi8(
9250 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9251 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9252 );
9253 let ones8 = _mm256_set1_epi8(1);
9254 let ones16 = _mm256_set1_epi16(1);
9255 let bp = bytes.as_ptr().add(r * gpr * 4);
9256 let mut acc = [0f32; 4];
9257 let mut gi = 0;
9258 while gi < gpr {
9259 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9260 let mut d = [[_mm256_setzero_si256(); 4]; 4];
9262 for t in 0..4 {
9263 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9264 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9265 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9266 for (k, xq) in xs.iter().enumerate() {
9267 let x = _mm256_loadu_si256(
9268 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9269 );
9270 let sel = _mm256_and_si256(x, mask);
9271 d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9272 }
9273 }
9274 for k in 0..4 {
9275 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9276 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9277 }
9278 gi += 4;
9279 }
9280 acc
9281 }
9282}
9283
9284#[cfg(target_arch = "x86_64")]
9286#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9287unsafe fn dot_q1g_row_1x4_vnni(
9288 bytes: &[u8],
9289 sc_off: usize,
9290 r: usize,
9291 gpr: usize,
9292 xs: [&[i8]; 4],
9293 gsums: [&[i32]; 4],
9294) -> [f32; 4] {
9295 unsafe {
9297 use core::arch::x86_64::*;
9298 let expand = _mm256_setr_epi8(
9299 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,
9300 3, 3, 3,
9301 );
9302 let bitsel = _mm256_setr_epi8(
9303 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9304 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9305 );
9306 let ones8 = _mm256_set1_epi8(1);
9307 let bp = bytes.as_ptr().add(r * gpr * 4);
9308 let mut acc = [0f32; 4];
9309 let mut gi = 0;
9310 while gi < gpr {
9311 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9312 let mut d = [[_mm256_setzero_si256(); 4]; 4];
9313 for t in 0..4 {
9314 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9315 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9316 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9317 for (k, xq) in xs.iter().enumerate() {
9318 let x = _mm256_loadu_si256(
9319 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9320 );
9321 d[k][t] = _mm256_dpbusd_epi32(
9322 _mm256_setzero_si256(),
9323 ones8,
9324 _mm256_and_si256(x, mask),
9325 );
9326 }
9327 }
9328 for k in 0..4 {
9329 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9330 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9331 }
9332 gi += 4;
9333 }
9334 acc
9335 }
9336}
9337
9338#[cfg(target_arch = "aarch64")]
9349#[target_feature(enable = "neon,dotprod")]
9350#[allow(clippy::too_many_arguments)]
9351unsafe fn dot_q1g_row_sdot(
9352 bytes: &[u8],
9353 sc_off: usize,
9354 rows: usize,
9355 r: usize,
9356 gpr: usize,
9357 xq: &[i8],
9358 gsum: &[i32],
9359) -> f32 {
9360 unsafe {
9363 use core::arch::aarch64::*;
9364 use core::arch::asm;
9365 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9366 const IW: [[u8; 16]; 8] = [
9368 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
9369 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
9370 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
9371 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
9372 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
9373 [
9374 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9375 ],
9376 [
9377 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
9378 ],
9379 [
9380 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
9381 ],
9382 ];
9383 let m = vld1q_u8(MASKS.as_ptr());
9384 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
9385 macro_rules! tile_dot {
9386 ($w0:expr, $w1:expr, $x:expr) => {{
9387 let x0 = vld1q_s8($x);
9388 let x1 = vld1q_s8($x.add(16));
9389 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9390 asm!(
9391 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9392 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9393 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9394 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9395 options(pure, nomem, nostack),
9396 );
9397 vaddq_s32(a0, a1)
9398 }};
9399 }
9400 macro_rules! tbl_mask {
9401 ($ld:expr, $i:expr) => {
9402 vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
9403 };
9404 }
9405 let bp = bytes.as_ptr().add(r * gpr * 4);
9406 let xp = xq.as_ptr();
9407 let gp = gsum.as_ptr();
9408 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
9409 let mut accv = vdupq_n_f32(0.0);
9410 let mut gi = 0;
9411 while gi + 4 <= vec_end {
9412 let ld = vld1q_u8(bp.add(gi * 4));
9413 let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
9414 let d1 = tile_dot!(
9415 tbl_mask!(ld, 2),
9416 tbl_mask!(ld, 3),
9417 xp.add((gi + 1) * GROUP_SIZE)
9418 );
9419 let d2 = tile_dot!(
9420 tbl_mask!(ld, 4),
9421 tbl_mask!(ld, 5),
9422 xp.add((gi + 2) * GROUP_SIZE)
9423 );
9424 let d3 = tile_dot!(
9425 tbl_mask!(ld, 6),
9426 tbl_mask!(ld, 7),
9427 xp.add((gi + 3) * GROUP_SIZE)
9428 );
9429 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9430 let g = vld1q_s32(gp.add(gi));
9431 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9432 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9433 let scf: float32x4_t;
9434 asm!(
9435 "fcvtl {o:v}.4s, {i:v}.4h",
9436 o = out(vreg) scf, i = in(vreg) h,
9437 options(pure, nomem, nostack),
9438 );
9439 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
9440 gi += 4;
9441 }
9442 let mut acc = vaddvq_f32(accv);
9443 while gi < gpr {
9444 let t = bp.add(gi * 4);
9445 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9446 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
9447 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
9448 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9449 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9450 let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
9451 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
9452 gi += 1;
9453 }
9454 acc
9455 }
9456}
9457
9458#[cfg(target_arch = "aarch64")]
9462#[target_feature(enable = "neon,dotprod")]
9463#[allow(clippy::too_many_arguments)]
9464unsafe fn dot_q1g_row_1x4_sdot(
9465 bytes: &[u8],
9466 sc_off: usize,
9467 rows: usize,
9468 r: usize,
9469 gpr: usize,
9470 xs: [&[i8]; 4],
9471 gs: [&[i32]; 4],
9472) -> [f32; 4] {
9473 unsafe {
9475 use core::arch::aarch64::*;
9476 use core::arch::asm;
9477 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9478 const IW: [[u8; 16]; 8] = [
9479 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
9480 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
9481 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
9482 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
9483 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
9484 [
9485 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9486 ],
9487 [
9488 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
9489 ],
9490 [
9491 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
9492 ],
9493 ];
9494 let m = vld1q_u8(MASKS.as_ptr());
9495 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
9496 macro_rules! sdot2 {
9497 ($w0:expr, $w1:expr, $x:expr) => {{
9498 let x0 = vld1q_s8($x);
9499 let x1 = vld1q_s8($x.add(16));
9500 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9501 asm!(
9502 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9503 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9504 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9505 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9506 options(pure, nomem, nostack),
9507 );
9508 vaddq_s32(a0, a1)
9509 }};
9510 }
9511 let bp = bytes.as_ptr().add(r * gpr * 4);
9512 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
9513 let mut accv = [vdupq_n_f32(0.0); 4];
9514 let mut gi = 0;
9515 while gi + 4 <= vec_end {
9516 let ld = vld1q_u8(bp.add(gi * 4));
9517 let w: [int8x16_t; 8] =
9519 core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
9520 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9521 let scf: float32x4_t;
9522 asm!(
9523 "fcvtl {o:v}.4s, {i:v}.4h",
9524 o = out(vreg) scf, i = in(vreg) h,
9525 options(pure, nomem, nostack),
9526 );
9527 for k in 0..4 {
9528 let xp = xs[k].as_ptr();
9529 let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
9530 let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
9531 let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
9532 let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
9533 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9534 let g = vld1q_s32(gs[k].as_ptr().add(gi));
9535 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9536 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
9537 }
9538 gi += 4;
9539 }
9540 let mut acc = [
9541 vaddvq_f32(accv[0]),
9542 vaddvq_f32(accv[1]),
9543 vaddvq_f32(accv[2]),
9544 vaddvq_f32(accv[3]),
9545 ];
9546 while gi < gpr {
9547 let t = bp.add(gi * 4);
9548 let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9549 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
9550 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
9551 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9552 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9553 for k in 0..4 {
9554 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
9555 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
9556 }
9557 gi += 1;
9558 }
9559 acc
9560 }
9561}
9562
9563#[inline]
9565fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
9566 let t = r * gpr + j / GROUP_SIZE;
9567 let k = j % GROUP_SIZE;
9568 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9569 let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
9570 ((bit as i32 * 2 - 1) as f32, s)
9571}
9572
9573#[inline]
9575fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
9576 let mut acc = 0f32;
9577 for gi in 0..gpr {
9578 let t = r * gpr + gi;
9579 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9580 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9581 let mut ga = 0f32;
9582 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
9583 for k in 0..8 {
9584 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
9585 }
9586 }
9587 acc += ga * s;
9588 }
9589 acc
9590}
9591
9592fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
9597 let entries = base_len + (rows + 1) * 4;
9598 (base_len, entries, entries <= bytes.len())
9599}
9600
9601#[inline]
9603fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
9604 let o = rp_off + r * 4;
9605 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
9606}
9607
9608const SIGN5: [[f32; 5]; 256] = {
9612 let mut lut = [[0.0f32; 5]; 256];
9613 let pow3 = [1u16, 3, 9, 27, 81];
9614 let mut byte = 0usize;
9615 while byte < 256 {
9616 let mut i = 0usize;
9617 while i < 5 {
9618 let code = (byte as u16 / pow3[i]) % 3;
9619 lut[byte][i] = if code == 1 {
9620 1.0
9621 } else if code == 2 {
9622 -1.0
9623 } else {
9624 0.0
9625 };
9626 i += 1;
9627 }
9628 byte += 1;
9629 }
9630 lut
9631};
9632
9633const SIGN5_I8: [[i8; 5]; 256] = {
9635 let mut lut = [[0i8; 5]; 256];
9636 let pow3 = [1u16, 3, 9, 27, 81];
9637 let mut byte = 0usize;
9638 while byte < 256 {
9639 let mut i = 0usize;
9640 while i < 5 {
9641 let code = (byte as u16 / pow3[i]) % 3;
9642 lut[byte][i] = if code == 1 {
9643 1
9644 } else if code == 2 {
9645 -1
9646 } else {
9647 0
9648 };
9649 i += 1;
9650 }
9651 byte += 1;
9652 }
9653 lut
9654};
9655
9656const SIGN5_U64: [u64; 256] = {
9662 let mut lut = [0u64; 256];
9663 let pow3 = [1u16, 3, 9, 27, 81];
9664 let mut byte = 0usize;
9665 while byte < 256 {
9666 let mut v = 0u64;
9667 let mut i = 0usize;
9668 while i < 5 {
9669 let code = (byte as u16 / pow3[i]) % 3;
9670 let s: u8 = if code == 1 {
9671 1
9672 } else if code == 2 {
9673 0xFF
9674 } else {
9675 0
9676 };
9677 v |= (s as u64) << (i * 8);
9678 i += 1;
9679 }
9680 lut[byte] = v;
9681 byte += 1;
9682 }
9683 lut
9684};
9685
9686#[inline]
9691fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
9692 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9693 let off = (r * gpr + j / GROUP_SIZE) * TILE;
9694 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9695 let within = j % GROUP_SIZE;
9696 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
9697}
9698
9699#[cfg(target_arch = "aarch64")]
9702#[target_feature(enable = "neon,dotprod")]
9703#[inline]
9704unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
9705 unsafe {
9707 use core::arch::aarch64::*;
9708 use core::arch::asm;
9709 let w0 = vld1q_s8(w);
9710 let w1 = vld1q_s8(w.add(16));
9711 let x0 = vld1q_s8(x);
9712 let x1 = vld1q_s8(x.add(16));
9713 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9714 asm!(
9715 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9716 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9717 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9718 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9719 options(pure, nomem, nostack),
9720 );
9721 vaddvq_s32(vaddq_s32(a0, a1))
9722 }
9723}
9724
9725#[cfg(target_arch = "x86_64")]
9728#[target_feature(enable = "avx2")]
9729#[inline]
9730unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
9731 unsafe {
9733 use core::arch::x86_64::*;
9734 let wv = _mm256_loadu_si256(w as *const __m256i);
9735 let xv = _mm256_loadu_si256(x as *const __m256i);
9736 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9737 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
9738 let hi128 = _mm256_extracti128_si256::<1>(d);
9739 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9740 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9741 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9742 _mm_cvtsi128_si32(s32)
9743 }
9744}
9745
9746#[inline]
9751fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
9752 debug_assert!(dst.len() >= 40);
9753 unsafe {
9756 let p = dst.as_mut_ptr();
9757 for bi in 0..7 {
9758 core::ptr::write_unaligned(
9759 p.add(bi * 5) as *mut u64,
9760 SIGN5_U64[*codes.add(bi) as usize],
9761 );
9762 }
9763 }
9764}
9765
9766#[inline]
9771fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
9772 #[cfg(target_arch = "aarch64")]
9773 unsafe {
9774 return sdot32_i8(w, x);
9775 }
9776 #[cfg(target_arch = "x86_64")]
9777 unsafe {
9778 return i8dot32_avx2(w, x);
9779 }
9780 #[allow(unreachable_code)]
9781 unsafe {
9782 let mut s = 0i32;
9783 for k in 0..GROUP_SIZE {
9784 s += *w.add(k) as i32 * *x.add(k) as i32;
9785 }
9786 s
9787 }
9788}
9789
9790#[inline]
9791unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
9792 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
9793 (
9794 SIGN5_U64[*codes as usize],
9795 SIGN5_U64[*codes.add(1) as usize],
9796 SIGN5_U64[*codes.add(2) as usize],
9797 SIGN5_U64[*codes.add(3) as usize],
9798 SIGN5_U64[*codes.add(4) as usize],
9799 SIGN5_U64[*codes.add(5) as usize],
9800 SIGN5_U64[*codes.add(6) as usize],
9801 )
9802 };
9803
9804 let u0 = s0 | (s1 << 40);
9805 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
9806 let u2 = (s3 >> 8) | (s4 << 32);
9807 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
9808
9809 (u0, u1, u2, u3)
9810}
9811
9812#[cfg(target_arch = "aarch64")]
9816#[target_feature(enable = "neon,dotprod")]
9817unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9818 use core::arch::aarch64::*;
9819 use core::arch::asm;
9820 unsafe {
9821 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9822 let mut acc = 0f32;
9823 let bytes_ptr = bytes.as_ptr();
9824 let xq_ptr = xq.as_ptr();
9825 let row_off = r * gpr * TILE;
9826
9827 let gpr2 = gpr & !1;
9828 let mut gi = 0;
9829 while gi < gpr2 {
9830 let off0 = row_off + gi * TILE;
9831 let off1 = off0 + TILE;
9832 let s0 = f16_to_f32(u16::from_le_bytes([
9833 *bytes_ptr.add(off0),
9834 *bytes_ptr.add(off0 + 1),
9835 ]));
9836 let s1 = f16_to_f32(u16::from_le_bytes([
9837 *bytes_ptr.add(off1),
9838 *bytes_ptr.add(off1 + 1),
9839 ]));
9840
9841 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9842 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9843
9844 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9845 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9846 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9847 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9848
9849 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
9850 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
9851 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
9852 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
9853
9854 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
9855 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9856 asm!(
9857 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
9858 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
9859 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
9860 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
9861 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
9862 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
9863 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
9864 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
9865 options(pure, nomem, nostack),
9866 );
9867 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
9868 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
9869 acc += d0 as f32 * s0 + d1 as f32 * s1;
9870 gi += 2;
9871 }
9872
9873 if gi < gpr {
9874 let off = row_off + gi * TILE;
9875 let s = f16_to_f32(u16::from_le_bytes([
9876 *bytes_ptr.add(off),
9877 *bytes_ptr.add(off + 1),
9878 ]));
9879 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9880 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9881 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9882 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
9883 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
9884 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9885 asm!(
9886 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9887 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9888 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9889 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9890 options(pure, nomem, nostack),
9891 );
9892 let d = vaddvq_s32(vaddq_s32(a0, a1));
9893 acc += d as f32 * s;
9894 }
9895 acc
9896 }
9897}
9898
9899#[cfg(target_arch = "x86_64")]
9901#[target_feature(enable = "avx2")]
9902unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9903 use core::arch::x86_64::*;
9904 unsafe {
9905 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9906 let mut acc = 0f32;
9907 let bytes_ptr = bytes.as_ptr();
9908 let xq_ptr = xq.as_ptr();
9909 let row_off = r * gpr * TILE;
9910
9911 let ones = _mm256_set1_epi16(1);
9912 for gi in 0..gpr {
9913 let off = row_off + gi * TILE;
9914 let s = f16_to_f32(u16::from_le_bytes([
9915 *bytes_ptr.add(off),
9916 *bytes_ptr.add(off + 1),
9917 ]));
9918 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9919 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9920 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9921 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9922 let d256 = _mm256_madd_epi16(p16, ones);
9923 let d128 = _mm_add_epi32(
9924 _mm256_castsi256_si128(d256),
9925 _mm256_extracti128_si256(d256, 1),
9926 );
9927 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
9928 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
9929 acc += d32 as f32 * s;
9930 }
9931 acc
9932 }
9933}
9934
9935#[cfg(target_arch = "x86_64")]
9937#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9938unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9939 use core::arch::x86_64::*;
9940 unsafe {
9942 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9943 let mut acc = 0f32;
9944 let bytes_ptr = bytes.as_ptr();
9945 let xq_ptr = xq.as_ptr();
9946 let row_off = r * gpr * TILE;
9947 for gi in 0..gpr {
9948 let off = row_off + gi * TILE;
9949 let s = f16_to_f32(u16::from_le_bytes([
9950 *bytes_ptr.add(off),
9951 *bytes_ptr.add(off + 1),
9952 ]));
9953 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9954 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9955 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9956 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9957 acc += d as f32 * s;
9958 }
9959 acc
9960 }
9961}
9962
9963#[inline]
9967fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9968 #[cfg(target_arch = "aarch64")]
9969 unsafe {
9970 return q1t_dot_row_sdot(bytes, r, gpr, xq);
9971 }
9972 #[cfg(target_arch = "x86_64")]
9973 unsafe {
9974 if vnni_tiles_enabled() {
9975 return q1t_dot_row_vnni(bytes, r, gpr, xq);
9976 }
9977 return q1t_dot_row_avx2(bytes, r, gpr, xq);
9978 }
9979 #[allow(unreachable_code)]
9980 {
9981 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9982 let mut acc = 0f32;
9983 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
9985 let off = (r * gpr + gi) * TILE;
9986 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9987 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
9988 let mut d = 0i32;
9989 for k in 0..GROUP_SIZE {
9990 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
9991 }
9992 acc += d as f32 * s;
9993 }
9994 acc
9995 }
9996}
9997
9998fn q1t_row_outlier_correction(
10005 bytes: &[u8],
10006 r: usize,
10007 rp_off: usize,
10008 entries_off: usize,
10009 has_ov: bool,
10010 x: &[f32],
10011) -> f32 {
10012 if !has_ov {
10013 return 0.0;
10014 }
10015 let (c0, c1) = (
10016 q1t_rowptr(bytes, rp_off, r),
10017 q1t_rowptr(bytes, rp_off, r + 1),
10018 );
10019 let mut corr = 0f32;
10020 for p in c0..c1 {
10021 let e = entries_off + p * 4;
10022 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10023 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10024 corr += val * x[col];
10025 }
10026 corr
10027}
10028
10029fn q1t_dequant_row(
10033 bytes: &[u8],
10034 r: usize,
10035 gpr: usize,
10036 rp_off: usize,
10037 entries_off: usize,
10038 has_ov: bool,
10039 buf: &mut [f32],
10040) {
10041 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10042 for g in 0..gpr {
10043 let off = (r * gpr + g) * TILE;
10044 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10045 let codes = &bytes[off + 2..off + TILE];
10046 let bc = g * GROUP_SIZE;
10047 for bi in 0..6 {
10049 let lut = &SIGN5[codes[bi] as usize];
10050 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
10051 for i in 0..5 {
10052 d[i] = lut[i] * s;
10053 }
10054 }
10055 let lut = &SIGN5[codes[6] as usize];
10056 buf[bc + 30] = lut[0] * s;
10057 buf[bc + 31] = lut[1] * s;
10058 }
10059 if !has_ov {
10060 return;
10061 }
10062 let (c0, c1) = (
10063 q1t_rowptr(bytes, rp_off, r),
10064 q1t_rowptr(bytes, rp_off, r + 1),
10065 );
10066 for p in c0..c1 {
10067 let e = entries_off + p * 4;
10068 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10069 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10070 }
10071}
10072
10073fn q1t_add_overlay(
10077 bytes: &[u8],
10078 x: &[f32],
10079 rows: usize,
10080 cols: usize,
10081 out: &mut [f32],
10082 pool: Option<&Pool>,
10083) {
10084 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10085 let gpr = cols / GROUP_SIZE;
10086 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10087 if !has_ov {
10088 return;
10089 }
10090 let out_addr = SendMut(out.as_mut_ptr());
10091 let run = move |start: usize, end: usize| {
10092 for r in start..end {
10093 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10094 unsafe { *out_addr.at(r) += corr };
10096 }
10097 };
10098 dispatch_rows(pool, rows, &run);
10099}
10100
10101#[allow(clippy::too_many_arguments)]
10104fn q1t_range_a8w8(
10105 bytes: &[u8],
10106 gpr: usize,
10107 rp_off: usize,
10108 ent_off: usize,
10109 has_ov: bool,
10110 act: &SplitAct,
10111 x: &[f32],
10112 out: SendMut,
10113 start: usize,
10114 end: usize,
10115) {
10116 for r in start..end {
10117 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10118 for &(j, xv) in &act.outliers {
10119 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10120 }
10121 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10122 unsafe { *out.at(r) = acc };
10124 }
10125}
10126
10127#[allow(clippy::too_many_arguments)]
10130fn q1t_range_f32_batch(
10131 bytes: &[u8],
10132 gpr: usize,
10133 rp_off: usize,
10134 ent_off: usize,
10135 has_ov: bool,
10136 x: &[f32],
10137 out: SendMut,
10138 start: usize,
10139 end: usize,
10140) {
10141 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10142 let mut sg = [0f32; GROUP_SIZE];
10143 for r in start..end {
10144 let mut acc = 0f32;
10145 for g in 0..gpr {
10146 let off = (r * gpr + g) * TILE;
10147 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10148 let codes = &bytes[off + 2..off + TILE];
10149 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10150 for bi in 0..6 {
10151 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10152 }
10153 let lut = &SIGN5[codes[6] as usize];
10154 sg[30] = lut[0];
10155 sg[31] = lut[1];
10156 let mut gsum = 0f32;
10157 for k in 0..GROUP_SIZE {
10158 gsum += sg[k] * xg[k];
10159 }
10160 acc += s * gsum;
10161 }
10162 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10163 unsafe { *out.at(r) = acc };
10165 }
10166}
10167
10168fn q1t_matvec(
10172 bytes: &[u8],
10173 x: &[f32],
10174 rows: usize,
10175 cols: usize,
10176 out: &mut [f32],
10177 pool: Option<&Pool>,
10178) {
10179 debug_assert_eq!(out.len(), rows);
10180 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10181 let gpr = cols / GROUP_SIZE;
10182 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10183 let out_addr = SendMut(out.as_mut_ptr());
10184 if a8w8_enabled() {
10188 let act = split_act(x);
10189 let act = &act;
10190 let run = move |start: usize, end: usize| {
10191 for r in start..end {
10192 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10193 for &(j, xv) in &act.outliers {
10194 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10195 }
10196 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10197 unsafe { *out_addr.at(r) = acc };
10199 }
10200 };
10201 dispatch_rows(pool, rows, &run);
10202 return;
10203 }
10204 let run = move |start: usize, end: usize| {
10205 let mut sg = [0f32; GROUP_SIZE];
10209 for r in start..end {
10210 let mut acc = 0f32;
10211 for g in 0..gpr {
10212 let off = (r * gpr + g) * TILE;
10213 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10214 let codes = &bytes[off + 2..off + TILE];
10215 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10216 for bi in 0..6 {
10217 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10218 }
10219 let lut = &SIGN5[codes[6] as usize];
10220 sg[30] = lut[0];
10221 sg[31] = lut[1];
10222 let mut gsum = 0f32;
10223 for k in 0..GROUP_SIZE {
10224 gsum += sg[k] * xg[k];
10225 }
10226 acc += s * gsum;
10227 }
10228 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10229 unsafe { *out_addr.at(r) = acc };
10230 }
10231 };
10232 dispatch_rows(pool, rows, &run);
10233}
10234
10235#[cfg(target_arch = "aarch64")]
10241#[target_feature(enable = "neon,dotprod")]
10242unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
10243 use core::arch::aarch64::*;
10244 use core::arch::asm;
10245 unsafe {
10247 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10248 let bytes_ptr = bytes.as_ptr();
10249 let row_off = r * gpr * TILE;
10250 let xp = [xa.as_ptr(), xb.as_ptr()];
10251 let mut acc = [0f32; 2];
10252 macro_rules! sdot2 {
10253 ($w0:expr, $w1:expr, $x:expr) => {{
10254 let x0 = vld1q_s8($x);
10255 let x1 = vld1q_s8($x.add(16));
10256 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10257 asm!(
10258 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10259 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10260 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10261 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10262 options(pure, nomem, nostack),
10263 );
10264 vaddvq_s32(vaddq_s32(a0, a1))
10265 }};
10266 }
10267 let gpr2 = gpr & !1;
10268 let mut gi = 0;
10269 while gi < gpr2 {
10270 let off0 = row_off + gi * TILE;
10271 let off1 = off0 + TILE;
10272 let s0 = f16_to_f32(u16::from_le_bytes([
10273 *bytes_ptr.add(off0),
10274 *bytes_ptr.add(off0 + 1),
10275 ]));
10276 let s1 = f16_to_f32(u16::from_le_bytes([
10277 *bytes_ptr.add(off1),
10278 *bytes_ptr.add(off1 + 1),
10279 ]));
10280 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10281 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10282 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10283 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10284 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10285 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10286 for k in 0..2 {
10287 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
10288 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
10289 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
10290 }
10291 gi += 2;
10292 }
10293 if gi < gpr {
10294 let off = row_off + gi * TILE;
10295 let s = f16_to_f32(u16::from_le_bytes([
10296 *bytes_ptr.add(off),
10297 *bytes_ptr.add(off + 1),
10298 ]));
10299 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10300 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10301 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10302 for k in 0..2 {
10303 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
10304 acc[k] += d as f32 * s;
10305 }
10306 }
10307 acc
10308 }
10309}
10310
10311fn q1t_matvec2(
10317 bytes: &[u8],
10318 x1: &[f32],
10319 x2: &[f32],
10320 rows: usize,
10321 cols: usize,
10322 o1: &mut [f32],
10323 o2: &mut [f32],
10324 pool: Option<&Pool>,
10325) {
10326 debug_assert_eq!(o1.len(), rows);
10327 debug_assert_eq!(o2.len(), rows);
10328 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10329 let gpr = cols / GROUP_SIZE;
10330 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10331 let out1 = SendMut(o1.as_mut_ptr());
10332 let out2 = SendMut(o2.as_mut_ptr());
10333 if a8w8_enabled() {
10334 let a1 = split_act(x1);
10335 let a2 = split_act(x2);
10336 let (a1, a2) = (&a1, &a2);
10337 let run = move |start: usize, end: usize| {
10338 for r in start..end {
10339 #[cfg(target_arch = "aarch64")]
10340 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
10343 #[cfg(not(target_arch = "aarch64"))]
10344 let ds = [
10345 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
10346 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
10347 ];
10348 let mut acc1 = ds[0] * a1.sx;
10349 for &(j, xv) in &a1.outliers {
10350 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
10351 }
10352 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
10353 let mut acc2 = ds[1] * a2.sx;
10354 for &(j, xv) in &a2.outliers {
10355 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
10356 }
10357 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
10358 unsafe {
10360 *out1.at(r) = acc1;
10361 *out2.at(r) = acc2;
10362 }
10363 }
10364 };
10365 dispatch_rows(pool, rows, &run);
10366 return;
10367 }
10368 let run = move |start: usize, end: usize| {
10369 let mut sg = [0f32; GROUP_SIZE];
10372 for r in start..end {
10373 let mut acc1 = 0f32;
10374 let mut acc2 = 0f32;
10375 for g in 0..gpr {
10376 let off = (r * gpr + g) * TILE;
10377 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10378 let codes = &bytes[off + 2..off + TILE];
10379 for bi in 0..6 {
10380 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10381 }
10382 let lut = &SIGN5[codes[6] as usize];
10383 sg[30] = lut[0];
10384 sg[31] = lut[1];
10385 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10386 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10387 let mut gsum1 = 0f32;
10388 for k in 0..GROUP_SIZE {
10389 gsum1 += sg[k] * xg1[k];
10390 }
10391 acc1 += s * gsum1;
10392 let mut gsum2 = 0f32;
10393 for k in 0..GROUP_SIZE {
10394 gsum2 += sg[k] * xg2[k];
10395 }
10396 acc2 += s * gsum2;
10397 }
10398 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
10399 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
10400 unsafe {
10402 *out1.at(r) = acc1;
10403 *out2.at(r) = acc2;
10404 }
10405 }
10406 };
10407 dispatch_rows(pool, rows, &run);
10408}
10409
10410fn q1t_matmat(
10413 bytes: &[u8],
10414 xs: &[f32],
10415 b: usize,
10416 rows: usize,
10417 cols: usize,
10418 out: &mut [f32],
10419 pool: Option<&Pool>,
10420) {
10421 debug_assert_eq!(out.len(), b * rows);
10422 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10423 let gpr = cols / GROUP_SIZE;
10424 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10425 let out_addr = SendMut(out.as_mut_ptr());
10426 if a8w8_enabled() {
10430 let acts: Vec<SplitAct> = (0..b)
10431 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
10432 .collect();
10433 let acts = &acts;
10434 let run = move |start: usize, end: usize| {
10435 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
10439 for g in 0..gpr {
10440 let off = (r * gpr + g) * TILE;
10441 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10442 q1t_unpack_group_i8(
10443 bytes.as_ptr().wrapping_add(off + 2),
10444 &mut sg[g * GROUP_SIZE..],
10445 );
10446 }
10447 for bi in 0..b {
10448 let act = &acts[bi];
10449 let mut isum = 0f32;
10450 for g in 0..gpr {
10451 let d = q1t_i8dot32(
10452 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
10453 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
10454 );
10455 isum += d as f32 * sc[g];
10456 }
10457 let mut acc = isum * act.sx;
10458 for &(j, xv) in &act.outliers {
10459 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10460 }
10461 accs[bi] = acc;
10462 }
10463 if has_ov {
10467 let (c0, c1) = (
10468 q1t_rowptr(bytes, rp_off, r),
10469 q1t_rowptr(bytes, rp_off, r + 1),
10470 );
10471 for p in c0..c1 {
10472 let e = ent_off + p * 4;
10473 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10474 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10475 for bi in 0..b {
10476 accs[bi] += val * xs[bi * cols + col];
10477 }
10478 }
10479 }
10480 for bi in 0..b {
10481 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
10482 }
10483 }
10484 };
10485 dispatch_rows(pool, rows, &run);
10486 return;
10487 }
10488 let run = move |start: usize, end: usize| {
10489 let mut buf = vec![0f32; cols];
10490 for r in start..end {
10491 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
10492 for bi in 0..b {
10493 let xr = &xs[bi * cols..(bi + 1) * cols];
10494 let mut acc = 0f32;
10495 for j in 0..cols {
10496 acc += buf[j] * xr[j];
10497 }
10498 unsafe { *out_addr.at(bi * rows + r) = acc };
10499 }
10500 }
10501 };
10502 dispatch_rows(pool, rows, &run);
10503}
10504
10505#[cfg(test)]
10508fn q1_matvec(
10509 bytes: &[u8],
10510 x: &[f32],
10511 rows: usize,
10512 cols: usize,
10513 out: &mut [f32],
10514 pool: Option<&Pool>,
10515) {
10516 bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
10517}
10518
10519fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
10522 debug_assert_eq!(out.len(), rows);
10523 let gpr = cols / GROUP_SIZE;
10524 let out_addr = SendMut(out.as_mut_ptr());
10525 if a8w8_enabled() {
10526 let act = split_act(x);
10527 let gsum = q1_group_sums(&act.xq, gpr);
10528 let (act, gsum) = (&act, &gsum);
10529 let run =
10530 move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
10531 dispatch_rows(pool, rows, &run);
10532 return;
10533 }
10534 let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
10535 dispatch_rows(pool, rows, &run);
10536}
10537
10538#[cfg(test)]
10540#[allow(clippy::too_many_arguments)]
10541fn q1_matvec2(
10542 bytes: &[u8],
10543 x1: &[f32],
10544 x2: &[f32],
10545 rows: usize,
10546 cols: usize,
10547 o1: &mut [f32],
10548 o2: &mut [f32],
10549 pool: Option<&Pool>,
10550) {
10551 bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
10552}
10553
10554#[allow(clippy::too_many_arguments)]
10556fn bin_matvec2(
10557 w: BinW,
10558 x1: &[f32],
10559 x2: &[f32],
10560 rows: usize,
10561 cols: usize,
10562 o1: &mut [f32],
10563 o2: &mut [f32],
10564 pool: Option<&Pool>,
10565) {
10566 let gpr = cols / GROUP_SIZE;
10567 let p1 = SendMut(o1.as_mut_ptr());
10568 let p2 = SendMut(o2.as_mut_ptr());
10569 if a8w8_enabled() {
10570 let a1 = split_act(x1);
10571 let a2 = split_act(x2);
10572 let g1 = q1_group_sums(&a1.xq, gpr);
10573 let g2 = q1_group_sums(&a2.xq, gpr);
10574 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
10575 let run = move |start: usize, end: usize| {
10576 for r in start..end {
10577 let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
10578 let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
10579 for &(j, xv) in &a1.outliers {
10580 let (wv, s) = w.outlier(r, gpr, j);
10581 v1 += wv * s * xv;
10582 }
10583 for &(j, xv) in &a2.outliers {
10584 let (wv, s) = w.outlier(r, gpr, j);
10585 v2 += wv * s * xv;
10586 }
10587 unsafe {
10589 *p1.at(r) = v1;
10590 *p2.at(r) = v2;
10591 }
10592 }
10593 };
10594 dispatch_rows(pool, rows, &run);
10595 return;
10596 }
10597 let run = move |start: usize, end: usize| {
10598 for r in start..end {
10599 unsafe {
10601 *p1.at(r) = w.row_exact(r, gpr, x1);
10602 *p2.at(r) = w.row_exact(r, gpr, x2);
10603 }
10604 }
10605 };
10606 dispatch_rows(pool, rows, &run);
10607}
10608
10609#[cfg(test)]
10611#[allow(clippy::too_many_arguments)]
10612fn q1_matmat(
10613 bytes: &[u8],
10614 xs_all: &[f32],
10615 b: usize,
10616 rows: usize,
10617 cols: usize,
10618 out: &mut [f32],
10619 pool: Option<&Pool>,
10620) {
10621 bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
10622}
10623
10624#[allow(clippy::too_many_arguments)]
10626fn bin_matmat(
10627 w: BinW,
10628 xs_all: &[f32],
10629 b: usize,
10630 rows: usize,
10631 cols: usize,
10632 out: &mut [f32],
10633 pool: Option<&Pool>,
10634) {
10635 debug_assert_eq!(out.len(), b * rows);
10636 let gpr = cols / GROUP_SIZE;
10637 let out_addr = SendMut(out.as_mut_ptr());
10638 if a8w8_enabled() {
10639 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
10640 .map(|bi| {
10641 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
10642 let gsum = q1_group_sums(&act.xq, gpr);
10643 (act, gsum)
10644 })
10645 .collect();
10646 let acts = &acts;
10647 #[cfg(target_arch = "x86_64")]
10648 let blocked_ok = avx2_enabled() && blocked_enabled();
10649 #[cfg(target_arch = "aarch64")]
10650 let blocked_ok = sdot_enabled() && blocked_enabled();
10651 let run = move |start: usize, end: usize| {
10652 for r in start..end {
10653 let mut bi = 0usize;
10654 #[cfg(target_arch = "aarch64")]
10657 if blocked_ok {
10658 while bi + 4 <= acts.len() {
10659 let xs = [
10660 acts[bi].0.xq.as_slice(),
10661 acts[bi + 1].0.xq.as_slice(),
10662 acts[bi + 2].0.xq.as_slice(),
10663 acts[bi + 3].0.xq.as_slice(),
10664 ];
10665 let gs = [
10666 acts[bi].1.as_slice(),
10667 acts[bi + 1].1.as_slice(),
10668 acts[bi + 2].1.as_slice(),
10669 acts[bi + 3].1.as_slice(),
10670 ];
10671 let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
10673 for k in 0..4 {
10674 let (act, _) = &acts[bi + k];
10675 let mut acc = d[k] * act.sx;
10676 for &(j, xv) in &act.outliers {
10677 let (wv, sc) = w.outlier(r, gpr, j);
10678 acc += wv * sc * xv;
10679 }
10680 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10682 }
10683 bi += 4;
10684 }
10685 }
10686 #[cfg(target_arch = "x86_64")]
10687 if blocked_ok {
10688 while bi + 4 <= acts.len() {
10689 let xs = [
10690 acts[bi].0.xq.as_slice(),
10691 acts[bi + 1].0.xq.as_slice(),
10692 acts[bi + 2].0.xq.as_slice(),
10693 acts[bi + 3].0.xq.as_slice(),
10694 ];
10695 let gs = [
10696 acts[bi].1.as_slice(),
10697 acts[bi + 1].1.as_slice(),
10698 acts[bi + 2].1.as_slice(),
10699 acts[bi + 3].1.as_slice(),
10700 ];
10701 let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
10703 for k in 0..4 {
10704 let (act, _) = &acts[bi + k];
10705 let mut acc = d[k] * act.sx;
10706 for &(j, xv) in &act.outliers {
10707 let (wv, sc) = w.outlier(r, gpr, j);
10708 acc += wv * sc * xv;
10709 }
10710 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10712 }
10713 bi += 4;
10714 }
10715 }
10716 while bi < acts.len() {
10717 let (act, gsum) = &acts[bi];
10718 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
10719 for &(j, xv) in &act.outliers {
10720 let (wv, s) = w.outlier(r, gpr, j);
10721 acc += wv * s * xv;
10722 }
10723 unsafe { *out_addr.at(bi * rows + r) = acc };
10725 bi += 1;
10726 }
10727 }
10728 };
10729 dispatch_rows(pool, rows, &run);
10730 return;
10731 }
10732 let run = move |start: usize, end: usize| {
10733 for r in start..end {
10734 for bi in 0..b {
10735 let x = &xs_all[bi * cols..(bi + 1) * cols];
10736 unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
10738 }
10739 }
10740 };
10741 dispatch_rows(pool, rows, &run);
10742}
10743
10744fn q4matvec(
10750 bytes: &[u8],
10751 x: &[f32],
10752 rows: usize,
10753 cols: usize,
10754 out: &mut [f32],
10755 pool: Option<&Pool>,
10756) {
10757 debug_assert_eq!(out.len(), rows);
10758 let (packed, scales) = q4_split(bytes, rows, cols);
10759 let gpr = cols / GROUP_SIZE;
10760 let out_addr = SendMut(out.as_mut_ptr());
10761
10762 if a8w8_enabled() {
10763 let act = split_act(x);
10764 let run = move |start: usize, end: usize| {
10765 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
10766 };
10767 dispatch_rows(pool, rows, &run);
10768 return;
10769 }
10770
10771 let run =
10772 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
10773 dispatch_rows(pool, rows, &run);
10774}
10775
10776#[inline]
10779#[allow(unreachable_code)]
10780#[cfg(target_arch = "x86_64")]
10785#[target_feature(enable = "avx2")]
10786unsafe fn dot_q4b_row_1x4_avx2(
10787 buf: &[u8],
10788 scales: &[u8],
10789 g0: usize,
10790 gpr: usize,
10791 xs: [&[i8]; 4],
10792) -> [f32; 4] {
10793 unsafe {
10795 use core::arch::x86_64::*;
10796 let ones = _mm256_set1_epi16(1);
10797 let mut acc = [0f32; 4];
10798 for gi in 0..gpr {
10799 let s = f16_to_f32(u16::from_le_bytes([
10800 scales[(g0 + gi) * 2],
10801 scales[(g0 + gi) * 2 + 1],
10802 ]));
10803 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10804 let aw = _mm256_abs_epi8(w);
10805 for (k, xq) in xs.iter().enumerate() {
10806 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10807 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
10808 let d = _mm256_madd_epi16(p16, ones);
10809 let hi128 = _mm256_extracti128_si256::<1>(d);
10810 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10811 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10812 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10813 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
10814 }
10815 }
10816 acc
10817 }
10818}
10819
10820#[cfg(target_arch = "x86_64")]
10822#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10823unsafe fn dot_q4b_row_1x4_vnni(
10824 buf: &[u8],
10825 scales: &[u8],
10826 g0: usize,
10827 gpr: usize,
10828 xs: [&[i8]; 4],
10829) -> [f32; 4] {
10830 unsafe {
10832 use core::arch::x86_64::*;
10833 let mut acc = [0f32; 4];
10834 for gi in 0..gpr {
10835 let s = f16_to_f32(u16::from_le_bytes([
10836 scales[(g0 + gi) * 2],
10837 scales[(g0 + gi) * 2 + 1],
10838 ]));
10839 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10840 let aw = _mm256_abs_epi8(w);
10841 for (k, xq) in xs.iter().enumerate() {
10842 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10843 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
10844 acc[k] += d as f32 * s;
10845 }
10846 }
10847 acc
10848 }
10849}
10850
10851#[cfg(target_arch = "x86_64")]
10857#[target_feature(enable = "avx2")]
10858unsafe fn dot_q4b_row_1x4_sx_avx2(
10859 buf: &[u8],
10860 scales: &[u8],
10861 g0: usize,
10862 gpr: usize,
10863 xs: [&[i8]; 4],
10864 sxs: [f32; 4],
10865) -> [f32; 4] {
10866 unsafe {
10868 use core::arch::x86_64::*;
10869 let ones = _mm256_set1_epi16(1);
10870 let mut acc = [0f32; 4];
10871 for gi in 0..gpr {
10872 let s = f16_to_f32(u16::from_le_bytes([
10873 scales[(g0 + gi) * 2],
10874 scales[(g0 + gi) * 2 + 1],
10875 ]));
10876 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10877 let aw = _mm256_abs_epi8(w);
10878 for (k, xq) in xs.iter().enumerate() {
10879 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10880 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
10881 let d = _mm256_madd_epi16(p16, ones);
10882 let hi128 = _mm256_extracti128_si256::<1>(d);
10883 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10884 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10885 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10886 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
10887 }
10888 }
10889 acc
10890 }
10891}
10892
10893#[cfg(target_arch = "x86_64")]
10896#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10897unsafe fn dot_q4b_row_1x4_sx_vnni(
10898 buf: &[u8],
10899 scales: &[u8],
10900 g0: usize,
10901 gpr: usize,
10902 xs: [&[i8]; 4],
10903 sxs: [f32; 4],
10904) -> [f32; 4] {
10905 unsafe {
10907 use core::arch::x86_64::*;
10908 let mut acc = [0f32; 4];
10909 for gi in 0..gpr {
10910 let s = f16_to_f32(u16::from_le_bytes([
10911 scales[(g0 + gi) * 2],
10912 scales[(g0 + gi) * 2 + 1],
10913 ]));
10914 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10915 let aw = _mm256_abs_epi8(w);
10916 for (k, xq) in xs.iter().enumerate() {
10917 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10918 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
10919 acc[k] += (d as f32 * sxs[k]) * s;
10920 }
10921 }
10922 acc
10923 }
10924}
10925
10926#[allow(unreachable_code)]
10927fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10928 #[cfg(target_arch = "aarch64")]
10929 unsafe {
10930 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
10931 }
10932 #[cfg(target_arch = "x86_64")]
10933 unsafe {
10934 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
10935 }
10936 let mut acc = 0f32;
10937 for gi in 0..gpr {
10938 let g = g0 + gi;
10939 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10940 let mut d = 0i32;
10941 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10942 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
10943 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
10944 }
10945 acc += d as f32 * s;
10946 }
10947 acc
10948}
10949
10950#[inline]
10952#[allow(unreachable_code)]
10953fn dot_q4_row_i8_2(
10954 packed: &[u8],
10955 scales: &[u8],
10956 g0: usize,
10957 gpr: usize,
10958 xq1: &[i8],
10959 xq2: &[i8],
10960) -> (f32, f32) {
10961 #[cfg(target_arch = "aarch64")]
10962 unsafe {
10963 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
10964 }
10965 #[cfg(target_arch = "x86_64")]
10966 unsafe {
10967 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
10968 }
10969 (
10970 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
10971 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
10972 )
10973}
10974
10975#[allow(clippy::too_many_arguments)]
10978fn q4_range_a8w8(
10979 packed: &[u8],
10980 scales: &[u8],
10981 gpr: usize,
10982 cols: usize,
10983 act: &SplitAct,
10984 out: SendMut,
10985 start: usize,
10986 end: usize,
10987) {
10988 for r in start..end {
10989 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
10990 for &(j, xv) in &act.outliers {
10992 let flat = r * cols + j;
10993 let byte = packed[flat / 2];
10994 let nib = if flat & 1 == 0 {
10995 byte & 0x0F
10996 } else {
10997 byte >> 4
10998 };
10999 let s = f16_to_f32(u16::from_le_bytes([
11000 scales[(flat / GROUP_SIZE) * 2],
11001 scales[(flat / GROUP_SIZE) * 2 + 1],
11002 ]));
11003 acc += ((nib as i32 - 8) as f32) * s * xv;
11004 }
11005 unsafe { *out.at(r) = acc };
11007 }
11008}
11009
11010#[allow(clippy::too_many_arguments)]
11013fn q4_range2_a8w8(
11014 packed: &[u8],
11015 scales: &[u8],
11016 gpr: usize,
11017 cols: usize,
11018 a1: &SplitAct,
11019 a2: &SplitAct,
11020 p1: SendMut,
11021 p2: SendMut,
11022 start: usize,
11023 end: usize,
11024) {
11025 for r in start..end {
11026 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
11027 let mut acc1 = s1 * a1.sx;
11028 let mut acc2 = s2 * a2.sx;
11029 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
11031 for &(j, xv) in outliers {
11032 let flat = r * cols + j;
11033 let byte = packed[flat / 2];
11034 let nib = if flat & 1 == 0 {
11035 byte & 0x0F
11036 } else {
11037 byte >> 4
11038 };
11039 let s = f16_to_f32(u16::from_le_bytes([
11040 scales[(flat / GROUP_SIZE) * 2],
11041 scales[(flat / GROUP_SIZE) * 2 + 1],
11042 ]));
11043 *acc += ((nib as i32 - 8) as f32) * s * xv;
11044 }
11045 };
11046 fix(&a1.outliers, &mut acc1);
11047 fix(&a2.outliers, &mut acc2);
11048 unsafe {
11050 *p1.at(r) = acc1;
11051 *p2.at(r) = acc2;
11052 }
11053 }
11054}
11055
11056fn q4_range_f32(
11058 packed: &[u8],
11059 scales: &[u8],
11060 gpr: usize,
11061 x: &[f32],
11062 out: SendMut,
11063 start: usize,
11064 end: usize,
11065) {
11066 for r in start..end {
11067 let mut acc = 0f32;
11068 for gi in 0..gpr {
11069 let g = r * gpr + gi;
11070 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11071 let pk = &packed[g * 16..(g + 1) * 16];
11072 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11073 let mut ga = 0f32;
11074 for (k, &b) in pk.iter().enumerate() {
11075 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
11076 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
11077 }
11078 acc += ga * s;
11079 }
11080 unsafe { *out.at(r) = acc };
11082 }
11083}
11084
11085#[allow(clippy::too_many_arguments)]
11089fn q4matvec2(
11090 bytes: &[u8],
11091 x1: &[f32],
11092 x2: &[f32],
11093 rows: usize,
11094 cols: usize,
11095 o1: &mut [f32],
11096 o2: &mut [f32],
11097 pool: Option<&Pool>,
11098) {
11099 debug_assert_eq!(o1.len(), rows);
11100 debug_assert_eq!(o2.len(), rows);
11101 let (packed, scales) = q4_split(bytes, rows, cols);
11102 let gpr = cols / GROUP_SIZE;
11103
11104 if a8w8_enabled() {
11105 let a1 = split_act(x1);
11106 let a2 = split_act(x2);
11107 let p1 = SendMut(o1.as_mut_ptr());
11108 let p2 = SendMut(o2.as_mut_ptr());
11109 let run = move |start: usize, end: usize| {
11110 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
11111 };
11112 dispatch_rows(pool, rows, &run);
11113 return;
11114 }
11115
11116 let p1 = SendMut(o1.as_mut_ptr());
11117 let p2 = SendMut(o2.as_mut_ptr());
11118 let run = move |start: usize, end: usize| {
11119 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
11120 };
11121 dispatch_rows(pool, rows, &run);
11122}
11123
11124#[allow(clippy::too_many_arguments)]
11126fn q4_range2_f32(
11127 packed: &[u8],
11128 scales: &[u8],
11129 gpr: usize,
11130 x1: &[f32],
11131 x2: &[f32],
11132 p1: SendMut,
11133 p2: SendMut,
11134 start: usize,
11135 end: usize,
11136) {
11137 for r in start..end {
11138 let (mut acc1, mut acc2) = (0f32, 0f32);
11139 for gi in 0..gpr {
11140 let g = r * gpr + gi;
11141 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11142 let pk = &packed[g * 16..(g + 1) * 16];
11143 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11144 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11145 let (mut g1, mut g2) = (0f32, 0f32);
11146 for (k, &b) in pk.iter().enumerate() {
11147 let wl = (b & 0x0F) as f32 - 8.0;
11148 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
11149 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
11150 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
11151 }
11152 acc1 += g1 * s;
11153 acc2 += g2 * s;
11154 }
11155 unsafe {
11157 *p1.at(r) = acc1;
11158 *p2.at(r) = acc2;
11159 }
11160 }
11161}
11162
11163thread_local! {
11164 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
11167 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11168}
11169
11170#[allow(clippy::too_many_arguments)]
11176fn q4matmat(
11177 bytes: &[u8],
11178 xs_all: &[f32],
11179 b: usize,
11180 rows: usize,
11181 cols: usize,
11182 out: &mut [f32],
11183 pool: Option<&Pool>,
11184) {
11185 debug_assert_eq!(xs_all.len(), b * cols);
11186 debug_assert_eq!(out.len(), b * rows);
11187 let (packed, scales) = q4_split(bytes, rows, cols);
11188 let gpr = cols / GROUP_SIZE;
11189 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11190
11191 if a8w8_enabled() {
11192 let acts: Vec<SplitAct> = (0..b)
11193 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11194 .collect();
11195 let acts = &acts;
11196 let out_addr = SendMut(out.as_mut_ptr());
11197 let run = move |start: usize, end: usize| {
11198 ROW_I8.with(|rb| {
11199 let mut buf = rb.borrow_mut();
11200 buf.resize(cols, 0);
11201 for r in start..end {
11202 for gi in 0..gpr {
11206 let g = r * gpr + gi;
11207 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11208 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
11209 buf[gi * GROUP_SIZE + k * 2 + 1] =
11210 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
11211 }
11212 }
11213 let mut bi = 0usize;
11214 #[cfg(target_arch = "x86_64")]
11215 if avx2_enabled() && blocked_enabled() {
11216 while bi + 4 <= acts.len() {
11217 let xs = [
11218 acts[bi].xq.as_slice(),
11219 acts[bi + 1].xq.as_slice(),
11220 acts[bi + 2].xq.as_slice(),
11221 acts[bi + 3].xq.as_slice(),
11222 ];
11223 let d = unsafe {
11224 if vnni_tiles_enabled() {
11225 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
11226 } else {
11227 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
11228 }
11229 };
11230 for k in 0..4 {
11231 let act = &acts[bi + k];
11232 let mut acc = d[k] * act.sx;
11233 for &(j, xv) in &act.outliers {
11234 acc += (buf[j] as i8) as f32
11235 * gscale((r * cols + j) / GROUP_SIZE)
11236 * xv;
11237 }
11238 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11240 }
11241 bi += 4;
11242 }
11243 }
11244 while bi < acts.len() {
11245 let act = &acts[bi];
11246 let mut acc = 0f32;
11247 for gi in 0..gpr {
11248 let d = dot_i8_i8(
11249 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11250 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11251 );
11252 acc += d as f32 * gscale(r * gpr + gi);
11253 }
11254 acc *= act.sx;
11255 for &(j, xv) in &act.outliers {
11257 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
11258 }
11259 unsafe { *out_addr.at(bi * rows + r) = acc };
11261 bi += 1;
11262 }
11263 }
11264 })
11265 };
11266 dispatch_rows(pool, rows, &run);
11267 return;
11268 }
11269
11270 let out_addr = SendMut(out.as_mut_ptr());
11271 let run = move |start: usize, end: usize| {
11272 ROW_F32.with(|rb| {
11273 let mut buf = rb.borrow_mut();
11274 buf.resize(cols, 0.0);
11275 for r in start..end {
11276 for gi in 0..gpr {
11279 let g = r * gpr + gi;
11280 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11281 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
11282 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
11283 }
11284 }
11285 for bi in 0..b {
11286 let x = &xs_all[bi * cols..(bi + 1) * cols];
11287 let mut acc = 0f32;
11288 for gi in 0..gpr {
11289 let mut ga = 0f32;
11290 for k in 0..GROUP_SIZE / 2 {
11295 let e = gi * GROUP_SIZE + k * 2;
11296 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
11297 }
11298 acc += ga * gscale(r * gpr + gi);
11299 }
11300 unsafe { *out_addr.at(bi * rows + r) = acc };
11302 }
11303 }
11304 })
11305 };
11306 dispatch_rows(pool, rows, &run);
11307}
11308
11309#[allow(clippy::too_many_arguments)]
11314fn vbitmatmat(
11315 bytes: &[u8],
11316 offsets: &[usize],
11317 xs_all: &[f32],
11318 b: usize,
11319 rows: usize,
11320 cols: usize,
11321 out: &mut [f32],
11322 pool: Option<&Pool>,
11323) {
11324 debug_assert_eq!(xs_all.len(), b * cols);
11325 debug_assert_eq!(out.len(), b * rows);
11326 debug_assert_eq!(offsets.len(), rows + 1);
11327 let ng = cols / GROUP_SIZE;
11328 let bits = &bytes[..rows];
11329 let sc_off = rows;
11330 let gscale = |r: usize, g: usize| {
11331 let so = (r * ng + g) * 2;
11332 f16_to_f32(u16::from_le_bytes([
11333 bytes[sc_off + so],
11334 bytes[sc_off + so + 1],
11335 ]))
11336 };
11337
11338 let decode_f32 = |r: usize, dst: &mut [f32]| {
11340 let bw = bits[r] as usize;
11341 let l = ((1i32 << (bw - 1)) - 1) as f32;
11342 let data = &bytes[offsets[r]..offsets[r + 1]];
11343 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
11344 for d in dst.iter_mut() {
11345 while nbits < bw {
11346 acc = (acc << 8) | data[idx] as u64;
11347 idx += 1;
11348 nbits += 8;
11349 }
11350 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
11351 nbits -= bw;
11352 *d = u - l;
11353 }
11354 };
11355
11356 if a8w8_enabled() {
11357 let acts: Vec<SplitAct> = (0..b)
11358 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11359 .collect();
11360 let acts = &acts;
11361 let out_addr = SendMut(out.as_mut_ptr());
11362 let run = move |start: usize, end: usize| {
11363 for r in start..end {
11364 let bw = bits[r] as usize;
11365 if bw == 8 {
11366 ROW_F32.with(|rb| {
11369 let mut buf = rb.borrow_mut();
11370 buf.resize(cols, 0.0);
11371 decode_f32(r, &mut buf);
11372 for bi in 0..b {
11373 let x = &xs_all[bi * cols..(bi + 1) * cols];
11374 let mut dot = 0f32;
11375 for g in 0..ng {
11376 let mut gd = 0f32;
11377 for k in 0..GROUP_SIZE {
11378 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
11379 }
11380 dot += gd * gscale(r, g);
11381 }
11382 unsafe { *out_addr.at(bi * rows + r) = dot };
11384 }
11385 });
11386 continue;
11387 }
11388 let l = (1i32 << (bw - 1)) - 1;
11389 let data = &bytes[offsets[r]..offsets[r + 1]];
11390 ROW_I8.with(|rb| {
11391 let mut buf = rb.borrow_mut();
11392 buf.resize(cols, 0);
11393 #[inline(always)]
11394 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
11395 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
11396 let u = unpack8::<B>(&data[blk * B..]);
11397 for k in 0..8 {
11398 chunk[k] = (u[k] - l) as i8 as u8;
11399 }
11400 }
11401 }
11402 match bw {
11403 3 => fill::<3>(data, l, &mut buf),
11404 4 => vbit_fill4(data, &mut buf),
11405 5 => fill::<5>(data, l, &mut buf),
11406 6 => fill::<6>(data, l, &mut buf),
11407 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
11408 }
11409 let mut bi = 0usize;
11410 #[cfg(target_arch = "x86_64")]
11414 if avx2_enabled() && blocked_enabled() {
11415 while bi + 4 <= acts.len() {
11416 let xs = [
11417 acts[bi].xq.as_slice(),
11418 acts[bi + 1].xq.as_slice(),
11419 acts[bi + 2].xq.as_slice(),
11420 acts[bi + 3].xq.as_slice(),
11421 ];
11422 let sxs = [
11423 acts[bi].sx,
11424 acts[bi + 1].sx,
11425 acts[bi + 2].sx,
11426 acts[bi + 3].sx,
11427 ];
11428 let d = unsafe {
11429 if vnni_tiles_enabled() {
11430 dot_q4b_row_1x4_sx_vnni(
11431 &buf,
11432 &bytes[sc_off..],
11433 r * ng,
11434 ng,
11435 xs,
11436 sxs,
11437 )
11438 } else {
11439 dot_q4b_row_1x4_sx_avx2(
11440 &buf,
11441 &bytes[sc_off..],
11442 r * ng,
11443 ng,
11444 xs,
11445 sxs,
11446 )
11447 }
11448 };
11449 for k in 0..4 {
11450 let act = &acts[bi + k];
11451 let mut dot = d[k];
11452 for &(j, xv) in &act.outliers {
11453 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
11454 }
11455 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
11457 }
11458 bi += 4;
11459 }
11460 }
11461 while bi < acts.len() {
11462 let act = &acts[bi];
11463 let mut dot = 0f32;
11464 for g in 0..ng {
11465 let d = dot_i8_i8(
11466 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
11467 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
11468 ) as f32
11469 * act.sx;
11470 dot += d * gscale(r, g);
11471 }
11472 for &(j, xv) in &act.outliers {
11473 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
11474 }
11475 unsafe { *out_addr.at(bi * rows + r) = dot };
11477 bi += 1;
11478 }
11479 });
11480 }
11481 };
11482 dispatch_rows(pool, rows, &run);
11483 return;
11484 }
11485
11486 let out_addr = SendMut(out.as_mut_ptr());
11487 let run = move |start: usize, end: usize| {
11488 ROW_F32.with(|rb| {
11489 let mut buf = rb.borrow_mut();
11490 buf.resize(cols, 0.0);
11491 for r in start..end {
11492 decode_f32(r, &mut buf);
11493 for bi in 0..b {
11494 let x = &xs_all[bi * cols..(bi + 1) * cols];
11495 let mut dot = 0f32;
11496 for g in 0..ng {
11497 let mut gd = 0f32;
11498 for k in 0..GROUP_SIZE {
11499 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
11500 }
11501 dot += gd * gscale(r, g);
11502 }
11503 unsafe { *out_addr.at(bi * rows + r) = dot };
11505 }
11506 }
11507 })
11508 };
11509 dispatch_rows(pool, rows, &run);
11510}
11511
11512pub(crate) fn gpu_batch_job<'a>(
11516 t: &'a QTensor,
11517 x: &[f32],
11518) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
11519 match t {
11520 QTensor::Mapped {
11521 model,
11522 idx,
11523 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
11524 rows,
11525 cols,
11526 row_scale,
11527 col_field,
11528 ..
11529 } => Some((
11530 model.clone(),
11531 crate::gpu::BatchJob {
11532 idx: *idx,
11533 rows: *rows,
11534 cols: *cols,
11535 row_scale,
11536 xs: prescale(x, col_field, *dt).into_owned(),
11537 layout: crate::gpu::BatchLayout::Q8,
11538 },
11539 )),
11540 QTensor::Mapped {
11542 model,
11543 idx,
11544 dtype: TensorDtype::Q1,
11545 rows,
11546 cols,
11547 ..
11548 } => Some((
11549 model.clone(),
11550 crate::gpu::BatchJob {
11551 idx: *idx,
11552 rows: *rows,
11553 cols: *cols,
11554 row_scale: &[],
11555 xs: x.to_vec(),
11556 layout: crate::gpu::BatchLayout::Q1,
11557 },
11558 )),
11559 QTensor::Mapped {
11561 model,
11562 idx,
11563 dtype: TensorDtype::Q1G,
11564 rows,
11565 cols,
11566 ..
11567 } => Some((
11568 model.clone(),
11569 crate::gpu::BatchJob {
11570 idx: *idx,
11571 rows: *rows,
11572 cols: *cols,
11573 row_scale: &[],
11574 xs: x.to_vec(),
11575 layout: crate::gpu::BatchLayout::Q1g,
11576 },
11577 )),
11578 QTensor::Mapped {
11583 model,
11584 idx,
11585 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
11586 rows,
11587 cols,
11588 ..
11589 } => Some((
11590 model.clone(),
11591 crate::gpu::BatchJob {
11592 idx: *idx,
11593 rows: *rows,
11594 cols: *cols,
11595 row_scale: &[],
11596 xs: x.to_vec(),
11597 layout: if *dt == TensorDtype::Q4Tiled {
11598 crate::gpu::BatchLayout::Q4t
11599 } else {
11600 crate::gpu::BatchLayout::Q4tp
11601 },
11602 },
11603 )),
11604 _ => None,
11605 }
11606}
11607
11608thread_local! {
11609 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11610 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11611}
11612
11613pub(crate) fn prescale<'a>(
11614 x: &'a [f32],
11615 col_field: &[f32],
11616 dtype: TensorDtype,
11617) -> std::borrow::Cow<'a, [f32]> {
11618 if dtype == TensorDtype::Q8_2f {
11619 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
11620 } else {
11621 std::borrow::Cow::Borrowed(x)
11622 }
11623}
11624
11625pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
11628 x: &[f32],
11629 col_field: &[f32],
11630 dtype: TensorDtype,
11631 buf_id: u8,
11632 f: F,
11633) -> R {
11634 if dtype == TensorDtype::Q8_2f {
11635 if buf_id == 1 {
11636 PRESCALE_BUF1.with(|b| {
11637 let mut buf = b.borrow_mut();
11638 buf.clear();
11639 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
11640 f(&buf)
11641 })
11642 } else {
11643 PRESCALE_BUF2.with(|b| {
11644 let mut buf = b.borrow_mut();
11645 buf.clear();
11646 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
11647 f(&buf)
11648 })
11649 }
11650 } else {
11651 f(x)
11652 }
11653}
11654
11655#[cfg(target_arch = "x86_64")]
11660pub(crate) fn avx2_enabled() -> bool {
11661 use std::sync::OnceLock;
11662 static ON: OnceLock<bool> = OnceLock::new();
11663 *ON.get_or_init(|| {
11664 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
11665 && std::arch::is_x86_feature_detected!("avx2")
11666 && std::arch::is_x86_feature_detected!("fma")
11667 })
11668}
11669
11670#[cfg(target_arch = "x86_64")]
11675fn avx2_a8w8_enabled() -> bool {
11676 if FLOAT_ACTIVATIONS.get() {
11677 return false;
11678 }
11679 use std::sync::OnceLock;
11680 static ON: OnceLock<bool> = OnceLock::new();
11681 *ON.get_or_init(|| {
11682 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
11683 })
11684}
11685
11686thread_local! {
11687 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
11688}
11689
11690pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
11693 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
11694 impl Drop for Restore {
11695 fn drop(&mut self) {
11696 FULL_GPU_Q8.set(self.0);
11697 }
11698 }
11699 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
11700}
11701
11702thread_local! {
11706 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
11707}
11708
11709fn flint_exact_cpu_default_enabled() -> bool {
11716 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11717 *ON.get_or_init(|| std::env::var_os("CMF_SDOT").is_none())
11718}
11719
11720impl QTensor {
11721 #[inline]
11722 fn flint_exact_cpu_default(&self) -> bool {
11723 !FLOAT_ACTIVATIONS.get()
11724 && flint_exact_cpu_default_enabled()
11725 && matches!(
11726 self,
11727 Self::Mapped { model, .. } if model.header.arch.arch_name == "flint"
11728 )
11729 }
11730
11731 #[inline]
11732 fn any_flint_exact_cpu_default(ts: &[&QTensor]) -> bool {
11733 ts.iter().any(|t| t.flint_exact_cpu_default())
11734 }
11735
11736 #[inline]
11737 fn any_flint_exact_cpu_default_pairs(pairs: &[(&QTensor, &QTensor)]) -> bool {
11738 pairs
11739 .iter()
11740 .any(|(gate, up)| gate.flint_exact_cpu_default() || up.flint_exact_cpu_default())
11741 }
11742}
11743
11744pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
11749 float_activations_scope(f)
11750}
11751
11752#[inline]
11756pub(crate) fn float_activations_active() -> bool {
11757 FLOAT_ACTIVATIONS.get()
11758}
11759
11760pub(crate) struct FloatActivationsScope {
11764 previous: bool,
11765 _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
11766}
11767
11768impl Drop for FloatActivationsScope {
11769 fn drop(&mut self) {
11770 FLOAT_ACTIVATIONS.set(self.previous);
11771 }
11772}
11773
11774#[inline]
11775pub(crate) fn enter_float_activations_scope() -> FloatActivationsScope {
11776 FloatActivationsScope {
11777 previous: FLOAT_ACTIVATIONS.replace(true),
11778 _not_send: std::marker::PhantomData,
11779 }
11780}
11781
11782#[doc(hidden)]
11788pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
11789 let _scope = enter_float_activations_scope();
11790 f()
11791}
11792
11793static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
11805
11806pub(crate) fn row_exact() -> bool {
11807 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
11808}
11809
11810fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
11811 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
11812 impl Drop for Restore<'_> {
11813 fn drop(&mut self) {
11814 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
11815 }
11816 }
11817 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
11818 let _restore = Restore(active);
11819 f()
11820}
11821
11822pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
11824 counted_row_exact_scope(&ROW_EXACT, f)
11825}
11826
11827#[inline]
11831pub(crate) fn a8w8_enabled() -> bool {
11832 #[cfg(target_arch = "aarch64")]
11833 {
11834 sdot_enabled()
11835 }
11836 #[cfg(target_arch = "x86_64")]
11837 {
11838 avx2_a8w8_enabled()
11839 }
11840 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
11841 {
11842 false
11843 }
11844}
11845
11846#[inline]
11849#[allow(unreachable_code)]
11850fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
11851 #[cfg(target_arch = "aarch64")]
11852 unsafe {
11853 return dot_i8_sdot(w, xq);
11854 }
11855 #[cfg(target_arch = "x86_64")]
11856 unsafe {
11857 if avx512vnni_enabled() {
11858 return dot_i8_i8_vnni(w, xq);
11859 }
11860 return dot_i8_i8_avx2(w, xq);
11861 }
11862 w.iter()
11863 .zip(xq)
11864 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
11865 .sum()
11866}
11867
11868#[cfg(target_arch = "x86_64")]
11872fn avx512vnni_enabled() -> bool {
11873 use std::sync::OnceLock;
11874 static ON: OnceLock<bool> = OnceLock::new();
11875 *ON.get_or_init(|| {
11876 std::env::var("CMF_AVX512")
11877 .map(|v| v != "0")
11878 .unwrap_or(true)
11879 && std::arch::is_x86_feature_detected!("avx512f")
11880 && std::arch::is_x86_feature_detected!("avx512bw")
11881 && std::arch::is_x86_feature_detected!("avx512vl")
11882 && std::arch::is_x86_feature_detected!("avx512vnni")
11883 })
11884}
11885
11886#[cfg(target_arch = "x86_64")]
11894fn vnni_tiles_enabled() -> bool {
11895 use std::sync::OnceLock;
11896 static ON: OnceLock<bool> = OnceLock::new();
11897 *ON.get_or_init(|| {
11898 std::env::var("CMF_VNNI_TILES")
11899 .map(|v| v != "0")
11900 .unwrap_or(true)
11901 && avx512vnni_enabled()
11902 })
11903}
11904
11905#[cfg(target_arch = "x86_64")]
11910#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11911#[inline]
11912unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
11913 unsafe {
11915 use core::arch::x86_64::*;
11916 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
11917 let hi128 = _mm256_extracti128_si256::<1>(d);
11918 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11919 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11920 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11921 _mm_cvtsi128_si32(s32)
11922 }
11923}
11924
11925#[cfg(target_arch = "x86_64")]
11930#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11931unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11932 unsafe {
11934 use core::arch::x86_64::*;
11935 let n = w.len();
11936 let mut j = 0usize;
11937 let mut total: i32;
11938 {
11943 #[inline(always)]
11944 unsafe fn step(
11945 w: *const u8,
11946 x: *const i8,
11947 acc: core::arch::x86_64::__m512i,
11948 ) -> core::arch::x86_64::__m512i {
11949 unsafe {
11950 use core::arch::x86_64::*;
11951 let wv = _mm512_loadu_si512(w as *const _);
11952 let xv = _mm512_loadu_si512(x as *const _);
11953 let aw = _mm512_abs_epi8(wv);
11954 let neg = _mm512_movepi8_mask(wv);
11955 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
11956 _mm512_dpbusd_epi32(acc, aw, sx)
11957 }
11958 }
11959 let (mut a0, mut a1, mut a2, mut a3) = (
11960 _mm512_setzero_si512(),
11961 _mm512_setzero_si512(),
11962 _mm512_setzero_si512(),
11963 _mm512_setzero_si512(),
11964 );
11965 while j + 256 <= n {
11966 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
11967 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
11968 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
11969 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
11970 j += 256;
11971 }
11972 while j + 64 <= n {
11973 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
11974 j += 64;
11975 }
11976 let s01 = _mm512_add_epi32(a0, a1);
11977 let s23 = _mm512_add_epi32(a2, a3);
11978 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
11979 }
11980 if j + 32 <= n {
11982 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
11983 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
11984 let d = _mm256_dpbusd_epi32(
11985 _mm256_setzero_si256(),
11986 _mm256_abs_epi8(wv),
11987 _mm256_sign_epi8(xv, wv),
11988 );
11989 let hi128 = _mm256_extracti128_si256::<1>(d);
11990 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11991 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11992 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11993 total += _mm_cvtsi128_si32(s32);
11994 j += 32;
11995 }
11996 while j < n {
11997 total += (w[j] as i8) as i32 * xq[j] as i32;
11998 j += 1;
11999 }
12000 total
12001 }
12002}
12003
12004#[cfg(target_arch = "x86_64")]
12006#[target_feature(enable = "avx2,fma")]
12007unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
12008 unsafe {
12010 use core::arch::x86_64::*;
12011 let n = x.len();
12012 let wp = w.as_ptr();
12013 let xp = x.as_ptr();
12014 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
12015 let mut j = 0usize;
12016 while j + 16 <= n {
12017 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
12018 let lo = _mm256_cvtepi8_epi32(wb);
12019 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
12020 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
12021 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
12022 j += 16;
12023 }
12024 let acc = _mm256_add_ps(a0, a1);
12025 let hi128 = _mm256_extractf128_ps::<1>(acc);
12026 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
12027 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
12028 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
12029 let mut sum = _mm_cvtss_f32(s32);
12030 while j < n {
12031 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
12032 j += 1;
12033 }
12034 sum
12035 }
12036}
12037
12038#[cfg(target_arch = "x86_64")]
12043#[target_feature(enable = "avx2")]
12044unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
12045 unsafe {
12047 use core::arch::x86_64::*;
12048 let n = w.len();
12049 let ones = _mm256_set1_epi16(1);
12050 let mut acc = _mm256_setzero_si256();
12051 let mut j = 0usize;
12052 while j + 32 <= n {
12053 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12054 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12055 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
12056 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
12057 j += 32;
12058 }
12059 let hi128 = _mm256_extracti128_si256::<1>(acc);
12060 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
12061 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12062 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12063 let mut s = _mm_cvtsi128_si32(s32);
12064 while j < n {
12065 s += (w[j] as i8) as i32 * xq[j] as i32;
12066 j += 1;
12067 }
12068 s
12069 }
12070}
12071
12072#[cfg(target_arch = "aarch64")]
12076#[target_feature(enable = "neon,i8mm")]
12077unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12078 unsafe {
12080 use core::arch::aarch64::*;
12081 use core::arch::asm;
12082 let n = w0.len();
12083 let w0p = w0.as_ptr() as *const i8;
12084 let w1p = w1.as_ptr() as *const i8;
12085 let mut acc01 = vdupq_n_s32(0);
12088 let mut acc23 = vdupq_n_s32(0);
12089 let mut i = 0usize;
12090 while i + 8 <= n {
12091 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
12092 let xb01 = vcombine_s8(
12093 vld1_s8(xs[0].as_ptr().add(i)),
12094 vld1_s8(xs[1].as_ptr().add(i)),
12095 );
12096 let xb23 = vcombine_s8(
12097 vld1_s8(xs[2].as_ptr().add(i)),
12098 vld1_s8(xs[3].as_ptr().add(i)),
12099 );
12100 asm!(
12101 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
12102 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
12103 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
12104 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
12105 options(pure, nomem, nostack),
12106 );
12107 i += 8;
12108 }
12109 let mut out = [[0i32; 4]; 2];
12110 let a01: [i32; 4] = core::mem::transmute(acc01);
12111 let a23: [i32; 4] = core::mem::transmute(acc23);
12112 out[0][0] = a01[0];
12113 out[0][1] = a01[1];
12114 out[1][0] = a01[2];
12115 out[1][1] = a01[3];
12116 out[0][2] = a23[0];
12117 out[0][3] = a23[1];
12118 out[1][2] = a23[2];
12119 out[1][3] = a23[3];
12120 if i < n {
12121 for (k, x) in xs.iter().enumerate() {
12122 for j in i..n {
12123 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12124 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12125 }
12126 }
12127 }
12128 out
12129 }
12130}
12131
12132#[cfg(target_arch = "aarch64")]
12136#[target_feature(enable = "neon,dotprod")]
12137unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12138 unsafe {
12140 use core::arch::aarch64::*;
12141 use core::arch::asm;
12142 let n = w0.len();
12143 let w0p = w0.as_ptr() as *const i8;
12144 let w1p = w1.as_ptr() as *const i8;
12145 let mut acc = [[vdupq_n_s32(0); 4]; 2];
12146 let mut i = 0usize;
12147 while i + 16 <= n {
12148 let wv0 = vld1q_s8(w0p.add(i));
12149 let wv1 = vld1q_s8(w1p.add(i));
12150 for (k, x) in xs.iter().enumerate() {
12151 let xv = vld1q_s8(x.as_ptr().add(i));
12152 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
12153 asm!(
12154 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
12155 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
12156 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12157 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
12158 options(pure, nomem, nostack),
12159 );
12160 acc[0][k] = a0;
12161 acc[1][k] = a1;
12162 }
12163 i += 16;
12164 }
12165 let mut out = [[0i32; 4]; 2];
12166 for r in 0..2 {
12167 for k in 0..4 {
12168 out[r][k] = vaddvq_s32(acc[r][k]);
12169 }
12170 }
12171 if i < n {
12172 for (k, x) in xs.iter().enumerate() {
12173 for j in i..n {
12174 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12175 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12176 }
12177 }
12178 }
12179 out
12180 }
12181}
12182
12183#[cfg(target_arch = "x86_64")]
12189#[target_feature(enable = "avx2")]
12190unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12191 unsafe {
12193 use core::arch::x86_64::*;
12194 let n = w0.len();
12195 let ones = _mm256_set1_epi16(1);
12196 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
12197 let mut j = 0usize;
12198 while j + 32 <= n {
12199 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
12200 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
12201 let aw0 = _mm256_abs_epi8(wv0);
12202 let aw1 = _mm256_abs_epi8(wv1);
12203 for (k, x) in xs.iter().enumerate() {
12204 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
12205 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
12206 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
12207 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
12208 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
12209 }
12210 j += 32;
12211 }
12212 let mut out = [[0i32; 4]; 2];
12213 for r in 0..2 {
12214 for k in 0..4 {
12215 let a = acc[r][k];
12216 let hi128 = _mm256_extracti128_si256::<1>(a);
12217 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
12218 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12219 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12220 out[r][k] = _mm_cvtsi128_si32(s32);
12221 }
12222 }
12223 if j < n {
12224 for (k, x) in xs.iter().enumerate() {
12225 for i in j..n {
12226 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
12227 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
12228 }
12229 }
12230 }
12231 out
12232 }
12233}
12234
12235#[cfg(target_arch = "x86_64")]
12240#[inline]
12241fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
12242 let dot = if avx512vnni_enabled() && row.len() >= 64 {
12243 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
12244 } else {
12245 unsafe { dot_i8_i8_avx2(row, &act.xq) }
12246 };
12247 let mut acc = dot as f32 * act.sx;
12248 for &(j, xv) in &act.outliers {
12249 acc += (row[j] as i8) as f32 * xv;
12250 }
12251 acc
12252}
12253
12254#[cfg(target_arch = "x86_64")]
12258#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12259unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12260 unsafe {
12262 use core::arch::x86_64::*;
12263 let n = w.len();
12264 let flip = _mm512_set1_epi8(-128); #[inline(always)]
12266 unsafe fn step(
12267 w: *const u8,
12268 x: *const i8,
12269 flip: core::arch::x86_64::__m512i,
12270 acc: core::arch::x86_64::__m512i,
12271 ) -> core::arch::x86_64::__m512i {
12272 unsafe {
12273 use core::arch::x86_64::*;
12274 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
12275 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
12276 }
12277 }
12278 let (mut a0, mut a1, mut a2, mut a3) = (
12279 _mm512_setzero_si512(),
12280 _mm512_setzero_si512(),
12281 _mm512_setzero_si512(),
12282 _mm512_setzero_si512(),
12283 );
12284 let mut j = 0usize;
12285 while j + 256 <= n {
12286 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12287 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
12288 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
12289 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
12290 j += 256;
12291 }
12292 while j + 64 <= n {
12293 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12294 j += 64;
12295 }
12296 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
12297 _mm512_add_epi32(a0, a1),
12298 _mm512_add_epi32(a2, a3),
12299 ));
12300 while j < n {
12302 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
12303 j += 1;
12304 }
12305 total
12306 }
12307}
12308
12309#[cfg(target_arch = "x86_64")]
12315#[target_feature(enable = "avx2")]
12316unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12317 unsafe {
12320 use core::arch::x86_64::*;
12321 let lomask = _mm_set1_epi8(0x0F);
12322 let eight = _mm256_set1_epi8(8);
12323 let ones = _mm256_set1_epi16(1);
12324 let mut acc = 0f32;
12325 for gi in 0..gpr {
12326 let g = g0 + gi;
12327 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12328 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
12329 let lo = _mm_and_si128(b, lomask);
12330 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
12331 let w = _mm256_sub_epi8(
12332 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
12333 eight,
12334 );
12335 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
12336 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
12337 let d = _mm256_madd_epi16(p16, ones);
12338 let hi128 = _mm256_extracti128_si256::<1>(d);
12339 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12340 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12341 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12342 acc += _mm_cvtsi128_si32(s32) as f32 * s;
12343 }
12344 acc
12345 }
12346}
12347
12348#[cfg(target_arch = "x86_64")]
12351#[target_feature(enable = "avx2")]
12352unsafe fn dot_q4_row_avx2_2(
12353 packed: &[u8],
12354 scales: &[u8],
12355 g0: usize,
12356 gpr: usize,
12357 xq1: &[i8],
12358 xq2: &[i8],
12359) -> (f32, f32) {
12360 unsafe {
12362 use core::arch::x86_64::*;
12363 let lomask = _mm_set1_epi8(0x0F);
12364 let eight = _mm256_set1_epi8(8);
12365 let ones = _mm256_set1_epi16(1);
12366 let (mut acc1, mut acc2) = (0f32, 0f32);
12367 #[inline(always)]
12368 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
12369 unsafe {
12370 use core::arch::x86_64::*;
12371 let hi128 = _mm256_extracti128_si256::<1>(d);
12372 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12373 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12374 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12375 _mm_cvtsi128_si32(s32)
12376 }
12377 }
12378 for gi in 0..gpr {
12379 let g = g0 + gi;
12380 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12381 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
12382 let lo = _mm_and_si128(b, lomask);
12383 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
12384 let w = _mm256_sub_epi8(
12385 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
12386 eight,
12387 );
12388 let aw = _mm256_abs_epi8(w);
12389 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
12390 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
12391 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
12392 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
12393 acc1 += hsum(d1) as f32 * s;
12394 acc2 += hsum(d2) as f32 * s;
12395 }
12396 (acc1, acc2)
12397 }
12398}
12399
12400#[cfg(target_arch = "x86_64")]
12402fn q8_range_avx2(
12403 q: &[u8],
12404 row_scale: &[f32],
12405 act: &SplitAct,
12406 cols: usize,
12407 out_addr: SendMut,
12408 start: usize,
12409 end: usize,
12410) {
12411 for o in start..end {
12412 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12413 unsafe { *out_addr.at(o) = v };
12415 }
12416}
12417
12418#[cfg(target_arch = "x86_64")]
12420#[allow(clippy::too_many_arguments)]
12421fn q8_range2_avx2(
12422 q: &[u8],
12423 row_scale: &[f32],
12424 a1: &SplitAct,
12425 a2: &SplitAct,
12426 cols: usize,
12427 p1: SendMut,
12428 p2: SendMut,
12429 start: usize,
12430 end: usize,
12431) {
12432 for o in start..end {
12433 let row = &q[o * cols..(o + 1) * cols];
12434 unsafe {
12436 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
12437 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
12438 }
12439 }
12440}
12441
12442#[cfg(target_arch = "aarch64")]
12453fn i8mm_enabled() -> bool {
12454 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12455 *ON.get_or_init(|| {
12456 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
12457 && std::arch::is_aarch64_feature_detected!("i8mm")
12458 })
12459}
12460
12461#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
12465fn sdot_enabled() -> bool {
12466 if FLOAT_ACTIVATIONS.get() {
12467 return false;
12468 }
12469 use std::sync::OnceLock;
12470 static ON: OnceLock<bool> = OnceLock::new();
12471 *ON.get_or_init(|| {
12472 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
12473 if !want {
12474 return false;
12475 }
12476
12477 #[cfg(target_arch = "aarch64")]
12478 {
12479 if std::arch::is_aarch64_feature_detected!("dotprod") {
12480 return true;
12481 }
12482 #[cfg(target_os = "android")]
12483 {
12484 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
12485 if cpuinfo.lines().any(|l| {
12486 (l.starts_with("Features") || l.starts_with("features"))
12487 && l.contains("asimddp")
12488 }) {
12489 return true;
12490 }
12491 }
12492 }
12493 false
12494 }
12495 #[cfg(not(target_arch = "aarch64"))]
12496 {
12497 false
12498 }
12499 })
12500}
12501
12502struct SplitAct {
12507 xq: Vec<i8>,
12508 sx: f32,
12509 outliers: Vec<(usize, f32)>,
12510 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
12513 xsum: i32,
12514 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
12524 xg8: Vec<f32>,
12525}
12526
12527thread_local! {
12528 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
12531 const { std::cell::RefCell::new(Vec::new()) };
12532 static XG8_FREE: std::cell::RefCell<Vec<Vec<f32>>> =
12534 const { std::cell::RefCell::new(Vec::new()) };
12535}
12536
12537fn xg8_take() -> Vec<f32> {
12539 let mut v = XG8_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
12540 v.clear();
12541 v
12542}
12543
12544impl Drop for SplitAct {
12545 fn drop(&mut self) {
12546 let buf = std::mem::take(&mut self.xq);
12547 if buf.capacity() > 0 {
12548 XQ_FREE.with(|f| {
12549 let mut f = f.borrow_mut();
12550 if f.len() < 16 {
12551 f.push(buf);
12552 }
12553 });
12554 }
12555 let buf = std::mem::take(&mut self.xg8);
12556 if buf.capacity() > 0 {
12557 XG8_FREE.with(|f| {
12558 let mut f = f.borrow_mut();
12559 if f.len() < 16 {
12560 f.push(buf);
12561 }
12562 });
12563 }
12564 }
12565}
12566
12567fn xg8_fill(xq: &[i8], out: &mut Vec<f32>) {
12571 out.clear();
12572 out.extend(xq.chunks_exact(GROUP_SIZE).map(|c| {
12573 let mut s = 0i32;
12574 for &v in c {
12575 s += v as i32;
12576 }
12577 s as f32
12578 }));
12579}
12580
12581thread_local! {
12582 static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
12589 static KROWS: UnsafeCell<[Vec<f32>; 2]> =
12594 const { UnsafeCell::new([Vec::new(), Vec::new()]) };
12595}
12596
12597#[inline]
12600fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
12601 KROW.with(|s| {
12602 let b = unsafe { &mut *s.get() };
12605 if b.len() < n {
12606 b.resize(n, 0.0);
12607 }
12608 f(&mut b[..n])
12609 })
12610}
12611
12612#[inline(always)]
12623fn q8_round(t: f32) -> i8 {
12624 let t = t.clamp(-127.0, 127.0);
12625 let i = t as i32;
12626 let f = t - i as f32;
12627 let r = if f >= 0.5 {
12628 i + 1
12629 } else if f <= -0.5 {
12630 i - 1
12631 } else {
12632 i
12633 };
12634 r as i8
12635}
12636
12637#[inline]
12638fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
12639 KROWS.with(|s| {
12640 let b = unsafe { &mut *s.get() };
12643 for row in b.iter_mut() {
12644 if row.len() < n {
12645 row.resize(n, 0.0);
12646 }
12647 }
12648 let (a, b) = b.split_at_mut(1);
12649 f(&mut a[0][..n], &mut b[0][..n])
12650 })
12651}
12652
12653#[inline]
12654fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
12655 if limit > 0.0 {
12656 up = up.clamp(-limit, limit);
12657 gate = gate.min(limit);
12658 }
12659 gate / (1.0 + (-gate).exp()) * up
12660}
12661
12662fn split_act(x: &[f32]) -> SplitAct {
12663 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
12664 let n = x.len();
12665 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
12666 let thr = 8.0 * rms;
12667 let mut outliers: Vec<(usize, f32)> = Vec::new();
12671 let mut amax = 0f32;
12672 for (j, &v) in x.iter().enumerate() {
12673 let a = v.abs();
12674 if a > thr {
12675 outliers.push((j, v));
12676 } else if a > amax {
12677 amax = a;
12678 }
12679 }
12680 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
12681 let inv = 1.0 / sx;
12682 crate::cpuprof::count_outliers(outliers.len() as u64);
12683 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
12684 xq.clear();
12685 xq.reserve(n);
12686 if outliers.is_empty() {
12687 xq.extend(
12688 x.iter()
12689 .map(|&v| q8_round(v * inv)),
12690 );
12691 } else {
12692 xq.extend(x.iter().map(|&v| {
12694 if v.abs() > thr {
12695 0
12696 } else {
12697 q8_round(v * inv)
12698 }
12699 }));
12700 }
12701 let xsum = xq.iter().map(|&v| v as i32).sum();
12702 let mut xg8 = xg8_take();
12703 xg8_fill(&xq, &mut xg8);
12704 SplitAct {
12705 xq,
12706 sx,
12707 outliers,
12708 xsum,
12709 xg8,
12710 }
12711}
12712
12713fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
12714 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
12715 let n = x.len();
12716 let rms = (x
12717 .iter()
12718 .zip(col)
12719 .map(|(&a, &c)| {
12720 let v = a * c;
12721 (v * v) as f64
12722 })
12723 .sum::<f64>()
12724 / n.max(1) as f64)
12725 .sqrt() as f32;
12726 let thr = 8.0 * rms;
12727
12728 let mut outliers = Vec::new();
12729 let mut amax = 0f32;
12730 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
12731 let v = a * c;
12732 let s = v.abs();
12733 if s > thr {
12734 outliers.push((j, v));
12735 } else if s > amax {
12736 amax = s;
12737 }
12738 }
12739
12740 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
12741 let inv = 1.0 / sx;
12742 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
12743 xq.clear();
12744 xq.reserve(n);
12745 if outliers.is_empty() {
12746 xq.extend(
12747 x.iter()
12748 .zip(col)
12749 .map(|(&a, &c)| q8_round((a * c) * inv)),
12750 );
12751 } else {
12752 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
12753 let v = a * c;
12754 if v.abs() > thr {
12755 0
12756 } else {
12757 q8_round(v * inv)
12758 }
12759 }));
12760 }
12761 let xsum = xq.iter().map(|&v| v as i32).sum();
12762 let mut xg8 = xg8_take();
12763 xg8_fill(&xq, &mut xg8);
12764 SplitAct {
12765 xq,
12766 sx,
12767 outliers,
12768 xsum,
12769 xg8,
12770 }
12771}
12772
12773#[cfg(target_arch = "aarch64")]
12776#[target_feature(enable = "neon,dotprod")]
12777unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
12778 unsafe {
12780 use core::arch::aarch64::*;
12781 use core::arch::asm;
12782 let wp = w.as_ptr() as *const i8;
12783 let n = w.len();
12784 let (mut a0, mut a1, mut a2, mut a3) = (
12785 vdupq_n_s32(0),
12786 vdupq_n_s32(0),
12787 vdupq_n_s32(0),
12788 vdupq_n_s32(0),
12789 );
12790 let mut i = 0;
12791 while i + 64 <= n {
12792 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
12793 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
12794 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
12795 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
12796 asm!(
12797 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
12798 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
12799 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
12800 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
12801 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
12802 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
12803 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
12804 options(pure, nomem, nostack),
12805 );
12806 i += 64;
12807 }
12808 while i + 16 <= n {
12809 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
12810 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
12811 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
12812 i += 16;
12813 }
12814 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
12815 while i < n {
12816 s += (*wp.add(i)) as i32 * xq[i] as i32;
12817 i += 1;
12818 }
12819 s
12820 }
12821}
12822
12823#[cfg(target_arch = "aarch64")]
12827#[target_feature(enable = "neon,dotprod")]
12828unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
12829 unsafe {
12831 use core::arch::aarch64::*;
12832 use core::arch::asm;
12833 let n = xq.len();
12834 let px = xq.as_ptr();
12835 let (p0, p1, p2, p3) = (
12836 w0.as_ptr() as *const i8,
12837 w1.as_ptr() as *const i8,
12838 w2.as_ptr() as *const i8,
12839 w3.as_ptr() as *const i8,
12840 );
12841 let (mut a0, mut a1, mut a2, mut a3) = (
12842 vdupq_n_s32(0),
12843 vdupq_n_s32(0),
12844 vdupq_n_s32(0),
12845 vdupq_n_s32(0),
12846 );
12847 let mut i = 0;
12848 while i + 16 <= n {
12849 let x = vld1q_s8(px.add(i));
12850 let v0 = vld1q_s8(p0.add(i));
12851 let v1 = vld1q_s8(p1.add(i));
12852 let v2 = vld1q_s8(p2.add(i));
12853 let v3 = vld1q_s8(p3.add(i));
12854 asm!(
12855 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
12856 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
12857 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
12858 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
12859 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
12860 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
12861 options(pure, nomem, nostack),
12862 );
12863 i += 16;
12864 }
12865 let mut r = [
12866 vaddvq_s32(a0),
12867 vaddvq_s32(a1),
12868 vaddvq_s32(a2),
12869 vaddvq_s32(a3),
12870 ];
12871 while i < n {
12872 let xi = *px.add(i) as i32;
12873 r[0] += (*p0.add(i)) as i32 * xi;
12874 r[1] += (*p1.add(i)) as i32 * xi;
12875 r[2] += (*p2.add(i)) as i32 * xi;
12876 r[3] += (*p3.add(i)) as i32 * xi;
12877 i += 1;
12878 }
12879 r
12880 }
12881}
12882
12883#[cfg(target_arch = "aarch64")]
12890#[target_feature(enable = "neon,dotprod")]
12891unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
12892 unsafe {
12895 use core::arch::aarch64::*;
12896 use core::arch::asm;
12897 let n = xq.len();
12898 let px = xq.as_ptr();
12899 let pg = g.as_ptr() as *const i8;
12900 let (mut a0, mut a1, mut a2, mut a3) = (
12901 vdupq_n_s32(0),
12902 vdupq_n_s32(0),
12903 vdupq_n_s32(0),
12904 vdupq_n_s32(0),
12905 );
12906 let mut i = 0;
12907 while i + 16 <= n {
12908 let x = vld1q_s8(px.add(i));
12909 let base = pg.add(4 * i);
12910 let v0 = vld1q_s8(base);
12911 let v1 = vld1q_s8(base.add(16));
12912 let v2 = vld1q_s8(base.add(32));
12913 let v3 = vld1q_s8(base.add(48));
12914 asm!(
12915 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
12916 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
12917 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
12918 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
12919 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
12920 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
12921 options(pure, nomem, nostack),
12922 );
12923 i += 16;
12924 }
12925 [
12926 vaddvq_s32(a0),
12927 vaddvq_s32(a1),
12928 vaddvq_s32(a2),
12929 vaddvq_s32(a3),
12930 ]
12931 }
12932}
12933
12934#[cfg(target_arch = "aarch64")]
12940fn q8_range_sdot(
12941 q: &[u8],
12942 rep: &[u8],
12943 row_scale: &[f32],
12944 act: &SplitAct,
12945 cols: usize,
12946 out_addr: SendMut,
12947 start: usize,
12948 end: usize,
12949) {
12950 let mut o = start;
12951 if !rep.is_empty() {
12954 while o < end && o % 4 != 0 {
12955 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12956 unsafe { *out_addr.at(o) = v };
12957 o += 1;
12958 }
12959 }
12960 while o + 4 <= end {
12961 let r = if rep.is_empty() {
12962 unsafe {
12963 dot_i8_sdot_4rows(
12964 &q[o * cols..(o + 1) * cols],
12965 &q[(o + 1) * cols..(o + 2) * cols],
12966 &q[(o + 2) * cols..(o + 3) * cols],
12967 &q[(o + 3) * cols..(o + 4) * cols],
12968 &act.xq,
12969 )
12970 }
12971 } else {
12972 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
12973 };
12974 for k in 0..4 {
12975 let mut acc = r[k] as f32 * act.sx;
12976 for &(j, xv) in &act.outliers {
12977 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
12978 }
12979 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
12981 }
12982 o += 4;
12983 }
12984 while o < end {
12985 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12986 unsafe { *out_addr.at(o) = v };
12987 o += 1;
12988 }
12989}
12990
12991#[cfg(target_arch = "aarch64")]
12994#[allow(clippy::too_many_arguments)]
12995fn q8_range2_sdot(
12996 q: &[u8],
12997 row_scale: &[f32],
12998 a1: &SplitAct,
12999 a2: &SplitAct,
13000 cols: usize,
13001 p1: SendMut,
13002 p2: SendMut,
13003 start: usize,
13004 end: usize,
13005) {
13006 for o in start..end {
13007 let row = &q[o * cols..(o + 1) * cols];
13008 unsafe {
13010 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
13011 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
13012 }
13013 }
13014}
13015
13016#[allow(clippy::too_many_arguments)]
13018fn q8_range2_f32(
13019 q: &[u8],
13020 row_scale: &[f32],
13021 x1: &[f32],
13022 x2: &[f32],
13023 cols: usize,
13024 p1: SendMut,
13025 p2: SendMut,
13026 start: usize,
13027 end: usize,
13028) {
13029 for o in start..end {
13030 let row = &q[o * cols..(o + 1) * cols];
13031 unsafe {
13033 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
13034 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
13035 }
13036 }
13037}
13038
13039fn q8_range_f32(
13041 q: &[u8],
13042 row_scale: &[f32],
13043 xs: &[f32],
13044 cols: usize,
13045 out_addr: SendMut,
13046 start: usize,
13047 end: usize,
13048) {
13049 for o in start..end {
13050 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
13051 unsafe { *out_addr.at(o) = v };
13053 }
13054}
13055
13056#[inline]
13060fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
13061 #[cfg(target_arch = "aarch64")]
13062 return row_dot_sdot(row, act);
13063 #[cfg(target_arch = "x86_64")]
13064 return row_dot_avx2(row, act);
13065 #[allow(unreachable_code)]
13066 q8_row_dot_scalar(row, act)
13067}
13068
13069#[allow(dead_code)]
13070fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
13071 let mut acc = 0i32;
13072 for (k, &b) in row.iter().enumerate() {
13073 acc += (b as i8) as i32 * act.xq[k] as i32;
13074 }
13075 let mut acc = acc as f32 * act.sx;
13076 for &(j, xv) in &act.outliers {
13077 acc += (row[j] as i8) as f32 * xv;
13078 }
13079 acc
13080}
13081
13082#[cfg(target_arch = "aarch64")]
13085#[inline]
13086fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
13087 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
13088 for &(j, xv) in &act.outliers {
13089 acc += (row[j] as i8) as f32 * xv;
13090 }
13091 acc
13092}
13093
13094#[cfg(target_arch = "aarch64")]
13102#[target_feature(enable = "neon,dotprod")]
13103unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
13104 unsafe {
13107 use core::arch::aarch64::*;
13108 use core::arch::asm;
13109 let lomask = vdupq_n_u8(0x0F);
13110 let eight = vdupq_n_s8(8);
13111 let mut acc = 0f32;
13112 for gi in 0..gpr {
13113 let g = g0 + gi;
13114 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13115 let b = vld1q_u8(packed.as_ptr().add(g * 16));
13116 let lo = vandq_u8(b, lomask);
13117 let hi = vshrq_n_u8::<4>(b);
13118 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13119 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13120 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
13121 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
13122 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
13123 asm!(
13124 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
13125 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
13126 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13127 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
13128 options(pure, nomem, nostack),
13129 );
13130 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13131 }
13132 acc
13133 }
13134}
13135
13136#[cfg(target_arch = "aarch64")]
13141#[target_feature(enable = "neon,dotprod")]
13142unsafe fn dot_q4_row_sdot2(
13143 packed: &[u8],
13144 scales: &[u8],
13145 g0: usize,
13146 gpr: usize,
13147 xq1: &[i8],
13148 xq2: &[i8],
13149) -> (f32, f32) {
13150 unsafe {
13153 use core::arch::aarch64::*;
13154 use core::arch::asm;
13155 let lomask = vdupq_n_u8(0x0F);
13156 let eight = vdupq_n_s8(8);
13157 let (mut acc1, mut acc2) = (0f32, 0f32);
13158 for gi in 0..gpr {
13159 let g = g0 + gi;
13160 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13161 let b = vld1q_u8(packed.as_ptr().add(g * 16));
13162 let lo = vandq_u8(b, lomask);
13163 let hi = vshrq_n_u8::<4>(b);
13164 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13165 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13166 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
13167 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
13168 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
13169 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
13170 let (mut a0, mut a1, mut b0, mut b1) = (
13171 vdupq_n_s32(0),
13172 vdupq_n_s32(0),
13173 vdupq_n_s32(0),
13174 vdupq_n_s32(0),
13175 );
13176 asm!(
13177 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
13178 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
13179 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
13180 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
13181 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13182 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
13183 e0 = in(vreg) e0, e1 = in(vreg) e1,
13184 x10 = in(vreg) x10, x11 = in(vreg) x11,
13185 x20 = in(vreg) x20, x21 = in(vreg) x21,
13186 options(pure, nomem, nostack),
13187 );
13188 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13189 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
13190 }
13191 (acc1, acc2)
13192 }
13193}
13194
13195#[inline]
13200pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
13201 #[cfg(target_arch = "aarch64")]
13202 unsafe {
13203 return axpy_i8_f32_neon(acc, row, w);
13204 }
13205 #[cfg(target_arch = "x86_64")]
13206 if avx2_enabled() {
13207 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
13208 }
13209 #[allow(unreachable_code)]
13210 {
13211 for (a, &b) in acc.iter_mut().zip(row) {
13212 *a += w * b as f32;
13213 }
13214 }
13215}
13216
13217#[cfg(target_arch = "x86_64")]
13219#[target_feature(enable = "avx2,fma")]
13220unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
13221 unsafe {
13223 use core::arch::x86_64::*;
13224 let n = acc.len().min(row.len());
13225 let ap = acc.as_mut_ptr();
13226 let rp = row.as_ptr();
13227 let wv = _mm256_set1_ps(w);
13228 let mut j = 0usize;
13229 while j + 16 <= n {
13230 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
13231 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
13232 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
13233 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
13234 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
13235 _mm256_storeu_ps(ap.add(j), v0);
13236 _mm256_storeu_ps(ap.add(j + 8), v1);
13237 j += 16;
13238 }
13239 while j < n {
13240 *ap.add(j) += w * (*rp.add(j)) as f32;
13241 j += 1;
13242 }
13243 }
13244}
13245
13246#[cfg(target_arch = "aarch64")]
13247#[target_feature(enable = "neon")]
13248unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
13249 unsafe {
13251 use core::arch::aarch64::*;
13252 let n = acc.len().min(row.len());
13253 let ap = acc.as_mut_ptr();
13254 let rp = row.as_ptr();
13255 let wv = vdupq_n_f32(w);
13256 let mut j = 0usize;
13257 while j + 16 <= n {
13258 let rb = vld1q_s8(rp.add(j));
13259 let lo = vmovl_s8(vget_low_s8(rb));
13260 let hi = vmovl_s8(vget_high_s8(rb));
13261 for (off, half) in [(0, lo), (8, hi)] {
13262 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
13263 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
13264 let o = j + off;
13265 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
13266 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
13267 }
13268 j += 16;
13269 }
13270 while j < n {
13271 *ap.add(j) += w * (*rp.add(j)) as f32;
13272 j += 1;
13273 }
13274 }
13275}
13276
13277#[inline]
13280pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
13281 #[cfg(target_arch = "aarch64")]
13282 unsafe {
13283 return dot_i8_f32_neon(w, x);
13284 }
13285 #[cfg(target_arch = "x86_64")]
13286 if avx2_enabled() {
13287 return unsafe { dot_i8_f32_avx2(w, x) };
13288 }
13289 #[allow(unreachable_code)]
13290 {
13291 let mut sum = 0.0f32;
13292 for (j, &b) in w.iter().enumerate() {
13293 sum += (b as i8) as f32 * x[j];
13294 }
13295 sum
13296 }
13297}
13298
13299#[inline]
13303fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
13304 #[cfg(target_arch = "aarch64")]
13305 unsafe {
13306 return dot_i8_col_f32_neon(w, x, col);
13307 }
13308 #[allow(unreachable_code)]
13309 {
13310 let mut sum = 0.0f32;
13311 for (j, &b) in w.iter().enumerate() {
13312 sum += (b as i8) as f32 * x[j] * col[j];
13313 }
13314 sum
13315 }
13316}
13317
13318#[cfg(target_arch = "aarch64")]
13319#[target_feature(enable = "neon")]
13320unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
13321 unsafe {
13323 use core::arch::aarch64::*;
13324 let n = x.len();
13325 let wp = w.as_ptr() as *const i8;
13326 let xp = x.as_ptr();
13327 let cp = col.as_ptr();
13328 let (mut a0, mut a1, mut a2, mut a3) = (
13329 vdupq_n_f32(0.0),
13330 vdupq_n_f32(0.0),
13331 vdupq_n_f32(0.0),
13332 vdupq_n_f32(0.0),
13333 );
13334 let mut j = 0usize;
13335 while j + 16 <= n {
13336 let wb = vld1q_s8(wp.add(j));
13337 let lo = vmovl_s8(vget_low_s8(wb));
13338 let hi = vmovl_s8(vget_high_s8(wb));
13339 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
13340 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
13341 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
13342 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
13343 a0 = vfmaq_f32(
13344 a0,
13345 w0,
13346 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
13347 );
13348 a1 = vfmaq_f32(
13349 a1,
13350 w1,
13351 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
13352 );
13353 a2 = vfmaq_f32(
13354 a2,
13355 w2,
13356 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
13357 );
13358 a3 = vfmaq_f32(
13359 a3,
13360 w3,
13361 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
13362 );
13363 j += 16;
13364 }
13365 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
13366 while j < n {
13367 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
13368 j += 1;
13369 }
13370 sum
13371 }
13372}
13373
13374#[cfg(target_arch = "aarch64")]
13375#[target_feature(enable = "neon")]
13376unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
13377 unsafe {
13379 use core::arch::aarch64::*;
13380 let n = x.len();
13381 let wp = w.as_ptr() as *const i8;
13382 let xp = x.as_ptr();
13383 let (mut a0, mut a1, mut a2, mut a3) = (
13384 vdupq_n_f32(0.0),
13385 vdupq_n_f32(0.0),
13386 vdupq_n_f32(0.0),
13387 vdupq_n_f32(0.0),
13388 );
13389 let mut j = 0usize;
13390 while j + 16 <= n {
13391 let wb = vld1q_s8(wp.add(j));
13392 let lo = vmovl_s8(vget_low_s8(wb));
13393 let hi = vmovl_s8(vget_high_s8(wb));
13394 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
13395 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
13396 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
13397 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
13398 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
13399 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
13400 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
13401 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
13402 j += 16;
13403 }
13404 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
13405 while j < n {
13406 sum += (*wp.add(j)) as f32 * *xp.add(j);
13407 j += 1;
13408 }
13409 sum
13410 }
13411}
13412
13413#[allow(clippy::too_many_arguments)]
13414fn qmatvec(
13415 q: &[u8],
13416 rep: &[u8],
13417 row_scale: &[f32],
13418 x: &[f32],
13419 col_field: &[f32],
13420 dtype: TensorDtype,
13421 rows: usize,
13422 cols: usize,
13423 out: &mut [f32],
13424 pool: Option<&Pool>,
13425) {
13426 debug_assert_eq!(out.len(), rows);
13427 #[cfg(not(target_arch = "aarch64"))]
13428 let _ = rep;
13429
13430 #[cfg(target_arch = "aarch64")]
13431 if sdot_enabled() {
13432 let act = if dtype == TensorDtype::Q8_2f {
13433 split_act_q8_2f(x, col_field)
13434 } else {
13435 split_act(x)
13436 };
13437 let out_addr = SendMut(out.as_mut_ptr());
13438 let run_range = |start: usize, end: usize| {
13439 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
13440 };
13441 match pool {
13442 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13443 _ => run_range(0, rows),
13444 }
13445 return;
13446 }
13447 #[cfg(target_arch = "x86_64")]
13450 if avx2_a8w8_enabled() {
13451 let act = if dtype == TensorDtype::Q8_2f {
13452 split_act_q8_2f(x, col_field)
13453 } else {
13454 split_act(x)
13455 };
13456 let out_addr = SendMut(out.as_mut_ptr());
13457 let run_range = |start: usize, end: usize| {
13458 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
13459 };
13460 match pool {
13461 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13462 _ => run_range(0, rows),
13463 }
13464 return;
13465 }
13466
13467 prescale_with(x, col_field, dtype, 1, |xs| {
13468 let out_addr = SendMut(out.as_mut_ptr());
13469 let run_range = move |start: usize, end: usize| {
13470 for o in start..end {
13471 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
13472 unsafe { *out_addr.at(o) = v };
13474 }
13475 };
13476 match pool {
13477 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13478 _ => run_range(0, rows),
13479 }
13480 });
13481}
13482
13483#[allow(clippy::too_many_arguments)]
13484fn qmatvec2(
13485 q: &[u8],
13486 row_scale: &[f32],
13487 x1: &[f32],
13488 x2: &[f32],
13489 col_field: &[f32],
13490 dtype: TensorDtype,
13491 rows: usize,
13492 cols: usize,
13493 o1: &mut [f32],
13494 o2: &mut [f32],
13495 pool: Option<&Pool>,
13496) {
13497 #[cfg(target_arch = "aarch64")]
13498 if sdot_enabled() {
13499 let a1s = if dtype == TensorDtype::Q8_2f {
13500 split_act_q8_2f(x1, col_field)
13501 } else {
13502 split_act(x1)
13503 };
13504 let a2s = if dtype == TensorDtype::Q8_2f {
13505 split_act_q8_2f(x2, col_field)
13506 } else {
13507 split_act(x2)
13508 };
13509 let p1 = SendMut(o1.as_mut_ptr());
13510 let p2 = SendMut(o2.as_mut_ptr());
13511 let run_range = |start: usize, end: usize| {
13512 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
13513 };
13514 match pool {
13515 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13516 _ => run_range(0, rows),
13517 }
13518 return;
13519 }
13520 #[cfg(target_arch = "x86_64")]
13521 if avx2_a8w8_enabled() {
13522 let a1s = if dtype == TensorDtype::Q8_2f {
13523 split_act_q8_2f(x1, col_field)
13524 } else {
13525 split_act(x1)
13526 };
13527 let a2s = if dtype == TensorDtype::Q8_2f {
13528 split_act_q8_2f(x2, col_field)
13529 } else {
13530 split_act(x2)
13531 };
13532 let p1 = SendMut(o1.as_mut_ptr());
13533 let p2 = SendMut(o2.as_mut_ptr());
13534 let run_range = |start: usize, end: usize| {
13535 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
13536 };
13537 match pool {
13538 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13539 _ => run_range(0, rows),
13540 }
13541 return;
13542 }
13543
13544 prescale_with(x1, col_field, dtype, 1, |x1s| {
13545 prescale_with(x2, col_field, dtype, 2, |x2s| {
13546 let p1 = SendMut(o1.as_mut_ptr());
13547 let p2 = SendMut(o2.as_mut_ptr());
13548 let run_range = move |start: usize, end: usize| {
13549 for o in start..end {
13550 let row = &q[o * cols..(o + 1) * cols];
13551 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
13552 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
13553 unsafe {
13555 *p1.at(o) = s1;
13556 *p2.at(o) = s2;
13557 }
13558 }
13559 };
13560 match pool {
13561 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13562 _ => run_range(0, rows),
13563 }
13564 });
13565 });
13566}
13567
13568#[derive(Clone, Copy)]
13569struct SendMut(*mut f32);
13570unsafe impl Send for SendMut {}
13571unsafe impl Sync for SendMut {}
13572
13573impl SendMut {
13574 #[inline]
13575 fn at(self, i: usize) -> *mut f32 {
13576 unsafe { self.0.add(i) }
13577 }
13578}
13579
13580#[cfg(test)]
13581mod tests {
13582
13583 #[test]
13587 fn q8_round_is_round_clamp() {
13588 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
13589 let mut probe = vec![
13590 0.0f32,
13591 -0.0,
13592 f32::NAN,
13593 f32::INFINITY,
13594 f32::NEG_INFINITY,
13595 f32::MAX,
13596 f32::MIN,
13597 1e30,
13598 -1e30,
13599 f32::MIN_POSITIVE,
13600 -f32::MIN_POSITIVE,
13601 ];
13602 for k in -300i32..=300 {
13603 let h = k as f32 * 0.5;
13604 let up = f32::from_bits(h.to_bits() + 1);
13605 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
13606 for t in [h, up, down] {
13607 probe.push(t);
13608 probe.push(-t);
13609 }
13610 }
13611 let mut t = -140.0f32;
13612 while t < 140.0 {
13613 probe.push(t);
13614 t += 0.000_731;
13615 }
13616 for t in probe {
13617 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
13618 }
13619 }
13620
13621 use super::*;
13622
13623 #[test]
13627 fn f32_matmat_equals_matvec_bitwise() {
13628 let (rows, cols, b) = (5usize, 37usize, 7usize);
13629 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
13630 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
13631 let t = QTensor::from_f32(w, rows, cols);
13632 let mut out = vec![0f32; b * rows];
13633 t.matmat(&xs, b, &mut out, None);
13634 for bi in 0..b {
13635 let mut one = vec![0f32; rows];
13636 t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
13637 for o in 0..rows {
13638 assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
13639 }
13640 }
13641 }
13642
13643 #[test]
13644 fn q2tp_i8_dot_matches_exact_on_grid() {
13645 let (rows, cols) = (5, 64);
13649 let gpr = cols / GROUP_SIZE;
13650 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
13654 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
13655 .collect();
13656 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
13657 let x: Vec<f32> = (0..cols)
13658 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
13659 .collect();
13660 let act = split_act(&x);
13661 assert!(
13662 act.outliers.is_empty(),
13663 "on-grid input must have no outliers"
13664 );
13665 let gsum = q1_group_sums(&act.xq, gpr);
13666 for r in 0..rows {
13667 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
13668 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
13669 assert!(
13670 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
13671 "row {r}: exact {exact} vs i8 {fast}"
13672 );
13673 }
13674 }
13675
13676 #[cfg(target_arch = "x86_64")]
13677 #[test]
13678 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
13679 if !std::arch::is_x86_feature_detected!("avx2") {
13684 return;
13685 }
13686 let mut seed = 0x9e3779b9u32;
13687 let mut next = || {
13688 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
13689 seed
13690 };
13691 for _ in 0..20_000 {
13692 let mut ch = [0u8; Q2TP_CHUNK];
13693 let mut x = [0i8; GROUP_SIZE];
13694 for b in &mut ch {
13695 *b = next() as u8;
13696 }
13697 for v in &mut x {
13698 *v = (next() >> 24) as i8;
13699 }
13700 let mut reference = 0i32;
13701 for (k, &b) in ch.iter().enumerate() {
13702 reference += (b & 3) as i32 * x[k * 4] as i32;
13703 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
13704 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
13705 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
13706 }
13707 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
13710 assert_eq!(got, reference, "packed q2 lane mismatch");
13711 }
13712 }
13713
13714 #[test]
13715 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
13716 let (rows, cols) = (1usize, GROUP_SIZE);
13717 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
13718 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
13722 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
13723 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
13725 let mut raw = vec![0.0f32; rows];
13726 let mut affine = vec![0.0f32; rows];
13727 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
13728 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
13729 assert_eq!(raw, vec![-24.0]);
13730 assert_eq!(affine, vec![-8.0]);
13731 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
13732 }
13733
13734 #[test]
13735 fn q8_row_dot_fast_matches_scalar() {
13736 let cols = 96;
13739 let row: Vec<u8> = (0..cols)
13740 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
13741 .collect();
13742 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
13743 let act = split_act(&x);
13744 let fast = q8_row_dot(&row, &act);
13745 let scalar = q8_row_dot_scalar(&row, &act);
13746 assert!(
13747 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
13748 "fast {fast} vs scalar {scalar}"
13749 );
13750 }
13751
13752 #[test]
13753 fn f32_matvec_matches_matvec_rows_bitexact() {
13754 let (rows, cols) = (300, 40);
13755 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
13756 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
13757 let qt = QTensor::from_f32(w.clone(), rows, cols);
13758
13759 let mut a = vec![0.0f32; rows];
13760 matvec_rows(None, &w, &x, &mut a);
13761 let mut b = vec![0.0f32; rows];
13762 qt.matvec(&x, &mut b, None);
13763 assert_eq!(a, b);
13764 }
13765
13766 #[test]
13767 fn sdot_kernel_exact_on_grid() {
13768 eprintln!("sdot_enabled = {}", sdot_enabled());
13773 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
13775 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13776 .collect();
13777 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
13778 let x: Vec<f32> = (0..cols)
13779 .map(|i| match i % 3 {
13780 0 => 1.0,
13781 1 => -1.0,
13782 _ => 0.0,
13783 })
13784 .collect();
13785 let mut a = vec![0.0f32; rows];
13786 qmatvec(
13787 &w,
13788 &[],
13789 &scales,
13790 &x,
13791 &[],
13792 TensorDtype::Q8Row,
13793 rows,
13794 cols,
13795 &mut a,
13796 None,
13797 );
13798 for o in 0..rows {
13799 let mut acc = 0.0f32;
13800 for j in 0..cols {
13801 acc += (w[o * cols + j] as i8) as f32 * x[j];
13802 }
13803 let expect = acc * scales[o];
13804 assert!(
13805 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
13806 "row {o}: {} vs {expect}",
13807 a[o]
13808 );
13809 }
13810 }
13811
13812 #[test]
13813 fn q1_tbl_fast_path_matches_reference() {
13814 let (rows, cols) = (5, 256);
13819 let gpr = cols / GROUP_SIZE;
13820 let mut bytes = Vec::new();
13821 for t in 0..rows * gpr {
13822 let s = 0.007 + (t % 11) as f32 * 0.004;
13823 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13824 for j in 0..4 {
13825 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
13826 }
13827 }
13828 let x: Vec<f32> = (0..cols)
13829 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
13830 .collect();
13831 let mut w = vec![0.0f32; rows * cols];
13832 cortiq_core::quant::dequant_q1(&bytes, &mut w);
13833 let mut got = vec![0.0f32; rows];
13834 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
13835 for o in 0..rows {
13836 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
13837 assert!(
13838 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
13839 "row {o}: {} vs {expect}",
13840 got[o]
13841 );
13842 }
13843 let b = 5usize;
13846 let mut xs_all = Vec::new();
13847 for bi in 0..b {
13848 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
13849 }
13850 let mut mm = vec![0.0f32; b * rows];
13851 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
13852 for bi in 0..b {
13853 let mut single = vec![0.0f32; rows];
13854 q1_matvec(
13855 &bytes,
13856 &xs_all[bi * cols..(bi + 1) * cols],
13857 rows,
13858 cols,
13859 &mut single,
13860 None,
13861 );
13862 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
13863 }
13864 }
13865
13866 #[test]
13867 fn q1_kernels_match_exact_reference() {
13868 let (rows, cols) = (7, 96);
13870 let gpr = cols / GROUP_SIZE;
13871 let mut bytes = Vec::new();
13872 for t in 0..rows * gpr {
13873 let s = 0.01 + (t % 13) as f32 * 0.003;
13874 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13875 for j in 0..4 {
13876 bytes.push(((t * 31 + j * 97) % 251) as u8);
13877 }
13878 }
13879 let x: Vec<f32> = (0..cols)
13881 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
13882 .collect();
13883 let mut w = vec![0.0f32; rows * cols];
13885 cortiq_core::quant::dequant_q1(&bytes, &mut w);
13886 let mut expect = vec![0.0f32; rows];
13887 for o in 0..rows {
13888 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
13889 }
13890 let mut got = vec![0.0f32; rows];
13891 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
13892 for o in 0..rows {
13893 assert!(
13894 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
13895 "row {o}: {} vs {}",
13896 got[o],
13897 expect[o]
13898 );
13899 }
13900 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
13902 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
13903 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
13904 assert_eq!(a1, got);
13905 let mut xs = x.clone();
13906 xs.extend_from_slice(&x2);
13907 let mut mm = vec![0.0f32; 2 * rows];
13908 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
13909 assert_eq!(&mm[..rows], got.as_slice());
13910 assert_eq!(&mm[rows..], a2.as_slice());
13911 }
13912
13913 fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
13916 let tiles = rows * cols / GROUP_SIZE;
13917 let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
13918 for t in 0..tiles {
13919 let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
13920 q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13921 let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
13922 q1.extend_from_slice(&bits.to_le_bytes());
13923 }
13924 let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
13925 assert_eq!(rp.differing, 0, "uniform scales repack exactly");
13926 (q1, rp.data)
13927 }
13928
13929 fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
13932 let mut x: Vec<f32> = (0..cols)
13933 .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
13934 .collect();
13935 x[3 % cols] = 41.0;
13936 x[(cols / 2 + 5) % cols] = -37.5;
13937 x
13938 }
13939
13940 #[test]
13947 fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
13948 for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
13949 let gpr = cols / GROUP_SIZE;
13950 let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
13951 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
13952 let x = q1g_test_x(cols, rows);
13953 let act = split_act(&x);
13954 assert!(!act.outliers.is_empty(), "test wants the outlier path");
13955 let gsum = q1_group_sums(&act.xq, gpr);
13956 for r in 0..rows {
13957 let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
13958 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
13959 let (a, b) = (
13960 wa.dot_i8(r, gpr, &act.xq, &gsum),
13961 wb.dot_i8(r, gpr, &act.xq, &gsum),
13962 );
13963 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
13964 for j in [0, 31, 32, 127, cols - 1] {
13965 let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
13966 assert_eq!(
13967 (a.0.to_bits(), a.1.to_bits()),
13968 (b.0.to_bits(), b.1.to_bits())
13969 );
13970 }
13971 #[cfg(target_arch = "x86_64")]
13972 if avx2_enabled() {
13973 let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
13974 unsafe {
13976 let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
13977 let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
13978 assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
13979 }
13980 if std::is_x86_feature_detected!("avx512vnni")
13981 && std::is_x86_feature_detected!("avx512bw")
13982 && std::is_x86_feature_detected!("avx512vl")
13983 {
13984 unsafe {
13986 let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
13987 let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
13988 assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
13989 }
13990 }
13991 }
13992 }
13993 let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
13995 bin_matvec(wa, &x, rows, cols, &mut ga, None);
13996 bin_matvec(wb, &x, rows, cols, &mut gb, None);
13997 assert_eq!(ga, gb, "{rows}x{cols} matvec");
13998 let x2 = q1g_test_x(cols, rows + 1);
13999 let (mut a1, mut a2, mut b1, mut b2) = (
14000 vec![0f32; rows],
14001 vec![0f32; rows],
14002 vec![0f32; rows],
14003 vec![0f32; rows],
14004 );
14005 bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14006 bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
14007 assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
14008 let b = 5usize;
14009 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
14010 let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14011 bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
14012 bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
14013 assert_eq!(ma, mb, "{rows}x{cols} matmat");
14014 let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
14017 .map(|k| {
14018 let a = split_act(&xs[k * cols..(k + 1) * cols]);
14019 let g = q1_group_sums(&a.xq, gpr);
14020 (a, g)
14021 })
14022 .collect();
14023 let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
14024 let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
14025 for r in 0..rows {
14026 #[cfg(target_arch = "x86_64")]
14027 if avx2_enabled() {
14028 let (a, b) = unsafe {
14030 (
14031 wa.dot_1x4_x86(r, gpr, xq4, gs4),
14032 wb.dot_1x4_x86(r, gpr, xq4, gs4),
14033 )
14034 };
14035 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14036 }
14037 #[cfg(target_arch = "aarch64")]
14038 if sdot_enabled() {
14039 let (a, b) = unsafe {
14041 (
14042 wa.dot_1x4_sdot(r, gpr, xq4, gs4),
14043 wb.dot_1x4_sdot(r, gpr, xq4, gs4),
14044 )
14045 };
14046 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14047 }
14048 let _ = (r, &xq4, &gs4);
14049 }
14050 }
14051 }
14052
14053 #[test]
14059 #[ignore]
14060 fn q1g_cpu_matvec_speed() {
14061 use crate::pool::Pool;
14062 let threads = std::env::var("CMF_THREADS")
14063 .ok()
14064 .and_then(|v| v.parse().ok())
14065 .unwrap_or(8usize);
14066 let pool = Pool::new(threads);
14067 for &(name, rows, cols, reps) in &[
14068 ("gate/up_proj", 6144usize, 2048usize, 41usize),
14069 ("down_proj", 2048, 6144, 41),
14070 ("lm_head", 151_669, 2048, 9),
14071 ] {
14072 let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
14073 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14074 let x = q1g_test_x(cols, 1);
14075 let mut out = vec![0f32; rows];
14076 let mut time = |w: BinW| {
14077 let t0 = std::time::Instant::now();
14078 bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
14079 t0.elapsed().as_secs_f64()
14080 };
14081 let _ = (time(wa), time(wb));
14082 let (mut ta, mut tb) = (Vec::new(), Vec::new());
14083 for r in 0..reps {
14084 if r % 2 == 0 {
14085 ta.push(time(wa));
14086 tb.push(time(wb));
14087 } else {
14088 tb.push(time(wb));
14089 ta.push(time(wa));
14090 }
14091 }
14092 ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
14093 tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
14094 let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
14095 eprintln!(
14096 "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
14097 a8w8_enabled(),
14098 ma * 1e6,
14099 mb * 1e6,
14100 mb / ma
14101 );
14102 }
14103 }
14104
14105 #[test]
14110 fn q1g_mapped_tensor_paths_equal_q1() {
14111 use crate::pool::Pool;
14112 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14113 let (rows, cols) = (300usize, 256usize);
14114 let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
14115 let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
14116 let header: CmfHeader = serde_json::from_value(serde_json::json!({
14117 "format": "cmf", "version": cortiq_core::CMF_VERSION,
14118 "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
14119 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14120 "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
14121 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14122 "max_position_embeddings": 8},
14123 "quant_type": "VBIT"
14124 }))
14125 .unwrap();
14126 let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
14127 name: name.into(),
14128 dtype,
14129 shape: vec![rows, cols],
14130 data: data.clone(),
14131 };
14132 let specs = [
14133 spec("g1", TensorDtype::Q1, &gq1),
14134 spec("u1", TensorDtype::Q1, &uq1),
14135 spec("gg", TensorDtype::Q1G, &gq1g),
14136 spec("ug", TensorDtype::Q1G, &uq1g),
14137 ];
14138 let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
14139 std::fs::create_dir_all(&dir).unwrap();
14140 let path = dir.join("m.cmf");
14141 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14142 let model = Arc::new(CmfModel::open(&path).unwrap());
14143 assert!(model.verify().is_empty(), "{:?}", model.verify());
14144 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14145 let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
14146 assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
14147 assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
14148 let pool = Pool::new(3);
14149 let x = q1g_test_x(cols, 4);
14150 let run = |a: &QTensor| {
14151 let mut o = vec![0f32; rows];
14152 a.matvec(&x, &mut o, Some(&pool));
14153 o
14154 };
14155 assert_eq!(run(&g1), run(&gg), "matvec");
14156 let many = |a: &QTensor, b: &QTensor| {
14157 let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
14158 QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
14159 (oa, ob)
14160 };
14161 let want = many(&g1, &u1);
14162 assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
14163 assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
14164 assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
14165 let fused = |g: &QTensor, u: &QTensor| {
14166 let mut o = vec![0f32; rows];
14167 let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
14168 (ok, o)
14169 };
14170 let (ok_a, fa) = fused(&g1, &u1);
14171 let (ok_b, fb) = fused(&gg, &ug);
14172 assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
14173 assert_eq!(fa, fb, "fused gate/up");
14174 let b = 6usize;
14175 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
14176 let mm = |a: &QTensor| {
14177 let mut o = vec![0f32; b * rows];
14178 a.matmat(&xs, b, &mut o, Some(&pool));
14179 o
14180 };
14181 assert_eq!(mm(&g1), mm(&gg), "matmat");
14182 let x2 = q1g_test_x(cols, 77);
14183 let pair = |a: &QTensor| {
14184 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14185 a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14186 (o1, o2)
14187 };
14188 assert_eq!(pair(&g1), pair(&gg), "matvec2");
14189 for r in [0, 1, rows / 2, rows - 1] {
14190 let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
14191 g1.row_f32(r, &mut ra);
14192 gg.row_f32(r, &mut rb);
14193 assert_eq!(ra, rb, "embedding row {r}");
14194 }
14195 drop((g1, u1, gg, ug, model));
14196 let _ = std::fs::remove_dir_all(&dir);
14197 }
14198
14199 #[test]
14206 fn blr_mapped_tensor_paths_equal_dense() {
14207 use crate::pool::Pool;
14208 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14209 let (rows, cols, r) = (96usize, 128usize, 16usize);
14210 let mut rng = crate::blr::tests::Rng(41);
14211 let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
14212 let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
14213 let mut q8: Vec<u8> = (0..rows * cols)
14215 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14216 .collect();
14217 for o in 0..rows {
14218 q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
14219 }
14220 let header: CmfHeader = serde_json::from_value(serde_json::json!({
14221 "format": "cmf", "version": cortiq_core::CMF_VERSION,
14222 "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
14223 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14224 "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
14225 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14226 "max_position_embeddings": 8},
14227 "quant_type": "Q8_2F"
14228 }))
14229 .unwrap();
14230 let specs = [
14231 TensorSpec {
14232 name: "g".into(),
14233 dtype: TensorDtype::BinaryLowRank,
14234 shape: vec![rows, cols],
14235 data: gate.clone(),
14236 },
14237 TensorSpec {
14238 name: "u".into(),
14239 dtype: TensorDtype::BinaryLowRank,
14240 shape: vec![rows, cols],
14241 data: up.clone(),
14242 },
14243 TensorSpec {
14244 name: "q".into(),
14245 dtype: TensorDtype::Q8Row,
14246 shape: vec![rows, cols],
14247 data: q8,
14248 },
14249 ];
14250 let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
14251 std::fs::create_dir_all(&dir).unwrap();
14252 let path = dir.join("m.cmf");
14253 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14254 let model = Arc::new(CmfModel::open(&path).unwrap());
14255 assert!(model.verify().is_empty(), "{:?}", model.verify());
14256 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14257 let (g, u, q) = (t("g"), t("u"), t("q"));
14258 assert!(g.is_blr() && u.is_blr() && !q.is_blr());
14259 assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
14260 assert!(g.as_f32().is_none() && g.model_idx().is_none());
14261 assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
14262 assert!(!g.sparse_col_ok());
14263 assert_eq!((g.rows(), g.cols()), (rows, cols));
14264 let deq = |bytes: &[u8]| {
14265 let mut w = vec![0f32; rows * cols];
14266 cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
14267 w
14268 };
14269 let (wg, wu) = (deq(&gate), deq(&up));
14270 let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
14271 (0..rows)
14272 .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
14273 .collect()
14274 };
14275 let close = |a: &[f32], b: &[f32], what: &str| {
14276 let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
14277 let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
14278 assert!(e / s < 1e-5, "{what}: rel {}", e / s);
14279 };
14280 let pool = Pool::new(3);
14281 let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
14282 let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
14283 let mut o = vec![0f32; rows];
14284 g.matvec(&x, &mut o, Some(&pool));
14285 close(&o, &mv(&wg, &x), "matvec");
14286 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14287 g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14288 close(&o1, &mv(&wg, &x), "matvec2 a");
14289 close(&o2, &mv(&wg, &x2), "matvec2 b");
14290 for b in [2usize, 5] {
14291 let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
14292 let mut om = vec![0f32; b * rows];
14293 g.matmat(&xs, b, &mut om, Some(&pool));
14294 for t in 0..b {
14295 close(
14296 &om[t * rows..(t + 1) * rows],
14297 &mv(&wg, &xs[t * cols..(t + 1) * cols]),
14298 "matmat",
14299 );
14300 }
14301 }
14302 let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
14305 QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
14306 close(&a, &mv(&wg, &x), "matvec_many blr gate");
14307 close(&c, &mv(&wu, &x), "matvec_many blr up");
14308 let mut qo = vec![0f32; rows];
14309 q.matvec(&x, &mut qo, Some(&pool));
14310 assert_eq!(bq, qo, "matvec_many q8_row sibling");
14311 let mut f = vec![0f32; rows];
14312 assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
14313 let mut row = vec![0f32; cols];
14314 g.row_f32(rows / 3, &mut row);
14315 close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
14316 let mut scratch = vec![0f32; cols];
14317 let d = g.row_dot(7, &x, &mut scratch);
14318 let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
14319 assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
14320 let mut col = vec![0f32; rows];
14321 g.add_col_scaled(cols / 2, 1.5, &mut col);
14322 let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
14323 close(&col, &want, "add_col_scaled");
14324 drop((g, u, q, model));
14325 let _ = std::fs::remove_dir_all(&dir);
14326 }
14327
14328 #[test]
14329 fn repack_is_bit_identical() {
14330 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
14336 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14337 .collect();
14338 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
14339 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
14340 let rep = q8_repack_layout(&w, rows, cols);
14341 for g in 0..rows / 4 {
14343 for c in 0..cols / 16 {
14344 for lane in 0..4 {
14345 assert_eq!(
14346 &rep[g * 4 * cols + c * 64 + lane * 16
14347 ..g * 4 * cols + c * 64 + lane * 16 + 16],
14348 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
14349 );
14350 }
14351 }
14352 }
14353 let mut a = vec![0.0f32; rows];
14354 qmatvec(
14355 &w,
14356 &[],
14357 &scales,
14358 &x,
14359 &[],
14360 TensorDtype::Q8Row,
14361 rows,
14362 cols,
14363 &mut a,
14364 None,
14365 );
14366 let mut b = vec![0.0f32; rows];
14367 qmatvec(
14368 &w,
14369 &rep,
14370 &scales,
14371 &x,
14372 &[],
14373 TensorDtype::Q8Row,
14374 rows,
14375 cols,
14376 &mut b,
14377 None,
14378 );
14379 assert_eq!(a, b, "full-range repack output diverged");
14380
14381 #[cfg(target_arch = "aarch64")]
14382 if sdot_enabled() {
14383 let act = split_act(&x);
14385 let mut c1 = vec![0.0f32; rows];
14386 let mut c2 = vec![0.0f32; rows];
14387 q8_range_sdot(
14388 &w,
14389 &[],
14390 &scales,
14391 &act,
14392 cols,
14393 SendMut(c1.as_mut_ptr()),
14394 3,
14395 rows - 2,
14396 );
14397 q8_range_sdot(
14398 &w,
14399 &rep,
14400 &scales,
14401 &act,
14402 cols,
14403 SendMut(c2.as_mut_ptr()),
14404 3,
14405 rows - 2,
14406 );
14407 assert_eq!(c1, c2, "unaligned-range repack output diverged");
14408 }
14409 }
14410
14411 #[test]
14412 fn sdot_a8w8_noise_is_bounded() {
14413 let (rows, cols) = (16, 512);
14417 let w: Vec<u8> = (0..rows * cols)
14418 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
14419 .collect();
14420 let scales = vec![0.01f32; rows];
14421 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
14422 let mut a = vec![0.0f32; rows];
14423 qmatvec(
14424 &w,
14425 &[],
14426 &scales,
14427 &x,
14428 &[],
14429 TensorDtype::Q8Row,
14430 rows,
14431 cols,
14432 &mut a,
14433 None,
14434 );
14435 let (mut num, mut den) = (0f64, 0f64);
14436 for o in 0..rows {
14437 let mut acc = 0.0f32;
14438 for j in 0..cols {
14439 acc += (w[o * cols + j] as i8) as f32 * x[j];
14440 }
14441 let expect = acc * scales[o];
14442 num += ((a[o] - expect) as f64).powi(2);
14443 den += (expect as f64).powi(2);
14444 }
14445 let rel = (num / den.max(1e-12)).sqrt();
14446 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
14447 }
14448
14449 #[test]
14450 fn i8_dot_neon_matches_scalar() {
14451 let n = 100;
14452 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
14453 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
14454 let mut scalar = 0.0f32;
14455 for j in 0..n {
14456 scalar += (w[j] as i8) as f32 * x[j];
14457 }
14458 let fast = dot_i8_f32(&w, &x);
14459 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
14460 }
14461
14462 #[test]
14464 fn vbitmatvec_matches_full_dequant() {
14465 let (rows, cols) = (6, 64);
14466 let ng = cols / GROUP_SIZE;
14467 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
14469 let mut bytes = bits.clone();
14470 for g in 0..rows * ng {
14471 let s = 0.02 + 0.001 * g as f32;
14472 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14473 }
14474 for r in 0..rows {
14475 let b = bits[r] as usize;
14476 let (mut acc, mut nb) = (0u64, 0usize);
14477 let mut rowbytes = Vec::new();
14478 for i in 0..cols {
14479 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
14480 acc = (acc << b) | v;
14481 nb += b;
14482 while nb >= 8 {
14483 nb -= 8;
14484 rowbytes.push(((acc >> nb) & 0xFF) as u8);
14485 }
14486 }
14487 if nb > 0 {
14488 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14489 }
14490 bytes.extend_from_slice(&rowbytes);
14491 }
14492 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
14493
14494 let mut reference = vec![0f32; rows * cols];
14495 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
14496 let mut expect = vec![0f32; rows];
14497 for r in 0..rows {
14498 expect[r] = reference[r * cols..(r + 1) * cols]
14499 .iter()
14500 .zip(&x)
14501 .map(|(w, xv)| w * xv)
14502 .sum();
14503 }
14504 let mut got = vec![0f32; rows];
14505 let offsets = vbit_row_offsets(&bytes, rows, cols);
14506 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
14507 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
14511 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
14512 for r in 0..rows {
14513 assert!(
14514 (got[r] - expect[r]).abs() < tol * scale,
14515 "row {r}: {} vs {}",
14516 got[r],
14517 expect[r]
14518 );
14519 }
14520 }
14521
14522 #[test]
14527 #[cfg(target_arch = "x86_64")]
14528 fn vbit_matmat_blocked_matches_per_row() {
14529 let (rows, cols, b) = (64usize, 128usize, 9usize);
14530 let ng = cols / GROUP_SIZE;
14531 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
14532 let mut bytes = bits.clone();
14533 for g in 0..rows * ng {
14534 let sc = 0.02 + 0.0005 * g as f32;
14535 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
14536 }
14537 for r in 0..rows {
14538 let bw = bits[r] as usize;
14539 let (mut acc, mut nb) = (0u64, 0usize);
14540 let mut rowbytes = Vec::new();
14541 for i in 0..cols {
14542 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14543 acc = (acc << bw) | v;
14544 nb += bw;
14545 while nb >= 8 {
14546 nb -= 8;
14547 rowbytes.push(((acc >> nb) & 0xFF) as u8);
14548 }
14549 }
14550 if nb > 0 {
14551 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14552 }
14553 bytes.extend_from_slice(&rowbytes);
14554 }
14555 let x: Vec<f32> = (0..b * cols)
14556 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14557 .collect();
14558 let offsets = vbit_row_offsets(&bytes, rows, cols);
14559 let mut y_a = vec![0f32; b * rows];
14560 let mut y_b = vec![0f32; b * rows];
14561 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
14562 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
14563 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
14564 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
14565 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
14566 let max_d = y_a
14567 .iter()
14568 .zip(&y_b)
14569 .map(|(p, q)| (p - q).abs())
14570 .fold(0.0f32, f32::max);
14571 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
14572 }
14573
14574 #[test]
14582 fn q4t_matmat_blocked_matches_per_row() {
14583 let (rows, cols, b) = (16usize, 64usize, 9usize);
14584 let gpr = cols / GROUP_SIZE;
14585 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
14586 for r in 0..rows {
14587 for g in 0..gpr {
14588 let t = (r * gpr + g) * Q4_TILE;
14589 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
14590 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
14591 for k in 0..16 {
14592 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
14593 }
14594 }
14595 }
14596 let x: Vec<f32> = (0..b * cols)
14597 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14598 .collect();
14599 let mut y_blk = vec![0f32; b * rows];
14600 let mut y_row = vec![0f32; b * rows];
14601 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
14602 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
14603 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
14604 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
14605 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
14606 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
14607 }
14608
14609 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
14616 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
14617 let gpr = cols / GROUP_SIZE;
14618 let stride = q4tp_code_stride(gpr);
14619 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
14620 let mut b = vec![0u8; codes_off + rows * stride];
14621 for r in 0..rows {
14622 for g in 0..gpr {
14623 let t = (r * gpr + g) * Q4TP_NIB;
14624 for k in 0..16 {
14625 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
14626 }
14627 }
14628 let lo = -6.0 - 0.03 * (r % 17) as f32;
14629 let step = 0.01 + 0.004 * (r % 11) as f32;
14630 let p = params_off + r * 4;
14631 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
14632 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
14633 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
14634 for g in 0..gpr {
14635 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
14636 }
14637 }
14638 b
14639 }
14640
14641 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
14645 let gpr = cols / GROUP_SIZE;
14646 let v = Q4tpView::new(bytes, rows, cols);
14647 let mut out = vec![0u8; rows * gpr * Q4_TILE];
14648 let mut sc = vec![0f32; gpr];
14649 for r in 0..rows {
14650 v.scales_into(r, gpr, &mut sc);
14651 for g in 0..gpr {
14652 let t = (r * gpr + g) * Q4_TILE;
14653 let s = sc[g];
14654 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14655 let src = (r * gpr + g) * Q4TP_NIB;
14656 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
14657 }
14658 }
14659 out
14660 }
14661
14662 #[test]
14668 fn q4tp_exact_path_matches_dequant_reference() {
14669 let (rows, cols) = (256usize, 512usize);
14670 let gpr = cols / GROUP_SIZE;
14671 let bytes = synth_q4tp(rows, cols);
14672 let mut w = vec![0f32; rows * cols];
14673 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
14674
14675 let x: Vec<f32> = (0..cols)
14676 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14677 .collect();
14678 let v = Q4tpView::new(&bytes, rows, cols);
14679 let mut sc = vec![0f32; gpr];
14680 for r in 0..rows {
14681 v.scales_into(r, gpr, &mut sc);
14682 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
14683 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
14684 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
14688 assert!(
14689 (got - want).abs() <= 1e-5 * mag,
14690 "row {r}: kernel {got} vs dequant {want}"
14691 );
14692 }
14693 }
14694
14695 #[test]
14701 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
14702 let (rows, cols) = (256usize, 512usize);
14703 let bytes = synth_q4tp(rows, cols);
14704 let twin = q4tp_as_q4t(&bytes, rows, cols);
14705 let x: Vec<f32> = (0..cols)
14706 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14707 .collect();
14708
14709 let mut got = vec![0f32; rows];
14710 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
14711 let mut want = vec![0f32; rows];
14712 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
14713
14714 let mut w = vec![0f32; rows * cols];
14717 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
14718 for r in 0..rows {
14719 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
14720 assert!(
14721 (got[r] - want[r]).abs() <= 1e-3 * mag,
14722 "row {r}: q4tp {} vs q4t {}",
14723 got[r],
14724 want[r]
14725 );
14726 }
14727 }
14728
14729 #[test]
14734 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
14735 let (rows, cols, b) = (256usize, 512usize, 5usize);
14736 let bytes = synth_q4tp(rows, cols);
14737 let twin = q4tp_as_q4t(&bytes, rows, cols);
14738 let xs: Vec<f32> = (0..b * cols)
14739 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
14740 .collect();
14741
14742 let mut got = vec![0f32; b * rows];
14743 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
14744 let mut want = vec![0f32; b * rows];
14745 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
14746
14747 let mut w = vec![0f32; rows * cols];
14748 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
14749 for t in 0..b {
14750 for r in 0..rows {
14751 let mag: f32 = (0..cols)
14752 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
14753 .sum();
14754 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
14755 assert!(
14756 (g - wa).abs() <= 1e-3 * mag,
14757 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
14758 );
14759 }
14760 }
14761 }
14762
14763 #[test]
14764 fn q4tp_matvec2_matches_the_single_stream_kernel() {
14765 let (rows, cols) = (128usize, 256usize);
14766 let gpr = cols / GROUP_SIZE;
14767 let bytes = synth_q4tp(rows, cols);
14768 let xs: Vec<f32> = (0..2 * cols)
14769 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
14770 .collect();
14771
14772 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14773 q4tp_matvec2(
14774 &bytes,
14775 &xs[..cols],
14776 &xs[cols..],
14777 rows,
14778 cols,
14779 &mut o1,
14780 &mut o2,
14781 None,
14782 );
14783
14784 let v = Q4tpView::new(&bytes, rows, cols);
14787 let mut sc = vec![0f32; gpr];
14788 for r in 0..rows {
14789 v.scales_into(r, gpr, &mut sc);
14790 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
14791 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
14792 }
14793 }
14794
14795 #[test]
14802 fn q4tp_matvec_keeps_pace_with_q4t() {
14803 let (rows, cols) = (4096usize, 3072usize);
14804 let bytes = synth_q4tp(rows, cols);
14805 let twin = q4tp_as_q4t(&bytes, rows, cols);
14806 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
14807 let mut o = vec![0f32; rows];
14808 let n = 12;
14809 let mut best = (f64::MAX, f64::MAX);
14810 for _ in 0..3 {
14813 let t0 = std::time::Instant::now();
14814 for _ in 0..n {
14815 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
14816 }
14817 best.0 = best.0.min(t0.elapsed().as_secs_f64());
14818 let t0 = std::time::Instant::now();
14819 for _ in 0..n {
14820 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
14821 }
14822 best.1 = best.1.min(t0.elapsed().as_secs_f64());
14823 }
14824 let ratio = best.1 / best.0;
14825 println!(
14826 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
14827 best.0 * 1e3 / n as f64,
14828 best.1 * 1e3 / n as f64
14829 );
14830 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
14831 }
14832
14833 #[cfg(target_os = "macos")]
14834 #[test]
14835 fn q4t_matmat_accel_matches_dequant_reference() {
14836 if !accel_gemm_enabled() {
14837 return; }
14839 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
14841 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
14842 for r in 0..rows {
14843 for g in 0..gpr {
14844 let t = (r * gpr + g) * Q4_TILE;
14845 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
14846 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
14847 for k in 0..16 {
14848 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
14849 }
14850 }
14851 }
14852 let x: Vec<f32> = (0..b * cols)
14853 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14854 .collect();
14855 let mut got = vec![0f32; b * rows];
14856 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
14857 let mut w = vec![0f32; rows * cols];
14859 for r in 0..rows {
14860 for g in 0..gpr {
14861 let t = (r * gpr + g) * Q4_TILE;
14862 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
14863 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
14864 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
14865 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
14866 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
14867 }
14868 }
14869 }
14870 for bi in 0..b {
14871 for r in 0..rows {
14872 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
14873 let d = (got[bi * rows + r] - want).abs();
14874 assert!(
14875 d <= want.abs().max(1.0) * 1e-4,
14876 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
14877 got[bi * rows + r]
14878 );
14879 }
14880 }
14881 }
14882
14883 #[test]
14884 fn q4matvec_matches_full_dequant() {
14885 let (rows, cols) = (8, 64);
14886 let groups = rows * cols / GROUP_SIZE;
14887 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14889 for i in 0..groups * 16 {
14890 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14891 }
14892 for g in 0..groups {
14893 let s = 0.01 + 0.003 * g as f32;
14894 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14895 }
14896 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14897
14898 let mut reference = vec![0.0f32; rows * cols];
14899 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14900 let mut expect = vec![0.0f32; rows];
14901 for r in 0..rows {
14902 expect[r] = reference[r * cols..(r + 1) * cols]
14903 .iter()
14904 .zip(&x)
14905 .map(|(w, xv)| w * xv)
14906 .sum();
14907 }
14908
14909 let mut got = vec![0.0f32; rows];
14910 q4matvec(&bytes, &x, rows, cols, &mut got, None);
14911 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
14915 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14916 for r in 0..rows {
14917 assert!(
14918 (got[r] - expect[r]).abs() < tol * scale,
14919 "row {r}: {} vs {}",
14920 got[r],
14921 expect[r]
14922 );
14923 }
14924 }
14925
14926 #[test]
14929 fn vbitmatvec2_equals_two_singles() {
14930 let (rows, cols) = (6, 64);
14931 let ng = cols / GROUP_SIZE;
14932 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
14933 let mut bytes = bits.clone();
14934 for g in 0..rows * ng {
14935 let s = 0.02 + 0.001 * g as f32;
14936 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14937 }
14938 for r in 0..rows {
14939 let b = bits[r] as usize;
14940 let (mut acc, mut nb) = (0u64, 0usize);
14941 let mut rowbytes = Vec::new();
14942 for i in 0..cols {
14943 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
14944 acc = (acc << b) | v;
14945 nb += b;
14946 while nb >= 8 {
14947 nb -= 8;
14948 rowbytes.push(((acc >> nb) & 0xFF) as u8);
14949 }
14950 }
14951 if nb > 0 {
14952 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14953 }
14954 bytes.extend_from_slice(&rowbytes);
14955 }
14956 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
14957 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
14958 let offsets = vbit_row_offsets(&bytes, rows, cols);
14959
14960 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14961 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
14962 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
14963 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
14964 vbitmatvec2(
14965 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
14966 );
14967 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
14968 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
14969 }
14970
14971 #[test]
14973 fn q4matvec2_equals_two_singles() {
14974 let (rows, cols) = (8, 128);
14975 let groups = rows * cols / GROUP_SIZE;
14976 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14977 for i in 0..groups * 16 {
14978 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14979 }
14980 for g in 0..groups {
14981 let s = 0.01 + 0.003 * g as f32;
14982 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14983 }
14984 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14987 x1[9] = 250.0;
14988 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14989
14990 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14991 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
14992 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
14993 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
14994 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
14995 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
14996 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
14997 }
14998
14999 #[test]
15002 fn matvec_many_equals_separate_matvecs() {
15003 use crate::pool::Pool;
15004 let (r1, r2, cols) = (300, 200, 64);
15005 let mk = |salt: usize, rows: usize| {
15006 QTensor::from_f32(
15007 (0..rows * cols)
15008 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
15009 .collect(),
15010 rows,
15011 cols,
15012 )
15013 };
15014 let (a, b) = (mk(1, r1), mk(5, r2));
15015 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
15016 let pool = Pool::new(3);
15017
15018 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
15019 a.matvec(&x, &mut ea, Some(&pool));
15020 b.matvec(&x, &mut eb, Some(&pool));
15021 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
15022 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15023 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
15024 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
15025 }
15026
15027 #[test]
15032 fn q4tp_matvec_many_equals_separate_matvecs() {
15033 use crate::pool::Pool;
15034 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15035
15036 let (r1, r2, cols) = (300usize, 200usize, 64usize);
15037 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15038 "arch_name": "tiny-q4tp",
15039 "hidden_size": cols,
15040 "intermediate_size": cols * 2,
15041 "num_layers": 1,
15042 "num_attention_heads": 2,
15043 "num_kv_heads": 1,
15044 "head_dim": 32,
15045 "vocab_size": r1,
15046 "layer_types": ["FullAttention"],
15047 "rms_norm_eps": 1e-6,
15048 "max_position_embeddings": 8,
15049 "linear_conv_kernel_dim": 0,
15050 "linear_num_key_heads": 0,
15051 "linear_num_value_heads": 0
15052 }))
15053 .unwrap();
15054 let header = CmfHeader {
15055 format: "cmf".into(),
15056 version: CMF_VERSION,
15057 arch,
15058 quant_type: QuantType::Q4Block,
15059 provenance: None,
15060 tokenizer_config: None,
15061 section_hashes: None,
15062 skills: Vec::new(),
15063 shard: None,
15064 calibration: None,
15065 routing: None,
15066 genome: None,
15067 lineage: Vec::new(),
15068 router: None,
15069 segments: Vec::new(),
15070 };
15071 let specs = [
15072 TensorSpec {
15073 name: "q".into(),
15074 dtype: TensorDtype::Q4TiledP,
15075 shape: vec![r1, cols],
15076 data: synth_q4tp(r1, cols),
15077 },
15078 TensorSpec {
15079 name: "kv".into(),
15080 dtype: TensorDtype::Q4TiledP,
15081 shape: vec![r2, cols],
15082 data: synth_q4tp(r2, cols),
15083 },
15084 ];
15085 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
15086 std::fs::create_dir_all(&dir).unwrap();
15087 let path = dir.join("m.cmf");
15088 CmfModel::write(&path, &header, &specs, None, None).unwrap();
15089 let model = Arc::new(CmfModel::open(&path).unwrap());
15090 let (a, b) = (
15091 QTensor::from_model(&model, "q").unwrap(),
15092 QTensor::from_model(&model, "kv").unwrap(),
15093 );
15094 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
15095 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
15096 let x: Vec<f32> = (0..cols)
15097 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
15098 .collect();
15099 let pool = Pool::new(3);
15100 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15101 a.matvec(&x, &mut ea, Some(&pool));
15102 b.matvec(&x, &mut eb, Some(&pool));
15103 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15104 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15105 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
15106 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
15107 let _ = std::fs::remove_dir_all(&dir);
15108 }
15109
15110 #[test]
15116 fn multi_token_moe_rows_equal_single_token_decode() {
15117 use crate::pool::Pool;
15118 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15119
15120 let (h, inter, ne) = (64usize, 128usize, 3usize);
15121 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15122 "arch_name": "tiny-q4tp-moe",
15123 "hidden_size": h,
15124 "intermediate_size": inter,
15125 "num_layers": 1,
15126 "num_attention_heads": 2,
15127 "num_kv_heads": 1,
15128 "head_dim": 32,
15129 "vocab_size": 8,
15130 "layer_types": ["FullAttention"],
15131 "rms_norm_eps": 1e-6,
15132 "max_position_embeddings": 8,
15133 "linear_conv_kernel_dim": 0,
15134 "linear_num_key_heads": 0,
15135 "linear_num_value_heads": 0
15136 }))
15137 .unwrap();
15138 let header = CmfHeader {
15139 format: "cmf".into(),
15140 version: CMF_VERSION,
15141 arch,
15142 quant_type: QuantType::Q4Block,
15143 provenance: None,
15144 tokenizer_config: None,
15145 section_hashes: None,
15146 skills: Vec::new(),
15147 shard: None,
15148 calibration: None,
15149 routing: None,
15150 genome: None,
15151 lineage: Vec::new(),
15152 router: None,
15153 segments: Vec::new(),
15154 };
15155 let mut specs = Vec::new();
15156 for e in 0..ne {
15157 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
15158 .into_iter()
15159 .enumerate()
15160 {
15161 let mut data = synth_q4tp(r, c);
15164 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
15165 .iter_mut()
15166 .enumerate()
15167 {
15168 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
15169 }
15170 specs.push(TensorSpec {
15171 name: format!("{n}{e}"),
15172 dtype: TensorDtype::Q4TiledP,
15173 shape: vec![r, c],
15174 data,
15175 });
15176 }
15177 }
15178 let dir = std::env::temp_dir().join(format!(
15179 "cmf-moe-rows-{}-{}",
15180 std::process::id(),
15181 FLOAT_ACTIVATIONS.get()
15182 ));
15183 std::fs::create_dir_all(&dir).unwrap();
15184 let path = dir.join("m.cmf");
15185 CmfModel::write(&path, &header, &specs, None, None).unwrap();
15186 let model = Arc::new(CmfModel::open(&path).unwrap());
15187 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
15188 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
15189 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
15190 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
15191 let b = 4usize;
15192 let mut xs: Vec<f32> = (0..b * h)
15193 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
15194 .collect();
15195 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
15198 (vec![2, 0], vec![0.6, 0.4]),
15199 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
15200 (vec![1], vec![1.0]),
15201 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
15202 ];
15203 let pool = Pool::new(3);
15204 let mut want = vec![0f32; b * h];
15206 for (tk, (idx, w)) in routes.iter().enumerate() {
15207 let x = &xs[tk * h..(tk + 1) * h];
15208 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
15209 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
15210 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
15211 if FLOAT_ACTIVATIONS.get() {
15212 for (slot, &e) in idx.iter().enumerate() {
15213 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
15214 g[e].matvec(x, &mut gate, Some(&pool));
15215 u[e].matvec(x, &mut up, Some(&pool));
15216 for (v, u) in gate.iter_mut().zip(up) {
15217 *v = (*v / (1.0 + (-*v).exp())) * u;
15218 }
15219 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
15220 }
15221 }
15222 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
15223 assert!(QTensor::moe_down_many(
15224 &downs,
15225 &gs,
15226 w,
15227 &mut want[tk * h..(tk + 1) * h],
15228 Some(&pool)
15229 ));
15230 }
15231 if FLOAT_ACTIVATIONS.get() {
15232 for (tk, (idx, w)) in routes.iter().enumerate() {
15233 let mut scalar = vec![0.0; h];
15234 for (&e, &weight) in idx.iter().zip(w) {
15235 let (mut gate, mut up, mut down) =
15236 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
15237 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
15238 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
15239 for (v, u) in gate.iter_mut().zip(up) {
15240 *v = (*v / (1.0 + (-*v).exp())) * u;
15241 }
15242 d[e].matvec(&gate, &mut down, Some(&pool));
15243 for (v, d) in scalar.iter_mut().zip(down) {
15244 *v += weight * d;
15245 }
15246 }
15247 assert_eq!(
15248 &want[tk * h..(tk + 1) * h],
15249 scalar,
15250 "float many equals scalar experts"
15251 );
15252 }
15253 }
15254 let mut experts: Vec<usize> = Vec::new();
15256 let mut groups: Vec<Vec<usize>> = Vec::new();
15257 for (tk, (idx, _)) in routes.iter().enumerate() {
15258 for &e in idx {
15259 match experts.iter().position(|&x| x == e) {
15260 Some(k) => groups[k].push(tk),
15261 None => {
15262 experts.push(e);
15263 groups.push(vec![tk]);
15264 }
15265 }
15266 }
15267 }
15268 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
15269 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
15270 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
15271 assert!(QTensor::moe_gate_up_rows(
15272 &pairs,
15273 &groups,
15274 &xs,
15275 &mut gs,
15276 Some(&pool)
15277 ));
15278 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
15279 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
15280 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15281 assert!(QTensor::moe_down_rows(
15282 &downs,
15283 &lens,
15284 &gs,
15285 &mut ds,
15286 Some(&pool)
15287 ));
15288 let slot = |tk: usize, e: usize| {
15289 let k = experts.iter().position(|&x| x == e).unwrap();
15290 groups[..k].iter().map(|g| g.len()).sum::<usize>()
15291 + groups[k].iter().position(|&x| x == tk).unwrap()
15292 };
15293 let mut got = vec![0f32; b * h];
15294 for (tk, (idx, w)) in routes.iter().enumerate() {
15295 for i in 0..h {
15296 let mut acc = 0f32;
15297 for (&e, &we) in idx.iter().zip(w) {
15298 acc += we * ds[slot(tk, e)][i];
15299 }
15300 got[tk * h + i] = acc;
15301 }
15302 }
15303 assert!(want.iter().any(|v| *v != 0.0));
15304 assert_eq!(
15305 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
15306 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
15307 "multi-token MoE must equal decode bit for bit"
15308 );
15309
15310 let b5 = 5usize;
15313 let x5: Vec<f32> = (0..b5 * h)
15314 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
15315 .collect();
15316 let mut mm = vec![0f32; b5 * inter];
15317 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
15318 for tk in 0..b5 {
15319 let mut mv = vec![0f32; inter];
15320 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
15321 assert_eq!(
15322 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
15323 mm[tk * inter..(tk + 1) * inter]
15324 .iter()
15325 .map(|v| v.to_bits())
15326 .collect::<Vec<_>>(),
15327 "row-exact matmat token {tk}"
15328 );
15329 }
15330 let _ = std::fs::remove_dir_all(&dir);
15333 }
15334
15335 #[test]
15343 fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
15344 use crate::pool::Pool;
15345 use std::sync::atomic::Ordering::Relaxed;
15346 let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
15347 Q4TP_ALT.store(2, Relaxed);
15350 let pool = Pool::new(3);
15351 for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
15355 let bytes = synth_q4tp(rows, cols);
15356 let mut xs: Vec<f32> = (0..b * cols)
15357 .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
15358 .collect();
15359 xs[3] = 7.5; let run = |exact: bool| {
15361 let mut out = vec![0f32; b * rows];
15362 q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
15363 out
15364 };
15365 let (fast, exact) = (run(false), run(true));
15366 #[cfg(not(target_arch = "aarch64"))]
15367 let _ = fast;
15368 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
15369 let mut matvecs = vec![0f32; b * rows];
15370 for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
15371 q4tp_matvec(
15372 &bytes,
15373 &xs[bi * cols..(bi + 1) * cols],
15374 rows,
15375 cols,
15376 o,
15377 Some(&pool),
15378 );
15379 }
15380 assert!(matvecs.iter().any(|v| *v != 0.0));
15381 assert_eq!(
15382 bits(&exact),
15383 bits(&matvecs),
15384 "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
15385 );
15386 #[cfg(target_arch = "aarch64")]
15387 {
15388 let gpr = cols / GROUP_SIZE;
15390 let v = Q4tpView::new(&bytes, rows, cols);
15391 let mut old = vec![0f32; b * rows];
15392 let mut sc = vec![0f32; gpr];
15393 let a8w8 = a8w8_enabled();
15394 let blocked = sdot_enabled() && blocked_enabled();
15395 let acts: Vec<SplitAct> = (0..b)
15396 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
15397 .collect();
15398 for r in 0..rows {
15399 v.scales_into(r, gpr, &mut sc);
15400 if !a8w8 {
15401 for bi in 0..b {
15402 let x = &xs[bi * cols..(bi + 1) * cols];
15403 old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
15404 }
15405 continue;
15406 }
15407 let finish = |d: f32, act: &SplitAct| {
15408 let mut acc = d * act.sx;
15409 for &(j, xv) in &act.outliers {
15410 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
15411 acc += w * s * xv;
15412 }
15413 acc
15414 };
15415 let mut bi = 0usize;
15416 while blocked && bi + 4 <= b {
15417 let xs4 = [
15418 acts[bi].xq.as_slice(),
15419 acts[bi + 1].xq.as_slice(),
15420 acts[bi + 2].xq.as_slice(),
15421 acts[bi + 3].xq.as_slice(),
15422 ];
15423 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
15424 for k in 0..4 {
15425 old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
15426 }
15427 bi += 4;
15428 }
15429 for (bi, act) in acts.iter().enumerate().skip(bi) {
15430 let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
15431 old[bi * rows + r] = finish(d, act);
15432 }
15433 }
15434 assert_eq!(
15435 bits(&fast),
15436 bits(&old),
15437 "{rows}x{cols} b={b}: the fast path changed outside row_exact"
15438 );
15439 if blocked {
15442 assert_ne!(
15443 bits(&fast),
15444 bits(&matvecs),
15445 "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
15446 );
15447 }
15448 }
15449 }
15450 Q4TP_ALT.store(0, Relaxed);
15451 }
15452
15453 #[test]
15454 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
15455 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
15456 let active = AtomicUsize::new(0);
15459 counted_row_exact_scope(&active, || {
15460 assert_eq!(active.load(Ordering::Acquire), 1);
15461 counted_row_exact_scope(&active, || {
15462 assert_eq!(active.load(Ordering::Acquire), 2);
15463 });
15464 assert_eq!(active.load(Ordering::Acquire), 1);
15465 });
15466 assert_eq!(active.load(Ordering::Acquire), 0);
15467
15468 let both_entered = Barrier::new(2);
15469 let release_last = Barrier::new(2);
15470 std::thread::scope(|s| {
15471 let first = s.spawn(|| counted_row_exact_scope(&active, || {
15472 both_entered.wait();
15473 }));
15474 let last = s.spawn(|| counted_row_exact_scope(&active, || {
15475 both_entered.wait();
15476 release_last.wait();
15477 }));
15478 first.join().unwrap();
15479 let after_first = active.load(Ordering::Acquire);
15480 release_last.wait();
15481 last.join().unwrap();
15482 assert_eq!(after_first, 1, "second request must remain exact");
15483 });
15484 assert_eq!(active.load(Ordering::Acquire), 0);
15485 let panic = std::panic::catch_unwind(|| {
15486 counted_row_exact_scope(&active, || panic!("scope unwind"));
15487 });
15488 assert!(panic.is_err());
15489 assert_eq!(active.load(Ordering::Acquire), 0);
15490 }
15491
15492 #[test]
15493 #[cfg(target_arch = "x86_64")]
15494 fn q4tp_float_avx2_is_bitwise_scalar() {
15495 if !avx2_enabled() {
15496 return;
15497 }
15498 for cols in [32, 64, 96, 2048, 4096] {
15499 let rows = 9;
15500 let bytes = synth_q4tp(rows, cols);
15501 let v = Q4tpView::new(&bytes, rows, cols);
15502 let gpr = cols / GROUP_SIZE;
15503 let mut sc = vec![0.0; gpr];
15504 for seed in 1..=5 {
15505 let xs: Vec<f32> = (0..cols)
15506 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
15507 .collect();
15508 for r in 0..rows {
15509 v.scales_into(r, gpr, &mut sc);
15510 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
15511 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
15512 assert_eq!(
15513 scalar.to_bits(),
15514 vector.to_bits(),
15515 "cols={cols} row={r} seed={seed}"
15516 );
15517 }
15518 }
15519 }
15520 }
15521
15522 #[test]
15523 fn multi_token_moe_rows_float_equal_single_token_decode() {
15524 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
15525 }
15526
15527 #[test]
15528 fn full_gpu_q8_scope_is_nested_and_thread_local() {
15529 assert!(!FULL_GPU_Q8.get());
15530 let before = gpu_split_frac();
15531 {
15532 let _guard = enter_full_gpu_q8_scope();
15533 assert_eq!(gpu_split_frac(), 1.0);
15534 {
15535 let _nested = enter_full_gpu_q8_scope();
15536 }
15537 assert_eq!(gpu_split_frac(), 1.0);
15538 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
15539 }
15540 assert!(!FULL_GPU_Q8.get());
15541 assert_eq!(gpu_split_frac(), before);
15542 }
15543
15544 #[test]
15545 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
15546 assert!(!FLOAT_ACTIVATIONS.get());
15547 let before = a8w8_enabled();
15548 float_activations_scope(|| {
15549 assert!(!a8w8_enabled());
15550 float_activations_scope(|| assert!(!a8w8_enabled()));
15551 assert!(FLOAT_ACTIVATIONS.get());
15552 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
15553 .join()
15554 .unwrap();
15555 });
15556 assert!(!FLOAT_ACTIVATIONS.get());
15557 assert_eq!(a8w8_enabled(), before);
15558 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
15559 assert!(!FLOAT_ACTIVATIONS.get());
15560 }
15561
15562 #[test]
15565 fn batched_matmat_equals_per_position_matvec() {
15566 let (rows, cols, b) = (8, 64, 5);
15567 let groups = rows * cols / GROUP_SIZE;
15569 let mut q4 = Vec::new();
15570 for i in 0..groups * 16 {
15571 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15572 }
15573 for g in 0..groups {
15574 q4.extend_from_slice(
15575 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
15576 );
15577 }
15578 let ng = cols / GROUP_SIZE;
15580 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
15581 let mut vb = bits.clone();
15582 for g in 0..rows * ng {
15583 vb.extend_from_slice(
15584 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
15585 );
15586 }
15587 for r in 0..rows {
15588 let bw = bits[r] as usize;
15589 let (mut acc, mut nb) = (0u64, 0usize);
15590 let mut rowbytes = Vec::new();
15591 for i in 0..cols {
15592 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
15593 acc = (acc << bw) | v;
15594 nb += bw;
15595 while nb >= 8 {
15596 nb -= 8;
15597 rowbytes.push(((acc >> nb) & 0xFF) as u8);
15598 }
15599 }
15600 if nb > 0 {
15601 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15602 }
15603 vb.extend_from_slice(&rowbytes);
15604 }
15605 let offsets = vbit_row_offsets(&vb, rows, cols);
15606
15607 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
15608
15609 let mut got = vec![0f32; b * rows];
15611 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
15612 for bi in 0..b {
15613 let mut expect = vec![0f32; rows];
15614 q4matvec(
15615 &q4,
15616 &xs[bi * cols..(bi + 1) * cols],
15617 rows,
15618 cols,
15619 &mut expect,
15620 None,
15621 );
15622 assert_eq!(
15623 &got[bi * rows..(bi + 1) * rows],
15624 &expect[..],
15625 "q4 batch pos {bi}"
15626 );
15627 }
15628
15629 let mut got = vec![0f32; b * rows];
15631 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
15632 for bi in 0..b {
15633 let mut expect = vec![0f32; rows];
15634 vbitmatvec(
15635 &vb,
15636 &offsets,
15637 &xs[bi * cols..(bi + 1) * cols],
15638 rows,
15639 cols,
15640 &mut expect,
15641 None,
15642 );
15643 assert_eq!(
15644 &got[bi * rows..(bi + 1) * rows],
15645 &expect[..],
15646 "vbit batch pos {bi}"
15647 );
15648 }
15649 }
15650
15651 #[test]
15655 fn q4_tiled_matches_q4_block_bitexact() {
15656 let (rows, cols, b) = (8usize, 128usize, 3usize);
15657 let groups = rows * cols / GROUP_SIZE;
15658 let mut split = Vec::with_capacity(groups * 18);
15659 for i in 0..groups * 16 {
15660 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15661 }
15662 for g in 0..groups {
15663 split.extend_from_slice(
15664 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
15665 );
15666 }
15667 let (packed, scales) = split.split_at(groups * 16);
15669 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
15670 for g in 0..groups {
15671 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
15672 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
15673 }
15674
15675 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15676 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
15678
15679 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
15680 q4matvec(&split, &x1, rows, cols, &mut a, None);
15681 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
15682 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
15683
15684 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15685 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
15686 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
15687 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
15688 assert_eq!(a1, t1);
15689 assert_eq!(a2, t2);
15690
15691 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
15692 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
15693 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
15694 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
15695 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
15696 }
15697
15698 #[test]
15705 fn q4matvec_sdot_outlier_exact() {
15706 let (rows, cols) = (4, 128);
15707 let groups = rows * cols / GROUP_SIZE;
15708 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15709 for i in 0..groups * 16 {
15710 bytes.push(((i * 11 + 5) % 256) as u8);
15711 }
15712 for g in 0..groups {
15713 let s = 0.02 + 0.002 * g as f32;
15714 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15715 }
15716 let mut x: Vec<f32> = (0..cols)
15717 .map(|i| match i % 3 {
15718 0 => 1.0,
15719 1 => -1.0,
15720 _ => 0.0,
15721 })
15722 .collect();
15723 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
15726 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
15727 let mut expect = vec![0.0f32; rows];
15728 for r in 0..rows {
15729 expect[r] = reference[r * cols..(r + 1) * cols]
15730 .iter()
15731 .zip(&x)
15732 .map(|(w, xv)| w * xv)
15733 .sum();
15734 }
15735 let mut got = vec![0.0f32; rows];
15736 q4matvec(&bytes, &x, rows, cols, &mut got, None);
15737 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
15738 for r in 0..rows {
15739 assert!(
15740 (got[r] - expect[r]).abs() < 2e-3 * scale,
15741 "row {r}: {} vs {} (outlier term must be exact)",
15742 got[r],
15743 expect[r]
15744 );
15745 }
15746 }
15747
15748 #[test]
15752 fn q1t_matvec_matches_reference() {
15753 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
15754 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
15756 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
15757 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
15759 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
15760 let mut bytes = Vec::new();
15761 for r in 0..rows {
15762 for g in 0..gpr {
15763 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
15764 let mut c = [0u8; 7];
15765 for k in 0..GROUP_SIZE {
15766 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
15768 0
15769 } else {
15770 ((k + r * 3 + g) % 3) as u8 };
15772 cortiq_core::quant::q1t_pack(&mut c, k, code);
15773 }
15774 bytes.extend_from_slice(&c);
15775 }
15776 }
15777 let mut row_ptr = vec![0u32; rows + 1];
15780 for &(idx, _) in &outliers {
15781 row_ptr[idx as usize / cols + 1] += 1;
15782 }
15783 for r in 0..rows {
15784 row_ptr[r + 1] += row_ptr[r];
15785 }
15786 for &p in &row_ptr {
15787 bytes.extend_from_slice(&p.to_le_bytes());
15788 }
15789 for &(idx, v) in &outliers {
15790 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
15791 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
15792 }
15793
15794 let mut refw = vec![0f32; rows * cols];
15795 dequant_q1t(&bytes, rows, cols, &mut refw);
15796 let x: Vec<f32> = (0..cols)
15799 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
15800 .collect();
15801 let mut expect = vec![0f32; rows];
15802 for r in 0..rows {
15803 let mut a = 0.0f32;
15804 for j in 0..cols {
15805 a += refw[r * cols + j] * x[j];
15806 }
15807 expect[r] = a;
15808 }
15809 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
15810 let mut got = vec![0f32; rows];
15811 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
15812 for r in 0..rows {
15813 assert!(
15814 (got[r] - expect[r]).abs() < tol(expect[r]),
15815 "row {r}: {} vs {}",
15816 got[r],
15817 expect[r]
15818 );
15819 }
15820 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
15822 let mut gm = vec![0f32; 2 * rows];
15823 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
15824 for r in 0..rows {
15825 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
15826 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
15827 }
15828 let xb: Vec<f32> = (0..cols)
15832 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
15833 .collect();
15834 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
15835 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
15836 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
15837 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
15838 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
15839 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
15840 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
15841 }
15842
15843 #[test]
15846 fn q1t_matvec2_odd_gpr_matches_singles() {
15847 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
15848 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
15850 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
15851 for r in 0..rows {
15852 for g in 0..gpr {
15853 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
15854 let mut c = [0u8; 7];
15855 for k in 0..GROUP_SIZE {
15856 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
15857 }
15858 bytes.extend_from_slice(&c);
15859 }
15860 }
15861 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
15862 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
15863 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
15864 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
15865 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
15866 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
15867 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
15868 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
15869 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
15870 }
15871
15872 #[test]
15876 #[ignore]
15877 fn q1t_matvec2_speed() {
15878 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
15879 use std::time::Instant;
15880 let (rows, cols) = (8192usize, 4096usize);
15881 let gpr = cols / GROUP_SIZE;
15882 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
15883 for r in 0..rows {
15884 for g in 0..gpr {
15885 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
15886 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
15887 let mut c = [0u8; 7];
15888 for k in 0..GROUP_SIZE {
15889 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
15890 }
15891 bytes.extend_from_slice(&c);
15892 }
15893 }
15894 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
15895 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
15896 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
15897 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
15898 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
15900 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
15901 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
15902 for _ in 0..8 {
15903 let t0 = Instant::now();
15904 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
15905 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
15906 let t1 = Instant::now();
15907 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
15908 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
15909 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
15910 }
15911 assert_eq!(p1, s1);
15912 assert_eq!(p2, s2);
15913 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
15914 }
15915
15916 #[test]
15920 #[ignore]
15921 fn q1t_matvec_speed() {
15922 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
15923 use std::time::Instant;
15924 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
15926 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
15927 for r in 0..rows {
15928 for g in 0..gpr {
15929 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
15930 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
15931 let mut c = [0u8; 7];
15932 for k in 0..GROUP_SIZE {
15933 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
15934 }
15935 bytes.extend_from_slice(&c);
15936 }
15937 }
15938 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
15940 let mut idx = 0usize;
15941 while idx < n {
15942 row_ptr[idx / cols + 1] += 1;
15943 idx += stride;
15944 }
15945 for r in 0..rows {
15946 row_ptr[r + 1] += row_ptr[r];
15947 }
15948 for &p in &row_ptr {
15949 bytes.extend_from_slice(&p.to_le_bytes());
15950 }
15951 let mut idx = 0usize;
15952 while idx < n {
15953 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
15954 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
15955 idx += stride;
15956 }
15957 let x: Vec<f32> = (0..cols)
15960 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
15961 .collect();
15962 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
15963
15964 let slow = |out: &mut [f32]| {
15966 let mut buf = vec![0f32; cols];
15967 for r in 0..rows {
15968 for g in 0..gpr {
15969 let off = (r * gpr + g) * Q1T_TILE;
15970 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
15971 let codes = &bytes[off + 2..off + Q1T_TILE];
15972 for k in 0..GROUP_SIZE {
15973 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
15974 1 => s,
15975 2 => -s,
15976 _ => 0.0,
15977 };
15978 }
15979 }
15980 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
15981 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
15982 }
15983 };
15984 let iters = 5;
15985 let mut a = vec![0f32; rows];
15986 slow(&mut a); let t = Instant::now();
15988 for _ in 0..iters {
15989 slow(&mut a);
15990 }
15991 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
15992
15993 let mut b = vec![0f32; rows];
15994 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
15996 for _ in 0..iters {
15997 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
15998 }
15999 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16000
16001 for r in 0..rows {
16002 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
16003 }
16004 println!(
16005 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
16006 slow_ms / fast_ms
16007 );
16008 }
16009}
16010
16011#[cfg(test)]
16012mod gemm_bench {
16013 #[test]
16024 #[ignore]
16025 fn q4tp_matmat_throughput() {
16026 let _alt = super::Q4TP_ALT_TEST_LOCK
16027 .lock()
16028 .unwrap_or_else(|e| e.into_inner());
16029 let b: usize = std::env::var("CMF_BENCH_B")
16033 .ok()
16034 .and_then(|v| v.parse().ok())
16035 .unwrap_or(296);
16036 let (rows, cols) = (9216usize, 2304usize);
16037 let (_, _, _) = (rows, cols, b);
16038 let total =
16039 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
16040 .unwrap();
16041 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16046 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16047 let lo = cortiq_core::quant::f32_to_f16(-4.0);
16048 let step = cortiq_core::quant::f32_to_f16(0.1);
16049 for r in 0..rows {
16050 let o = params_off + r * 4;
16051 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16052 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16053 }
16054 let _ = codes_off;
16055 let xs: Vec<f32> = (0..b * cols)
16056 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
16057 .collect();
16058 let mut out = vec![0f32; b * rows];
16059 let pool = crate::pool::Pool::from_env();
16060 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16066 let reps: usize = std::env::var("CMF_BENCH_REPS")
16067 .ok()
16068 .and_then(|v| v.parse().ok())
16069 .unwrap_or(10);
16070 let mut best = [f64::MAX; 2];
16071 let mut sums = [0f32; 2];
16072 for _ in 0..reps {
16073 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
16074 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
16075 let t = std::time::Instant::now();
16076 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16077 best[k] = best[k].min(t.elapsed().as_secs_f64());
16078 sums[k] = out.iter().take(64).sum::<f32>();
16079 }
16080 }
16081 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
16082 for (k, name) in ["previous", "tuned "].iter().enumerate() {
16083 println!(
16084 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
16085 best[k] * 1e3,
16086 flops / best[k] / 1e9,
16087 sums[k]
16088 );
16089 }
16090 assert!(
16091 (sums[0] - sums[1]).abs() < 1e-2,
16092 "the tuned kernel changed the result: {} vs {}",
16093 sums[0],
16094 sums[1]
16095 );
16096 }
16097
16098 #[test]
16104 fn q4tp_matmat_blocked_matches_scalar() {
16105 use std::sync::atomic::Ordering::Relaxed;
16106 let _alt = super::Q4TP_ALT_TEST_LOCK
16107 .lock()
16108 .unwrap_or_else(|e| e.into_inner());
16109 for &(rows, cols, b) in &[
16117 (64usize, 128usize, 7usize),
16118 (33, 96, 4),
16119 (16, 256, 9),
16120 (192, 2304, 37),
16121 ] {
16122 let total = cortiq_core::quant::expected_nbytes(
16123 cortiq_core::TensorDtype::Q4TiledP,
16124 &[rows, cols],
16125 )
16126 .unwrap();
16127 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16128 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
16129 let lo = cortiq_core::quant::f32_to_f16(-4.0);
16130 let step = cortiq_core::quant::f32_to_f16(0.1);
16131 for r in 0..rows {
16132 let o = params_off + r * 4;
16133 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16134 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16135 }
16136 let xs: Vec<f32> = (0..b * cols)
16137 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
16138 .collect();
16139 let mut got = vec![0f32; b * rows];
16140 let mut want = vec![0f32; b * rows];
16141 let gpr = cols / 32;
16142 let view = super::Q4tpView::new(&bytes, rows, cols);
16143 let pool = crate::pool::Pool::from_env();
16144 super::Q4TP_ALT.store(2, Relaxed);
16145 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
16146 super::Q4TP_ALT.store(1, Relaxed);
16147 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
16148 super::Q4TP_ALT.store(0, Relaxed);
16149 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
16156 let (mut worst, mut at) = (0f32, 0usize);
16157 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
16158 if (g - w).abs() > worst {
16159 worst = (g - w).abs();
16160 at = i;
16161 }
16162 }
16163 assert!(
16164 worst <= 1e-4 * scale,
16165 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
16166 (scale {scale:.3e}) at cell {at}: {} vs {}",
16167 got[at],
16168 want[at]
16169 );
16170
16171 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
16178 for bi in 0..b {
16179 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
16180 for r in 0..rows {
16181 let mut sc = vec![0f32; gpr];
16182 view.scales_into(r, gpr, &mut sc);
16183 let mut exact = 0f64;
16184 for j in 0..cols {
16185 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
16186 exact += w as f64 * sq as f64 * act.xq[j] as f64;
16187 }
16188 exact *= act.sx as f64;
16189 for &(j, xv) in &act.outliers {
16190 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
16191 exact += w as f64 * sq as f64 * xv as f64;
16192 }
16193 let i = bi * rows + r;
16194 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
16195 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
16196 }
16197 }
16198 println!(
16199 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
16200 per-column {e_scalar:.3e}"
16201 );
16202 assert!(
16207 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
16208 "{rows}x{cols} b={b}: error against f64 too large — blocked \
16209 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
16210 );
16211 }
16212 }
16213
16214 #[test]
16220 #[ignore]
16221 fn q4tp_matmat_row_exact_cost() {
16222 let _alt = super::Q4TP_ALT_TEST_LOCK
16223 .lock()
16224 .unwrap_or_else(|e| e.into_inner());
16225 let pool = crate::pool::Pool::from_env();
16226 let reps: usize = std::env::var("CMF_BENCH_REPS")
16227 .ok()
16228 .and_then(|v| v.parse().ok())
16229 .unwrap_or(30);
16230 for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
16231 let total = cortiq_core::quant::expected_nbytes(
16232 cortiq_core::TensorDtype::Q4TiledP,
16233 &[rows, cols],
16234 )
16235 .unwrap();
16236 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16237 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16238 let lo = cortiq_core::quant::f32_to_f16(-4.0);
16239 let step = cortiq_core::quant::f32_to_f16(0.1);
16240 for r in 0..rows {
16241 let o = params_off + r * 4;
16242 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16243 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16244 }
16245 for &b in &[2usize, 4, 5, 8] {
16246 let xs: Vec<f32> = (0..b * cols)
16247 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
16248 .collect();
16249 let mut out = vec![0f32; b * rows];
16250 let mut best = [f64::MAX; 3];
16251 for _ in 0..reps {
16252 for (k, best_k) in best.iter_mut().enumerate() {
16253 let t = std::time::Instant::now();
16254 match k {
16255 0 | 1 => super::q4tp_matmat_with(
16256 &bytes,
16257 &xs,
16258 b,
16259 rows,
16260 cols,
16261 &mut out,
16262 pool.as_deref(),
16263 k == 1,
16264 ),
16265 _ => {
16266 for (bi, o) in out.chunks_mut(rows).enumerate() {
16267 super::q4tp_matvec(
16268 &bytes,
16269 &xs[bi * cols..(bi + 1) * cols],
16270 rows,
16271 cols,
16272 o,
16273 pool.as_deref(),
16274 );
16275 }
16276 }
16277 }
16278 *best_k = best_k.min(t.elapsed().as_secs_f64());
16279 }
16280 }
16281 println!(
16282 "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
16283 {b} matvecs {:.3} ms",
16284 best[0] * 1e3,
16285 best[1] * 1e3,
16286 best[2] * 1e3
16287 );
16288 }
16289 }
16290 }
16291
16292 #[test]
16296 #[ignore]
16297 fn q4t_matmat_throughput() {
16298 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
16299 let total =
16300 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
16301 .unwrap();
16302 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16305 let sc = cortiq_core::quant::f32_to_f16(0.02);
16306 for t in bytes.chunks_mut(super::Q4_TILE) {
16307 t[..2].copy_from_slice(&sc.to_le_bytes());
16308 }
16309 let xs: Vec<f32> = (0..b * cols)
16310 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
16311 .collect();
16312 let mut out = vec![0f32; b * rows];
16313 let pool = crate::pool::Pool::from_env();
16314 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16315 let reps: usize = std::env::var("CMF_BENCH_REPS")
16316 .ok()
16317 .and_then(|v| v.parse().ok())
16318 .unwrap_or(10);
16319 let mut best = f64::MAX;
16320 for _ in 0..reps {
16321 let t = std::time::Instant::now();
16322 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16323 best = best.min(t.elapsed().as_secs_f64());
16324 }
16325 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
16326 println!(
16327 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
16328 best * 1e3,
16329 flops / best / 1e9,
16330 out.iter().take(64).sum::<f32>()
16331 );
16332 }
16333}