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
27#[cfg(target_arch = "x86_64")]
30mod moe_fast;
31
32pub static MATMAT_GATE: AtomicU64 = AtomicU64::new(0);
35pub static MATMAT_GPU_HIT: AtomicU64 = AtomicU64::new(0);
36
37pub enum QTensor {
38 F32 {
39 data: Vec<f32>,
40 rows: usize,
41 cols: usize,
42 },
43 Mapped {
44 model: Arc<CmfModel>,
45 idx: usize,
47 dtype: TensorDtype,
48 rows: usize,
49 cols: usize,
50 row_scale: Vec<f32>,
52 col_field: Vec<f32>,
54 vbit_offsets: Vec<usize>,
58 repack: Vec<u8>,
66 },
67 Blr(Box<crate::blr::BlrTensor>),
73}
74
75fn repack_enabled() -> bool {
82 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
83 *ON.get_or_init(|| {
84 std::env::var("CMF_REPACK")
85 .map(|v| v == "1")
86 .unwrap_or(cfg!(target_os = "android"))
87 })
88}
89
90fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
94 #[cfg(target_arch = "aarch64")]
95 let arch_ok = sdot_enabled();
96 #[cfg(not(target_arch = "aarch64"))]
97 let arch_ok = false;
98 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
99 return Vec::new();
100 }
101 q8_repack_layout(bytes, rows, cols)
102}
103
104fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
107 let groups = rows / 4;
108 let mut rep = vec![0u8; groups * 4 * cols];
109 for g in 0..groups {
110 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
111 for c in 0..cols / 16 {
112 for lane in 0..4 {
113 let src = (g * 4 + lane) * cols + c * 16;
114 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
115 .copy_from_slice(&bytes[src..src + 16]);
116 }
117 }
118 }
119 rep
120}
121
122fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
125 let ng = cols / GROUP_SIZE;
126 let bits = &bytes[..rows];
127 let mut offsets = Vec::with_capacity(rows + 1);
128 let mut off = rows + rows * ng * 2;
129 for r in 0..rows {
130 offsets.push(off);
131 off += (cols * bits[r] as usize).div_ceil(8);
132 }
133 offsets.push(off);
134 offsets
135}
136
137fn blocked_enabled() -> bool {
143 use std::sync::atomic::Ordering::Relaxed;
144 match BLOCKED_OVERRIDE.load(Relaxed) {
145 1 => false,
146 2 => true,
147 _ => {
148 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
149 *ON.get_or_init(|| {
150 std::env::var("CMF_X86_BLOCKED")
151 .map(|v| v != "0")
152 .unwrap_or(true)
153 })
154 }
155 }
156}
157
158static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
159
160pub fn set_blocked_override(on: Option<bool>) {
167 let v = match on {
168 None => 0,
169 Some(false) => 1,
170 Some(true) => 2,
171 };
172 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
173}
174
175fn gpu_lmhead_enabled() -> bool {
176 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
177 *ON.get_or_init(|| {
178 std::env::var("CMF_GPU_LMHEAD")
179 .map(|v| v != "0")
180 .unwrap_or(true)
181 })
182}
183
184fn gpu_split_frac() -> f32 {
185 if FULL_GPU_Q8.get() {
188 return 1.0;
189 }
190 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
191 *F.get_or_init(|| {
192 std::env::var("CMF_GPU_SPLIT")
193 .ok()
194 .and_then(|v| v.parse::<f32>().ok())
195 .unwrap_or(0.5)
196 .clamp(0.0, 1.0)
197 })
198}
199
200impl QTensor {
201 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
202 debug_assert_eq!(data.len(), rows * cols);
203 Self::F32 { data, rows, cols }
204 }
205
206 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
209 let idx = model
212 .tensor_index(name)
213 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
214 let entry = &model.tensors[idx];
215 if entry.shape.len() != 2 {
216 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
217 }
218 let (rows, cols) = (entry.shape[0], entry.shape[1]);
219 let bytes = model.entry_bytes(entry);
220
221 match entry.dtype {
222 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
223 let n = rows * cols;
224 let scales_off = n;
225 let row_scale: Vec<f32> = (0..rows)
226 .map(|o| {
227 f16_to_f32(u16::from_le_bytes([
228 bytes[scales_off + o * 2],
229 bytes[scales_off + o * 2 + 1],
230 ]))
231 })
232 .collect();
233 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
234 let col_off = n + rows * 2;
235 (0..cols)
236 .map(|i| {
237 f16_to_f32(u16::from_le_bytes([
238 bytes[col_off + i * 2],
239 bytes[col_off + i * 2 + 1],
240 ]))
241 })
242 .collect()
243 } else {
244 Vec::new()
245 };
246 Ok(Self::Mapped {
247 model: model.clone(),
248 idx,
249 dtype: entry.dtype,
250 rows,
251 cols,
252 row_scale,
253 col_field,
254 vbit_offsets: Vec::new(),
255 repack: q8_repack(bytes, rows, cols),
256 })
257 }
258 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
260 model: model.clone(),
261 idx,
262 dtype: entry.dtype,
263 rows,
264 cols,
265 row_scale: Vec::new(),
266 col_field: Vec::new(),
267 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
268 repack: Vec::new(),
269 }),
270 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
274 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
275 let offsets: Vec<usize> = (0..=rows)
276 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
277 .collect();
278 Ok(Self::Mapped {
279 model: model.clone(),
280 idx,
281 dtype: entry.dtype,
282 rows,
283 cols,
284 row_scale: Vec::new(),
285 col_field: Vec::new(),
286 vbit_offsets: offsets,
287 repack: Vec::new(),
288 })
289 }
290 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
296 model: model.clone(),
297 idx,
298 dtype: entry.dtype,
299 rows,
300 cols,
301 row_scale: Vec::new(),
302 col_field: Vec::new(),
303 vbit_offsets: Vec::new(),
304 repack: Vec::new(),
305 }),
306 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
309 model: model.clone(),
310 idx,
311 dtype: entry.dtype,
312 rows,
313 cols,
314 row_scale: Vec::new(),
315 col_field: Vec::new(),
316 vbit_offsets: Vec::new(),
317 repack: Vec::new(),
318 }),
319 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
321 model: model.clone(),
322 idx,
323 dtype: entry.dtype,
324 rows,
325 cols,
326 row_scale: Vec::new(),
327 col_field: Vec::new(),
328 vbit_offsets: Vec::new(),
329 repack: Vec::new(),
330 }),
331 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
332 model: model.clone(),
333 idx,
334 dtype: entry.dtype,
335 rows,
336 cols,
337 row_scale: Vec::new(),
338 col_field: Vec::new(),
339 vbit_offsets: Vec::new(),
340 repack: Vec::new(),
341 }),
342 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
345 model: model.clone(),
346 idx,
347 dtype: entry.dtype,
348 rows,
349 cols,
350 row_scale: Vec::new(),
351 col_field: Vec::new(),
352 vbit_offsets: Vec::new(),
353 repack: Vec::new(),
354 }),
355 TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
356 model: model.clone(),
357 idx,
358 dtype: entry.dtype,
359 rows,
360 cols,
361 row_scale: Vec::new(),
362 col_field: Vec::new(),
363 vbit_offsets: Vec::new(),
364 repack: Vec::new(),
365 }),
366 TensorDtype::Mxfp4 if cols % cortiq_core::mxfp4::GROUP == 0 => Ok(Self::Mapped {
369 model: model.clone(),
370 idx,
371 dtype: entry.dtype,
372 rows,
373 cols,
374 row_scale: Vec::new(),
375 col_field: Vec::new(),
376 vbit_offsets: Vec::new(),
377 repack: Vec::new(),
378 }),
379 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
383 model: model.clone(),
384 idx,
385 dtype: entry.dtype,
386 rows,
387 cols,
388 row_scale: Vec::new(),
389 col_field: Vec::new(),
390 vbit_offsets: Vec::new(),
391 repack: Vec::new(),
392 }),
393 TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
396 Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
397 model, idx,
398 )?)))
399 }
400 _ => {
402 let mut data = vec![0.0f32; rows * cols];
403 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
404 Ok(Self::from_f32(data, rows, cols))
405 }
406 }
407 }
408
409 pub fn is_blr(&self) -> bool {
412 matches!(self, Self::Blr(_))
413 }
414
415 pub fn is_mxfp4(&self) -> bool {
418 matches!(
419 self,
420 Self::Mapped {
421 dtype: TensorDtype::Mxfp4,
422 ..
423 }
424 )
425 }
426
427 fn per_tensor_only(&self) -> bool {
430 self.is_blr() || self.is_mxfp4()
431 }
432
433 pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
439 match self {
440 Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
441 _ => None,
442 }
443 }
444
445 pub(crate) fn graph_weight_or_blr(
449 &self,
450 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
451 self.graph_weight()
452 .or_else(|| self.graph_weight_descriptor())
453 .or_else(|| {
454 self.graph_blr()
455 .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
456 })
457 }
458
459 pub(crate) fn is_q1(&self) -> bool {
463 matches!(
464 self,
465 Self::Mapped {
466 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
467 ..
468 }
469 )
470 }
471
472 fn bin_w(&self) -> Option<BinW<'_>> {
475 match self {
476 Self::Mapped {
477 dtype: TensorDtype::Q1,
478 ..
479 } => Some(BinW::Q1(self.quant_bytes())),
480 Self::Mapped {
481 dtype: TensorDtype::Q1G,
482 rows,
483 cols,
484 ..
485 } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
486 _ => None,
487 }
488 }
489
490 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
493 match self {
494 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
495 _ => None,
496 }
497 }
498
499 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
506 if self.has_prism_contract() {
507 return None;
508 }
509 match self {
510 #[cfg(target_os = "macos")]
511 Self::Mapped {
512 dtype: TensorDtype::Q1T,
513 ..
514 } if !crate::gpu::metal_q1t_enabled() => None,
515 Self::Mapped {
516 idx,
517 dtype:
518 TensorDtype::Q1
519 | TensorDtype::Q1G
520 | TensorDtype::Q1T
521 | TensorDtype::Q4Block
522 | TensorDtype::Q4Tiled
523 | TensorDtype::Q4TiledP
527 | TensorDtype::Q8Row
528 | TensorDtype::Q8_2f,
529 rows,
530 cols,
531 ..
532 } => Some((*idx, *rows, *cols)),
533 _ => None,
534 }
535 }
536
537 #[cfg(target_os = "macos")]
544 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
545 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
546 let name = &model.tensors[idx].name;
547 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
548 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
549 if forward && affine {
550 let e = model.tensors.get(idx)?;
551 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
552 }
553 }
554 self.q1_parts()
555 }
556
557 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
563 if self.has_prism_contract() {
564 return None;
565 }
566 match self {
567 Self::Mapped {
568 idx,
569 dtype: TensorDtype::Q4Tiled,
570 rows,
571 cols,
572 ..
573 } => Some((*idx, *rows, *cols)),
574 _ => None,
575 }
576 }
577
578 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
582 if self.has_prism_contract() {
583 return None;
584 }
585 match self {
586 Self::Mapped {
587 idx,
588 dtype: TensorDtype::Q4TiledP,
589 rows,
590 cols,
591 ..
592 } => Some((*idx, *rows, *cols)),
593 _ => None,
594 }
595 }
596
597 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
602 if self.has_prism_contract() {
603 return None;
604 }
605 match self {
606 Self::Mapped {
607 idx,
608 dtype: TensorDtype::Q8Row,
609 rows,
610 cols,
611 row_scale,
612 col_field,
613 ..
614 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
615 _ => None,
616 }
617 }
618
619 pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
622 if self.has_prism_contract() {
623 return None;
624 }
625 match self {
626 Self::Mapped {
627 idx,
628 dtype: TensorDtype::Q8_2f,
629 rows,
630 cols,
631 row_scale,
632 col_field,
633 ..
634 } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
635 _ => None,
636 }
637 }
638
639 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
643 match self {
644 Self::Mapped { dtype, .. } => Some(*dtype),
645 _ => None,
646 }
647 }
648
649 pub fn model_idx(&self) -> Option<usize> {
654 match self {
655 Self::Mapped { idx, .. } => Some(*idx),
656 _ => None,
657 }
658 }
659
660 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
665 match self {
666 Self::Mapped { model, .. } => Some(model.clone()),
667 _ => None,
668 }
669 }
670
671 pub fn row_spans(&self, r: usize, out: &mut Vec<(usize, usize)>) -> bool {
677 let Self::Mapped { dtype, .. } = self else {
678 return false;
679 };
680 let Some(spans) = row_span_offsets(*dtype, self.rows(), self.cols(), r) else {
681 return false;
682 };
683 let base = self.quant_bytes().as_ptr() as usize;
684 out.extend(spans.iter().filter(|s| s.1 > 0).map(|&(o, n)| (base + o, n)));
685 true
686 }
687
688 pub(crate) fn has_prism_contract(&self) -> bool {
693 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
694 }
695
696 pub fn rows(&self) -> usize {
697 match self {
698 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
699 Self::Blr(t) => t.rows,
700 }
701 }
702
703 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
706 if self.has_prism_contract() {
707 return None;
708 }
709 match self {
710 Self::Mapped {
711 model,
712 idx,
713 dtype: TensorDtype::Q4Tiled,
714 ..
715 } => Some((model, *idx)),
716 _ => None,
717 }
718 }
719
720 pub fn mapped_q4tp(&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,
731 ..
732 } => Some((model, *idx)),
733 _ => None,
734 }
735 }
736
737 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
745 if self.has_prism_contract() {
746 return None;
747 }
748 match self {
749 Self::Mapped {
750 model,
751 idx,
752 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
753 ..
754 } => Some((model, *idx)),
755 _ => None,
756 }
757 }
758
759 pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
762 if self.has_prism_contract() {
763 return None;
764 }
765 match self {
766 Self::Mapped {
767 model,
768 idx,
769 dtype: TensorDtype::Q8_2f,
770 ..
771 } => Some((model, *idx)),
772 _ => None,
773 }
774 }
775
776 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
779 if self.has_prism_contract() {
780 return None;
781 }
782 match self {
783 Self::Mapped {
784 model,
785 idx,
786 dtype: TensorDtype::Q2TiledP,
787 ..
788 } => Some((model, *idx)),
789 _ => None,
790 }
791 }
792
793 pub fn cols(&self) -> usize {
794 match self {
795 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
796 Self::Blr(t) => t.cols,
797 }
798 }
799
800 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
803 if self.has_prism_contract() {
804 return None;
805 }
806 match self {
807 Self::Mapped {
808 model,
809 idx,
810 dtype: TensorDtype::Q1,
811 ..
812 } => Some((model, *idx)),
813 _ => None,
814 }
815 }
816
817 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
828 if self.has_prism_contract() {
829 return None;
830 }
831 self.graph_weight_descriptor()
832 }
833
834 pub(crate) fn graph_weight_descriptor(
839 &self,
840 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
841 match self {
842 Self::Mapped {
843 model,
844 idx,
845 dtype: TensorDtype::Q8Row,
846 row_scale,
847 ..
848 } => Some((model, *idx, 0, row_scale.as_slice())),
849 Self::Mapped {
850 model,
851 idx,
852 dtype: TensorDtype::Q1,
853 ..
854 } => Some((model, *idx, 1, &[])),
855 Self::Mapped {
858 model,
859 idx,
860 dtype: TensorDtype::Q1G,
861 ..
862 } => Some((model, *idx, 10, &[])),
863 Self::Mapped {
868 model,
869 idx,
870 dtype: TensorDtype::Q4Tiled,
871 ..
872 } => Some((model, *idx, 5, &[])),
873 Self::Mapped {
877 model,
878 idx,
879 dtype: TensorDtype::Q4TiledP,
880 ..
881 } => Some((model, *idx, 6, &[])),
882 Self::Mapped {
883 model,
884 idx,
885 dtype: TensorDtype::Q4Block,
886 ..
887 } => Some((model, *idx, 2, &[])),
888 Self::Mapped {
893 model,
894 idx,
895 dtype: TensorDtype::Q8_2f,
896 ..
897 } => Some((model, *idx, 7, &[])),
898 Self::Mapped {
899 model,
900 idx,
901 dtype: TensorDtype::Q1T,
902 ..
903 } => Some((model, *idx, 3, &[])),
904 Self::Mapped {
908 model,
909 idx,
910 dtype: TensorDtype::Q2TiledP,
911 ..
912 } => Some((model, *idx, 9, &[])),
913 _ => None,
914 }
915 }
916
917 pub fn q8_2f_codes(&self) -> Option<(&[u8], &[f32], &[f32])> {
921 match self {
922 Self::Mapped { dtype: TensorDtype::Q8_2f, row_scale, col_field, .. } => {
923 Some((self.quant_bytes(), row_scale, col_field))
924 }
925 _ => None,
926 }
927 }
928
929 pub fn as_f32(&self) -> Option<&[f32]> {
930 match self {
931 Self::F32 { data, .. } => Some(data),
932 Self::Mapped { .. } | Self::Blr(_) => None,
933 }
934 }
935
936 pub fn mapped_bytes(&self) -> Option<&[u8]> {
938 match self {
939 Self::Mapped { .. } => Some(self.quant_bytes()),
940 _ => None,
941 }
942 }
943
944 fn quant_bytes(&self) -> &[u8] {
945 match self {
946 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
947 Self::Blr(t) => t.bytes(),
948 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
949 }
950 }
951
952 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
954 let cols = self.cols();
955 debug_assert_eq!(dst.len(), cols);
956 match self {
957 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
958 Self::Blr(t) => t.row_f32(r, dst),
959 Self::Mapped {
960 model,
961 idx,
962 dtype,
963 row_scale,
964 col_field,
965 vbit_offsets,
966 ..
967 } => {
968 if *dtype == TensorDtype::Q4Tiled {
969 let bytes = self.quant_bytes();
970 let gpr = cols / GROUP_SIZE;
971 for gi in 0..gpr {
972 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
973 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
974 for (k, &b) in tile[2..].iter().enumerate() {
975 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
976 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
977 }
978 }
979 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
980 crate::prism::inverse_embedding(model, dst);
981 }
982 return;
983 }
984 if *dtype == TensorDtype::Q4TiledP {
985 let bytes = self.quant_bytes();
986 let gpr = cols / GROUP_SIZE;
987 let v = Q4tpView::new(bytes, self.rows(), cols);
988 let mut sc = vec![0f32; gpr];
989 v.scales_into(r, gpr, &mut sc);
990 for gi in 0..gpr {
991 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
992 let s = sc[gi];
993 for (k, &b) in tile.iter().enumerate() {
994 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
995 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
996 }
997 }
998 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
999 crate::prism::inverse_embedding(model, dst);
1000 }
1001 return;
1002 }
1003 if *dtype == TensorDtype::Q2TiledP {
1004 let bytes = self.quant_bytes();
1005 let gpr = cols / GROUP_SIZE;
1006 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
1007 let mut sc = vec![0f32; gpr];
1008 v.scales_into(r, gpr, &mut sc);
1009 for gi in 0..gpr {
1010 let ch =
1011 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
1012 let s = sc[gi];
1013 for (k, &b) in ch.iter().enumerate() {
1014 for j in 0..4 {
1015 let center = if crate::prism::is_affine_target(
1016 model,
1017 &model.tensors[*idx].name,
1018 ) {
1019 1.0
1020 } else {
1021 1.5
1022 };
1023 dst[gi * GROUP_SIZE + k * 4 + j] =
1024 (((b >> (2 * j)) & 3) as f32 - center) * s;
1025 }
1026 }
1027 }
1028 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1029 crate::prism::inverse_embedding(model, dst);
1030 }
1031 return;
1032 }
1033 if *dtype == TensorDtype::Q4Block {
1034 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
1035 let gpr = cols / GROUP_SIZE;
1036 for gi in 0..gpr {
1037 let g = r * gpr + gi;
1038 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1039 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1040 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
1041 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
1042 }
1043 }
1044 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1045 crate::prism::inverse_embedding(model, dst);
1046 }
1047 return;
1048 }
1049 if *dtype == TensorDtype::Mxfp4 {
1050 mxfp4::Mx::new(self.quant_bytes(), self.rows(), cols).row_f32(r, dst);
1051 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1052 crate::prism::inverse_embedding(model, dst);
1053 }
1054 return;
1055 }
1056 if *dtype == TensorDtype::Q1 {
1057 let bytes = self.quant_bytes();
1058 let gpr = cols / GROUP_SIZE;
1059 for gi in 0..gpr {
1060 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
1061 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1062 for (j, &b) in tile[2..].iter().enumerate() {
1063 for k in 0..8 {
1064 dst[gi * GROUP_SIZE + j * 8 + k] =
1065 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1066 }
1067 }
1068 }
1069 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1070 crate::prism::inverse_embedding(model, dst);
1071 }
1072 return;
1073 }
1074 if *dtype == TensorDtype::Q1G {
1075 let bytes = self.quant_bytes();
1078 let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
1079 let gpr = cols / GROUP_SIZE;
1080 for gi in 0..gpr {
1081 let t = r * gpr + gi;
1082 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
1083 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
1084 for k in 0..8 {
1085 dst[gi * GROUP_SIZE + j * 8 + k] =
1086 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1087 }
1088 }
1089 }
1090 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1091 crate::prism::inverse_embedding(model, dst);
1092 }
1093 return;
1094 }
1095 if *dtype == TensorDtype::Q1T {
1096 let bytes = self.quant_bytes();
1097 let gpr = cols / GROUP_SIZE;
1098 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1099 for gi in 0..gpr {
1100 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1101 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1102 bytes[off],
1103 bytes[off + 1],
1104 ]));
1105 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1106 for k in 0..GROUP_SIZE {
1107 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1108 {
1109 1 => s,
1110 2 => -s,
1111 _ => 0.0,
1112 };
1113 }
1114 }
1115 let rows = self.rows();
1117 let entries = base_len + (rows + 1) * 4;
1118 if entries <= bytes.len() {
1119 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1120 let r0 = u32::from_le_bytes([
1121 ptrs[r * 4],
1122 ptrs[r * 4 + 1],
1123 ptrs[r * 4 + 2],
1124 ptrs[r * 4 + 3],
1125 ]) as usize;
1126 let r1 = u32::from_le_bytes([
1127 ptrs[(r + 1) * 4],
1128 ptrs[(r + 1) * 4 + 1],
1129 ptrs[(r + 1) * 4 + 2],
1130 ptrs[(r + 1) * 4 + 3],
1131 ]) as usize;
1132 let off = entries + r0 * 4;
1133 for i in 0..r1 - r0 {
1134 let item = &bytes[off + i * 4..off + i * 4 + 4];
1135 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1136 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1137 item[2], item[3],
1138 ]));
1139 if c < cols {
1140 dst[c] = v;
1141 }
1142 }
1143 }
1144 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1145 crate::prism::inverse_embedding(model, dst);
1146 }
1147 return;
1148 }
1149 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1150 let bytes = self.quant_bytes();
1151 let rows = self.rows();
1152 let ng = cols / GROUP_SIZE;
1153 let bits = &bytes[..rows];
1154 let sc_off = rows;
1155 let off = vbit_offsets[r];
1158 let b = bits[r] as usize;
1159 let l = ((1usize << (b - 1)) - 1) as f32;
1160 let data = &bytes[off..];
1161 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1162 for (i, d) in dst.iter_mut().enumerate() {
1163 while nbits < b {
1164 acc = (acc << 8) | data[byte_idx] as u64;
1165 byte_idx += 1;
1166 nbits += 8;
1167 }
1168 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1169 nbits -= b;
1170 let so = (r * ng + i / GROUP_SIZE) * 2;
1171 let sv = f16_to_f32(u16::from_le_bytes([
1172 bytes[sc_off + so],
1173 bytes[sc_off + so + 1],
1174 ]));
1175 *d = (u - l) * sv;
1176 }
1177 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1178 crate::prism::inverse_embedding(model, dst);
1179 }
1180 return;
1181 }
1182 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1183 let s = row_scale[r];
1184 match dtype {
1185 TensorDtype::Q8Row => {
1186 for (d, &b) in dst.iter_mut().zip(q) {
1187 *d = (b as i8) as f32 * s;
1188 }
1189 }
1190 TensorDtype::Q8_2f => {
1191 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1192 *d = (b as i8) as f32 * s * col_field[i];
1193 }
1194 }
1195 _ => unreachable!(),
1196 }
1197 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1198 crate::prism::inverse_embedding(model, dst);
1199 }
1200 }
1201 }
1202 }
1203
1204 pub fn sparse_col_ok(&self) -> bool {
1209 match self {
1210 Self::F32 { .. } => true,
1211 Self::Blr(_) => false,
1212 Self::Mapped { dtype, .. } => {
1213 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1214 }
1215 }
1216 }
1217
1218 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1222 let inter = self.cols();
1223 let hidden = self.rows();
1224 debug_assert_eq!(out.len(), hidden);
1225 match self {
1226 Self::F32 { data, .. } => {
1227 for (k, o) in out.iter_mut().enumerate() {
1228 *o += w * data[k * inter + c];
1229 }
1230 }
1231 Self::Blr(t) => t.add_col_scaled(c, w, out),
1232 Self::Mapped {
1233 dtype,
1234 row_scale,
1235 col_field,
1236 ..
1237 } => {
1238 let q = self.quant_bytes();
1239 let colf = if *dtype == TensorDtype::Q8_2f {
1240 col_field[c]
1241 } else {
1242 1.0
1243 };
1244 let wc = w * colf;
1245 for (k, o) in out.iter_mut().enumerate() {
1246 let b = q[k * inter + c] as i8 as f32;
1247 *o += wc * b * row_scale[k];
1248 }
1249 }
1250 }
1251 }
1252
1253 #[inline]
1262 pub fn prefetch_row(&self, r: usize) {
1263 let Self::Mapped { dtype, .. } = self else {
1264 return;
1265 };
1266 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1267 return;
1268 }
1269 let cols = self.cols();
1270 let q = self.quant_bytes();
1271 let (a, b) = (r * cols, (r + 1) * cols);
1272 if b > q.len() {
1273 return;
1274 }
1275 let mut j = a;
1276 while j < b {
1277 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1278 j += 512;
1279 }
1280 }
1281
1282 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1291 let cols = self.cols();
1292 debug_assert_eq!(out.len(), cols);
1293 match self {
1294 Self::F32 { data, .. } => {
1295 let row = &data[r * cols..(r + 1) * cols];
1296 for (o, v) in out.iter_mut().zip(row) {
1297 *o += w * v;
1298 }
1299 }
1300 Self::Blr(t) => {
1301 let tmp = &mut scratch[..cols];
1302 t.row_f32(r, tmp);
1303 for (o, v) in out.iter_mut().zip(tmp.iter()) {
1304 *o += w * v;
1305 }
1306 }
1307 Self::Mapped {
1308 dtype,
1309 row_scale,
1310 col_field,
1311 ..
1312 } => match dtype {
1313 TensorDtype::Q8Row => {
1314 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1315 let ws = w * row_scale[r];
1316 let row: &[i8] =
1317 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1318 axpy_i8_f32(out, row, ws);
1319 }
1320 TensorDtype::Q8_2f => {
1321 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1322 let ws = w * row_scale[r];
1323 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1324 *o += ws * c * (*b as i8 as f32);
1325 }
1326 }
1327 _ => {
1328 self.row_f32(r, scratch);
1329 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1330 *o += w * v;
1331 }
1332 }
1333 },
1334 }
1335 }
1336
1337 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1341 let cols = self.cols();
1342 match self {
1343 Self::F32 { data, .. } => {
1344 let row = &data[r * cols..(r + 1) * cols];
1345 row.iter().zip(x).map(|(w, v)| w * v).sum()
1346 }
1347 Self::Blr(t) => {
1348 let tmp = &mut scratch[..cols];
1349 t.row_f32(r, tmp);
1350 tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1351 }
1352 Self::Mapped {
1353 model,
1354 idx,
1355 dtype,
1356 row_scale,
1357 col_field,
1358 ..
1359 } => {
1360 let prism_forward =
1361 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1362 if prism_forward {
1363 let transformed = crate::prism::forward(model, &x[..cols]);
1364 let gpr = cols / GROUP_SIZE;
1365 match dtype {
1366 TensorDtype::Q2TiledP => {
1367 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1368 let mut sc = vec![0f32; gpr];
1369 v.scales_into(r, gpr, &mut sc);
1370 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1371 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1372 }
1373 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1374 }
1375 _ => {
1376 self.row_f32(r, scratch);
1377 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1378 }
1379 }
1380 }
1381 match dtype {
1382 TensorDtype::Q8Row => {
1383 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1384 dot_i8_f32(q, x) * row_scale[r]
1385 }
1386 TensorDtype::Q8_2f => {
1387 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1388 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1389 }
1390 _ => {
1391 self.row_f32(r, scratch);
1392 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1393 }
1394 }
1395 }
1396 }
1397 }
1398
1399 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1402 if self.exact_cpu_default() {
1403 return float_activations_scope(|| self.matvec(x, out, pool));
1404 }
1405 match self {
1406 Self::F32 { data, .. } => {
1410 if !crate::f32_backend::matvec(data, x, out) {
1411 matvec_rows(pool, data, x, out);
1412 }
1413 }
1414 Self::Blr(t) => {
1415 if crate::gpu::enabled_here()
1416 && crate::blr::gpu_wanted()
1417 && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1418 {
1419 return;
1420 }
1421 t.matvec(x, out, pool)
1422 }
1423 Self::Mapped {
1424 model,
1425 idx,
1426 dtype,
1427 rows,
1428 cols,
1429 row_scale,
1430 col_field,
1431 vbit_offsets,
1432 repack,
1433 } => {
1434 let _ = (model, idx);
1435 assert!(
1445 out.len() >= *rows && x.len() >= *cols,
1446 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1447 out.len(),
1448 x.len(),
1449 );
1450 let prism_forward =
1451 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1452 if *dtype == TensorDtype::Q2TiledP
1453 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1454 {
1455 use std::sync::atomic::{AtomicUsize, Ordering};
1456 static N: AtomicUsize = AtomicUsize::new(0);
1457 let n = N.fetch_add(1, Ordering::Relaxed);
1458 if n < 128 {
1459 eprintln!(
1460 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1461 model.tensors[*idx].name,
1462 prism_forward,
1463 rows,
1464 cols,
1465 crate::gpu::enabled_here(),
1466 crate::gpu::q2tp_gpu_opt_in(),
1467 crate::gpu::cur_layer(),
1468 );
1469 }
1470 }
1471 if prism_forward {
1476 let transformed = crate::prism::forward(model, &x[..*cols]);
1477 match dtype {
1478 TensorDtype::Q4Block => {
1479 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1480 }
1481 TensorDtype::Q4Tiled => {
1482 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1483 }
1484 TensorDtype::Q4TiledP => {
1485 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1486 }
1487 TensorDtype::Q2TiledP => {
1488 let affine =
1489 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1490 if *rows * *cols >= 8_388_608
1491 && crate::gpu::enabled_here()
1492 && crate::gpu::q2tp_gpu_opt_in()
1493 {
1494 let gpu_ok = if affine {
1495 crate::gpu::q2tp_affine_matvec(
1496 model,
1497 *idx,
1498 &transformed,
1499 *rows,
1500 *cols,
1501 out,
1502 )
1503 } else {
1504 crate::gpu::q2tp_matvec(
1505 model,
1506 *idx,
1507 &transformed,
1508 *rows,
1509 *cols,
1510 out,
1511 )
1512 };
1513 if gpu_ok {
1514 return;
1515 }
1516 }
1517 if affine {
1518 q2tp_affine_matvec(
1519 self.quant_bytes(),
1520 &transformed,
1521 *rows,
1522 *cols,
1523 out,
1524 pool,
1525 )
1526 } else {
1527 q2tp_matvec(
1528 self.quant_bytes(),
1529 &transformed,
1530 *rows,
1531 *cols,
1532 out,
1533 pool,
1534 )
1535 }
1536 }
1537 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1538 self.bin_w().expect("binary dtype"),
1539 &transformed,
1540 *rows,
1541 *cols,
1542 out,
1543 pool,
1544 ),
1545 TensorDtype::Q1T => {
1546 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1547 }
1548 TensorDtype::Mxfp4 => {
1549 mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1550 }
1551 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1552 self.quant_bytes(),
1553 vbit_offsets,
1554 &transformed,
1555 *rows,
1556 *cols,
1557 out,
1558 pool,
1559 ),
1560 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1561 self.quant_bytes(),
1562 repack,
1563 row_scale,
1564 &transformed,
1565 col_field,
1566 *dtype,
1567 *rows,
1568 *cols,
1569 out,
1570 pool,
1571 ),
1572 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1573 }
1574 return;
1575 }
1576 if *dtype == TensorDtype::Q4Block {
1577 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1581 let t0 = std::time::Instant::now();
1582 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1583 crate::gpu::ProbeArm::Gpu => {
1584 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1585 crate::gpu::probe_record(
1586 crate::gpu::OpClass::Matvec,
1587 true,
1588 t0.elapsed(),
1589 );
1590 return;
1591 }
1592 }
1593 crate::gpu::ProbeArm::CpuTimed => {
1594 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1595 crate::gpu::probe_record(
1596 crate::gpu::OpClass::Matvec,
1597 false,
1598 t0.elapsed(),
1599 );
1600 return;
1601 }
1602 crate::gpu::ProbeArm::Cpu => {}
1603 }
1604 }
1605 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1606 return;
1607 }
1608 if *dtype == TensorDtype::Q4Tiled {
1609 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1614 let t0 = std::time::Instant::now();
1615 let cls = crate::gpu::matvec_class(*rows, *cols);
1616 match crate::gpu::probe_arm(cls) {
1617 crate::gpu::ProbeArm::Gpu => {
1618 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1619 crate::gpu::probe_record(cls, true, t0.elapsed());
1620 return;
1621 }
1622 }
1623 crate::gpu::ProbeArm::CpuTimed => {
1624 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1625 crate::gpu::probe_record(cls, false, t0.elapsed());
1626 return;
1627 }
1628 crate::gpu::ProbeArm::Cpu => {}
1629 }
1630 }
1631 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1632 return;
1633 }
1634 if *dtype == TensorDtype::Q4TiledP {
1635 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1641 let t0 = std::time::Instant::now();
1642 let cls = crate::gpu::matvec_class(*rows, *cols);
1643 match crate::gpu::probe_arm(cls) {
1644 crate::gpu::ProbeArm::Gpu => {
1645 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1646 crate::gpu::probe_record(cls, true, t0.elapsed());
1647 return;
1648 }
1649 }
1650 crate::gpu::ProbeArm::CpuTimed => {
1651 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1652 crate::gpu::probe_record(cls, false, t0.elapsed());
1653 return;
1654 }
1655 crate::gpu::ProbeArm::Cpu => {}
1656 }
1657 }
1658 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1659 return;
1660 }
1661 if *dtype == TensorDtype::Q2TiledP {
1662 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1663 return;
1664 }
1665 if *dtype == TensorDtype::Mxfp4 {
1666 self.mxfp4_matvec(x, out, pool);
1667 return;
1668 }
1669 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1670 let w = self.bin_w().expect("binary dtype");
1675 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1676 let t0 = std::time::Instant::now();
1677 let arm = if crate::gpu::q1_force() {
1678 crate::gpu::ProbeArm::Gpu
1679 } else {
1680 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1681 };
1682 match arm {
1683 crate::gpu::ProbeArm::Gpu => {
1684 let ok = if *dtype == TensorDtype::Q1G {
1685 crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1686 } else {
1687 crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1688 };
1689 if ok {
1690 crate::gpu::probe_record(
1691 crate::gpu::OpClass::Matvec,
1692 true,
1693 t0.elapsed(),
1694 );
1695 return;
1696 }
1697 }
1698 crate::gpu::ProbeArm::CpuTimed => {
1699 bin_matvec(w, x, *rows, *cols, out, pool);
1700 crate::gpu::probe_record(
1701 crate::gpu::OpClass::Matvec,
1702 false,
1703 t0.elapsed(),
1704 );
1705 return;
1706 }
1707 crate::gpu::ProbeArm::Cpu => {}
1708 }
1709 }
1710 bin_matvec(w, x, *rows, *cols, out, pool);
1711 return;
1712 }
1713 if *dtype == TensorDtype::Q1T {
1714 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1718 let t0 = std::time::Instant::now();
1719 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1720 crate::gpu::ProbeArm::Gpu => {
1721 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1722 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1723 crate::gpu::probe_record(
1724 crate::gpu::OpClass::Matvec,
1725 true,
1726 t0.elapsed(),
1727 );
1728 return;
1729 }
1730 }
1731 crate::gpu::ProbeArm::CpuTimed => {
1732 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1733 crate::gpu::probe_record(
1734 crate::gpu::OpClass::Matvec,
1735 false,
1736 t0.elapsed(),
1737 );
1738 return;
1739 }
1740 crate::gpu::ProbeArm::Cpu => {}
1741 }
1742 }
1743 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1744 return;
1745 }
1746 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1747 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1748 return;
1749 }
1750 let xs = prescale(x, col_field, *dtype);
1751 if *rows >= crate::gpu::min_rows()
1756 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1757 && gpu_lmhead_enabled()
1758 && crate::gpu::enabled_here()
1759 {
1760 let t0 = std::time::Instant::now();
1763 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1764 crate::gpu::ProbeArm::Gpu => {}
1765 crate::gpu::ProbeArm::CpuTimed => {
1766 qmatvec(
1767 self.quant_bytes(),
1768 repack,
1769 row_scale,
1770 x,
1771 col_field,
1772 *dtype,
1773 *rows,
1774 *cols,
1775 out,
1776 pool,
1777 );
1778 crate::gpu::probe_record(
1779 crate::gpu::OpClass::Matvec,
1780 false,
1781 t0.elapsed(),
1782 );
1783 return;
1784 }
1785 crate::gpu::ProbeArm::Cpu => {
1786 qmatvec(
1787 self.quant_bytes(),
1788 repack,
1789 row_scale,
1790 x,
1791 col_field,
1792 *dtype,
1793 *rows,
1794 *cols,
1795 out,
1796 pool,
1797 );
1798 return;
1799 }
1800 }
1801 let frac = gpu_split_frac();
1802 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1803 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1804 let bytes = self.quant_bytes();
1805 let ok = std::thread::scope(|sc| {
1806 let g = sc.spawn(|| {
1807 crate::gpu::q8_matvec_range(
1808 model,
1809 *idx,
1810 cpu_rows,
1811 &row_scale[cpu_rows..],
1812 &xs,
1813 *rows - cpu_rows,
1814 *cols,
1815 out_gpu,
1816 )
1817 });
1818 if cpu_rows > 0 {
1819 let rep_cpu = if repack.is_empty() {
1822 &[][..]
1823 } else {
1824 &repack[..(cpu_rows / 4) * 4 * *cols]
1825 };
1826 qmatvec(
1827 &bytes[..cpu_rows * *cols],
1828 rep_cpu,
1829 &row_scale[..cpu_rows],
1830 x,
1831 col_field,
1832 *dtype,
1833 cpu_rows,
1834 *cols,
1835 out_cpu,
1836 pool,
1837 );
1838 }
1839 g.join().unwrap_or(false)
1840 });
1841 if ok {
1842 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1843 return;
1844 }
1845 qmatvec(
1848 &bytes[cpu_rows * *cols..(*rows) * *cols],
1849 &[],
1850 &row_scale[cpu_rows..],
1851 x,
1852 col_field,
1853 *dtype,
1854 *rows - cpu_rows,
1855 *cols,
1856 out_gpu,
1857 pool,
1858 );
1859 return;
1860 }
1861 qmatvec(
1862 self.quant_bytes(),
1863 repack,
1864 row_scale,
1865 x,
1866 col_field,
1867 *dtype,
1868 *rows,
1869 *cols,
1870 out,
1871 pool,
1872 );
1873 }
1874 }
1875 }
1876
1877 pub fn matvec2(
1879 &self,
1880 x1: &[f32],
1881 x2: &[f32],
1882 o1: &mut [f32],
1883 o2: &mut [f32],
1884 pool: Option<&Pool>,
1885 ) {
1886 if self.exact_cpu_default() {
1887 return float_activations_scope(|| self.matvec2(x1, x2, o1, o2, pool));
1888 }
1889 match self {
1890 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1891 Self::Blr(_) => {
1892 self.matvec(x1, o1, pool);
1893 self.matvec(x2, o2, pool);
1894 }
1895 Self::Mapped {
1896 model,
1897 idx,
1898 dtype,
1899 rows,
1900 cols,
1901 row_scale,
1902 col_field,
1903 vbit_offsets,
1904 ..
1905 } => {
1906 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1907 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1908 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1909 match dtype {
1910 TensorDtype::Q4Block => {
1911 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1912 }
1913 TensorDtype::Q4Tiled => {
1914 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1915 }
1916 TensorDtype::Q4TiledP => {
1917 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1918 }
1919 TensorDtype::Q2TiledP => {
1920 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1921 q2tp_affine_matvec2(
1922 self.quant_bytes(),
1923 &tx1,
1924 &tx2,
1925 *rows,
1926 *cols,
1927 o1,
1928 o2,
1929 pool,
1930 )
1931 } else {
1932 q2tp_matvec2(
1933 self.quant_bytes(),
1934 &tx1,
1935 &tx2,
1936 *rows,
1937 *cols,
1938 o1,
1939 o2,
1940 pool,
1941 )
1942 }
1943 }
1944 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1945 self.bin_w().expect("binary dtype"),
1946 &tx1,
1947 &tx2,
1948 *rows,
1949 *cols,
1950 o1,
1951 o2,
1952 pool,
1953 ),
1954 TensorDtype::Q1T => {
1955 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1956 }
1957 TensorDtype::Mxfp4 => mxfp4::matvec2(
1958 self.quant_bytes(),
1959 &tx1,
1960 &tx2,
1961 *rows,
1962 *cols,
1963 o1,
1964 o2,
1965 pool,
1966 ),
1967 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1968 self.quant_bytes(),
1969 vbit_offsets,
1970 &tx1,
1971 &tx2,
1972 *rows,
1973 *cols,
1974 o1,
1975 o2,
1976 pool,
1977 ),
1978 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1979 self.quant_bytes(),
1980 row_scale,
1981 &tx1,
1982 &tx2,
1983 col_field,
1984 *dtype,
1985 *rows,
1986 *cols,
1987 o1,
1988 o2,
1989 pool,
1990 ),
1991 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1992 }
1993 return;
1994 }
1995 if *dtype == TensorDtype::Q4Block {
1996 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1997 return;
1998 }
1999 if *dtype == TensorDtype::Q4Tiled {
2000 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2001 return;
2002 }
2003 if *dtype == TensorDtype::Q4TiledP {
2004 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2005 return;
2006 }
2007 if *dtype == TensorDtype::Q2TiledP {
2008 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2009 return;
2010 }
2011 if *dtype == TensorDtype::Mxfp4 {
2012 mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2013 return;
2014 }
2015 if let Some(w) = self.bin_w() {
2016 bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
2017 return;
2018 }
2019 if *dtype == TensorDtype::Q1T {
2020 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2026 return;
2027 }
2028 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2029 vbitmatvec2(
2030 self.quant_bytes(),
2031 vbit_offsets,
2032 x1,
2033 x2,
2034 *rows,
2035 *cols,
2036 o1,
2037 o2,
2038 pool,
2039 );
2040 return;
2041 }
2042 qmatvec2(
2043 self.quant_bytes(),
2044 row_scale,
2045 x1,
2046 x2,
2047 col_field,
2048 *dtype,
2049 *rows,
2050 *cols,
2051 o1,
2052 o2,
2053 pool,
2054 );
2055 }
2056 }
2057 }
2058}
2059
2060impl QTensor {
2061 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2069 if self.has_prism_contract() {
2070 return None;
2071 }
2072 match self {
2073 Self::Mapped {
2074 model, idx, dtype, ..
2075 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2076 _ => None,
2077 }
2078 }
2079
2080 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2081 if self.exact_cpu_default() {
2082 return float_activations_scope(|| self.matmat(xs_all, b, out, pool));
2083 }
2084 let cols = self.cols();
2085 let rows = self.rows();
2086 debug_assert_eq!(xs_all.len(), b * cols);
2087 debug_assert_eq!(out.len(), b * rows);
2088 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2089 if crate::gptq_capture::capturing() {
2093 if let Self::Mapped { model, idx, .. } = self {
2094 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2095 }
2096 if let Self::Blr(t) = self {
2097 crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2098 }
2099 }
2100 match self {
2101 Self::Blr(t) => {
2102 if crate::gpu::enabled_here()
2103 && crate::blr::gpu_wanted()
2104 && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2105 {
2106 return;
2107 }
2108 t.matmat(xs_all, b, out, pool)
2109 }
2110 Self::F32 { data, .. } => {
2111 if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2112 return;
2113 }
2114 let out_addr = SendMut(out.as_mut_ptr());
2115 let kernel = |start: usize, end: usize, r0: usize, r1: usize| {
2122 let mut bi = start;
2123 while bi < end {
2124 let n = (end - bi).min(4);
2125 for o in r0..r1 {
2126 let row = &data[o * cols..(o + 1) * cols];
2127 if n == 4 {
2128 let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2129 let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2130 let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2131 for j in 0..cols {
2132 let w = row[j];
2133 a0 += w * x0[j];
2134 a1 += w * x1[j];
2135 a2 += w * x2[j];
2136 a3 += w * x3[j];
2137 }
2138 for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2139 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2140 }
2141 } else {
2142 for k in 0..n {
2143 let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2144 let mut acc = 0f32;
2145 for j in 0..cols {
2146 acc += row[j] * x[j];
2147 }
2148 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2149 }
2150 }
2151 }
2152 bi += n;
2153 }
2154 };
2155 match pool {
2156 Some(pool) if b < 256 && rows >= 256 => {
2162 pool.run_rows(rows, &|r0, r1| kernel(0, b, r0, r1))
2163 }
2164 _ => dispatch_rows(pool, b, &|s, e| kernel(s, e, 0, rows)),
2165 }
2166 }
2167 Self::Mapped {
2168 model,
2169 idx,
2170 dtype,
2171 row_scale,
2172 col_field,
2173 vbit_offsets,
2174 ..
2175 } => {
2176 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2177 let mut transformed = Vec::with_capacity(xs_all.len());
2178 for bi in 0..b {
2179 transformed.extend_from_slice(&crate::prism::forward(
2180 model,
2181 &xs_all[bi * cols..(bi + 1) * cols],
2182 ));
2183 }
2184 match dtype {
2185 TensorDtype::Q4Block => {
2186 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2187 }
2188 TensorDtype::Q4Tiled => {
2189 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2190 }
2191 TensorDtype::Q4TiledP => {
2192 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2193 }
2194 TensorDtype::Q2TiledP => {
2195 let affine =
2196 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2197 let gpu_batch_ok = if affine {
2203 b >= 2
2204 } else {
2205 b >= 32 && b * rows * cols >= 128_000_000
2206 };
2207 if gpu_batch_ok
2208 && cols % 32 == 0
2209 && crate::gpu::enabled_here()
2210 && crate::gpu::q2tp_gpu_opt_in()
2211 {
2212 let gpu_ok = if affine {
2213 crate::gpu::q2tp_affine_matmat(
2214 model,
2215 *idx,
2216 &transformed,
2217 b,
2218 rows,
2219 cols,
2220 out,
2221 )
2222 } else {
2223 crate::gpu::q2tp_matmat(
2224 model,
2225 *idx,
2226 &transformed,
2227 b,
2228 rows,
2229 cols,
2230 out,
2231 )
2232 };
2233 if gpu_ok {
2234 return;
2235 }
2236 }
2237 if affine
2241 && b == 1
2242 && cols % 32 == 0
2243 && crate::gpu::enabled_here()
2244 && crate::gpu::q2tp_gpu_opt_in()
2245 && crate::gpu::q2tp_affine_matvec(
2246 model,
2247 *idx,
2248 &transformed[..cols],
2249 rows,
2250 cols,
2251 &mut out[..rows],
2252 )
2253 {
2254 return;
2255 }
2256 if affine {
2257 q2tp_affine_matmat(
2258 self.quant_bytes(),
2259 &transformed,
2260 b,
2261 rows,
2262 cols,
2263 out,
2264 pool,
2265 )
2266 } else {
2267 q2tp_matmat(
2268 self.quant_bytes(),
2269 &transformed,
2270 b,
2271 rows,
2272 cols,
2273 out,
2274 pool,
2275 )
2276 }
2277 }
2278 TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2279 self.bin_w().expect("binary dtype"),
2280 &transformed,
2281 b,
2282 rows,
2283 cols,
2284 out,
2285 pool,
2286 ),
2287 TensorDtype::Q1T => {
2288 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2289 }
2290 TensorDtype::Mxfp4 => {
2291 mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2292 }
2293 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2294 self.quant_bytes(),
2295 vbit_offsets,
2296 &transformed,
2297 b,
2298 rows,
2299 cols,
2300 out,
2301 pool,
2302 ),
2303 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2304 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2305 .map(|bi| {
2306 prescale(
2307 &transformed[bi * cols..(bi + 1) * cols],
2308 col_field,
2309 *dtype,
2310 )
2311 })
2312 .collect();
2313 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2314 }
2315 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2316 }
2317 return;
2318 }
2319 if *dtype == TensorDtype::Q4Block {
2320 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2321 return;
2322 }
2323 if *dtype == TensorDtype::Mxfp4 {
2324 self.mxfp4_matmat(xs_all, b, out, pool);
2325 return;
2326 }
2327 if *dtype == TensorDtype::Q4TiledP {
2328 if crate::gpu::cuda_active()
2349 && b >= 2
2350 && rows * cols >= (1 << 20)
2351 && cols % 32 == 0
2352 && !row_exact()
2353 && crate::gpu::enabled_here()
2354 {
2355 if let Self::Mapped { model, idx, .. } = self {
2356 if crate::gpu::q4tp_matmat(model, *idx, xs_all, b, rows, cols, out) {
2357 return;
2358 }
2359 }
2360 }
2361 if b >= 32
2362 && b * rows * cols >= 128_000_000
2363 && cols % 32 == 0
2364 && !row_exact()
2365 && !crate::gpu::mm_killed()
2366 && crate::gpu::enabled_here()
2367 {
2368 let class = if b >= 128 {
2369 crate::gpu::OpClass::MatmatWide
2370 } else {
2371 crate::gpu::OpClass::Matmat
2372 };
2373 if let Self::Mapped { model, idx, .. } = self {
2374 if crate::mm_ab::on() {
2386 let mut g = vec![0f32; b * rows];
2387 let t = std::time::Instant::now();
2388 let took = crate::gpu::q4tp_matmat(
2389 model, *idx, xs_all, b, rows, cols, &mut g,
2390 );
2391 let dg = t.elapsed();
2392 let t = std::time::Instant::now();
2393 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2394 let dc = t.elapsed();
2395 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2396 return;
2397 }
2398 let t0 = std::time::Instant::now();
2399 let resident = crate::gpu::weight_is_resident(model, *idx);
2403 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2404 crate::gpu::ProbeArm::Gpu => {
2405 if crate::gpu::q4tp_matmat(
2406 model, *idx, xs_all, b, rows, cols, out,
2407 ) {
2408 let el = t0.elapsed();
2409 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2419 let budget = std::time::Duration::from_secs_f64(
2420 flops / 1.5e12 * 8.0 + 0.020,
2421 );
2422 crate::gpu::mm_budget_check(
2423 "q4tp matmat",
2424 el,
2425 budget,
2426 crate::gpu::probe_was_cold() || !resident,
2427 );
2428 crate::gpu::probe_record(class, true, el);
2429 return;
2430 }
2431 }
2432 crate::gpu::ProbeArm::CpuTimed => {
2433 q4tp_matmat(
2434 self.quant_bytes(),
2435 xs_all,
2436 b,
2437 rows,
2438 cols,
2439 out,
2440 pool,
2441 );
2442 crate::gpu::probe_record(class, false, t0.elapsed());
2443 return;
2444 }
2445 crate::gpu::ProbeArm::Cpu => {}
2446 }
2447 }
2448 }
2449 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2450 return;
2451 }
2452 if *dtype == TensorDtype::Q2TiledP {
2453 if b >= 32
2459 && b * rows * cols >= 128_000_000
2460 && cols % 32 == 0
2461 && !crate::gpu::mm_killed()
2462 && crate::gpu::enabled_here()
2463 {
2464 let class = if b >= 128 {
2465 crate::gpu::OpClass::MatmatWide
2466 } else {
2467 crate::gpu::OpClass::Matmat
2468 };
2469 if let Self::Mapped { model, idx, .. } = self {
2470 let t0 = std::time::Instant::now();
2471 match crate::gpu::probe_arm(class) {
2472 crate::gpu::ProbeArm::Gpu => {
2473 if crate::gpu::q2tp_matmat(
2474 model, *idx, xs_all, b, rows, cols, out,
2475 ) {
2476 crate::gpu::probe_record(class, true, t0.elapsed());
2477 return;
2478 }
2479 }
2480 crate::gpu::ProbeArm::CpuTimed => {
2481 q2tp_matmat(
2482 self.quant_bytes(),
2483 xs_all,
2484 b,
2485 rows,
2486 cols,
2487 out,
2488 pool,
2489 );
2490 crate::gpu::probe_record(class, false, t0.elapsed());
2491 return;
2492 }
2493 crate::gpu::ProbeArm::Cpu => {}
2494 }
2495 }
2496 }
2497 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2502 return;
2503 }
2504 if *dtype == TensorDtype::Q4Tiled {
2505 if b >= 32
2517 && b * rows * cols >= 128_000_000
2518 && cols % 32 == 0
2519 && !crate::gpu::mm_killed()
2520 && crate::gpu::enabled_here()
2521 {
2522 let class = if b >= 128 {
2523 crate::gpu::OpClass::MatmatWide
2524 } else {
2525 crate::gpu::OpClass::Matmat
2526 };
2527 if let Self::Mapped { model, idx, .. } = self {
2528 let t0 = std::time::Instant::now();
2529 match crate::gpu::probe_arm(class) {
2530 crate::gpu::ProbeArm::Gpu => {
2531 if crate::gpu::q4t_matmat(
2532 model, *idx, xs_all, b, rows, cols, out,
2533 ) {
2534 let el = t0.elapsed();
2535 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2545 let budget = std::time::Duration::from_secs_f64(
2546 flops / 1.5e12 * 8.0 + 0.020,
2547 );
2548 crate::gpu::mm_budget_check(
2549 "q4t matmat",
2550 el,
2551 budget,
2552 crate::gpu::probe_was_cold(),
2553 );
2554 crate::gpu::probe_record(class, true, el);
2555 return;
2556 }
2557 }
2558 crate::gpu::ProbeArm::CpuTimed => {
2559 q4t_matmat(
2560 self.quant_bytes(),
2561 xs_all,
2562 b,
2563 rows,
2564 cols,
2565 out,
2566 pool,
2567 );
2568 crate::gpu::probe_record(class, false, t0.elapsed());
2569 return;
2570 }
2571 crate::gpu::ProbeArm::Cpu => {}
2572 }
2573 }
2574 }
2575 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2576 return;
2577 }
2578 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2579 let w = self.bin_w().expect("binary dtype");
2583 if b >= 32
2584 && b * rows * cols >= 128_000_000
2585 && cols % 64 == 0
2586 && crate::gpu::enabled_here()
2587 {
2588 if let Self::Mapped { model, idx, .. } = self {
2589 let t0 = std::time::Instant::now();
2590 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2591 crate::gpu::ProbeArm::Gpu => {
2592 let ok = if *dtype == TensorDtype::Q1G {
2593 crate::gpu::q1g_matmat(
2594 model, *idx, xs_all, b, rows, cols, out,
2595 )
2596 } else {
2597 crate::gpu::q1_matmat(
2598 model, *idx, xs_all, b, rows, cols, out,
2599 )
2600 };
2601 if ok {
2602 crate::gpu::probe_record(
2603 crate::gpu::OpClass::Matmat,
2604 true,
2605 t0.elapsed(),
2606 );
2607 return;
2608 }
2609 }
2610 crate::gpu::ProbeArm::CpuTimed => {
2611 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2612 crate::gpu::probe_record(
2613 crate::gpu::OpClass::Matmat,
2614 false,
2615 t0.elapsed(),
2616 );
2617 return;
2618 }
2619 crate::gpu::ProbeArm::Cpu => {}
2620 }
2621 }
2622 }
2623 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2624 return;
2625 }
2626 if *dtype == TensorDtype::Q1T {
2627 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2630 if let Self::Mapped { model, idx, .. } = self {
2631 let t0 = std::time::Instant::now();
2632 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2633 crate::gpu::ProbeArm::Gpu => {
2634 if crate::gpu::q1t_matmat(
2635 model, *idx, xs_all, b, rows, cols, out,
2636 ) {
2637 crate::gpu::probe_record(
2638 crate::gpu::OpClass::Matmat,
2639 true,
2640 t0.elapsed(),
2641 );
2642 return;
2643 }
2644 }
2645 crate::gpu::ProbeArm::CpuTimed => {
2646 q1t_matmat(
2647 self.quant_bytes(),
2648 xs_all,
2649 b,
2650 rows,
2651 cols,
2652 out,
2653 pool,
2654 );
2655 crate::gpu::probe_record(
2656 crate::gpu::OpClass::Matmat,
2657 false,
2658 t0.elapsed(),
2659 );
2660 return;
2661 }
2662 crate::gpu::ProbeArm::Cpu => {}
2663 }
2664 }
2665 }
2666 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2667 return;
2668 }
2669 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2670 vbitmatmat(
2671 self.quant_bytes(),
2672 vbit_offsets,
2673 xs_all,
2674 b,
2675 rows,
2676 cols,
2677 out,
2678 pool,
2679 );
2680 return;
2681 }
2682 let _pmp = crate::cpuprof::spans::span("mm.prescale");
2683 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2684 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2685 .collect();
2686 drop(_pmp);
2687 if row_exact()
2693 && (1..=4).contains(&b)
2694 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2695 && crate::gpu::enabled_here()
2696 && crate::gpu::wgpu_active()
2697 {
2698 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2699 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2700 return;
2701 }
2702 }
2703 if b >= 8
2715 && b * rows * cols >= 128_000_000
2716 && crate::gpu::enabled_here()
2717 && std::env::var("CMF_MATMAT_GPU").as_deref() != Ok("0")
2718 {
2719 if let Self::Mapped { model, idx, .. } = self {
2720 MATMAT_GATE.fetch_add(1, Ordering::Relaxed);
2721 let t0 = std::time::Instant::now();
2722 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2723 crate::gpu::ProbeArm::Gpu
2724 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2725 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2726 {
2727 let q = self.quant_bytes();
2731 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2732 return;
2733 }
2734 crate::gpu::ProbeArm::Gpu => {
2735 if *dtype == TensorDtype::Q8_2f
2743 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2744 && crate::gpu::q8_matmat_2f(
2745 model, *idx, row_scale, col_field, xs_all, b, rows,
2746 cols, out,
2747 )
2748 {
2749 MATMAT_GPU_HIT.fetch_add(1, Ordering::Relaxed);
2750 crate::gpu::probe_record(
2751 crate::gpu::OpClass::Matmat,
2752 true,
2753 t0.elapsed(),
2754 );
2755 return;
2756 }
2757 let flat: Vec<f32> =
2758 pre.iter().flat_map(|v| v.iter().copied()).collect();
2759 if crate::gpu::q8_matmat(
2760 model, *idx, row_scale, &flat, b, rows, cols, out,
2761 ) {
2762 crate::gpu::probe_record(
2763 crate::gpu::OpClass::Matmat,
2764 true,
2765 t0.elapsed(),
2766 );
2767 return;
2768 }
2769 }
2770 crate::gpu::ProbeArm::CpuTimed => {
2771 let q = self.quant_bytes();
2772 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2773 crate::gpu::probe_record(
2774 crate::gpu::OpClass::Matmat,
2775 false,
2776 t0.elapsed(),
2777 );
2778 return;
2779 }
2780 crate::gpu::ProbeArm::Cpu => {}
2781 }
2782 }
2783 }
2784 let q = self.quant_bytes();
2785 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2786 }
2787 }
2788 }
2789}
2790
2791impl QTensor {
2792 fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2796 let (rows, cols) = (self.rows(), self.cols());
2797 if let Self::Mapped { model, idx, .. } = self {
2798 if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2799 let t0 = std::time::Instant::now();
2800 let cls = crate::gpu::matvec_class(rows, cols);
2801 match crate::gpu::probe_arm(cls) {
2802 crate::gpu::ProbeArm::Gpu => {
2803 if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2804 crate::gpu::probe_record(cls, true, t0.elapsed());
2805 return;
2806 }
2807 }
2808 crate::gpu::ProbeArm::CpuTimed => {
2809 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2810 crate::gpu::probe_record(cls, false, t0.elapsed());
2811 return;
2812 }
2813 crate::gpu::ProbeArm::Cpu => {}
2814 }
2815 }
2816 }
2817 mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2818 }
2819
2820 fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2825 let (rows, cols) = (self.rows(), self.cols());
2826 if let Self::Mapped { model, idx, .. } = self {
2827 if b >= 2
2828 && b * rows * cols >= 33_554_432
2829 && !row_exact()
2830 && !crate::gpu::mm_killed()
2831 && crate::gpu::enabled_here()
2832 {
2833 let class = if b >= 128 {
2834 crate::gpu::OpClass::MatmatWide
2835 } else {
2836 crate::gpu::OpClass::Matmat
2837 };
2838 let t0 = std::time::Instant::now();
2839 match crate::gpu::probe_arm(class) {
2840 crate::gpu::ProbeArm::Gpu => {
2841 if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2842 crate::gpu::probe_record(class, true, t0.elapsed());
2843 return;
2844 }
2845 }
2846 crate::gpu::ProbeArm::CpuTimed => {
2847 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2848 crate::gpu::probe_record(class, false, t0.elapsed());
2849 return;
2850 }
2851 crate::gpu::ProbeArm::Cpu => {}
2852 }
2853 }
2854 }
2855 mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2856 }
2857}
2858
2859impl QTensor {
2860 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2870 let (rows, cols) = (self.rows(), self.cols());
2871 let Self::Mapped {
2872 model,
2873 idx,
2874 dtype,
2875 row_scale,
2876 col_field,
2877 ..
2878 } = self
2879 else {
2880 return false;
2881 };
2882 if crate::prism::has_contract(model) {
2883 return false;
2884 }
2885 match *dtype {
2886 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2887 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2891 if *dtype == TensorDtype::Q8_2f
2894 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2895 && crate::gpu::q8_matmat_2f(
2896 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2897 )
2898 {
2899 return true;
2900 }
2901 let flat: Vec<f32> = (0..b)
2902 .flat_map(|bi| {
2903 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2904 })
2905 .collect();
2906 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2907 }
2908 _ => false,
2909 }
2910 }
2911
2912 pub fn matvec_many<const N: usize>(
2919 ts: [&QTensor; N],
2920 x: &[f32],
2921 mut outs: [&mut [f32]; N],
2922 pool: Option<&Pool>,
2923 ) {
2924 if Self::any_exact_cpu_default(&ts) {
2925 return float_activations_scope(|| Self::matvec_many(ts, x, outs, pool));
2926 }
2927 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2928 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2929 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2934 t.matvec(x, o, pool);
2935 }
2936 return;
2937 }
2938 let uniform_q8 = ts.iter().all(|t| {
2939 matches!(
2940 t,
2941 Self::Mapped {
2942 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2943 ..
2944 }
2945 )
2946 });
2947 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2948 if uniform_f32 && crate::f32_backend::active() {
2949 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2950 t.matvec(x, o, pool);
2951 }
2952 return;
2953 }
2954 let uniform_q4 = ts.iter().all(|t| {
2955 matches!(
2956 t,
2957 Self::Mapped {
2958 dtype: TensorDtype::Q4Block,
2959 ..
2960 }
2961 )
2962 });
2963 let uniform_vbit = ts.iter().all(|t| {
2964 matches!(
2965 t,
2966 Self::Mapped {
2967 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2968 ..
2969 }
2970 )
2971 });
2972 let uniform_q1 =
2975 ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2976 let uniform_q1t = ts.iter().all(|t| {
2977 matches!(
2978 t,
2979 Self::Mapped {
2980 dtype: TensorDtype::Q1T,
2981 ..
2982 }
2983 )
2984 });
2985 let uniform_q4tp = ts.iter().all(|t| {
2990 matches!(
2991 t,
2992 Self::Mapped {
2993 dtype: TensorDtype::Q4TiledP,
2994 ..
2995 }
2996 )
2997 }) && ts
2998 .iter()
2999 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
3000 let Some(pool) = pool else {
3001 for (t, o) in ts.iter().zip(outs.iter_mut()) {
3002 t.matvec(x, o, None);
3003 }
3004 return;
3005 };
3006 if total_rows < 256
3007 || !(uniform_q8
3008 || uniform_f32
3009 || uniform_q4
3010 || uniform_vbit
3011 || uniform_q1
3012 || uniform_q1t
3013 || uniform_q4tp)
3014 {
3015 for (t, o) in ts.iter().zip(outs.iter_mut()) {
3016 t.matvec(x, o, Some(pool));
3017 }
3018 return;
3019 }
3020
3021 if uniform_q4tp {
3022 let cols = ts[0].cols();
3028 let gpr = cols / GROUP_SIZE;
3029 let views: [Q4tpView; N] =
3030 std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
3031 let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
3032 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3033 let locate = |flat: usize| -> (usize, usize) {
3035 let mut acc = 0;
3036 for (i, &r) in rows_of.iter().enumerate() {
3037 if flat < acc + r {
3038 return (i, flat - acc);
3039 }
3040 acc += r;
3041 }
3042 (rows_of.len() - 1, 0)
3043 };
3044 let (views, outs_addr) = (&views, &outs_addr);
3045 if a8w8_enabled() {
3046 let act = split_act(x);
3047 let act = &act;
3048 #[cfg(target_arch = "x86_64")]
3049 let lanes = q4tp_lanes();
3050 let run = |start: usize, end: usize| {
3051 #[cfg(target_arch = "x86_64")]
3052 if lanes > 1 && act.xg8.len() == gpr {
3053 with_krow(gpr * lanes, |sc| {
3056 unsafe {
3059 match lanes {
3060 2 => q4tp_many_lanes::<2>(
3061 &views[..],
3062 &locate,
3063 start,
3064 end,
3065 gpr,
3066 &act.xq,
3067 &act.xg8,
3068 act.sx,
3069 &act.outliers,
3070 sc,
3071 &outs_addr[..],
3072 ),
3073 8 => q4tp_many_lanes::<8>(
3074 &views[..],
3075 &locate,
3076 start,
3077 end,
3078 gpr,
3079 &act.xq,
3080 &act.xg8,
3081 act.sx,
3082 &act.outliers,
3083 sc,
3084 &outs_addr[..],
3085 ),
3086 _ => q4tp_many_lanes::<4>(
3087 &views[..],
3088 &locate,
3089 start,
3090 end,
3091 gpr,
3092 &act.xq,
3093 &act.xg8,
3094 act.sx,
3095 &act.outliers,
3096 sc,
3097 &outs_addr[..],
3098 ),
3099 }
3100 }
3101 });
3102 return;
3103 }
3104 crate::cpuprof::count_plain_rows((end - start) as u64);
3105 with_krow(gpr, |sc| {
3106 for flat in start..end {
3107 let (t, r) = locate(flat);
3108 let v = &views[t];
3109 v.scales_into(r, gpr, sc);
3110 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
3111 for &(j, xv) in &act.outliers {
3112 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3113 acc += w * s * xv;
3114 }
3115 unsafe { *outs_addr[t].at(r) = acc };
3117 }
3118 });
3119 };
3120 pool.run_rows(total_rows, &run);
3121 } else {
3122 let run = |start: usize, end: usize| {
3123 with_krow(gpr, |sc| {
3124 for flat in start..end {
3125 let (t, r) = locate(flat);
3126 let v = &views[t];
3127 v.scales_into(r, gpr, sc);
3128 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
3130 }
3131 });
3132 };
3133 pool.run_rows(total_rows, &run);
3134 }
3135 return;
3136 }
3137
3138 if uniform_q1 {
3139 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3142 if a8w8_enabled() {
3143 let act = split_act(x);
3144 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
3145 let (act, gsum) = (&act, &gsum);
3146 let closures: [_; N] = std::array::from_fn(|i| {
3147 let (w, gpr, out) = (
3148 ts[i].bin_w().expect("uniform_q1 checked"),
3149 ts[i].cols() / GROUP_SIZE,
3150 outs_addr[i],
3151 );
3152 move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
3153 });
3154 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3155 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3156 pool.run_many(&parts);
3157 } else {
3158 let closures: [_; N] = std::array::from_fn(|i| {
3159 let (w, gpr, out) = (
3160 ts[i].bin_w().expect("uniform_q1 checked"),
3161 ts[i].cols() / GROUP_SIZE,
3162 outs_addr[i],
3163 );
3164 move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
3165 });
3166 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3167 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3168 pool.run_many(&parts);
3169 }
3170 return;
3171 }
3172
3173 if uniform_q1t {
3174 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3178 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3179 if a8w8_enabled() {
3180 let act = split_act(x);
3181 let act = &act;
3182 let x_ref = x;
3183 let closures: [_; N] = std::array::from_fn(|i| {
3184 let bytes = ts[i].quant_bytes();
3185 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3186 let gpr = cols / GROUP_SIZE;
3187 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3188 let out = outs_addr[i];
3189 move |s: usize, e: usize| {
3190 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
3191 }
3192 });
3193 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3194 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3195 pool.run_many(&parts);
3196 } else {
3197 let x_ref = x;
3198 let closures: [_; N] = std::array::from_fn(|i| {
3199 let bytes = ts[i].quant_bytes();
3200 let (rows, cols) = (ts[i].rows(), ts[i].cols());
3201 let gpr = cols / GROUP_SIZE;
3202 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3203 let out = outs_addr[i];
3204 move |s: usize, e: usize| {
3205 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3206 }
3207 });
3208 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3209 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3210 pool.run_many(&parts);
3211 }
3212 return;
3213 }
3214
3215 if uniform_q4 || uniform_vbit {
3216 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3217 if a8w8_enabled() {
3219 let act = split_act(x);
3220 let act = &act;
3221 if uniform_q4 {
3222 let closures: [_; N] = std::array::from_fn(|i| {
3223 let (packed, scales) =
3224 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3225 let (gpr, cols, out) =
3226 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3227 move |s: usize, e: usize| {
3228 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3229 }
3230 });
3231 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3232 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3233 pool.run_many(&parts);
3234 } else {
3235 let closures: [_; N] = std::array::from_fn(|i| {
3236 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3237 unreachable!()
3238 };
3239 let (bytes, rows, cols, out) = (
3240 ts[i].quant_bytes(),
3241 ts[i].rows(),
3242 ts[i].cols(),
3243 outs_addr[i],
3244 );
3245 move |s: usize, e: usize| {
3246 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3247 }
3248 });
3249 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3250 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3251 pool.run_many(&parts);
3252 }
3253 return;
3254 }
3255 if uniform_q4 {
3256 let closures: [_; N] = std::array::from_fn(|i| {
3257 let (packed, scales) =
3258 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3259 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3260 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3261 });
3262 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3263 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3264 pool.run_many(&parts);
3265 } else {
3266 let closures: [_; N] = std::array::from_fn(|i| {
3267 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3268 unreachable!()
3269 };
3270 let (bytes, rows, cols, out) = (
3271 ts[i].quant_bytes(),
3272 ts[i].rows(),
3273 ts[i].cols(),
3274 outs_addr[i],
3275 );
3276 move |s: usize, e: usize| {
3277 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3278 }
3279 });
3280 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3281 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3282 pool.run_many(&parts);
3283 }
3284 return;
3285 }
3286
3287 if uniform_f32 {
3288 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3289 let closures: [_; N] = std::array::from_fn(|i| {
3290 let Self::F32 { data, cols, .. } = ts[i] else {
3291 unreachable!()
3292 };
3293 let out = outs_addr[i];
3294 move |start: usize, end: usize| {
3295 for o in start..end {
3296 let row = &data[o * cols..(o + 1) * cols];
3297 let mut sum = 0.0f32;
3298 for j in 0..*cols {
3299 sum += row[j] * x[j];
3300 }
3301 unsafe { *out.at(o) = sum };
3303 }
3304 }
3305 });
3306 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3307 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3308 pool.run_many(&parts);
3309 return;
3310 }
3311
3312 struct Ctx<'a> {
3315 bytes: &'a [u8],
3316 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3317 rep: &'a [u8],
3318 row_scale: &'a [f32],
3319 cols: usize,
3320 xs: std::borrow::Cow<'a, [f32]>,
3321 }
3322 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3323 let Self::Mapped {
3324 dtype,
3325 cols,
3326 row_scale,
3327 col_field,
3328 repack,
3329 ..
3330 } = ts[i]
3331 else {
3332 unreachable!()
3333 };
3334 Ctx {
3335 bytes: ts[i].quant_bytes(),
3336 rep: repack,
3337 row_scale,
3338 cols: *cols,
3339 xs: prescale(x, col_field, *dtype),
3340 }
3341 });
3342 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3343 #[cfg(target_arch = "aarch64")]
3344 if sdot_enabled() {
3345 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3346 let closures: [_; N] = std::array::from_fn(|i| {
3347 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3348 move |start: usize, end: usize| {
3349 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3350 }
3351 });
3352 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3353 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3354 pool.run_many(&parts);
3355 return;
3356 }
3357 #[cfg(target_arch = "x86_64")]
3358 if avx2_a8w8_enabled() {
3359 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3360 let closures: [_; N] = std::array::from_fn(|i| {
3361 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3362 move |start: usize, end: usize| {
3363 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3364 }
3365 });
3366 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3367 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3368 pool.run_many(&parts);
3369 return;
3370 }
3371 let closures: [_; N] = std::array::from_fn(|i| {
3372 let (c, out) = (&ctxs[i], outs_addr[i]);
3373 move |start: usize, end: usize| {
3374 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3375 }
3376 });
3377 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3378 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3379 pool.run_many(&parts);
3380 }
3381}
3382
3383impl QTensor {
3384 #[allow(clippy::needless_range_loop)]
3389 pub fn matvec2_many<const N: usize>(
3390 ts: [&QTensor; N],
3391 x1: &[f32],
3392 x2: &[f32],
3393 mut o1s: [&mut [f32]; N],
3394 mut o2s: [&mut [f32]; N],
3395 pool: Option<&Pool>,
3396 ) {
3397 if Self::any_exact_cpu_default(&ts) {
3398 return float_activations_scope(|| {
3399 Self::matvec2_many(ts, x1, x2, o1s, o2s, pool)
3400 });
3401 }
3402 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3403 if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3404 for i in 0..N {
3405 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3406 }
3407 return;
3408 }
3409 let uniform_q8 = ts.iter().all(|t| {
3410 matches!(
3411 t,
3412 Self::Mapped {
3413 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3414 ..
3415 }
3416 )
3417 });
3418 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3419 let uniform_q4 = ts.iter().all(|t| {
3420 matches!(
3421 t,
3422 Self::Mapped {
3423 dtype: TensorDtype::Q4Block,
3424 ..
3425 }
3426 )
3427 });
3428 let uniform_vbit = ts.iter().all(|t| {
3429 matches!(
3430 t,
3431 Self::Mapped {
3432 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3433 ..
3434 }
3435 )
3436 });
3437 let fusable = pool.is_some()
3438 && total_rows >= 256
3439 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3440 if !fusable {
3441 for i in 0..N {
3442 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3443 }
3444 return;
3445 }
3446 let pool = pool.unwrap();
3447
3448 if uniform_q4 || uniform_vbit {
3449 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3450 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3451 if a8w8_enabled() {
3453 let a1 = split_act(x1);
3454 let a2 = split_act(x2);
3455 let (a1, a2) = (&a1, &a2);
3456 if uniform_q4 {
3457 let closures: [_; N] = std::array::from_fn(|i| {
3458 let (packed, scales) =
3459 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3460 let (gpr, cols, o1, o2) =
3461 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3462 move |s: usize, e: usize| {
3463 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3464 }
3465 });
3466 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3467 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3468 pool.run_many(&parts);
3469 } else {
3470 let closures: [_; N] = std::array::from_fn(|i| {
3471 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3472 unreachable!()
3473 };
3474 let (bytes, rows, cols, o1, o2) = (
3475 ts[i].quant_bytes(),
3476 ts[i].rows(),
3477 ts[i].cols(),
3478 p1[i],
3479 p2[i],
3480 );
3481 move |s: usize, e: usize| {
3482 vbit_range2_a8w8(
3483 bytes,
3484 vbit_offsets,
3485 x1,
3486 x2,
3487 a1,
3488 a2,
3489 rows,
3490 cols,
3491 o1,
3492 o2,
3493 s,
3494 e,
3495 )
3496 }
3497 });
3498 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3499 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3500 pool.run_many(&parts);
3501 }
3502 return;
3503 }
3504 if uniform_q4 {
3505 let closures: [_; N] = std::array::from_fn(|i| {
3506 let (packed, scales) =
3507 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3508 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3509 move |s: usize, e: usize| {
3510 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3511 }
3512 });
3513 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3514 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3515 pool.run_many(&parts);
3516 } else {
3517 let closures: [_; N] = std::array::from_fn(|i| {
3518 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3519 unreachable!()
3520 };
3521 let (bytes, rows, cols, o1, o2) = (
3522 ts[i].quant_bytes(),
3523 ts[i].rows(),
3524 ts[i].cols(),
3525 p1[i],
3526 p2[i],
3527 );
3528 move |s: usize, e: usize| {
3529 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3530 }
3531 });
3532 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3533 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3534 pool.run_many(&parts);
3535 }
3536 return;
3537 }
3538
3539 if uniform_f32 {
3540 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3541 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3542 let closures: [_; N] = std::array::from_fn(|i| {
3543 let Self::F32 { data, cols, .. } = ts[i] else {
3544 unreachable!()
3545 };
3546 let (o1, o2) = (p1[i], p2[i]);
3547 move |start: usize, end: usize| {
3548 for o in start..end {
3549 let row = &data[o * cols..(o + 1) * cols];
3550 let (mut s1, mut s2) = (0.0f32, 0.0f32);
3551 for j in 0..*cols {
3552 s1 += row[j] * x1[j];
3553 s2 += row[j] * x2[j];
3554 }
3555 unsafe {
3557 *o1.at(o) = s1;
3558 *o2.at(o) = s2;
3559 }
3560 }
3561 }
3562 });
3563 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3564 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3565 pool.run_many(&parts);
3566 return;
3567 }
3568
3569 struct Ctx<'a> {
3570 bytes: &'a [u8],
3571 row_scale: &'a [f32],
3572 cols: usize,
3573 xs1: std::borrow::Cow<'a, [f32]>,
3574 xs2: std::borrow::Cow<'a, [f32]>,
3575 }
3576 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3577 let Self::Mapped {
3578 dtype,
3579 cols,
3580 row_scale,
3581 col_field,
3582 ..
3583 } = ts[i]
3584 else {
3585 unreachable!()
3586 };
3587 Ctx {
3588 bytes: ts[i].quant_bytes(),
3589 row_scale,
3590 cols: *cols,
3591 xs1: prescale(x1, col_field, *dtype),
3592 xs2: prescale(x2, col_field, *dtype),
3593 }
3594 });
3595 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3596 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3597 #[cfg(target_arch = "aarch64")]
3598 if sdot_enabled() {
3599 let acts: [(SplitAct, SplitAct); N] =
3600 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3601 let closures: [_; N] = std::array::from_fn(|i| {
3602 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3603 move |start: usize, end: usize| {
3604 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3605 }
3606 });
3607 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3608 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3609 pool.run_many(&parts);
3610 return;
3611 }
3612 #[cfg(target_arch = "x86_64")]
3613 if avx2_a8w8_enabled() {
3614 let acts: [(SplitAct, SplitAct); N] =
3615 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3616 let closures: [_; N] = std::array::from_fn(|i| {
3617 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3618 move |start: usize, end: usize| {
3619 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3620 }
3621 });
3622 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3623 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3624 pool.run_many(&parts);
3625 return;
3626 }
3627 let closures: [_; N] = std::array::from_fn(|i| {
3628 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3629 move |start: usize, end: usize| {
3630 q8_range2_f32(
3631 c.bytes,
3632 c.row_scale,
3633 &c.xs1,
3634 &c.xs2,
3635 c.cols,
3636 o1,
3637 o2,
3638 start,
3639 end,
3640 )
3641 }
3642 });
3643 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3644 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3645 pool.run_many(&parts);
3646 }
3647
3648 pub fn matvec_silu_mul(
3653 gate: &QTensor,
3654 up: &QTensor,
3655 x: &[f32],
3656 out: &mut [f32],
3657 pool: Option<&Pool>,
3658 ) -> bool {
3659 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3660 }
3661
3662 pub fn matvec_silu_mul_limited(
3668 gate: &QTensor,
3669 up: &QTensor,
3670 x: &[f32],
3671 out: &mut [f32],
3672 limit: f32,
3673 pool: Option<&Pool>,
3674 ) -> bool {
3675 if gate.exact_cpu_default() || up.exact_cpu_default() {
3676 return float_activations_scope(|| {
3677 Self::matvec_silu_mul_limited(gate, up, x, out, limit, pool)
3678 });
3679 }
3680 if gate.has_prism_contract()
3681 || up.has_prism_contract()
3682 || gate.per_tensor_only()
3683 || up.per_tensor_only()
3684 {
3685 return false;
3689 }
3690 let inter = gate.rows();
3691 debug_assert_eq!(up.rows(), inter);
3692 debug_assert_eq!(out.len(), inter);
3693 debug_assert_eq!(gate.cols(), up.cols());
3694 if !a8w8_enabled() {
3695 return false;
3696 }
3697 let act = split_act(x);
3698 let act = &act;
3699 let x_ref = x;
3700 let out_addr = SendMut(out.as_mut_ptr());
3701
3702 match (gate, up) {
3703 (
3705 Self::Mapped {
3706 dtype: TensorDtype::Q4Block,
3707 ..
3708 },
3709 Self::Mapped {
3710 dtype: TensorDtype::Q4Block,
3711 ..
3712 },
3713 ) => {
3714 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3715 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3716 let gpr = gate.cols() / GROUP_SIZE;
3717 let cols = gate.cols();
3718 let run = move |start: usize, end: usize| {
3719 for r in start..end {
3720 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3721 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3722 for &(j, xv) in &act.outliers {
3723 let flat = r * cols + j;
3724 let gb = gp[flat / 2];
3725 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3726 let gsc = f16_to_f32(u16::from_le_bytes([
3727 gs[(flat / GROUP_SIZE) * 2],
3728 gs[(flat / GROUP_SIZE) * 2 + 1],
3729 ]));
3730 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3731 let ub = up_p[flat / 2];
3732 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3733 let usc = f16_to_f32(u16::from_le_bytes([
3734 up_s[(flat / GROUP_SIZE) * 2],
3735 up_s[(flat / GROUP_SIZE) * 2 + 1],
3736 ]));
3737 uv += ((un as i32 - 8) as f32) * usc * xv;
3738 }
3739 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3741 }
3742 };
3743 dispatch_rows(pool, inter, &run);
3744 true
3745 }
3746 (
3750 Self::Mapped {
3751 dtype: TensorDtype::Q4Tiled,
3752 ..
3753 },
3754 Self::Mapped {
3755 dtype: TensorDtype::Q4Tiled,
3756 ..
3757 },
3758 ) => {
3759 let g_bytes = gate.quant_bytes();
3760 let u_bytes = up.quant_bytes();
3761 let gpr = gate.cols() / GROUP_SIZE;
3762 let run = move |start: usize, end: usize| {
3763 for r in start..end {
3764 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3765 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3766 for &(j, xv) in &act.outliers {
3767 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3768 gv += w * s * xv;
3769 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3770 uv += w * s * xv;
3771 }
3772 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3774 }
3775 };
3776 dispatch_rows(pool, inter, &run);
3777 true
3778 }
3779 (
3782 Self::Mapped {
3783 dtype: TensorDtype::Q4TiledP,
3784 ..
3785 },
3786 Self::Mapped {
3787 dtype: TensorDtype::Q4TiledP,
3788 ..
3789 },
3790 ) => {
3791 let cols = gate.cols();
3792 let gpr = cols / GROUP_SIZE;
3793 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3794 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3795 #[cfg(target_arch = "x86_64")]
3796 let lanes = q4tp_lanes();
3797 let run = |start: usize, end: usize| {
3798 #[cfg(target_arch = "x86_64")]
3799 if lanes > 1 && act.xg8.len() == gpr {
3800 with_krows(gpr * lanes, |gsc, usc| {
3803 unsafe {
3805 match lanes {
3806 2 => q4tp_gateup_lanes::<2>(
3807 &gv_view, &uv_view, start, end, gpr, &act.xq,
3808 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3809 out_addr,
3810 ),
3811 8 => q4tp_gateup_lanes::<8>(
3812 &gv_view, &uv_view, start, end, gpr, &act.xq,
3813 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3814 out_addr,
3815 ),
3816 _ => q4tp_gateup_lanes::<4>(
3817 &gv_view, &uv_view, start, end, gpr, &act.xq,
3818 &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3819 out_addr,
3820 ),
3821 }
3822 }
3823 });
3824 return;
3825 }
3826 crate::cpuprof::count_plain_rows((end - start) as u64);
3827 with_krows(gpr, |gsc, usc| {
3828 for r in start..end {
3829 gv_view.scales_into(r, gpr, gsc);
3830 uv_view.scales_into(r, gpr, usc);
3831 let mut gv =
3832 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3833 let mut uv =
3834 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3835 for &(j, xv) in &act.outliers {
3836 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3837 gv += w * s * xv;
3838 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3839 uv += w * s * xv;
3840 }
3841 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3843 }
3844 });
3845 };
3846 dispatch_rows(pool, inter, &run);
3847 true
3848 }
3849 (
3856 Self::Mapped {
3857 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3858 ..
3859 },
3860 Self::Mapped {
3861 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3862 ..
3863 },
3864 ) => {
3865 let gw = gate.bin_w().expect("binary dtype matched");
3866 let uw = up.bin_w().expect("binary dtype matched");
3867 let gpr = gate.cols() / GROUP_SIZE;
3868 let gsum = q1_group_sums(&act.xq, gpr);
3869 let gsum = &gsum;
3870 let run = move |start: usize, end: usize| {
3871 for r in start..end {
3872 let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3873 let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3874 for &(j, xv) in &act.outliers {
3875 let (w, s) = gw.outlier(r, gpr, j);
3876 gv += w * s * xv;
3877 let (w, s) = uw.outlier(r, gpr, j);
3878 uv += w * s * xv;
3879 }
3880 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3882 }
3883 };
3884 dispatch_rows(pool, inter, &run);
3885 true
3886 }
3887 (
3891 Self::Mapped {
3892 dtype: TensorDtype::Q2TiledP,
3893 ..
3894 },
3895 Self::Mapped {
3896 dtype: TensorDtype::Q2TiledP,
3897 ..
3898 },
3899 ) => {
3900 let cols = gate.cols();
3901 let gpr = cols / GROUP_SIZE;
3902 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3903 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3904 let gsum = q1_group_sums(&act.xq, gpr);
3905 let gsum = &gsum;
3906 let run = move |start: usize, end: usize| {
3907 with_krows(gpr, |gsc, usc| {
3908 for r in start..end {
3909 gv_view.scales_into(r, gpr, gsc);
3910 uv_view.scales_into(r, gpr, usc);
3911 let mut gv =
3912 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3913 let mut uv =
3914 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3915 for &(j, xv) in &act.outliers {
3916 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3917 gv += w * s * xv;
3918 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3919 uv += w * s * xv;
3920 }
3921 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3923 }
3924 });
3925 };
3926 dispatch_rows(pool, inter, &run);
3927 true
3928 }
3929 (
3934 Self::Mapped {
3935 dtype: TensorDtype::Q8Row,
3936 row_scale: g_rs,
3937 ..
3938 },
3939 Self::Mapped {
3940 dtype: TensorDtype::Q8Row,
3941 row_scale: u_rs,
3942 ..
3943 },
3944 ) => {
3945 let g_bytes = gate.quant_bytes();
3946 let u_bytes = up.quant_bytes();
3947 let cols = gate.cols();
3948 let run = move |start: usize, end: usize| {
3949 for r in start..end {
3950 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3951 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3952 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3954 }
3955 };
3956 dispatch_rows(pool, inter, &run);
3957 true
3958 }
3959 (
3961 Self::Mapped {
3962 dtype: TensorDtype::Q1T,
3963 ..
3964 },
3965 Self::Mapped {
3966 dtype: TensorDtype::Q1T,
3967 ..
3968 },
3969 ) => {
3970 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3971 let g_bytes = gate.quant_bytes();
3972 let u_bytes = up.quant_bytes();
3973 let gpr = gate.cols() / GROUP_SIZE;
3974 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3975 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3976 let run = move |start: usize, end: usize| {
3977 for r in start..end {
3978 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3979 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3980 for &(j, xv) in &act.outliers {
3981 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3982 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3983 }
3984 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3985 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3986 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3988 }
3989 };
3990 dispatch_rows(pool, inter, &run);
3991 true
3992 }
3993 _ => false,
3994 }
3995 }
3996
3997 pub fn moe_gate_up_many(
4012 pairs: &[(&QTensor, &QTensor)],
4013 x: &[f32],
4014 outs: &mut [Vec<f32>],
4015 pool: Option<&Pool>,
4016 ) -> bool {
4017 Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
4018 }
4019
4020 pub fn moe_gate_up_many_limited(
4024 pairs: &[(&QTensor, &QTensor)],
4025 x: &[f32],
4026 outs: &mut [Vec<f32>],
4027 limit: f32,
4028 pool: Option<&Pool>,
4029 ) -> bool {
4030 if Self::any_exact_cpu_default_pairs(pairs) {
4031 return float_activations_scope(|| {
4032 Self::moe_gate_up_many_limited(pairs, x, outs, limit, pool)
4033 });
4034 }
4035 if pairs.is_empty() || pairs.len() != outs.len() {
4036 return false;
4037 }
4038 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4039 return false;
4040 }
4041 if !a8w8_enabled() {
4042 if limit > 0.0 {
4043 return false;
4046 }
4047 let groups = vec![vec![0]; pairs.len()];
4048 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
4049 }
4050 let inter = pairs[0].0.rows();
4051 let cols = pairs[0].0.cols();
4052 if cols % GROUP_SIZE != 0 {
4053 return false;
4054 }
4055 let gpr = cols / GROUP_SIZE;
4056 let q2 = matches!(
4059 pairs[0].0,
4060 Self::Mapped {
4061 dtype: TensorDtype::Q2TiledP,
4062 ..
4063 }
4064 );
4065 let want = if q2 {
4066 TensorDtype::Q2TiledP
4067 } else {
4068 TensorDtype::Q4TiledP
4069 };
4070 let mut views = Vec::with_capacity(pairs.len() * 2);
4071 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
4072 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
4073 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
4074 if !both
4075 || g.rows() != inter
4076 || u.rows() != inter
4077 || g.cols() != cols
4078 || u.cols() != cols
4079 || o.len() != inter
4080 {
4081 return false;
4082 }
4083 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
4084 views.push(mk(g.quant_bytes(), inter, cols));
4085 views.push(mk(u.quant_bytes(), inter, cols));
4086 }
4087 let act = split_act(x);
4088 let gsum = if q2 {
4089 q1_group_sums(&act.xq, gpr)
4090 } else {
4091 Vec::new()
4092 };
4093 let (act, gsum) = (&act, &gsum);
4094 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4095 let (views, ptrs) = (&views, &ptrs);
4096 let run = |start: usize, end: usize| {
4097 with_krows(gpr, |gsc, usc| {
4098 for flat in start..end {
4099 let (e, r) = (flat / inter, flat % inter);
4100 let gv_view = &views[e * 2];
4101 let uv_view = &views[e * 2 + 1];
4102 gv_view.scales_into(r, gpr, gsc);
4103 uv_view.scales_into(r, gpr, usc);
4104 let (mut gv, mut uv) = if q2 {
4105 (
4106 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
4107 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
4108 )
4109 } else {
4110 (
4111 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
4112 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
4113 )
4114 };
4115 for &(j, xv) in &act.outliers {
4116 let (og, ou) = if q2 {
4117 (
4118 q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
4119 q2tp_outlier(uv_view.nib, r, gpr, j, usc),
4120 )
4121 } else {
4122 (
4123 q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
4124 q4tp_outlier(uv_view.nib, r, gpr, j, usc),
4125 )
4126 };
4127 gv += og.0 * og.1 * xv;
4128 uv += ou.0 * ou.1 * xv;
4129 }
4130 unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
4132 }
4133 });
4134 };
4135 dispatch_rows(pool, pairs.len() * inter, &run);
4136 true
4137 }
4138
4139 pub fn moe_down_many(
4148 downs: &[&QTensor],
4149 gs: &[Vec<f32>],
4150 weights: &[f32],
4151 out: &mut [f32],
4152 pool: Option<&Pool>,
4153 ) -> bool {
4154 if Self::any_exact_cpu_default(downs) {
4155 return float_activations_scope(|| {
4156 Self::moe_down_many(downs, gs, weights, out, pool)
4157 });
4158 }
4159 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
4160 return false;
4161 }
4162 if downs.iter().any(|d| d.per_tensor_only()) {
4163 return false;
4164 }
4165 if !a8w8_enabled() {
4166 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
4167 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
4168 return false;
4169 }
4170 out.fill(0.0);
4171 for (row, &w) in terms.iter().zip(weights) {
4172 for (o, &v) in out.iter_mut().zip(row) {
4173 *o += w * v;
4174 }
4175 }
4176 return true;
4177 }
4178 let rows = out.len();
4179 let cols = downs[0].cols();
4180 if cols % GROUP_SIZE != 0 {
4181 return false;
4182 }
4183 let gpr = cols / GROUP_SIZE;
4184 let mut views = Vec::with_capacity(downs.len());
4185 for (d, g) in downs.iter().zip(gs.iter()) {
4186 if !matches!(
4187 d,
4188 Self::Mapped {
4189 dtype: TensorDtype::Q4TiledP,
4190 ..
4191 }
4192 ) || d.rows() != rows
4193 || d.cols() != cols
4194 || g.len() != cols
4195 {
4196 return false;
4197 }
4198 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4199 }
4200 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
4202 let out_addr = SendMut(out.as_mut_ptr());
4210 let (views, acts, weights) = (&views, &acts, &weights);
4211 let run = |start: usize, end: usize| {
4212 with_krow(gpr, |sc| {
4213 for r in start..end {
4214 let mut acc = 0f32;
4215 for (e, v) in views.iter().enumerate() {
4216 v.scales_into(r, gpr, sc);
4217 let a = &acts[e];
4218 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4219 for &(j, xv) in &a.outliers {
4220 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4221 d += w * s * xv;
4222 }
4223 acc += weights[e] * d;
4224 }
4225 unsafe { *out_addr.at(r) = acc };
4227 }
4228 });
4229 };
4230 dispatch_rows(pool, rows, &run);
4231 true
4232 }
4233
4234 pub fn moe_gate_up_rows(
4246 pairs: &[(&QTensor, &QTensor)],
4247 groups: &[Vec<usize>],
4248 xs: &[f32],
4249 outs: &mut [Vec<f32>],
4250 pool: Option<&Pool>,
4251 ) -> bool {
4252 if Self::any_exact_cpu_default_pairs(pairs) {
4253 return float_activations_scope(|| {
4254 Self::moe_gate_up_rows(pairs, groups, xs, outs, pool)
4255 });
4256 }
4257 if pairs.is_empty() || pairs.len() != groups.len() {
4258 return false;
4259 }
4260 if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4261 return false;
4262 }
4263 let inter = pairs[0].0.rows();
4264 let cols = pairs[0].0.cols();
4265 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4266 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4267 return false;
4268 }
4269 let b = xs.len() / cols;
4270 let gpr = cols / GROUP_SIZE;
4271 let mut views = Vec::with_capacity(pairs.len() * 2);
4272 for (g, u) in pairs {
4273 let q4tp = |t: &QTensor| {
4274 matches!(
4275 t,
4276 Self::Mapped {
4277 dtype: TensorDtype::Q4TiledP,
4278 ..
4279 }
4280 )
4281 };
4282 if g.has_prism_contract()
4283 || u.has_prism_contract()
4284 || !q4tp(g)
4285 || !q4tp(u)
4286 || g.rows() != inter
4287 || u.rows() != inter
4288 || g.cols() != cols
4289 || u.cols() != cols
4290 {
4291 return false;
4292 }
4293 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4294 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4295 }
4296 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4297 return false;
4298 }
4299 let quantized = a8w8_enabled();
4300 let acts: Vec<SplitAct> = if quantized {
4301 (0..b)
4302 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4303 .collect()
4304 } else {
4305 Vec::new()
4306 };
4307 let mut offs = Vec::with_capacity(groups.len());
4308 let mut o = 0usize;
4309 for g in groups {
4310 offs.push(o);
4311 o += g.len();
4312 }
4313 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4314 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4315 let run = |start: usize, end: usize| {
4316 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4317 for flat in start..end {
4318 let (e, r) = (flat / inter, flat % inter);
4319 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4320 gv_view.scales_into(r, gpr, &mut gsc);
4321 uv_view.scales_into(r, gpr, &mut usc);
4322 for (k, &t) in groups[e].iter().enumerate() {
4323 if !quantized {
4324 let x = &xs[t * cols..(t + 1) * cols];
4325 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4326 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4327 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4328 continue;
4329 }
4330 let act = &acts[t];
4331 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4332 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4333 for &(j, xv) in &act.outliers {
4334 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4335 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4336 gv += og.0 * og.1 * xv;
4337 uv += ou.0 * ou.1 * xv;
4338 }
4339 let silu_g = gv / (1.0 + (-gv).exp());
4340 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4343 }
4344 }
4345 };
4346 dispatch_rows(pool, pairs.len() * inter, &run);
4347 true
4348 }
4349
4350 pub fn moe_down_rows(
4357 downs: &[&QTensor],
4358 group_lens: &[usize],
4359 gs: &[Vec<f32>],
4360 outs: &mut [Vec<f32>],
4361 pool: Option<&Pool>,
4362 ) -> bool {
4363 if Self::any_exact_cpu_default(downs) {
4364 return float_activations_scope(|| {
4365 Self::moe_down_rows(downs, group_lens, gs, outs, pool)
4366 });
4367 }
4368 if downs.is_empty() || downs.len() != group_lens.len() {
4369 return false;
4370 }
4371 if downs.iter().any(|d| d.per_tensor_only()) {
4372 return false;
4373 }
4374 let rows = downs[0].rows();
4375 let cols = downs[0].cols();
4376 let n_pairs: usize = group_lens.iter().sum();
4377 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4378 return false;
4379 }
4380 let gpr = cols / GROUP_SIZE;
4381 let mut views = Vec::with_capacity(downs.len());
4382 for d in downs {
4383 if d.has_prism_contract()
4384 || !matches!(
4385 d,
4386 Self::Mapped {
4387 dtype: TensorDtype::Q4TiledP,
4388 ..
4389 }
4390 )
4391 || d.rows() != rows
4392 || d.cols() != cols
4393 {
4394 return false;
4395 }
4396 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4397 }
4398 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4399 return false;
4400 }
4401 let quantized = a8w8_enabled();
4402 let acts: Vec<SplitAct> = if quantized {
4403 gs.iter().map(|g| split_act(g)).collect()
4404 } else {
4405 Vec::new()
4406 };
4407 let mut offs = Vec::with_capacity(group_lens.len());
4408 let mut o = 0usize;
4409 for &l in group_lens {
4410 offs.push(o);
4411 o += l;
4412 }
4413 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4414 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4415 let run = |start: usize, end: usize| {
4416 let mut sc = vec![0f32; gpr];
4417 for flat in start..end {
4418 let (e, r) = (flat / rows, flat % rows);
4419 let v = &views[e];
4420 v.scales_into(r, gpr, &mut sc);
4421 for k in 0..group_lens[e] {
4422 if !quantized {
4423 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4424 unsafe { *ptrs[offs[e] + k].at(r) = d };
4425 continue;
4426 }
4427 let a = &acts[offs[e] + k];
4428 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4429 for &(j, xv) in &a.outliers {
4430 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4431 d += w * s * xv;
4432 }
4433 unsafe { *ptrs[offs[e] + k].at(r) = d };
4435 }
4436 }
4437 };
4438 dispatch_rows(pool, downs.len() * rows, &run);
4439 true
4440 }
4441
4442 pub fn moe_expert_rows_fast(
4455 experts: &[(&QTensor, &QTensor, &QTensor)],
4456 groups: &[Vec<usize>],
4457 xs: &[f32],
4458 out: &mut [f32],
4459 pool: Option<&Pool>,
4460 ) -> bool {
4461 #[cfg(target_arch = "x86_64")]
4462 {
4463 moe_fast::run(experts, groups, xs, out, pool)
4464 }
4465 #[cfg(not(target_arch = "x86_64"))]
4466 {
4467 let _ = (experts, groups, xs, out, pool);
4468 false
4469 }
4470 }
4471
4472 pub fn moe_expert_rows_fast_available() -> bool {
4475 #[cfg(target_arch = "x86_64")]
4476 {
4477 avx512vnni_enabled() && a8w8_enabled()
4478 }
4479 #[cfg(not(target_arch = "x86_64"))]
4480 {
4481 false
4482 }
4483 }
4484}
4485
4486#[cfg(target_os = "macos")]
4491mod accel_blas {
4492 #[link(name = "Accelerate", kind = "framework")]
4493 unsafe extern "C" {
4494 pub fn cblas_sgemm(
4495 order: i32,
4496 trans_a: i32,
4497 trans_b: i32,
4498 m: i32,
4499 n: i32,
4500 k: i32,
4501 alpha: f32,
4502 a: *const f32,
4503 lda: i32,
4504 b: *const f32,
4505 ldb: i32,
4506 beta: f32,
4507 c: *mut f32,
4508 ldc: i32,
4509 );
4510 }
4511}
4512
4513#[cfg(target_os = "macos")]
4514pub(crate) fn accel_gemm_enabled() -> bool {
4515 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4516 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4517}
4518
4519#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4522pub(crate) fn accel_gemm_enabled() -> bool {
4523 true
4524}
4525
4526#[cfg(all(target_arch = "x86_64", not(target_os = "macos")))]
4530pub(crate) fn accel_gemm_enabled() -> bool {
4531 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4532 *ON.get_or_init(|| {
4533 std::is_x86_feature_detected!("avx2")
4534 && std::is_x86_feature_detected!("fma")
4535 && std::env::var("CMF_CPU_ATTEND").map(|v| v != "0").unwrap_or(true)
4536 })
4537}
4538
4539#[cfg(not(any(target_os = "macos", target_arch = "aarch64", target_arch = "x86_64")))]
4541#[allow(dead_code)]
4542pub(crate) fn accel_gemm_enabled() -> bool {
4543 false
4544}
4545
4546#[cfg(target_arch = "aarch64")]
4553#[allow(clippy::too_many_arguments)]
4554pub(crate) fn neon_gemm_rm(
4555 m: usize,
4556 n: usize,
4557 k: usize,
4558 alpha: f32,
4559 a: &[f32],
4560 lda: usize,
4561 b_mat: &[f32],
4562 ldb: usize,
4563 b_rows_are_n: bool,
4564 c: &mut [f32],
4565 ldc: usize,
4566) {
4567 debug_assert!(a.len() >= (m - 1) * lda + k);
4568 debug_assert!(c.len() >= (m - 1) * ldc + n);
4569 unsafe {
4571 use core::arch::aarch64::*;
4572 let mut i = 0usize;
4573 while i < m {
4574 let mi = (m - i).min(4);
4575 let mut j = 0usize;
4576 while j < n {
4577 let nj = (n - j).min(8);
4578 if mi == 4 && nj == 8 {
4579 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4580 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4581 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4582 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4583 for p in 0..k {
4584 let (b0, b1) = if b_rows_are_n {
4585 let base = b_mat.as_ptr().add(j * ldb + p);
4588 let g = |o: usize| *base.add(o * ldb);
4589 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4590 } else {
4591 let base = b_mat.as_ptr().add(p * ldb + j);
4592 (
4593 [*base, *base.add(1), *base.add(2), *base.add(3)],
4594 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4595 )
4596 };
4597 let bv0 = vld1q_f32(b0.as_ptr());
4598 let bv1 = vld1q_f32(b1.as_ptr());
4599 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4600 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4601 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4602 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4603 c0a = vfmaq_f32(c0a, a0, bv0);
4604 c0b = vfmaq_f32(c0b, a0, bv1);
4605 c1a = vfmaq_f32(c1a, a1, bv0);
4606 c1b = vfmaq_f32(c1b, a1, bv1);
4607 c2a = vfmaq_f32(c2a, a2, bv0);
4608 c2b = vfmaq_f32(c2b, a2, bv1);
4609 c3a = vfmaq_f32(c3a, a3, bv0);
4610 c3b = vfmaq_f32(c3b, a3, bv1);
4611 }
4612 let al = vdupq_n_f32(alpha);
4613 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4614 .iter()
4615 .enumerate()
4616 {
4617 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4618 vst1q_f32(dst, vmulq_f32(*ca, al));
4619 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4620 }
4621 } else {
4622 for r in 0..mi {
4623 for q in 0..nj {
4624 let mut acc = 0f32;
4625 for p in 0..k {
4626 let bv = if b_rows_are_n {
4627 b_mat[(j + q) * ldb + p]
4628 } else {
4629 b_mat[p * ldb + j + q]
4630 };
4631 acc += a[(i + r) * lda + p] * bv;
4632 }
4633 c[(i + r) * ldc + j + q] = acc * alpha;
4634 }
4635 }
4636 }
4637 j += nj;
4638 }
4639 i += mi;
4640 }
4641 }
4642}
4643
4644#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4646#[allow(clippy::too_many_arguments)]
4647pub(crate) fn sgemm_rm(
4648 m: usize,
4649 n: usize,
4650 k: usize,
4651 alpha: f32,
4652 a: &[f32],
4653 lda: usize,
4654 b_mat: &[f32],
4655 ldb: usize,
4656 b_rows_are_n: bool,
4657 c: &mut [f32],
4658 ldc: usize,
4659) {
4660 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4661}
4662
4663#[cfg(target_arch = "aarch64")]
4666#[allow(clippy::too_many_arguments)]
4667pub(crate) fn micro_gemm_rm(
4668 m: usize,
4669 n: usize,
4670 k: usize,
4671 alpha: f32,
4672 a: &[f32],
4673 lda: usize,
4674 b_mat: &[f32],
4675 ldb: usize,
4676 b_rows_are_n: bool,
4677 c: &mut [f32],
4678 ldc: usize,
4679) {
4680 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4681}
4682
4683#[cfg(target_arch = "x86_64")]
4685#[allow(clippy::too_many_arguments)]
4686pub(crate) fn micro_gemm_rm(
4687 m: usize,
4688 n: usize,
4689 k: usize,
4690 alpha: f32,
4691 a: &[f32],
4692 lda: usize,
4693 b_mat: &[f32],
4694 ldb: usize,
4695 b_rows_are_n: bool,
4696 c: &mut [f32],
4697 ldc: usize,
4698) {
4699 x86_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4700}
4701
4702#[cfg(all(target_arch = "x86_64", not(target_os = "macos")))]
4704#[allow(clippy::too_many_arguments)]
4705pub(crate) fn sgemm_rm(
4706 m: usize,
4707 n: usize,
4708 k: usize,
4709 alpha: f32,
4710 a: &[f32],
4711 lda: usize,
4712 b_mat: &[f32],
4713 ldb: usize,
4714 b_rows_are_n: bool,
4715 c: &mut [f32],
4716 ldc: usize,
4717) {
4718 x86_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4719}
4720
4721#[cfg(target_arch = "x86_64")]
4734#[allow(clippy::too_many_arguments)]
4735pub(crate) fn x86_gemm_rm(
4736 m: usize,
4737 n: usize,
4738 k: usize,
4739 alpha: f32,
4740 a: &[f32],
4741 lda: usize,
4742 b_mat: &[f32],
4743 ldb: usize,
4744 b_rows_are_n: bool,
4745 c: &mut [f32],
4746 ldc: usize,
4747) {
4748 if m == 0 || n == 0 {
4749 return;
4750 }
4751 debug_assert!(a.len() >= (m - 1) * lda + k);
4752 debug_assert!(c.len() >= (m - 1) * ldc + n);
4753 if b_rows_are_n {
4754 debug_assert!(k == 0 || b_mat.len() >= (n - 1) * ldb + k);
4755 let mut bt = vec![0f32; k * n];
4756 for j in 0..n {
4757 for (p, &v) in b_mat[j * ldb..j * ldb + k].iter().enumerate() {
4758 bt[p * n + j] = v;
4759 }
4760 }
4761 return x86_gemm_rm(m, n, k, alpha, a, lda, &bt, n, false, c, ldc);
4762 }
4763 if k == 0 {
4764 for r in 0..m {
4765 c[r * ldc..r * ldc + n].fill(0.0);
4766 }
4767 return;
4768 }
4769 debug_assert!(b_mat.len() >= (k - 1) * ldb + n);
4770 if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
4771 unsafe { x86_gemm_avx2(m, n, k, alpha, a, lda, b_mat, ldb, c, ldc) }
4774 } else {
4775 for r in 0..m {
4776 for j in 0..n {
4777 let mut acc = 0f32;
4778 for p in 0..k {
4779 acc += a[r * lda + p] * b_mat[p * ldb + j];
4780 }
4781 c[r * ldc + j] = acc * alpha;
4782 }
4783 }
4784 }
4785}
4786
4787#[cfg(target_arch = "x86_64")]
4796#[target_feature(enable = "avx2,fma")]
4797#[allow(clippy::too_many_arguments)]
4798unsafe fn x86_gemm_avx2(
4799 m: usize,
4800 n: usize,
4801 k: usize,
4802 alpha: f32,
4803 a: &[f32],
4804 lda: usize,
4805 b: &[f32],
4806 ldb: usize,
4807 c: &mut [f32],
4808 ldc: usize,
4809) {
4810 use core::arch::x86_64::*;
4811 const NC: usize = 128;
4812 const KC: usize = 256;
4813 unsafe {
4814 let ap = a.as_ptr();
4815 let bp = b.as_ptr();
4816 let cp = c.as_mut_ptr();
4817 let al = _mm256_set1_ps(alpha);
4818 let mut jc = 0usize;
4819 while jc < n {
4820 let jend = (jc + NC).min(n);
4821 let mut pc = 0usize;
4822 while pc < k {
4823 let pend = (pc + KC).min(k);
4824 let (first, last) = (pc == 0, pend == k);
4825 let mut i = 0usize;
4826 while i < m {
4827 let mi = (m - i).min(4);
4828 let mut j = jc;
4829 while j < jend {
4830 let nj = (jend - j).min(16);
4831 if mi == 4 && nj == 16 {
4832 let a0 = ap.add(i * lda);
4833 let a1 = ap.add((i + 1) * lda);
4834 let a2 = ap.add((i + 2) * lda);
4835 let a3 = ap.add((i + 3) * lda);
4836 let c0 = cp.add(i * ldc + j);
4837 let c1 = cp.add((i + 1) * ldc + j);
4838 let c2 = cp.add((i + 2) * ldc + j);
4839 let c3 = cp.add((i + 3) * ldc + j);
4840 let (mut c00, mut c01, mut c10, mut c11, mut c20, mut c21, mut c30, mut c31) =
4841 if first {
4842 let z = _mm256_setzero_ps();
4843 (z, z, z, z, z, z, z, z)
4844 } else {
4845 (
4846 _mm256_loadu_ps(c0),
4847 _mm256_loadu_ps(c0.add(8)),
4848 _mm256_loadu_ps(c1),
4849 _mm256_loadu_ps(c1.add(8)),
4850 _mm256_loadu_ps(c2),
4851 _mm256_loadu_ps(c2.add(8)),
4852 _mm256_loadu_ps(c3),
4853 _mm256_loadu_ps(c3.add(8)),
4854 )
4855 };
4856 for p in pc..pend {
4857 let bb = bp.add(p * ldb + j);
4858 let b0 = _mm256_loadu_ps(bb);
4859 let b1 = _mm256_loadu_ps(bb.add(8));
4860 let x0 = _mm256_broadcast_ss(&*a0.add(p));
4861 c00 = _mm256_fmadd_ps(x0, b0, c00);
4862 c01 = _mm256_fmadd_ps(x0, b1, c01);
4863 let x1 = _mm256_broadcast_ss(&*a1.add(p));
4864 c10 = _mm256_fmadd_ps(x1, b0, c10);
4865 c11 = _mm256_fmadd_ps(x1, b1, c11);
4866 let x2 = _mm256_broadcast_ss(&*a2.add(p));
4867 c20 = _mm256_fmadd_ps(x2, b0, c20);
4868 c21 = _mm256_fmadd_ps(x2, b1, c21);
4869 let x3 = _mm256_broadcast_ss(&*a3.add(p));
4870 c30 = _mm256_fmadd_ps(x3, b0, c30);
4871 c31 = _mm256_fmadd_ps(x3, b1, c31);
4872 }
4873 if last {
4874 c00 = _mm256_mul_ps(c00, al);
4875 c01 = _mm256_mul_ps(c01, al);
4876 c10 = _mm256_mul_ps(c10, al);
4877 c11 = _mm256_mul_ps(c11, al);
4878 c20 = _mm256_mul_ps(c20, al);
4879 c21 = _mm256_mul_ps(c21, al);
4880 c30 = _mm256_mul_ps(c30, al);
4881 c31 = _mm256_mul_ps(c31, al);
4882 }
4883 _mm256_storeu_ps(c0, c00);
4884 _mm256_storeu_ps(c0.add(8), c01);
4885 _mm256_storeu_ps(c1, c10);
4886 _mm256_storeu_ps(c1.add(8), c11);
4887 _mm256_storeu_ps(c2, c20);
4888 _mm256_storeu_ps(c2.add(8), c21);
4889 _mm256_storeu_ps(c3, c30);
4890 _mm256_storeu_ps(c3.add(8), c31);
4891 } else {
4892 for r in 0..mi {
4895 let ar = ap.add((i + r) * lda);
4896 let cr = cp.add((i + r) * ldc + j);
4897 let mut q = 0usize;
4898 while q + 8 <= nj {
4899 let mut acc = if first {
4900 _mm256_setzero_ps()
4901 } else {
4902 _mm256_loadu_ps(cr.add(q))
4903 };
4904 for p in pc..pend {
4905 acc = _mm256_fmadd_ps(
4906 _mm256_broadcast_ss(&*ar.add(p)),
4907 _mm256_loadu_ps(bp.add(p * ldb + j + q)),
4908 acc,
4909 );
4910 }
4911 if last {
4912 acc = _mm256_mul_ps(acc, al);
4913 }
4914 _mm256_storeu_ps(cr.add(q), acc);
4915 q += 8;
4916 }
4917 while q < nj {
4918 let mut acc = if first { 0.0 } else { *cr.add(q) };
4919 for p in pc..pend {
4920 acc += *ar.add(p) * *bp.add(p * ldb + j + q);
4921 }
4922 if last {
4923 acc *= alpha;
4924 }
4925 *cr.add(q) = acc;
4926 q += 1;
4927 }
4928 }
4929 }
4930 j += nj;
4931 }
4932 i += mi;
4933 }
4934 pc = pend;
4935 }
4936 jc = jend;
4937 }
4938 }
4939}
4940
4941#[allow(clippy::too_many_arguments)]
4945pub fn sgemm_public(
4946 m: usize,
4947 n: usize,
4948 k: usize,
4949 alpha: f32,
4950 a: &[f32],
4951 lda: usize,
4952 b_mat: &[f32],
4953 ldb: usize,
4954 b_rows_are_n: bool,
4955 c: &mut [f32],
4956 ldc: usize,
4957) {
4958 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4959 {
4960 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4961 }
4962 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4967 {
4968 if m == 0 || n == 0 {
4969 return;
4970 }
4971 assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4972 let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4973 assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4974 unsafe {
4976 matrixmultiply::sgemm(
4977 m,
4978 k,
4979 n,
4980 alpha,
4981 a.as_ptr(),
4982 lda as isize,
4983 1,
4984 b_mat.as_ptr(),
4985 rsb,
4986 csb,
4987 0.0,
4988 c.as_mut_ptr(),
4989 ldc as isize,
4990 1,
4991 );
4992 }
4993 }
4994}
4995
4996#[cfg(target_os = "macos")]
4999#[allow(clippy::too_many_arguments)]
5000pub(crate) fn sgemm_rm(
5001 m: usize,
5002 n: usize,
5003 k: usize,
5004 alpha: f32,
5005 a: &[f32],
5006 lda: usize,
5007 b_mat: &[f32],
5008 ldb: usize,
5009 b_rows_are_n: bool,
5010 c: &mut [f32],
5011 ldc: usize,
5012) {
5013 debug_assert!(a.len() >= (m - 1) * lda + k);
5014 debug_assert!(c.len() >= (m - 1) * ldc + n);
5015 #[cfg(target_arch = "aarch64")]
5020 if std::env::var("CMF_FORCE_NEON_GEMM")
5021 .map(|v| v == "1")
5022 .unwrap_or(false)
5023 {
5024 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
5025 }
5026 unsafe {
5027 accel_blas::cblas_sgemm(
5028 101, 111, if b_rows_are_n { 112 } else { 111 },
5031 m as i32,
5032 n as i32,
5033 k as i32,
5034 alpha,
5035 a.as_ptr(),
5036 lda as i32,
5037 b_mat.as_ptr(),
5038 ldb as i32,
5039 0.0,
5040 c.as_mut_ptr(),
5041 ldc as i32,
5042 );
5043 }
5044}
5045
5046#[cfg(target_os = "macos")]
5053fn qmatmat_accel(
5054 q: &[u8],
5055 row_scale: &[f32],
5056 pre: &[std::borrow::Cow<'_, [f32]>],
5057 rows: usize,
5058 cols: usize,
5059 out: &mut [f32],
5060 pool: Option<&Pool>,
5061) {
5062 const TR: usize = 2048;
5067 let b = pre.len();
5068 thread_local! {
5069 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5070 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5071 }
5072 XPANEL.with(|xp| {
5073 WTILE.with(|wt| {
5074 let mut xpanel = xp.borrow_mut();
5075 xpanel.clear();
5076 for x in pre {
5077 xpanel.extend_from_slice(x);
5078 }
5079 let mut wtile = wt.borrow_mut();
5080 wtile.resize(TR * cols, 0.0);
5081 let mut r0 = 0usize;
5082 while r0 < rows {
5083 let tr = TR.min(rows - r0);
5084 let wt_addr = SendMut(wtile.as_mut_ptr());
5086 let run = |start: usize, end: usize| {
5087 for r in start..end {
5088 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
5089 let s = row_scale[r0 + r];
5090 let dst =
5092 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5093 for (d, &v) in dst.iter_mut().zip(row) {
5094 *d = (v as i8) as f32 * s;
5095 }
5096 }
5097 };
5098 dispatch_rows(pool, tr, &run);
5099 unsafe {
5101 accel_blas::cblas_sgemm(
5102 101, 111, 112, b as i32,
5106 tr as i32,
5107 cols as i32,
5108 1.0,
5109 xpanel.as_ptr(),
5110 cols as i32,
5111 wtile.as_ptr(),
5112 cols as i32,
5113 0.0,
5114 out.as_mut_ptr().add(r0),
5115 rows as i32,
5116 );
5117 }
5118 r0 += tr;
5119 }
5120 })
5121 });
5122}
5123
5124fn qmatmat(
5125 q: &[u8],
5126 row_scale: &[f32],
5127 pre: &[std::borrow::Cow<'_, [f32]>],
5128 rows: usize,
5129 cols: usize,
5130 out: &mut [f32],
5131 pool: Option<&Pool>,
5132) {
5133 let b = pre.len();
5134 debug_assert_eq!(out.len(), b * rows);
5135 let _p = crate::cpuprof::spans::span("mm.qmatmat");
5136 #[cfg(target_os = "macos")]
5141 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5142 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
5143 return;
5144 }
5145 #[cfg(target_arch = "aarch64")]
5146 if sdot_enabled() {
5147 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
5148 let out_addr = SendMut(out.as_mut_ptr());
5149 let blocked_ok = blocked_enabled();
5152 let use_i8mm = i8mm_enabled();
5153 if blocked_ok {
5154 let run = |start: usize, end: usize| {
5155 let mut o = start;
5156 while o < end {
5157 if o + 2 <= end {
5158 let r0 = &q[o * cols..(o + 1) * cols];
5159 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
5160 let mut bi = 0usize;
5161 while bi + 4 <= acts.len() {
5162 let xs = [
5163 acts[bi].xq.as_slice(),
5164 acts[bi + 1].xq.as_slice(),
5165 acts[bi + 2].xq.as_slice(),
5166 acts[bi + 3].xq.as_slice(),
5167 ];
5168 let d = if use_i8mm {
5169 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
5170 } else {
5171 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
5172 };
5173 for (r, row) in [r0, r1].into_iter().enumerate() {
5174 for k in 0..4 {
5175 let act = &acts[bi + k];
5176 let mut v = d[r][k] as f32 * act.sx;
5177 for &(j, xv) in &act.outliers {
5178 v += (row[j] as i8) as f32 * xv;
5179 }
5180 unsafe {
5181 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
5182 };
5183 }
5184 }
5185 bi += 4;
5186 }
5187 while bi < acts.len() {
5188 for (r, row) in [r0, r1].into_iter().enumerate() {
5189 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
5190 unsafe { *out_addr.at(bi * rows + o + r) = v };
5191 }
5192 bi += 1;
5193 }
5194 o += 2;
5195 } else {
5196 let row = &q[o * cols..(o + 1) * cols];
5197 for (bi, act) in acts.iter().enumerate() {
5198 let v = row_dot_sdot(row, act) * row_scale[o];
5199 unsafe { *out_addr.at(bi * rows + o) = v };
5200 }
5201 o += 1;
5202 }
5203 }
5204 };
5205 dispatch_rows(pool, rows, &run);
5206 return;
5207 }
5208 let run = |start: usize, end: usize| {
5209 for o in start..end {
5210 let row = &q[o * cols..(o + 1) * cols];
5211 for (bi, act) in acts.iter().enumerate() {
5212 let v = row_dot_sdot(row, act) * row_scale[o];
5213 unsafe { *out_addr.at(bi * rows + o) = v };
5214 }
5215 }
5216 };
5217 dispatch_rows(pool, rows, &run);
5218 return;
5219 }
5220 #[cfg(target_arch = "x86_64")]
5225 if avx2_a8w8_enabled() {
5226 let acts = {
5227 let _psa = crate::cpuprof::spans::span("mm.split_act");
5228 let a: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
5229 drop(_psa);
5230 a
5231 };
5232 let out_addr = SendMut(out.as_mut_ptr());
5233 let blocked_ok = blocked_enabled();
5236 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
5237 let run = |start: usize, end: usize| {
5238 let mut o = start;
5239 while o < end {
5240 if o + 2 <= end {
5241 let r0 = &q[o * cols..(o + 1) * cols];
5242 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
5243 let mut bi = 0usize;
5244 while bi + 4 <= acts.len() {
5245 let xs = [
5246 acts[bi].xq.as_slice(),
5247 acts[bi + 1].xq.as_slice(),
5248 acts[bi + 2].xq.as_slice(),
5249 acts[bi + 3].xq.as_slice(),
5250 ];
5251 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
5252 for (r, row) in [r0, r1].into_iter().enumerate() {
5253 for k in 0..4 {
5254 let act = &acts[bi + k];
5255 let mut v = d[r][k] as f32 * act.sx;
5256 for &(j, xv) in &act.outliers {
5257 v += (row[j] as i8) as f32 * xv;
5258 }
5259 unsafe {
5260 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
5261 };
5262 }
5263 }
5264 bi += 4;
5265 }
5266 while bi < acts.len() {
5267 for (r, row) in [r0, r1].into_iter().enumerate() {
5268 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
5269 unsafe { *out_addr.at(bi * rows + o + r) = v };
5270 }
5271 bi += 1;
5272 }
5273 o += 2;
5274 } else {
5275 let row = &q[o * cols..(o + 1) * cols];
5276 for (bi, act) in acts.iter().enumerate() {
5277 let v = row_dot_avx2(row, act) * row_scale[o];
5278 unsafe { *out_addr.at(bi * rows + o) = v };
5279 }
5280 o += 1;
5281 }
5282 }
5283 };
5284 dispatch_rows(pool, rows, &run);
5285 return;
5286 }
5287 let run = |start: usize, end: usize| {
5288 for o in start..end {
5289 let row = &q[o * cols..(o + 1) * cols];
5290 for (bi, act) in acts.iter().enumerate() {
5291 let v = row_dot_avx2(row, act) * row_scale[o];
5292 unsafe { *out_addr.at(bi * rows + o) = v };
5293 }
5294 }
5295 };
5296 dispatch_rows(pool, rows, &run);
5297 return;
5298 }
5299 let out_addr = SendMut(out.as_mut_ptr());
5300 let run = |start: usize, end: usize| {
5301 for o in start..end {
5302 let row = &q[o * cols..(o + 1) * cols];
5303 for (bi, x) in pre.iter().enumerate() {
5304 let mut acc = 0f32;
5305 for j in 0..cols {
5306 acc += (row[j] as i8) as f32 * x[j];
5307 }
5308 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
5309 }
5310 }
5311 };
5312 dispatch_rows(pool, rows, &run);
5313}
5314
5315fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
5318 match pool {
5319 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
5320 _ => run(0, rows),
5321 }
5322}
5323
5324fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
5326 let groups = rows * cols / GROUP_SIZE;
5327 bytes.split_at(groups * 16)
5328}
5329
5330#[inline]
5335fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
5336 #[cfg(target_arch = "aarch64")]
5337 unsafe {
5338 return vbit_fill4_neon(data, buf);
5339 }
5340 #[cfg(target_arch = "x86_64")]
5341 if avx2_enabled() {
5342 return unsafe { vbit_fill4_avx2(data, buf) };
5343 }
5344 #[allow(unreachable_code)]
5345 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5346 let u = unpack8::<4>(&data[blk * 4..]);
5347 for k in 0..8 {
5348 chunk[k] = (u[k] - 7) as i8 as u8;
5349 }
5350 }
5351}
5352
5353#[cfg(target_arch = "aarch64")]
5354#[target_feature(enable = "neon")]
5355unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
5356 unsafe {
5359 use core::arch::aarch64::*;
5360 let n = buf.len();
5361 let mask = vdupq_n_u8(0x0F);
5362 let seven = vdupq_n_s8(7);
5363 let mut g = 0usize;
5364 while g * 32 + 32 <= n {
5365 let b = vld1q_u8(data.as_ptr().add(g * 16));
5366 let hi = vshrq_n_u8::<4>(b);
5367 let lo = vandq_u8(b, mask);
5368 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
5369 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
5370 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
5371 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
5372 g += 1;
5373 }
5374 }
5375}
5376
5377#[cfg(target_arch = "x86_64")]
5378#[target_feature(enable = "avx2")]
5379unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
5380 unsafe {
5382 use core::arch::x86_64::*;
5383 let n = buf.len();
5384 let mask = _mm_set1_epi8(0x0F);
5385 let seven = _mm256_set1_epi8(7);
5386 let mut g = 0usize;
5387 while g * 32 + 32 <= n {
5388 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
5389 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
5390 let lo = _mm_and_si128(b, mask);
5391 let z = _mm256_sub_epi8(
5392 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
5393 seven,
5394 );
5395 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
5396 g += 1;
5397 }
5398 }
5399}
5400
5401#[inline(always)]
5406fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
5407 let mut acc = 0u64;
5408 for i in 0..B {
5409 acc = (acc << 8) | data[i] as u64;
5410 }
5411 let mask = (1u64 << B) - 1;
5412 let mut out = [0i32; 8];
5413 for (k, o) in out.iter_mut().enumerate() {
5414 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
5415 }
5416 out
5417}
5418
5419#[allow(clippy::too_many_arguments)]
5425fn vbitmatvec(
5426 bytes: &[u8],
5427 offsets: &[usize],
5428 x: &[f32],
5429 rows: usize,
5430 cols: usize,
5431 out: &mut [f32],
5432 pool: Option<&Pool>,
5433) {
5434 debug_assert_eq!(out.len(), rows);
5435 debug_assert_eq!(offsets.len(), rows + 1);
5436
5437 if a8w8_enabled() {
5441 let act = split_act(x);
5442 let out_addr = SendMut(out.as_mut_ptr());
5443 let run = move |start: usize, end: usize| {
5444 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
5445 };
5446 dispatch_rows(pool, rows, &run);
5447 return;
5448 }
5449
5450 let out_addr = SendMut(out.as_mut_ptr());
5451 let run = move |start: usize, end: usize| {
5452 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
5453 };
5454 dispatch_rows(pool, rows, &run);
5455}
5456
5457#[allow(clippy::too_many_arguments)]
5461fn vbit_range_a8w8(
5462 bytes: &[u8],
5463 offsets: &[usize],
5464 x: &[f32],
5465 act: &SplitAct,
5466 rows: usize,
5467 cols: usize,
5468 out: SendMut,
5469 start: usize,
5470 end: usize,
5471) {
5472 let ng = cols / GROUP_SIZE;
5473 let bits = &bytes[..rows];
5474 let sc_off = rows;
5475 let row_dot = |r: usize| -> f32 {
5476 let b = bits[r] as usize;
5477 let l = (1i32 << (b - 1)) - 1;
5478 let mask = (1u64 << b) - 1;
5479 let data = &bytes[offsets[r]..offsets[r + 1]];
5480 if b == 8 {
5481 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5483 let mut dot = 0f32;
5484 for g in 0..ng {
5485 let so = (r * ng + g) * 2;
5486 let sgf = f16_to_f32(u16::from_le_bytes([
5487 bytes[sc_off + so],
5488 bytes[sc_off + so + 1],
5489 ]));
5490 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5491 let mut gd = 0f32;
5492 for &xv in xg.iter() {
5493 if nbits < 8 {
5494 acc = (acc << 8) | data[idx] as u64;
5495 idx += 1;
5496 nbits += 8;
5497 }
5498 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5499 nbits -= 8;
5500 gd += (u - l) as f32 * xv;
5501 }
5502 dot += gd * sgf;
5503 }
5504 return dot;
5505 }
5506 thread_local! {
5510 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
5511 const { std::cell::RefCell::new(Vec::new()) };
5512 }
5513 #[inline(always)]
5514 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5515 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5516 let u = unpack8::<B>(&data[blk * B..]);
5517 for k in 0..8 {
5518 chunk[k] = (u[k] - l) as i8 as u8;
5519 }
5520 }
5521 }
5522 let _ = mask;
5523 VBIT_SCRATCH.with(|scratch| {
5524 let mut buf = scratch.borrow_mut();
5525 buf.resize(cols, 0);
5526 match b {
5527 3 => fill::<3>(data, l, &mut buf),
5528 4 => vbit_fill4(data, &mut buf),
5529 5 => fill::<5>(data, l, &mut buf),
5530 6 => fill::<6>(data, l, &mut buf),
5531 _ => unreachable!(),
5532 }
5533 let mut dot = 0f32;
5534 for g in 0..ng {
5535 let so = (r * ng + g) * 2;
5536 let s = f16_to_f32(u16::from_le_bytes([
5537 bytes[sc_off + so],
5538 bytes[sc_off + so + 1],
5539 ]));
5540 let d = dot_i8_i8(
5541 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5542 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5543 ) as f32
5544 * act.sx;
5545 dot += d * s;
5546 }
5547 for &(j, xv) in &act.outliers {
5548 let so = (r * ng + j / GROUP_SIZE) * 2;
5549 let s = f16_to_f32(u16::from_le_bytes([
5550 bytes[sc_off + so],
5551 bytes[sc_off + so + 1],
5552 ]));
5553 dot += (buf[j] as i8) as f32 * s * xv;
5555 }
5556 dot
5557 })
5558 };
5559 for r in start..end {
5560 unsafe { *out.at(r) = row_dot(r) };
5562 }
5563}
5564
5565#[allow(clippy::too_many_arguments)]
5567fn vbit_range_f32(
5568 bytes: &[u8],
5569 offsets: &[usize],
5570 x: &[f32],
5571 rows: usize,
5572 cols: usize,
5573 out: SendMut,
5574 start: usize,
5575 end: usize,
5576) {
5577 let ng = cols / GROUP_SIZE;
5578 let bits = &bytes[..rows];
5579 let sc_off = rows;
5580 #[inline(always)]
5584 fn dot_row<const B: usize>(
5585 data: &[u8],
5586 bytes: &[u8],
5587 sc_off: usize,
5588 r: usize,
5589 ng: usize,
5590 x: &[f32],
5591 ) -> f32 {
5592 let l = ((1i32 << (B - 1)) - 1) as f32;
5593 let gbytes = GROUP_SIZE * B / 8;
5594 let mut dot = 0f32;
5595 for g in 0..ng {
5596 let so = (r * ng + g) * 2;
5597 let s = f16_to_f32(u16::from_le_bytes([
5598 bytes[sc_off + so],
5599 bytes[sc_off + so + 1],
5600 ]));
5601 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5602 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5603 let mut gd = 0f32;
5604 for blk in 0..GROUP_SIZE / 8 {
5605 let u = unpack8::<B>(&gd0[blk * B..]);
5606 let xb = &xg[blk * 8..blk * 8 + 8];
5607 for k in 0..8 {
5608 gd += (u[k] as f32 - l) * xb[k];
5609 }
5610 }
5611 dot += gd * s;
5612 }
5613 dot
5614 }
5615 for r in start..end {
5616 let data = &bytes[offsets[r]..offsets[r + 1]];
5617 let v = match bits[r] {
5618 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5619 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5620 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5621 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5622 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5623 b => unreachable!("vbit bit-width {b} (validated at load)"),
5624 };
5625 unsafe { *out.at(r) = v };
5627 }
5628}
5629
5630#[allow(clippy::too_many_arguments)]
5635fn vbitmatvec2(
5636 bytes: &[u8],
5637 offsets: &[usize],
5638 x1: &[f32],
5639 x2: &[f32],
5640 rows: usize,
5641 cols: usize,
5642 o1: &mut [f32],
5643 o2: &mut [f32],
5644 pool: Option<&Pool>,
5645) {
5646 debug_assert_eq!(o1.len(), rows);
5647 debug_assert_eq!(o2.len(), rows);
5648
5649 if a8w8_enabled() {
5650 let a1 = split_act(x1);
5651 let a2 = split_act(x2);
5652 let p1 = SendMut(o1.as_mut_ptr());
5653 let p2 = SendMut(o2.as_mut_ptr());
5654 let run = move |start: usize, end: usize| {
5655 vbit_range2_a8w8(
5656 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5657 )
5658 };
5659 dispatch_rows(pool, rows, &run);
5660 return;
5661 }
5662
5663 let p1 = SendMut(o1.as_mut_ptr());
5664 let p2 = SendMut(o2.as_mut_ptr());
5665 let run = move |start: usize, end: usize| {
5666 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5667 };
5668 dispatch_rows(pool, rows, &run);
5669}
5670
5671#[allow(clippy::too_many_arguments)]
5675fn vbit_range2_a8w8(
5676 bytes: &[u8],
5677 offsets: &[usize],
5678 x1: &[f32],
5679 x2: &[f32],
5680 a1: &SplitAct,
5681 a2: &SplitAct,
5682 rows: usize,
5683 cols: usize,
5684 p1: SendMut,
5685 p2: SendMut,
5686 start: usize,
5687 end: usize,
5688) {
5689 let ng = cols / GROUP_SIZE;
5690 let bits = &bytes[..rows];
5691 let sc_off = rows;
5692 let row_dots = |r: usize| -> (f32, f32) {
5693 let b = bits[r] as usize;
5694 let l = (1i32 << (b - 1)) - 1;
5695 let data = &bytes[offsets[r]..offsets[r + 1]];
5696 if b == 8 {
5697 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5700 let (mut d1, mut d2) = (0f32, 0f32);
5701 for g in 0..ng {
5702 let so = (r * ng + g) * 2;
5703 let sgf = f16_to_f32(u16::from_le_bytes([
5704 bytes[sc_off + so],
5705 bytes[sc_off + so + 1],
5706 ]));
5707 let (mut g1, mut g2) = (0f32, 0f32);
5708 for k in 0..GROUP_SIZE {
5709 if nbits < 8 {
5710 acc = (acc << 8) | data[idx] as u64;
5711 idx += 1;
5712 nbits += 8;
5713 }
5714 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5715 nbits -= 8;
5716 let w = (u - l) as f32;
5717 g1 += w * x1[g * GROUP_SIZE + k];
5718 g2 += w * x2[g * GROUP_SIZE + k];
5719 }
5720 d1 += g1 * sgf;
5721 d2 += g2 * sgf;
5722 }
5723 return (d1, d2);
5724 }
5725 thread_local! {
5726 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5727 const { std::cell::RefCell::new(Vec::new()) };
5728 }
5729 #[inline(always)]
5730 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5731 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5732 let u = unpack8::<B>(&data[blk * B..]);
5733 for k in 0..8 {
5734 chunk[k] = (u[k] - l) as i8 as u8;
5735 }
5736 }
5737 }
5738 VBIT_SCRATCH2.with(|scratch| {
5739 let mut buf = scratch.borrow_mut();
5740 buf.resize(cols, 0);
5741 match b {
5742 3 => fill::<3>(data, l, &mut buf),
5743 4 => vbit_fill4(data, &mut buf),
5744 5 => fill::<5>(data, l, &mut buf),
5745 6 => fill::<6>(data, l, &mut buf),
5746 _ => unreachable!(),
5747 }
5748 let (mut d1, mut d2) = (0f32, 0f32);
5749 for g in 0..ng {
5750 let so = (r * ng + g) * 2;
5751 let s = f16_to_f32(u16::from_le_bytes([
5752 bytes[sc_off + so],
5753 bytes[sc_off + so + 1],
5754 ]));
5755 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5756 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5757 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5758 d1 += v1 * s;
5759 d2 += v2 * s;
5760 }
5761 for &(j, xv) in &a1.outliers {
5762 let so = (r * ng + j / GROUP_SIZE) * 2;
5763 let s = f16_to_f32(u16::from_le_bytes([
5764 bytes[sc_off + so],
5765 bytes[sc_off + so + 1],
5766 ]));
5767 d1 += (buf[j] as i8) as f32 * s * xv;
5768 }
5769 for &(j, xv) in &a2.outliers {
5770 let so = (r * ng + j / GROUP_SIZE) * 2;
5771 let s = f16_to_f32(u16::from_le_bytes([
5772 bytes[sc_off + so],
5773 bytes[sc_off + so + 1],
5774 ]));
5775 d2 += (buf[j] as i8) as f32 * s * xv;
5776 }
5777 (d1, d2)
5778 })
5779 };
5780 for r in start..end {
5781 let (v1, v2) = row_dots(r);
5782 unsafe {
5784 *p1.at(r) = v1;
5785 *p2.at(r) = v2;
5786 }
5787 }
5788}
5789
5790#[allow(clippy::too_many_arguments)]
5794fn vbit_range2_f32(
5795 bytes: &[u8],
5796 offsets: &[usize],
5797 x1: &[f32],
5798 x2: &[f32],
5799 rows: usize,
5800 cols: usize,
5801 p1: SendMut,
5802 p2: SendMut,
5803 start: usize,
5804 end: usize,
5805) {
5806 let ng = cols / GROUP_SIZE;
5807 let bits = &bytes[..rows];
5808 let sc_off = rows;
5809 #[inline(always)]
5810 #[allow(clippy::too_many_arguments)]
5811 fn dot_row2<const B: usize>(
5812 data: &[u8],
5813 bytes: &[u8],
5814 sc_off: usize,
5815 r: usize,
5816 ng: usize,
5817 x1: &[f32],
5818 x2: &[f32],
5819 ) -> (f32, f32) {
5820 let l = ((1i32 << (B - 1)) - 1) as f32;
5821 let gbytes = GROUP_SIZE * B / 8;
5822 let (mut d1, mut d2) = (0f32, 0f32);
5823 for g in 0..ng {
5824 let so = (r * ng + g) * 2;
5825 let s = f16_to_f32(u16::from_le_bytes([
5826 bytes[sc_off + so],
5827 bytes[sc_off + so + 1],
5828 ]));
5829 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5830 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5831 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5832 let (mut g1, mut g2) = (0f32, 0f32);
5833 for blk in 0..GROUP_SIZE / 8 {
5834 let u = unpack8::<B>(&gd0[blk * B..]);
5835 for k in 0..8 {
5836 let w = u[k] as f32 - l;
5837 g1 += w * x1g[blk * 8 + k];
5838 g2 += w * x2g[blk * 8 + k];
5839 }
5840 }
5841 d1 += g1 * s;
5842 d2 += g2 * s;
5843 }
5844 (d1, d2)
5845 }
5846 for r in start..end {
5847 let data = &bytes[offsets[r]..offsets[r + 1]];
5848 let (v1, v2) = match bits[r] {
5849 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5850 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5851 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5852 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5853 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5854 b => unreachable!("vbit bit-width {b} (validated at load)"),
5855 };
5856 unsafe {
5858 *p1.at(r) = v1;
5859 *p2.at(r) = v2;
5860 }
5861 }
5862}
5863
5864#[inline]
5871#[allow(unreachable_code)]
5872fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5873 #[cfg(target_arch = "aarch64")]
5874 unsafe {
5875 return dot_q4t_row_sdot(bytes, r, gpr, xq);
5876 }
5877 #[cfg(target_arch = "x86_64")]
5878 unsafe {
5879 if vnni_tiles_enabled() {
5880 return dot_q4t_row_vnni(bytes, r, gpr, xq);
5881 }
5882 return dot_q4t_row_avx2(bytes, r, gpr, xq);
5883 }
5884 let mut acc = 0f32;
5885 for gi in 0..gpr {
5886 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5887 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5888 let mut d = 0i32;
5889 for (k, &b) in tile[2..].iter().enumerate() {
5890 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5891 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5892 }
5893 acc += d as f32 * s;
5894 }
5895 acc
5896}
5897
5898#[cfg(target_arch = "aarch64")]
5899#[target_feature(enable = "neon,dotprod")]
5900unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5901 unsafe {
5904 use core::arch::aarch64::*;
5905 use core::arch::asm;
5906 let lomask = vdupq_n_u8(0x0F);
5907 let eight = vdupq_n_s8(8);
5908 let mut acc = 0f32;
5909 for gi in 0..gpr {
5910 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5911 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5912 let b = vld1q_u8(t.add(2));
5913 let lo = vandq_u8(b, lomask);
5914 let hi = vshrq_n_u8::<4>(b);
5915 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5916 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5917 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5918 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5919 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5920 asm!(
5921 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5922 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5923 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5924 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5925 options(pure, nomem, nostack),
5926 );
5927 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5928 }
5929 acc
5930 }
5931}
5932
5933#[cfg(target_arch = "x86_64")]
5934#[target_feature(enable = "avx2")]
5935unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5936 unsafe {
5938 use core::arch::x86_64::*;
5939 let lomask = _mm_set1_epi8(0x0F);
5940 let eight = _mm256_set1_epi8(8);
5941 let ones = _mm256_set1_epi16(1);
5942 let mut acc = 0f32;
5943 for gi in 0..gpr {
5944 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5945 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5946 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5947 let lo = _mm_and_si128(b, lomask);
5948 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5949 let w = _mm256_sub_epi8(
5950 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5951 eight,
5952 );
5953 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5954 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5955 let d = _mm256_madd_epi16(p16, ones);
5956 let hi128 = _mm256_extracti128_si256::<1>(d);
5957 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5958 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5959 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5960 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5961 }
5962 acc
5963 }
5964}
5965
5966#[cfg(target_arch = "x86_64")]
5970#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5971unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5972 unsafe {
5974 use core::arch::x86_64::*;
5975 let lomask = _mm_set1_epi8(0x0F);
5976 let eight = _mm256_set1_epi8(8);
5977 let mut acc = 0f32;
5978 for gi in 0..gpr {
5979 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5980 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5981 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5982 let lo = _mm_and_si128(b, lomask);
5983 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5984 let w = _mm256_sub_epi8(
5985 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5986 eight,
5987 );
5988 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5989 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5990 acc += d as f32 * s;
5991 }
5992 acc
5993 }
5994}
5995
5996#[cfg(target_arch = "x86_64")]
6001#[target_feature(enable = "avx2,fma")]
6006unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6007 unsafe {
6009 use core::arch::x86_64::*;
6010 let lomask = _mm_set1_epi8(0x0F);
6011 let eight = _mm256_set1_epi8(8);
6012 let ones = _mm256_set1_epi16(1);
6013 let mut f0 = _mm256_setzero_ps();
6026 let mut f1 = _mm256_setzero_ps();
6027 let mut f2 = _mm256_setzero_ps();
6028 let mut f3 = _mm256_setzero_ps();
6029 for gi in 0..gpr {
6030 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6031 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6032 let sv = _mm256_set1_ps(s);
6033 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
6034 let lo = _mm_and_si128(bb, lomask);
6035 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
6036 let w = _mm256_sub_epi8(
6037 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6038 eight,
6039 );
6040 let aw = _mm256_abs_epi8(w);
6041 let off = gi * GROUP_SIZE;
6042 let dot = |xq: &[i8]| {
6043 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
6044 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6045 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
6046 };
6047 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
6048 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
6049 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
6050 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
6051 }
6052 [
6053 hsum256_ps(f0),
6054 hsum256_ps(f1),
6055 hsum256_ps(f2),
6056 hsum256_ps(f3),
6057 ]
6058 }
6059}
6060
6061#[cfg(target_arch = "x86_64")]
6064#[target_feature(enable = "avx2")]
6065#[inline]
6066unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
6067 unsafe {
6069 use core::arch::x86_64::*;
6070 let hi = _mm256_extractf128_ps::<1>(v);
6071 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
6072 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
6073 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
6074 _mm_cvtss_f32(s)
6075 }
6076}
6077
6078#[cfg(target_arch = "x86_64")]
6080#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
6081unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6082 unsafe {
6084 use core::arch::x86_64::*;
6085 let lomask = _mm_set1_epi8(0x0F);
6086 let eight = _mm256_set1_epi8(8);
6087 let mut f0 = _mm256_setzero_ps();
6090 let mut f1 = _mm256_setzero_ps();
6091 let mut f2 = _mm256_setzero_ps();
6092 let mut f3 = _mm256_setzero_ps();
6093 for gi in 0..gpr {
6094 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6095 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6096 let sv = _mm256_set1_ps(s);
6097 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
6098 let lo = _mm_and_si128(bb, lomask);
6099 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
6100 let w = _mm256_sub_epi8(
6101 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6102 eight,
6103 );
6104 let aw = _mm256_abs_epi8(w);
6105 let off = gi * GROUP_SIZE;
6106 let dot = |xq: &[i8]| {
6107 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
6108 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
6109 _mm256_setzero_si256(),
6110 aw,
6111 _mm256_sign_epi8(x, w),
6112 ))
6113 };
6114 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
6115 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
6116 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
6117 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
6118 }
6119 let acc = [
6120 hsum256_ps(f0),
6121 hsum256_ps(f1),
6122 hsum256_ps(f2),
6123 hsum256_ps(f3),
6124 ];
6125 acc
6126 }
6127}
6128
6129#[cfg(target_arch = "aarch64")]
6134#[target_feature(enable = "neon,dotprod")]
6135unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6136 unsafe {
6138 use core::arch::aarch64::*;
6139 use core::arch::asm;
6140 let lomask = vdupq_n_u8(0x0F);
6141 let eight = vdupq_n_s8(8);
6142 let mut acc = [0f32; 4];
6143 for gi in 0..gpr {
6144 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6145 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6146 let b = vld1q_u8(t.add(2));
6147 let lo = vandq_u8(b, lomask);
6148 let hi = vshrq_n_u8::<4>(b);
6149 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6150 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6151 for (k, xq) in xs.iter().enumerate() {
6152 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6153 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6154 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6155 asm!(
6156 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6157 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6158 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6159 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6160 options(pure, nomem, nostack),
6161 );
6162 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6163 }
6164 }
6165 acc
6166 }
6167}
6168
6169#[inline]
6171fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6172 let gi = j / GROUP_SIZE;
6173 let k = j % GROUP_SIZE;
6174 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6175 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6176 let byte = tile[2 + k / 2];
6177 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
6178 ((nib as i32 - 8) as f32, s)
6179}
6180
6181#[inline]
6184fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6185 let mut acc = 0f32;
6186 for gi in 0..gpr {
6187 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6188 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6189 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6190 let mut ga = 0f32;
6191 for (k, &b) in tile[2..].iter().enumerate() {
6192 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6193 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6194 }
6195 acc += ga * s;
6196 }
6197 acc
6198}
6199
6200fn row_span_offsets(
6208 dtype: TensorDtype,
6209 rows: usize,
6210 cols: usize,
6211 r: usize,
6212) -> Option<[(usize, usize); 3]> {
6213 if r >= rows || cols % GROUP_SIZE != 0 {
6214 return None;
6215 }
6216 let gpr = cols / GROUP_SIZE;
6217 match dtype {
6218 TensorDtype::Q4Tiled => Some([(r * gpr * Q4_TILE, gpr * Q4_TILE), (0, 0), (0, 0)]),
6219 TensorDtype::Q4TiledP => {
6220 let (params, codes, stride) = q4tp_sections(rows, cols);
6221 Some([
6222 (r * gpr * Q4TP_NIB, gpr * Q4TP_NIB),
6223 (params + r * 4, 4),
6224 (codes + r * stride, stride),
6225 ])
6226 }
6227 _ => None,
6228 }
6229}
6230
6231struct Q4tpView<'a> {
6232 nib: &'a [u8],
6233 params: &'a [u8],
6234 codes: &'a [u8],
6235 stride: usize,
6236 zero_rung: bool,
6238}
6239
6240impl<'a> Q4tpView<'a> {
6241 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
6242 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
6243 Self {
6244 nib: &bytes[..params_off],
6245 params: &bytes[params_off..codes_off],
6246 codes: &bytes[codes_off..],
6247 stride,
6248 zero_rung: false,
6249 }
6250 }
6251
6252 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
6254 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
6255 Self {
6256 nib: &bytes[..params_off],
6257 params: &bytes[params_off..codes_off],
6258 codes: &bytes[codes_off..],
6259 stride,
6260 zero_rung: true,
6261 }
6262 }
6263
6264 #[cfg(target_arch = "x86_64")]
6287 #[inline]
6288 fn ladder_via_avx2(&self) -> bool {
6289 avx2_enabled()
6290 }
6291 #[cfg(not(target_arch = "x86_64"))]
6292 #[inline]
6293 fn ladder_via_avx2(&self) -> bool {
6294 false
6295 }
6296
6297 #[inline]
6298 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
6299 let tab = if self.zero_rung {
6300 q2tp_ladder(self.params, r)
6301 } else if self.ladder_via_avx2() {
6302 #[cfg(target_arch = "x86_64")]
6309 let tab = {
6310 let p = &self.params[r * 4..r * 4 + 4];
6311 let lo = f16_to_f32(u16::from_le_bytes([p[0], p[1]]));
6312 let st = f16_to_f32(u16::from_le_bytes([p[2], p[3]]));
6313 let mut tabc = [0f32; 128];
6314 unsafe { ladder_4(&[lo; 4], &[st; 4], &mut tabc) };
6315 let mut t = [0f32; 32];
6316 for (g, o) in t.iter_mut().enumerate() {
6317 *o = tabc[g * 4];
6318 }
6319 t
6320 };
6321 #[cfg(not(target_arch = "x86_64"))]
6325 let tab = q4tp_ladder(self.params, r);
6326 tab
6327 } else {
6328 q4tp_ladder(self.params, r)
6329 };
6330 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
6331 let out = &mut out[..gpr];
6332 let mut chunks = out.chunks_exact_mut(8);
6333 let mut ci = 0usize;
6334 for c in &mut chunks {
6335 let w = u64::from(codes[ci])
6336 | u64::from(codes[ci + 1]) << 8
6337 | u64::from(codes[ci + 2]) << 16
6338 | u64::from(codes[ci + 3]) << 24
6339 | u64::from(codes[ci + 4]) << 32;
6340 for (k, o) in c.iter_mut().enumerate() {
6341 *o = tab[((w >> (5 * k)) & 31) as usize];
6342 }
6343 ci += 5;
6344 }
6345 let tail = &codes[ci..];
6348 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6349 *o = tab[q4tp_code(tail, k)];
6350 }
6351 }
6352}
6353
6354#[inline]
6355fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6356 #[cfg(target_arch = "aarch64")]
6357 unsafe {
6358 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
6359 }
6360 #[cfg(target_arch = "x86_64")]
6361 unsafe {
6362 if vnni_tiles_enabled() {
6363 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
6364 }
6365 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
6366 }
6367 #[allow(unreachable_code)]
6368 {
6369 let mut acc = 0f32;
6370 for gi in 0..gpr {
6371 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6372 let s = scales[gi];
6373 let mut d = 0i32;
6374 for (k, &b) in tile.iter().enumerate() {
6375 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6376 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6377 }
6378 acc += d as f32 * s;
6379 }
6380 acc
6381 }
6382}
6383
6384#[cfg(target_arch = "aarch64")]
6387#[target_feature(enable = "neon,dotprod")]
6388unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6389 unsafe {
6392 use core::arch::aarch64::*;
6393 use core::arch::asm;
6394 let lomask = vdupq_n_u8(0x0F);
6395 let eight = vdupq_n_s8(8);
6396 let mut acc = 0f32;
6397 for gi in 0..gpr {
6398 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6399 let s = *scales.get_unchecked(gi);
6400 let b = vld1q_u8(t);
6401 let lo = vandq_u8(b, lomask);
6402 let hi = vshrq_n_u8::<4>(b);
6403 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6404 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6405 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6406 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6407 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6408 asm!(
6409 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6410 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6411 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6412 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6413 options(pure, nomem, nostack),
6414 );
6415 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6416 }
6417 acc
6418 }
6419}
6420
6421#[cfg(target_arch = "x86_64")]
6422#[target_feature(enable = "avx2")]
6423unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6424 unsafe {
6426 use core::arch::x86_64::*;
6427 let lomask = _mm_set1_epi8(0x0F);
6428 let eight = _mm256_set1_epi8(8);
6429 let ones = _mm256_set1_epi16(1);
6430 let mut acc = 0f32;
6431 for gi in 0..gpr {
6432 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6433 let s = *scales.get_unchecked(gi);
6434 let b = _mm_loadu_si128(t as *const __m128i);
6435 let lo = _mm_and_si128(b, lomask);
6436 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6437 let w = _mm256_sub_epi8(
6438 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6439 eight,
6440 );
6441 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6442 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6443 let d = _mm256_madd_epi16(p16, ones);
6444 let hi128 = _mm256_extracti128_si256::<1>(d);
6445 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6446 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6447 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6448 acc += _mm_cvtsi128_si32(s32) as f32 * s;
6449 }
6450 acc
6451 }
6452}
6453
6454#[cfg(target_arch = "x86_64")]
6457#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6458unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6459 unsafe {
6461 use core::arch::x86_64::*;
6462 let lomask = _mm_set1_epi8(0x0F);
6463 let eight = _mm256_set1_epi8(8);
6464 let mut acc = 0f32;
6465 for gi in 0..gpr {
6466 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6467 let s = *scales.get_unchecked(gi);
6468 let b = _mm_loadu_si128(t as *const __m128i);
6469 let lo = _mm_and_si128(b, lomask);
6470 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6471 let w = _mm256_sub_epi8(
6472 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6473 eight,
6474 );
6475 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6476 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
6477 }
6478 acc
6479 }
6480}
6481
6482#[cfg(target_arch = "x86_64")]
6489#[target_feature(enable = "avx2")]
6490#[inline]
6491unsafe fn exp2_128(x: core::arch::x86_64::__m128) -> core::arch::x86_64::__m128 {
6492 unsafe {
6493 use core::arch::x86_64::*;
6494 let tr = _mm_cvttps_epi32(x);
6496 let below = _mm_and_si128(
6499 _mm_castps_si128(_mm_cmplt_ps(x, _mm_cvtepi32_ps(tr))),
6500 _mm_set1_epi32(1),
6501 );
6502 let n = _mm_sub_epi32(tr, below); let r = _mm_sub_ps(x, _mm_cvtepi32_ps(n)); #[allow(clippy::excessive_precision, clippy::approx_constant)]
6513 let a = [
6514 6.931_471_805_599_453e-1,
6515 2.402_265_069_591_007e-1,
6516 5.550_410_866_482_158e-2,
6517 9.618_129_107_628_477e-3,
6518 1.333_355_814_642_844e-3,
6519 1.540_353_039_338_161e-4,
6520 1.525_273_380_405_984e-5,
6521 1.321_548_679_014_431e-6,
6522 1.017_808_600_923_969_6e-7,
6523 ];
6524 let mut p = _mm_set1_ps(a[8]);
6525 for &c in a[..8].iter().rev() {
6526 p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(c));
6527 }
6528 p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(1.0));
6529 let ncl = _mm_min_epi32(
6535 _mm_max_epi32(n, _mm_set1_epi32(-126)),
6536 _mm_set1_epi32(127),
6537 );
6538 let mut bits = _mm_add_epi32(_mm_castps_si128(p), _mm_slli_epi32(ncl, 23));
6539 let over = _mm_cmpgt_epi32(n, _mm_set1_epi32(127));
6540 bits = _mm_or_si128(
6541 _mm_andnot_si128(over, bits),
6542 _mm_and_si128(over, _mm_set1_epi32(0x007F_8000)),
6543 );
6544 let zero = _mm_cmplt_ps(x, _mm_set1_ps(-149.0));
6545 _mm_andnot_ps(zero, _mm_castsi128_ps(bits))
6546 }
6547}
6548
6549#[cfg(target_arch = "x86_64")]
6551#[inline]
6552unsafe fn lanes4(v: core::arch::x86_64::__m128) -> [f32; 4] {
6553 use core::arch::x86_64::*;
6554 unsafe {
6555 [
6556 _mm_cvtss_f32(v),
6557 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0x55)),
6558 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xAA)),
6559 _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xFF)),
6560 ]
6561 }
6562}
6563
6564#[cfg(target_arch = "x86_64")]
6572#[target_feature(enable = "avx2")]
6573#[inline]
6574unsafe fn ladder_4(lo4: &[f32; 4], st4: &[f32; 4], tabc: &mut [f32; 128]) {
6575 unsafe {
6576 use core::arch::x86_64::*;
6577 let ratio = exp2_128(_mm_loadu_ps(st4.as_ptr()));
6578 let mut t = exp2_128(_mm_loadu_ps(lo4.as_ptr()));
6579 tabc[0..4].copy_from_slice(&lanes4(t));
6580 for c in 1..32 {
6581 t = _mm_mul_ps(t, ratio);
6582 tabc[c * 4..c * 4 + 4].copy_from_slice(&lanes4(t));
6583 }
6584 }
6585}
6586
6587#[cfg(target_arch = "x86_64")]
6602#[target_feature(enable = "avx2")]
6603unsafe fn q4tp_scales_lanes(
6604 vs: &[&Q4tpView<'_>],
6605 rows: &[usize],
6606 gpr: usize,
6607 out: &mut [f32],
6608) {
6609 use core::arch::x86_64::*;
6610 debug_assert_eq!(vs.len(), rows.len());
6611 debug_assert!(out.len() >= rows.len() * gpr);
6612 let mut done = 0usize;
6613 for (ci, chunk) in rows.chunks(4).enumerate() {
6614 let vchunk = &vs[ci * 4..ci * 4 + chunk.len()];
6615 if chunk.len() < 4 || vchunk.iter().any(|v| v.zero_rung) {
6616 for (l, &r) in chunk.iter().enumerate() {
6617 vchunk[l].scales_into(r, gpr, &mut out[(done + l) * gpr..(done + l + 1) * gpr]);
6618 }
6619 done += chunk.len();
6620 continue;
6621 }
6622 let (r0, r1, r2, r3) = (chunk[0], chunk[1], chunk[2], chunk[3]);
6623 let v0 = vchunk[0];
6624 let v1 = vchunk[1];
6625 let v2 = vchunk[2];
6626 let v3 = vchunk[3];
6627 let mut lo = [0f32; 4];
6629 let mut st = [0f32; 4];
6630 for (l, r) in [r0, r1, r2, r3].iter().enumerate() {
6631 let o = *r * 4;
6632 let vv = vchunk[l];
6633 lo[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6634 vv.params[o],
6635 vv.params[o + 1],
6636 ]));
6637 st[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6638 vv.params[o + 2],
6639 vv.params[o + 3],
6640 ]));
6641 }
6642 let lo_v = _mm_loadu_ps(lo.as_ptr());
6643 let st_v = _mm_loadu_ps(st.as_ptr());
6644 let ratio = exp2_128(st_v);
6645 let mut t = exp2_128(lo_v);
6646 let mut tabc = [0f32; 32 * 4];
6648 tabc[0..4].copy_from_slice(&[
6649 _mm_cvtss_f32(t),
6650 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6651 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6652 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6653 ]);
6654 for c in 1..32 {
6655 t = _mm_mul_ps(t, ratio);
6656 tabc[c * 4..c * 4 + 4].copy_from_slice(&[
6657 _mm_cvtss_f32(t),
6658 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6659 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6660 _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6661 ]);
6662 }
6663 let vs4 = [v0, v1, v2, v3];
6665 for (l, &r) in [r0, r1, r2, r3].iter().enumerate() {
6666 let v = vs4[l];
6667 let codes = &v.codes[r * v.stride..(r + 1) * v.stride];
6668 let dst = &mut out[(done + l) * gpr..(done + l + 1) * gpr];
6669 let mut chunks = dst.chunks_exact_mut(8);
6670 let mut ci = 0usize;
6671 for c in &mut chunks {
6672 let w = u64::from(codes[ci])
6673 | u64::from(codes[ci + 1]) << 8
6674 | u64::from(codes[ci + 2]) << 16
6675 | u64::from(codes[ci + 3]) << 24
6676 | u64::from(codes[ci + 4]) << 32;
6677 for (k, o) in c.iter_mut().enumerate() {
6678 *o = tabc[(((w >> (5 * k)) & 31) as usize) * 4 + l];
6679 }
6680 ci += 5;
6681 }
6682 let tail = &codes[ci..];
6683 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6684 *o = tabc[q4tp_code(tail, k) * 4 + l];
6685 }
6686 }
6687 done += 4;
6688 }
6689}
6690
6691fn q4tp_lanes() -> usize {
6699 #[cfg(not(target_arch = "x86_64"))]
6700 {
6701 1
6702 }
6703 #[cfg(target_arch = "x86_64")]
6704 {
6705 static ON: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6706 *ON.get_or_init(|| {
6707 let n: usize = std::env::var("CMF_Q4TP_LANES")
6708 .ok()
6709 .and_then(|v| v.parse().ok())
6710 .unwrap_or(4);
6711 match n {
6712 1 => 1,
6713 2 => 2,
6714 8 => 8,
6715 _ => 4,
6716 }
6717 })
6718 }
6719}
6720
6721#[cfg(target_arch = "x86_64")]
6732#[target_feature(enable = "avx2")]
6733#[inline]
6734unsafe fn q4tp_unpack_w(tile: *const u8) -> core::arch::x86_64::__m256i {
6735 unsafe {
6736 use core::arch::x86_64::*;
6737 let lomask = _mm_set1_epi8(0x0F);
6738 let b = _mm_loadu_si128(tile as *const __m128i);
6739 let lo = _mm_and_si128(b, lomask);
6740 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6741 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi))
6742 }
6743}
6744
6745#[cfg(target_arch = "x86_64")]
6752#[target_feature(enable = "avx2")]
6753#[inline]
6754unsafe fn q4tp_madd_wx(
6755 w: core::arch::x86_64::__m256i,
6756 xp: *const i8,
6757) -> core::arch::x86_64::__m256i {
6758 unsafe {
6759 use core::arch::x86_64::*;
6760 let ones = _mm256_set1_epi16(1);
6761 let x = _mm256_loadu_si256(xp as *const __m256i);
6762 _mm256_madd_epi16(_mm256_maddubs_epi16(w, x), ones)
6763 }
6764}
6765
6766#[cfg(target_arch = "x86_64")]
6771#[target_feature(enable = "avx2")]
6772#[inline]
6773unsafe fn q4tp_hsum256(acc: core::arch::x86_64::__m256) -> f32 {
6774 unsafe {
6775 use core::arch::x86_64::*;
6776 let a = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps::<1>(acc));
6777 let b = _mm_add_ps(a, _mm_movehl_ps(a, a));
6778 let c = _mm_add_ss(b, _mm_shuffle_ps(b, b, 0x55));
6779 _mm_cvtss_f32(c)
6780 }
6781}
6782
6783#[cfg(target_arch = "x86_64")]
6787#[inline]
6788fn q4tp_xg8_corr(xg8: &[f32], sc: &[f32]) -> f32 {
6789 let n = sc.len().min(xg8.len());
6790 8.0 * crate::attention::dot_f32(&xg8[..n], &sc[..n])
6791}
6792
6793#[cfg(target_arch = "x86_64")]
6801#[target_feature(enable = "avx2,fma")]
6802#[inline]
6803unsafe fn q4tp_cell_dot_v(
6804 nib: *const u8,
6805 r: usize,
6806 gpr: usize,
6807 xp: *const i8,
6808 sc: *const f32,
6809 xg8: *const f32,
6810) -> f32 {
6811 unsafe {
6812 use core::arch::x86_64::*;
6813 let mut acc = _mm256_setzero_ps();
6814 for gi in 0..gpr {
6815 let w = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
6816 let d = q4tp_madd_wx(w, xp.add(gi * GROUP_SIZE));
6817 acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), _mm256_set1_ps(*sc.add(gi)), acc);
6818 }
6819 q4tp_hsum256(acc)
6820 - q4tp_xg8_corr(
6821 std::slice::from_raw_parts(xg8, gpr),
6822 std::slice::from_raw_parts(sc, gpr),
6823 )
6824 }
6825}
6826
6827#[cfg(target_arch = "x86_64")]
6840#[target_feature(enable = "avx2,fma")]
6841unsafe fn q4tp_matmat_row_cols_v(
6842 v: &Q4tpView<'_>,
6843 r: usize,
6844 gpr: usize,
6845 acts: &[SplitAct],
6846 sc: &[f32],
6847 abase: usize,
6848 rows: usize,
6849 out: SendMut,
6850) {
6851 use core::arch::x86_64::*;
6852 debug_assert_eq!(sc.len(), gpr);
6853 let n = acts.len();
6854 let mut bi = 0usize;
6855 while bi + 4 <= n {
6856 let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
6857 let mut accs = [_mm256_setzero_ps(); 4];
6858 for gi in 0..gpr {
6859 let w = q4tp_unpack_w(v.nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB));
6860 let sb = _mm256_set1_ps(*sc.get_unchecked(gi));
6861 for (k, act) in a.iter().enumerate() {
6862 let xp = act.xq.as_ptr().add(gi * GROUP_SIZE);
6863 let d = q4tp_madd_wx(w, xp);
6864 accs[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[k]);
6865 }
6866 }
6867 for k in 0..4 {
6868 let act = a[k];
6869 let mut acc = (q4tp_hsum256(accs[k]) - q4tp_xg8_corr(&act.xg8, sc)) * act.sx;
6870 for &(j, xv) in &act.outliers {
6871 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6872 acc += w * s * xv;
6873 }
6874 unsafe { *out.at((abase + bi + k) * rows + r) = acc };
6876 }
6877 bi += 4;
6878 }
6879 while bi < n {
6881 let act = &acts[bi];
6882 let mut acc = q4tp_cell_dot_v(
6883 v.nib.as_ptr(),
6884 r,
6885 gpr,
6886 act.xq.as_ptr(),
6887 sc.as_ptr(),
6888 act.xg8.as_ptr(),
6889 ) * act.sx;
6890 for &(j, xv) in &act.outliers {
6891 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6892 acc += w * s * xv;
6893 }
6894 unsafe { *out.at((abase + bi) * rows + r) = acc };
6896 bi += 1;
6897 }
6898}
6899
6900fn q4tp_panel_cols(cols: usize) -> usize {
6906 static ENV: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
6907 if let Some(v) = *ENV.get_or_init(|| {
6908 std::env::var("CMF_Q4TP_PANEL")
6909 .ok()
6910 .and_then(|v| v.parse().ok())
6911 .filter(|v: &usize| *v > 0)
6912 }) {
6913 return v;
6914 }
6915 #[cfg(target_arch = "x86_64")]
6916 {
6917 ((l2_bytes() / 2) / cols.max(1)).clamp(8, 256) & !3
6918 }
6919 #[cfg(not(target_arch = "x86_64"))]
6920 {
6921 let _ = cols;
6922 256
6923 }
6924}
6925
6926#[cfg(target_arch = "x86_64")]
6928fn l2_bytes() -> usize {
6929 static L2: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6930 *L2.get_or_init(|| {
6931 let s = std::fs::read_to_string("/sys/devices/system/cpu/cpu0/cache/index2/size").unwrap_or_default();
6932 let s = s.trim();
6933 let (num, mul) = match s.strip_suffix('K') {
6934 Some(n) => (n, 1024),
6935 None => match s.strip_suffix('M') {
6936 Some(n) => (n, 1024 * 1024),
6937 None => (s, 1),
6938 },
6939 };
6940 num.parse::<usize>().ok().map(|n| n * mul).filter(|&b| b >= 64 * 1024).unwrap_or(1 << 20)
6941 })
6942}
6943
6944#[cfg(target_arch = "x86_64")]
6947fn q4tp_fma_x86() -> bool {
6948 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6949 *ON.get_or_init(|| {
6950 std::is_x86_feature_detected!("avx2")
6951 && std::is_x86_feature_detected!("fma")
6952 && std::env::var("CMF_Q4TP_FMA").map(|v| v != "0").unwrap_or(true)
6953 })
6954}
6955
6956#[cfg(target_arch = "x86_64")]
6972#[target_feature(enable = "avx2,fma")]
6973#[allow(clippy::too_many_arguments)]
6974unsafe fn q4tp_matmat_2rows_fma(
6975 v: &Q4tpView<'_>,
6976 r: usize,
6977 gpr: usize,
6978 acts: &[SplitAct],
6979 sc0: &[f32],
6980 sc1: &[f32],
6981 abase: usize,
6982 rows: usize,
6983 out: SendMut,
6984) {
6985 use core::arch::x86_64::*;
6986 debug_assert!(sc0.len() == gpr && sc1.len() == gpr);
6987 let n = acts.len();
6988 let nib = v.nib.as_ptr();
6989 let ones = _mm256_set1_epi16(1);
6990 let mut bi = 0usize;
6991 unsafe {
6992 while bi + 4 <= n {
6993 let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
6994 let xp: [*const i8; 4] = [
6995 a[0].xq.as_ptr(),
6996 a[1].xq.as_ptr(),
6997 a[2].xq.as_ptr(),
6998 a[3].xq.as_ptr(),
6999 ];
7000 let mut c0 = [_mm256_setzero_ps(); 4];
7001 let mut c1 = [_mm256_setzero_ps(); 4];
7002 for gi in 0..gpr {
7003 let w0 = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
7004 let w1 = q4tp_unpack_w(nib.add(((r + 1) * gpr + gi) * Q4TP_NIB));
7005 let s0 = _mm256_set1_ps(*sc0.get_unchecked(gi));
7006 let s1 = _mm256_set1_ps(*sc1.get_unchecked(gi));
7007 for k in 0..4 {
7008 let x = _mm256_loadu_si256(xp[k].add(gi * GROUP_SIZE) as *const __m256i);
7009 let d0 = _mm256_madd_epi16(_mm256_maddubs_epi16(w0, x), ones);
7010 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(w1, x), ones);
7011 c0[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d0), s0, c0[k]);
7012 c1[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d1), s1, c1[k]);
7013 }
7014 }
7015 for (row, cr, scr) in [(r, &c0, sc0), (r + 1, &c1, sc1)] {
7016 for k in 0..4 {
7017 let act = a[k];
7018 let mut acc = (q4tp_hsum256(cr[k]) - q4tp_xg8_corr(&act.xg8, scr)) * act.sx;
7019 for &(j, xv) in &act.outliers {
7020 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7021 acc += w * s * xv;
7022 }
7023 *out.at((abase + bi + k) * rows + row) = acc;
7025 }
7026 }
7027 bi += 4;
7028 }
7029 if bi < n {
7030 q4tp_matmat_row_cols_v(v, r, gpr, &acts[bi..], sc0, abase + bi, rows, out);
7031 q4tp_matmat_row_cols_v(v, r + 1, gpr, &acts[bi..], sc1, abase + bi, rows, out);
7032 }
7033 }
7034}
7035
7036#[cfg(target_arch = "x86_64")]
7045#[target_feature(enable = "avx2,fma")]
7046unsafe fn q4tp_matvec_lanes<const L: usize>(
7047 v: &Q4tpView<'_>,
7048 start: usize,
7049 end: usize,
7050 gpr: usize,
7051 xq: &[i8],
7052 xg8: &[f32],
7053 sx: f32,
7054 outliers: &[(usize, f32)],
7055 sc: &mut [f32],
7056 out: SendMut,
7057) {
7058 use core::arch::x86_64::*;
7059 debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
7060 crate::cpuprof::count_lane_rows((end - start) as u64);
7061 let mut base = start;
7062 while base + L <= end {
7063 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7064 let mut rows = [0usize; L];
7065 for l in 0..L {
7066 rows[l] = base + l;
7067 }
7068 let v4: [&Q4tpView; L] = [v; L];
7069 q4tp_scales_lanes(&v4, &rows, gpr, sc);
7070 drop(_t0);
7071 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7072 let mut accs = [_mm256_setzero_ps(); L];
7073 for gi in 0..gpr {
7074 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7075 for l in 0..L {
7076 let w = q4tp_unpack_w(v.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7077 let d = q4tp_madd_wx(w, xp);
7078 let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
7079 accs[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[l]);
7080 }
7081 }
7082 drop(_t1);
7083 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7084 for l in 0..L {
7085 let r = rows[l];
7086 let scl = &sc[l * gpr..(l + 1) * gpr];
7087 let mut acc = (q4tp_hsum256(accs[l]) - q4tp_xg8_corr(xg8, scl)) * sx;
7088 for &(j, xv) in outliers {
7089 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7090 acc += w * s * xv;
7091 }
7092 unsafe { *out.at(r) = acc };
7094 }
7095 drop(_t2);
7096 base += L;
7097 }
7098 for r in base..end {
7102 let scl = &mut sc[..gpr];
7103 v.scales_into(r, gpr, scl);
7104 let scl: &[f32] = scl;
7105 let mut acc = q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr())
7106 * sx;
7107 for &(j, xv) in outliers {
7108 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7109 acc += w * s * xv;
7110 }
7111 unsafe { *out.at(r) = acc };
7113 }
7114}
7115
7116#[cfg(target_arch = "x86_64")]
7124#[target_feature(enable = "avx2,fma")]
7125unsafe fn q4tp_gateup_lanes<const L: usize>(
7126 gv: &Q4tpView<'_>,
7127 uv: &Q4tpView<'_>,
7128 start: usize,
7129 end: usize,
7130 gpr: usize,
7131 xq: &[i8],
7132 xg8: &[f32],
7133 sx: f32,
7134 limit: f32,
7135 outliers: &[(usize, f32)],
7136 gsc: &mut [f32],
7137 usc: &mut [f32],
7138 out: SendMut,
7139) {
7140 use core::arch::x86_64::*;
7141 debug_assert!(
7142 L >= 2 && gsc.len() >= L * gpr && usc.len() >= L * gpr && xg8.len() == gpr
7143 );
7144 crate::cpuprof::count_lane_rows((end - start) as u64);
7145 let mut base = start;
7146 while base + L <= end {
7147 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7148 let mut rows = [0usize; L];
7149 for l in 0..L {
7150 rows[l] = base + l;
7151 }
7152 let gv4: [&Q4tpView; L] = [gv; L];
7153 let uv4: [&Q4tpView; L] = [uv; L];
7154 q4tp_scales_lanes(&gv4, &rows, gpr, gsc);
7155 q4tp_scales_lanes(&uv4, &rows, gpr, usc);
7156 drop(_t0);
7157 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7158 let mut ag = [_mm256_setzero_ps(); L];
7159 let mut au = [_mm256_setzero_ps(); L];
7160 for gi in 0..gpr {
7161 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7162 for l in 0..L {
7163 let w = q4tp_unpack_w(gv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7164 let d = q4tp_madd_wx(w, xp);
7165 let sb = _mm256_set1_ps(*gsc.get_unchecked(l * gpr + gi));
7166 ag[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, ag[l]);
7167 }
7168 for l in 0..L {
7169 let w = q4tp_unpack_w(uv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7170 let d = q4tp_madd_wx(w, xp);
7171 let sb = _mm256_set1_ps(*usc.get_unchecked(l * gpr + gi));
7172 au[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, au[l]);
7173 }
7174 }
7175 drop(_t1);
7176 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7177 for l in 0..L {
7178 let r = rows[l];
7179 let gs = &gsc[l * gpr..(l + 1) * gpr];
7180 let us = &usc[l * gpr..(l + 1) * gpr];
7181 let mut g = (q4tp_hsum256(ag[l]) - q4tp_xg8_corr(xg8, gs)) * sx;
7182 let mut u = (q4tp_hsum256(au[l]) - q4tp_xg8_corr(xg8, us)) * sx;
7183 for &(j, xv) in outliers {
7184 let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, gs);
7185 g += w * s * xv;
7186 let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, us);
7187 u += w * s * xv;
7188 }
7189 unsafe { *out.at(r) = silu_mul_limited(g, u, limit) };
7191 }
7192 drop(_t2);
7193 base += L;
7194 }
7195 for r in base..end {
7197 let g = &mut gsc[..gpr];
7198 let u = &mut usc[..gpr];
7199 gv.scales_into(r, gpr, g);
7200 uv.scales_into(r, gpr, u);
7201 let (g, u): (&[f32], &[f32]) = (g, u);
7202 let mut gv_ =
7203 q4tp_cell_dot_v(gv.nib.as_ptr(), r, gpr, xq.as_ptr(), g.as_ptr(), xg8.as_ptr()) * sx;
7204 let mut uv_ =
7205 q4tp_cell_dot_v(uv.nib.as_ptr(), r, gpr, xq.as_ptr(), u.as_ptr(), xg8.as_ptr()) * sx;
7206 for &(j, xv) in outliers {
7207 let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, g);
7208 gv_ += w * s * xv;
7209 let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, u);
7210 uv_ += w * s * xv;
7211 }
7212 unsafe { *out.at(r) = silu_mul_limited(gv_, uv_, limit) };
7214 }
7215}
7216
7217#[cfg(target_arch = "x86_64")]
7225#[target_feature(enable = "avx2,fma")]
7226unsafe fn q4tp_many_lanes<const L: usize>(
7227 views: &[Q4tpView<'_>],
7228 locate: &(dyn Fn(usize) -> (usize, usize)),
7229 start: usize,
7230 end: usize,
7231 gpr: usize,
7232 xq: &[i8],
7233 xg8: &[f32],
7234 sx: f32,
7235 outliers: &[(usize, f32)],
7236 sc: &mut [f32],
7237 outs: &[SendMut],
7238) {
7239 use core::arch::x86_64::*;
7240 debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
7241 crate::cpuprof::count_lane_rows((end - start) as u64);
7242 let mut base = start;
7243 while base + L <= end {
7244 let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7245 let mut tk = [0usize; L];
7246 let mut rows = [0usize; L];
7247 for l in 0..L {
7248 let (t, r) = locate(base + l);
7249 tk[l] = t;
7250 rows[l] = r;
7251 }
7252 let v4: [&Q4tpView; L] = std::array::from_fn(|l| &views[tk[l]]);
7253 q4tp_scales_lanes(&v4, &rows, gpr, sc);
7254 drop(_t0);
7255 let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7256 let mut accs = [_mm256_setzero_ps(); L];
7257 for gi in 0..gpr {
7258 let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7259 for l in 0..L {
7260 let w = q4tp_unpack_w(
7261 views[tk[l]].nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB),
7262 );
7263 let d = q4tp_madd_wx(w, xp);
7264 let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
7265 accs[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[l]);
7266 }
7267 }
7268 drop(_t1);
7269 let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7270 for l in 0..L {
7271 let (t, r) = (tk[l], rows[l]);
7272 let scl = &sc[l * gpr..(l + 1) * gpr];
7273 let mut acc = (q4tp_hsum256(accs[l]) - q4tp_xg8_corr(xg8, scl)) * sx;
7274 for &(j, xv) in outliers {
7275 let (w, s) = q4tp_outlier(views[t].nib, r, gpr, j, scl);
7276 acc += w * s * xv;
7277 }
7278 unsafe { *outs[t].at(r) = acc };
7280 }
7281 drop(_t2);
7282 base += L;
7283 }
7284 for flat in base..end {
7286 let (t, r) = locate(flat);
7287 let v = &views[t];
7288 let scl = &mut sc[..gpr];
7289 v.scales_into(r, gpr, scl);
7290 let scl: &[f32] = scl;
7291 let mut acc =
7292 q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr()) * sx;
7293 for &(j, xv) in outliers {
7294 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7295 acc += w * s * xv;
7296 }
7297 unsafe { *outs[t].at(r) = acc };
7299 }
7300}
7301
7302#[inline]
7305fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7306 #[cfg(target_arch = "x86_64")]
7307 if avx2_enabled() {
7308 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
7310 }
7311 q4tp_row_float_scalar(nib, r, gpr, x, scales)
7312}
7313
7314#[inline]
7315fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7316 let mut acc = 0f32;
7317 for gi in 0..gpr {
7318 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
7319 let s = scales[gi];
7320 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7321 let mut ga = 0f32;
7322 for (k, &b) in tile.iter().enumerate() {
7323 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7324 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7325 }
7326 acc += ga * s;
7327 }
7328 acc
7329}
7330
7331#[cfg(target_arch = "x86_64")]
7335#[target_feature(enable = "avx2")]
7336unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7337 unsafe {
7340 use core::arch::x86_64::*;
7341 let mask = _mm_set1_epi8(15);
7342 let eight = _mm_set1_epi8(8);
7343 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
7344 let mut acc = 0.0f32;
7345 for gi in 0..gpr {
7346 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
7347 let lo = _mm_and_si128(packed, mask);
7348 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
7349 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
7350 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
7351 let xp = x.as_ptr().add(gi * GROUP_SIZE);
7352 let a = _mm256_mul_ps(
7353 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
7354 _mm256_loadu_ps(xp),
7355 );
7356 let b = _mm256_mul_ps(
7357 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
7358 _mm256_loadu_ps(xp.add(8)),
7359 );
7360 let c = _mm256_mul_ps(
7361 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
7362 _mm256_loadu_ps(xp.add(16)),
7363 );
7364 let d = _mm256_mul_ps(
7365 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
7366 _mm256_loadu_ps(xp.add(24)),
7367 );
7368 let mut pairs = [0.0f32; 16];
7369 _mm256_storeu_ps(
7370 pairs.as_mut_ptr(),
7371 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
7372 );
7373 _mm256_storeu_ps(
7374 pairs.as_mut_ptr().add(8),
7375 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
7376 );
7377 let mut ga = 0.0f32;
7378 for v in pairs {
7379 ga += v;
7380 }
7381 acc += ga * scales[gi];
7382 }
7383 acc
7384 }
7385}
7386
7387#[inline]
7390fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
7391 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
7392 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
7393 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
7394 ((n as i32 - 8) as f32, scales[gi])
7395}
7396
7397fn q4tp_matvec(
7399 bytes: &[u8],
7400 x: &[f32],
7401 rows: usize,
7402 cols: usize,
7403 out: &mut [f32],
7404 pool: Option<&Pool>,
7405) {
7406 debug_assert_eq!(out.len(), rows);
7407 let gpr = cols / GROUP_SIZE;
7408 let v = Q4tpView::new(bytes, rows, cols);
7409 let out_addr = SendMut(out.as_mut_ptr());
7410 if a8w8_enabled() {
7411 let act = split_act(x);
7412 #[cfg(target_arch = "x86_64")]
7413 let lanes = q4tp_lanes();
7414 let run = |start: usize, end: usize| {
7415 #[cfg(target_arch = "x86_64")]
7416 if lanes > 1 && act.xg8.len() == gpr {
7417 with_krow(gpr * lanes, |sc| {
7420 unsafe {
7422 match lanes {
7423 2 => q4tp_matvec_lanes::<2>(
7424 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7425 &act.outliers, sc, out_addr,
7426 ),
7427 8 => q4tp_matvec_lanes::<8>(
7428 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7429 &act.outliers, sc, out_addr,
7430 ),
7431 _ => q4tp_matvec_lanes::<4>(
7432 &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7433 &act.outliers, sc, out_addr,
7434 ),
7435 }
7436 }
7437 });
7438 return;
7439 }
7440 crate::cpuprof::count_plain_rows((end - start) as u64);
7441 with_krow(gpr, |sc| {
7443 for r in start..end {
7444 v.scales_into(r, gpr, sc);
7445 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
7446 for &(j, xv) in &act.outliers {
7447 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
7448 acc += w * s * xv;
7449 }
7450 unsafe { *out_addr.at(r) = acc }
7452 }
7453 })
7454 };
7455 dispatch_rows(pool, rows, &run);
7456 return;
7457 }
7458 let run = |start: usize, end: usize| {
7459 with_krow(gpr, |sc| {
7460 for r in start..end {
7461 v.scales_into(r, gpr, sc);
7462 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
7464 }
7465 })
7466 };
7467 dispatch_rows(pool, rows, &run);
7468}
7469
7470#[allow(clippy::too_many_arguments)]
7473fn q4tp_matvec2(
7474 bytes: &[u8],
7475 x1: &[f32],
7476 x2: &[f32],
7477 rows: usize,
7478 cols: usize,
7479 o1: &mut [f32],
7480 o2: &mut [f32],
7481 pool: Option<&Pool>,
7482) {
7483 let gpr = cols / GROUP_SIZE;
7484 let v = Q4tpView::new(bytes, rows, cols);
7485 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7486 let run = |start: usize, end: usize| {
7487 let mut sc = vec![0f32; gpr];
7488 for r in start..end {
7489 v.scales_into(r, gpr, &mut sc);
7490 unsafe {
7492 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
7493 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
7494 }
7495 }
7496 };
7497 dispatch_rows(pool, rows, &run);
7498}
7499
7500#[inline]
7503fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
7504 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
7505 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
7506 let c = (byte >> (2 * (k % 4))) & 3;
7507 (c as f32 - 1.5, scales[gi])
7508}
7509
7510#[cfg(target_arch = "x86_64")]
7511const Q2TP_DECODE_U32: [u32; 256] = {
7512 let mut tab = [0u32; 256];
7513 let mut b = 0usize;
7514 while b < 256 {
7515 tab[b] = ((b as u32) & 3)
7516 | ((((b as u32) >> 2) & 3) << 8)
7517 | ((((b as u32) >> 4) & 3) << 16)
7518 | ((((b as u32) >> 6) & 3) << 24);
7519 b += 1;
7520 }
7521 tab
7522};
7523
7524#[cfg(target_arch = "x86_64")]
7528#[target_feature(enable = "avx2")]
7529unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
7530 use core::arch::x86_64::*;
7531 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
7532 let codes = _mm256_setr_epi32(
7533 Q2TP_DECODE_U32[ch[0] as usize] as i32,
7534 Q2TP_DECODE_U32[ch[1] as usize] as i32,
7535 Q2TP_DECODE_U32[ch[2] as usize] as i32,
7536 Q2TP_DECODE_U32[ch[3] as usize] as i32,
7537 Q2TP_DECODE_U32[ch[4] as usize] as i32,
7538 Q2TP_DECODE_U32[ch[5] as usize] as i32,
7539 Q2TP_DECODE_U32[ch[6] as usize] as i32,
7540 Q2TP_DECODE_U32[ch[7] as usize] as i32,
7541 );
7542 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
7543 let pair = _mm256_maddubs_epi16(codes, xv);
7544 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
7545 let sum128 = _mm_add_epi32(
7546 _mm256_castsi256_si128(quad),
7547 _mm256_extracti128_si256(quad, 1),
7548 );
7549 let sum64 = _mm_hadd_epi32(sum128, sum128);
7550 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
7551}
7552
7553#[cfg(target_arch = "x86_64")]
7554#[inline]
7555fn q2tp_avx2_enabled() -> bool {
7556 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7557 *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
7558}
7559
7560#[inline]
7567fn dot_q2tp_row_i8(
7568 chunks: &[u8],
7569 r: usize,
7570 gpr: usize,
7571 xq: &[i8],
7572 gsum: &[i32],
7573 scales: &[f32],
7574) -> f32 {
7575 let mut acc = 0f32;
7576 let base = r * gpr * Q2TP_CHUNK;
7577 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7578 let mut codes = [0i8; GROUP_SIZE];
7579 #[cfg(target_arch = "x86_64")]
7580 let avx2 = q2tp_avx2_enabled();
7585 for gi in 0..gpr {
7586 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
7587 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7588 #[cfg(target_arch = "aarch64")]
7589 let dot = unsafe {
7595 use core::arch::aarch64::*;
7596 let b = vld1_u8(ch.as_ptr());
7597 let three = vdup_n_u8(3);
7598 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
7599 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
7600 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
7601 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
7602 let x4 = vld4_s8(xg.as_ptr());
7603 let mut acc4 = vdupq_n_s32(0);
7604 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
7605 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
7606 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
7607 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
7608 vaddvq_s32(acc4)
7609 };
7610 #[cfg(target_arch = "x86_64")]
7611 let dot: i32 = if avx2 {
7612 unsafe { q2tp_code_dot_avx2(ch, xg) }
7615 } else {
7616 ch.iter()
7617 .enumerate()
7618 .map(|(k, &b)| {
7619 ((b & 3) as i32) * xg[k * 4] as i32
7620 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
7621 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
7622 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
7623 })
7624 .sum()
7625 };
7626 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7627 let dot: i32 = {
7628 for (k, &b) in ch.iter().enumerate() {
7629 codes[k * 4] = (b & 3) as i8;
7630 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
7631 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
7632 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
7633 }
7634 codes
7635 .iter()
7636 .zip(xg)
7637 .map(|(&c, &x)| c as i32 * x as i32)
7638 .sum()
7639 };
7640 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
7641 }
7642 acc
7643}
7644
7645fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7649 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
7650}
7651
7652#[inline]
7656fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7657 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
7658}
7659
7660#[inline]
7661fn q2tp_row_exact_center(
7662 chunks: &[u8],
7663 r: usize,
7664 gpr: usize,
7665 x: &[f32],
7666 scales: &[f32],
7667 center: f32,
7668) -> f32 {
7669 let mut acc = 0f32;
7670 for gi in 0..gpr {
7671 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
7672 let s = scales[gi];
7673 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7674 let mut g = 0f32;
7675 for (k, &b) in ch.iter().enumerate() {
7676 g += ((b & 3) as f32 - center) * xb[k * 4]
7677 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
7678 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
7679 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
7680 }
7681 acc += s * g;
7682 }
7683 acc
7684}
7685
7686fn q2tp_matvec(
7687 bytes: &[u8],
7688 x: &[f32],
7689 rows: usize,
7690 cols: usize,
7691 out: &mut [f32],
7692 pool: Option<&Pool>,
7693) {
7694 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
7695}
7696
7697fn q2tp_affine_matvec(
7698 bytes: &[u8],
7699 x: &[f32],
7700 rows: usize,
7701 cols: usize,
7702 out: &mut [f32],
7703 pool: Option<&Pool>,
7704) {
7705 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
7706}
7707
7708fn q2tp_matvec_mode(
7709 bytes: &[u8],
7710 x: &[f32],
7711 rows: usize,
7712 cols: usize,
7713 out: &mut [f32],
7714 pool: Option<&Pool>,
7715 affine: bool,
7716) {
7717 debug_assert_eq!(out.len(), rows);
7718 let gpr = cols / GROUP_SIZE;
7719 let v = Q4tpView::new_q2(bytes, rows, cols);
7720 let out_addr = SendMut(out.as_mut_ptr());
7721 if !affine && a8w8_enabled() {
7726 let act = split_act(x);
7727 let gsum = q1_group_sums(&act.xq, gpr);
7728 let (act, gsum) = (&act, &gsum);
7729 let run = move |start: usize, end: usize| {
7730 with_krow(gpr, |sc| {
7731 for r in start..end {
7732 v.scales_into(r, gpr, sc);
7733 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
7734 for &(j, xv) in &act.outliers {
7735 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
7736 acc += w * s * xv;
7737 }
7738 unsafe { *out_addr.at(r) = acc };
7740 }
7741 })
7742 };
7743 dispatch_rows(pool, rows, &run);
7744 return;
7745 }
7746 let run = |start: usize, end: usize| {
7747 with_krow(gpr, |sc| {
7748 for r in start..end {
7749 v.scales_into(r, gpr, sc);
7750 unsafe {
7752 *out_addr.at(r) = if affine {
7753 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
7754 } else {
7755 q2tp_row_exact(v.nib, r, gpr, x, sc)
7756 }
7757 };
7758 }
7759 })
7760 };
7761 dispatch_rows(pool, rows, &run);
7762}
7763
7764#[allow(clippy::too_many_arguments)]
7766fn q2tp_matvec2(
7767 bytes: &[u8],
7768 x1: &[f32],
7769 x2: &[f32],
7770 rows: usize,
7771 cols: usize,
7772 o1: &mut [f32],
7773 o2: &mut [f32],
7774 pool: Option<&Pool>,
7775) {
7776 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
7777}
7778
7779#[allow(clippy::too_many_arguments)]
7780fn q2tp_affine_matvec2(
7781 bytes: &[u8],
7782 x1: &[f32],
7783 x2: &[f32],
7784 rows: usize,
7785 cols: usize,
7786 o1: &mut [f32],
7787 o2: &mut [f32],
7788 pool: Option<&Pool>,
7789) {
7790 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
7791}
7792
7793#[allow(clippy::too_many_arguments)]
7794fn q2tp_matvec2_mode(
7795 bytes: &[u8],
7796 x1: &[f32],
7797 x2: &[f32],
7798 rows: usize,
7799 cols: usize,
7800 o1: &mut [f32],
7801 o2: &mut [f32],
7802 pool: Option<&Pool>,
7803 affine: bool,
7804) {
7805 let gpr = cols / GROUP_SIZE;
7806 let v = Q4tpView::new_q2(bytes, rows, cols);
7807 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7808 let run = |start: usize, end: usize| {
7809 let mut sc = vec![0f32; gpr];
7810 for r in start..end {
7811 v.scales_into(r, gpr, &mut sc);
7812 unsafe {
7814 *p1.at(r) = if affine {
7815 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
7816 } else {
7817 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
7818 };
7819 *p2.at(r) = if affine {
7820 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
7821 } else {
7822 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
7823 };
7824 }
7825 }
7826 };
7827 dispatch_rows(pool, rows, &run);
7828}
7829
7830pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
7837 let gpr = cols / GROUP_SIZE;
7842 let v = Q4tpView::new_q2(bytes, rows, cols);
7843 with_krow(gpr, |sc| {
7844 for r in 0..rows {
7845 v.scales_into(r, gpr, sc);
7846 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
7847 }
7848 });
7849}
7850
7851pub fn q2tp_affine_matvec_for_test(
7854 bytes: &[u8],
7855 x: &[f32],
7856 rows: usize,
7857 cols: usize,
7858 out: &mut [f32],
7859) {
7860 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
7861}
7862
7863pub fn q2tp_matmat_for_test(
7864 bytes: &[u8],
7865 xs_all: &[f32],
7866 b: usize,
7867 rows: usize,
7868 cols: usize,
7869 out: &mut [f32],
7870) {
7871 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
7872}
7873
7874fn q2tp_matmat(
7875 bytes: &[u8],
7876 xs_all: &[f32],
7877 b: usize,
7878 rows: usize,
7879 cols: usize,
7880 out: &mut [f32],
7881 pool: Option<&Pool>,
7882) {
7883 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
7884}
7885
7886fn q2tp_affine_matmat(
7887 bytes: &[u8],
7888 xs_all: &[f32],
7889 b: usize,
7890 rows: usize,
7891 cols: usize,
7892 out: &mut [f32],
7893 pool: Option<&Pool>,
7894) {
7895 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
7896}
7897
7898fn q2tp_matmat_mode(
7899 bytes: &[u8],
7900 xs_all: &[f32],
7901 b: usize,
7902 rows: usize,
7903 cols: usize,
7904 out: &mut [f32],
7905 pool: Option<&Pool>,
7906 affine: bool,
7907) {
7908 debug_assert_eq!(out.len(), b * rows);
7909 let gpr = cols / GROUP_SIZE;
7910 let v = Q4tpView::new_q2(bytes, rows, cols);
7911 let out_addr = SendMut(out.as_mut_ptr());
7912 let run = |start: usize, end: usize| {
7913 let mut sc = vec![0f32; gpr];
7914 for r in start..end {
7915 v.scales_into(r, gpr, &mut sc);
7916 for bi in 0..b {
7917 let x = &xs_all[bi * cols..(bi + 1) * cols];
7918 unsafe {
7920 *out_addr.at(bi * rows + r) = if affine {
7921 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
7922 } else {
7923 q2tp_row_exact(v.nib, r, gpr, x, &sc)
7924 }
7925 };
7926 }
7927 }
7928 };
7929 dispatch_rows(pool, rows, &run);
7930}
7931
7932#[cfg(target_arch = "aarch64")]
7940#[target_feature(enable = "neon,dotprod")]
7941unsafe fn dot_q4tp_row_1x4_sdot_v1(
7942 nib: &[u8],
7943 r: usize,
7944 gpr: usize,
7945 xs: [&[i8]; 4],
7946 scales: &[f32],
7947) -> [f32; 4] {
7948 unsafe {
7949 use core::arch::aarch64::*;
7950 use core::arch::asm;
7951 let lomask = vdupq_n_u8(0x0F);
7952 let eight = vdupq_n_s8(8);
7953 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
7954 for gi in 0..gpr {
7955 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7956 let s = *scales.get_unchecked(gi);
7957 let bb = vld1q_u8(t);
7958 let lo = vandq_u8(bb, lomask);
7959 let hi = vshrq_n_u8::<4>(bb);
7960 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7961 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7962 let mut d = [0f32; 4];
7963 for (k, dk) in d.iter_mut().enumerate() {
7964 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
7965 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
7966 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7967 asm!(
7968 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7969 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7970 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7971 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7972 options(pure, nomem, nostack),
7973 );
7974 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7975 }
7976 f0 += d[0];
7977 f1 += d[1];
7978 f2 += d[2];
7979 f3 += d[3];
7980 }
7981 [f0, f1, f2, f3]
7982 }
7983}
7984
7985#[allow(dead_code)]
7992static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
7993
7994#[cfg(test)]
7997static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
7998
7999#[cfg(target_arch = "x86_64")]
8005fn q4tp_blocked_x86() -> bool {
8006 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
8007 1 => false,
8008 2 => avx512vnni_enabled(),
8013 _ => blocked_enabled() && avx512vnni_enabled(),
8017 }
8018}
8019
8020#[cfg(target_arch = "aarch64")]
8022#[allow(dead_code)]
8023fn q4tp_v1() -> bool {
8024 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
8025 1 => true,
8026 2 => false,
8027 _ => {
8028 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8029 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
8030 }
8031 }
8032}
8033
8034#[cfg(target_arch = "x86_64")]
8046#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8047unsafe fn dot_q4tp_2x8_avx512(
8048 nib: &[u8],
8049 r0: usize,
8050 gpr: usize,
8051 xs: [&[i8]; 8],
8052 sc0: &[f32],
8053 sc1: &[f32],
8054) -> [[f32; 8]; 2] {
8055 unsafe {
8058 use core::arch::x86_64::*;
8059 let lomask = _mm256_set1_epi8(0x0F);
8060 let eight = _mm256_set1_epi8(8);
8061 let zero = _mm512_setzero_si512();
8062 let mut v0 = [_mm512_setzero_ps(); 8];
8063 let mut v1 = [_mm512_setzero_ps(); 8];
8064 let pairs = gpr / 2;
8065 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
8066 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8067 let bb = _mm256_loadu_si256(t as *const __m256i);
8068 let lo = _mm256_and_si256(bb, lomask);
8069 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8070 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8071 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8072 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8073 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8074 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
8075 };
8076 for gp in 0..pairs {
8077 let gi = gp * 2;
8078 let (wa0, neg0) = unpack(r0, gi);
8079 let (wa1, neg1) = unpack(r0 + 1, gi);
8080 let off = gi * GROUP_SIZE;
8081 let sv = |sc: &[f32]| {
8082 _mm512_insertf32x8::<1>(
8083 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
8084 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
8085 )
8086 };
8087 let s0 = sv(sc0);
8088 let s1 = sv(sc1);
8089 for k in 0..8 {
8090 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
8091 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
8092 zero,
8093 wa0,
8094 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
8095 ));
8096 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
8097 zero,
8098 wa1,
8099 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
8100 ));
8101 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
8102 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
8103 }
8104 }
8105 let mut acc = [[0f32; 8]; 2];
8106 for k in 0..8 {
8107 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
8108 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
8109 }
8110 if gpr % 2 == 1 {
8111 let off = (gpr - 1) * GROUP_SIZE;
8112 for j in off..off + GROUP_SIZE {
8113 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
8114 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
8115 for k in 0..8 {
8116 let x = *xs[k].get_unchecked(j) as f32;
8117 acc[0][k] += w0 * sa * x;
8118 acc[1][k] += w1 * sb * x;
8119 }
8120 }
8121 }
8122 acc
8123 }
8124}
8125
8126#[cfg(target_arch = "x86_64")]
8131#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8132unsafe fn dot_q4tp_row_1x8_avx512(
8133 nib: &[u8],
8134 r: usize,
8135 gpr: usize,
8136 xs: [&[i8]; 8],
8137 scales: &[f32],
8138) -> [f32; 8] {
8139 unsafe {
8141 use core::arch::x86_64::*;
8142 let lomask = _mm256_set1_epi8(0x0F);
8143 let eight = _mm256_set1_epi8(8);
8144 let zero = _mm512_setzero_si512();
8145 let (mut v0, mut v1, mut v2, mut v3) = (
8146 _mm512_setzero_ps(),
8147 _mm512_setzero_ps(),
8148 _mm512_setzero_ps(),
8149 _mm512_setzero_ps(),
8150 );
8151 let (mut v4, mut v5, mut v6, mut v7) = (
8152 _mm512_setzero_ps(),
8153 _mm512_setzero_ps(),
8154 _mm512_setzero_ps(),
8155 _mm512_setzero_ps(),
8156 );
8157 let pairs = gpr / 2;
8158 for gp in 0..pairs {
8159 let gi = gp * 2;
8160 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8161 let bb = _mm256_loadu_si256(t as *const __m256i);
8162 let lo = _mm256_and_si256(bb, lomask);
8163 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8164 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8169 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8170 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8171 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8172 let wabs = _mm512_abs_epi8(w);
8173 let neg = _mm512_movepi8_mask(w);
8174 let off = gi * GROUP_SIZE;
8175 let sv = _mm512_insertf32x8::<1>(
8176 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
8177 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
8178 );
8179 let dot = |x: &[i8]| -> __m512 {
8180 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
8181 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
8182 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
8183 };
8184 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
8185 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
8186 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
8187 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
8188 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
8189 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
8190 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
8191 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
8192 }
8193 let mut acc = [
8194 _mm512_reduce_add_ps(v0),
8195 _mm512_reduce_add_ps(v1),
8196 _mm512_reduce_add_ps(v2),
8197 _mm512_reduce_add_ps(v3),
8198 _mm512_reduce_add_ps(v4),
8199 _mm512_reduce_add_ps(v5),
8200 _mm512_reduce_add_ps(v6),
8201 _mm512_reduce_add_ps(v7),
8202 ];
8203 if gpr % 2 == 1 {
8206 let off = (gpr - 1) * GROUP_SIZE;
8207 for j in off..off + GROUP_SIZE {
8208 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
8209 let ws = w * s;
8210 for k in 0..8 {
8211 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
8212 }
8213 }
8214 }
8215 acc
8216 }
8217}
8218
8219#[cfg(target_arch = "x86_64")]
8232#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8233unsafe fn dot_q4tp_row_1x4_avx512(
8234 nib: &[u8],
8235 r: usize,
8236 gpr: usize,
8237 xs: [&[i8]; 4],
8238 scales: &[f32],
8239) -> [f32; 4] {
8240 unsafe {
8242 use core::arch::x86_64::*;
8243 let lomask = _mm256_set1_epi8(0x0F);
8244 let eight = _mm256_set1_epi8(8);
8245 let zero = _mm512_setzero_si512();
8246 let (mut v0, mut v1, mut v2, mut v3) = (
8247 _mm512_setzero_ps(),
8248 _mm512_setzero_ps(),
8249 _mm512_setzero_ps(),
8250 _mm512_setzero_ps(),
8251 );
8252 let pairs = gpr / 2;
8253 for gp in 0..pairs {
8254 let gi = gp * 2;
8255 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8256 let bb = _mm256_loadu_si256(t as *const __m256i);
8257 let lo = _mm256_and_si256(bb, lomask);
8258 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8259 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8264 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8265 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8266 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8267 let wabs = _mm512_abs_epi8(w);
8268 let neg = _mm512_movepi8_mask(w);
8269 let off = gi * GROUP_SIZE;
8270 let sv = _mm512_insertf32x8::<1>(
8271 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
8272 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
8273 );
8274 let dot = |x: &[i8]| -> __m512 {
8275 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
8276 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
8277 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
8278 };
8279 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
8280 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
8281 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
8282 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
8283 }
8284 let mut acc = [
8285 _mm512_reduce_add_ps(v0),
8286 _mm512_reduce_add_ps(v1),
8287 _mm512_reduce_add_ps(v2),
8288 _mm512_reduce_add_ps(v3),
8289 ];
8290 if gpr % 2 == 1 {
8293 let off = (gpr - 1) * GROUP_SIZE;
8294 for j in off..off + GROUP_SIZE {
8295 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
8296 let ws = w * s;
8297 for k in 0..4 {
8298 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
8299 }
8300 }
8301 }
8302 acc
8303 }
8304}
8305
8306#[cfg(target_arch = "aarch64")]
8310#[target_feature(enable = "neon,dotprod")]
8311unsafe fn dot_q4tp_row_1x4_sdot(
8312 nib: &[u8],
8313 r: usize,
8314 gpr: usize,
8315 xs: [&[i8]; 4],
8316 scales: &[f32],
8317) -> [f32; 4] {
8318 unsafe {
8320 use core::arch::aarch64::*;
8321 use core::arch::asm;
8322 let lomask = vdupq_n_u8(0x0F);
8323 let eight = vdupq_n_s8(8);
8324 let (mut v0, mut v1, mut v2, mut v3) = (
8340 vdupq_n_f32(0.0),
8341 vdupq_n_f32(0.0),
8342 vdupq_n_f32(0.0),
8343 vdupq_n_f32(0.0),
8344 );
8345 for gi in 0..gpr {
8346 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8347 let s = *scales.get_unchecked(gi);
8348 let bb = vld1q_u8(t);
8349 let lo = vandq_u8(bb, lomask);
8350 let hi = vshrq_n_u8::<4>(bb);
8351 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8352 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8353 let off = gi * GROUP_SIZE;
8354 let dot4 = |x: &[i8]| -> int32x4_t {
8355 let x0 = vld1q_s8(x.as_ptr().add(off));
8356 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
8357 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8358 asm!(
8359 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8360 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8361 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8362 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8363 options(pure, nomem, nostack),
8364 );
8365 vaddq_s32(a0, a1)
8366 };
8367 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
8368 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
8369 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
8370 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
8371 }
8372 [
8373 vaddvq_f32(v0),
8374 vaddvq_f32(v1),
8375 vaddvq_f32(v2),
8376 vaddvq_f32(v3),
8377 ]
8378 }
8379}
8380
8381fn q4tp_matmat(
8389 bytes: &[u8],
8390 xs_all: &[f32],
8391 b: usize,
8392 rows: usize,
8393 cols: usize,
8394 out: &mut [f32],
8395 pool: Option<&Pool>,
8396) {
8397 q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
8398}
8399
8400#[allow(clippy::too_many_arguments)]
8404fn q4tp_matmat_with(
8405 bytes: &[u8],
8406 xs_all: &[f32],
8407 b: usize,
8408 rows: usize,
8409 cols: usize,
8410 out: &mut [f32],
8411 pool: Option<&Pool>,
8412 exact: bool,
8413) {
8414 debug_assert_eq!(out.len(), b * rows);
8415 let gpr = cols / GROUP_SIZE;
8416 let v = Q4tpView::new(bytes, rows, cols);
8417
8418 #[cfg(target_os = "macos")]
8422 if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8423 dequant_matmat_accel(
8424 &|r, dst| {
8425 let mut sc = [0f32; 32];
8426 let mut scv;
8427 let s: &[f32] = if gpr <= 32 {
8428 v.scales_into(r, gpr, &mut sc);
8429 &sc[..gpr]
8430 } else {
8431 scv = vec![0f32; gpr];
8432 v.scales_into(r, gpr, &mut scv);
8433 &scv
8434 };
8435 for gi in 0..gpr {
8436 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
8437 for (k, &bb) in tile.iter().enumerate() {
8438 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
8439 dst[gi * GROUP_SIZE + k * 2 + 1] =
8440 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
8441 }
8442 }
8443 },
8444 xs_all,
8445 b,
8446 rows,
8447 cols,
8448 out,
8449 pool,
8450 );
8451 return;
8452 }
8453
8454 let out_addr = SendMut(out.as_mut_ptr());
8455 if a8w8_enabled() {
8456 let acts = split_acts_pooled(xs_all, b, cols, pool);
8457 let acts = &acts;
8458 #[cfg(target_arch = "aarch64")]
8465 let blocked_ok = sdot_enabled() && blocked_enabled();
8466 #[cfg(target_arch = "x86_64")]
8473 let blocked_ok = q4tp_blocked_x86() && !exact;
8474 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8475 let blocked_ok = {
8476 let _ = exact;
8477 false
8478 };
8479 let panel_cols = q4tp_panel_cols(cols);
8489 let run = |start: usize, end: usize| {
8490 for abase in (0..acts.len()).step_by(panel_cols) {
8491 let alen = (acts.len() - abase).min(panel_cols);
8492 let mut sc = vec![0f32; gpr];
8493 #[cfg(target_arch = "x86_64")]
8494 let mut r_lo = start;
8495 #[cfg(target_arch = "x86_64")]
8496 if blocked_ok && alen >= 8 {
8497 let mut sc1 = vec![0f32; gpr];
8498 while r_lo + 2 <= end {
8499 v.scales_into(r_lo, gpr, &mut sc);
8500 v.scales_into(r_lo + 1, gpr, &mut sc1);
8501 let mut bi = 0usize;
8502 while bi + 8 <= alen {
8503 let xs = [
8504 acts[abase + bi].xq.as_slice(),
8505 acts[abase + bi + 1].xq.as_slice(),
8506 acts[abase + bi + 2].xq.as_slice(),
8507 acts[abase + bi + 3].xq.as_slice(),
8508 acts[abase + bi + 4].xq.as_slice(),
8509 acts[abase + bi + 5].xq.as_slice(),
8510 acts[abase + bi + 6].xq.as_slice(),
8511 acts[abase + bi + 7].xq.as_slice(),
8512 ];
8513 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
8514 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
8515 for k in 0..8 {
8516 let act = &acts[abase + bi + k];
8517 let mut acc = dr[k] * act.sx;
8518 for &(j, xv) in &act.outliers {
8519 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
8520 acc += w * s * xv;
8521 }
8522 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
8524 }
8525 }
8526 bi += 8;
8527 }
8528 for row in [r_lo, r_lo + 1] {
8531 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
8532 for b2 in bi..alen {
8533 let act = &acts[abase + b2];
8534 let xs4 = [
8535 act.xq.as_slice(),
8536 act.xq.as_slice(),
8537 act.xq.as_slice(),
8538 act.xq.as_slice(),
8539 ];
8540 let d =
8541 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
8542 let mut acc = d[0] * act.sx;
8543 for &(j, xv) in &act.outliers {
8544 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
8545 acc += w * s * xv;
8546 }
8547 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
8549 }
8550 }
8551 r_lo += 2;
8552 }
8553 }
8554 #[cfg(target_arch = "x86_64")]
8559 if !blocked_ok
8560 && !exact
8561 && alen >= 4
8562 && q4tp_fma_x86()
8563 && acts[abase..abase + alen].iter().all(|a| a.xg8.len() == gpr)
8564 {
8565 let mut sc1 = vec![0f32; gpr];
8566 while r_lo + 2 <= end {
8567 v.scales_into(r_lo, gpr, &mut sc);
8568 v.scales_into(r_lo + 1, gpr, &mut sc1);
8569 unsafe {
8572 q4tp_matmat_2rows_fma(
8573 &v,
8574 r_lo,
8575 gpr,
8576 &acts[abase..abase + alen],
8577 &sc,
8578 &sc1,
8579 abase,
8580 rows,
8581 out_addr,
8582 )
8583 };
8584 r_lo += 2;
8585 }
8586 }
8587 #[cfg(target_arch = "x86_64")]
8588 let row_start = r_lo;
8589 #[cfg(not(target_arch = "x86_64"))]
8590 let row_start = start;
8591 for r in row_start..end {
8592 v.scales_into(r, gpr, &mut sc);
8593 let mut bi = 0usize;
8594 #[cfg(target_arch = "x86_64")]
8595 if blocked_ok {
8596 while bi + 8 <= alen {
8597 let xs = [
8598 acts[abase + bi].xq.as_slice(),
8599 acts[abase + bi + 1].xq.as_slice(),
8600 acts[abase + bi + 2].xq.as_slice(),
8601 acts[abase + bi + 3].xq.as_slice(),
8602 acts[abase + bi + 4].xq.as_slice(),
8603 acts[abase + bi + 5].xq.as_slice(),
8604 acts[abase + bi + 6].xq.as_slice(),
8605 acts[abase + bi + 7].xq.as_slice(),
8606 ];
8607 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
8608 for k in 0..8 {
8609 let act = &acts[abase + bi + k];
8610 let mut acc = d[k] * act.sx;
8611 for &(j, xv) in &act.outliers {
8612 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8613 acc += w * s * xv;
8614 }
8615 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8617 }
8618 bi += 8;
8619 }
8620 while bi + 4 <= alen {
8621 let xs = [
8622 acts[abase + bi].xq.as_slice(),
8623 acts[abase + bi + 1].xq.as_slice(),
8624 acts[abase + bi + 2].xq.as_slice(),
8625 acts[abase + bi + 3].xq.as_slice(),
8626 ];
8627 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
8628 for k in 0..4 {
8629 let act = &acts[abase + bi + k];
8630 let mut acc = d[k] * act.sx;
8631 for &(j, xv) in &act.outliers {
8632 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8633 acc += w * s * xv;
8634 }
8635 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8637 }
8638 bi += 4;
8639 }
8640 }
8641 #[cfg(target_arch = "aarch64")]
8642 if blocked_ok {
8643 while bi + 4 <= alen {
8644 let xs = [
8645 acts[abase + bi].xq.as_slice(),
8646 acts[abase + bi + 1].xq.as_slice(),
8647 acts[abase + bi + 2].xq.as_slice(),
8648 acts[abase + bi + 3].xq.as_slice(),
8649 ];
8650 let d = unsafe {
8651 if exact || q4tp_v1() {
8652 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
8653 } else {
8654 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
8655 }
8656 };
8657 for k in 0..4 {
8658 let act = &acts[abase + bi + k];
8659 let mut acc = d[k] * act.sx;
8660 for &(j, xv) in &act.outliers {
8661 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8662 acc += w * s * xv;
8663 }
8664 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8666 }
8667 bi += 4;
8668 }
8669 }
8670 let _ = blocked_ok;
8671 #[cfg(target_arch = "x86_64")]
8677 if acts[abase..abase + alen]
8678 .iter()
8679 .all(|a| a.xg8.len() == gpr)
8680 {
8681 unsafe {
8684 q4tp_matmat_row_cols_v(
8685 &v,
8686 r,
8687 gpr,
8688 &acts[abase + bi..abase + alen],
8689 &sc,
8690 abase + bi,
8691 rows,
8692 out_addr,
8693 )
8694 };
8695 bi = alen;
8696 }
8697 while bi < alen {
8698 let act = &acts[abase + bi];
8699 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
8700 for &(j, xv) in &act.outliers {
8701 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8702 acc += w * s * xv;
8703 }
8704 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
8706 bi += 1;
8707 }
8708 }
8709 }
8710 };
8711 dispatch_rows(pool, rows, &run);
8712 return;
8713 }
8714
8715 let run = |start: usize, end: usize| {
8716 let mut sc = vec![0f32; gpr];
8717 for r in start..end {
8718 v.scales_into(r, gpr, &mut sc);
8719 for bi in 0..b {
8720 let x = &xs_all[bi * cols..(bi + 1) * cols];
8721 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
8723 }
8724 }
8725 };
8726 dispatch_rows(pool, rows, &run);
8727}
8728
8729fn q4t_matvec(
8731 bytes: &[u8],
8732 x: &[f32],
8733 rows: usize,
8734 cols: usize,
8735 out: &mut [f32],
8736 pool: Option<&Pool>,
8737) {
8738 debug_assert_eq!(out.len(), rows);
8739 let gpr = cols / GROUP_SIZE;
8740 let out_addr = SendMut(out.as_mut_ptr());
8741 if a8w8_enabled() {
8742 let act = split_act(x);
8743 let run = move |start: usize, end: usize| {
8744 for r in start..end {
8745 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8746 for &(j, xv) in &act.outliers {
8747 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8748 acc += w * s * xv;
8749 }
8750 unsafe { *out_addr.at(r) = acc };
8752 }
8753 };
8754 dispatch_rows(pool, rows, &run);
8755 return;
8756 }
8757 let run = move |start: usize, end: usize| {
8758 for r in start..end {
8759 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
8761 }
8762 };
8763 dispatch_rows(pool, rows, &run);
8764}
8765
8766#[allow(clippy::too_many_arguments)]
8768fn q4t_matvec2(
8769 bytes: &[u8],
8770 x1: &[f32],
8771 x2: &[f32],
8772 rows: usize,
8773 cols: usize,
8774 o1: &mut [f32],
8775 o2: &mut [f32],
8776 pool: Option<&Pool>,
8777) {
8778 let gpr = cols / GROUP_SIZE;
8779 let p1 = SendMut(o1.as_mut_ptr());
8780 let p2 = SendMut(o2.as_mut_ptr());
8781 if a8w8_enabled() {
8782 let a1 = split_act(x1);
8783 let a2 = split_act(x2);
8784 let run = move |start: usize, end: usize| {
8785 for r in start..end {
8786 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
8787 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
8788 for &(j, xv) in &a1.outliers {
8789 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8790 v1 += w * s * xv;
8791 }
8792 for &(j, xv) in &a2.outliers {
8793 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8794 v2 += w * s * xv;
8795 }
8796 unsafe {
8798 *p1.at(r) = v1;
8799 *p2.at(r) = v2;
8800 }
8801 }
8802 };
8803 dispatch_rows(pool, rows, &run);
8804 return;
8805 }
8806 let run = move |start: usize, end: usize| {
8807 for r in start..end {
8808 unsafe {
8810 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
8811 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
8812 }
8813 }
8814 };
8815 dispatch_rows(pool, rows, &run);
8816}
8817
8818#[allow(clippy::too_many_arguments)]
8820#[cfg(target_os = "macos")]
8826fn dequant_matmat_accel(
8827 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
8828 xs_all: &[f32],
8829 b: usize,
8830 rows: usize,
8831 cols: usize,
8832 out: &mut [f32],
8833 pool: Option<&Pool>,
8834) {
8835 const TR: usize = 2048;
8836 thread_local! {
8837 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8838 }
8839 WTILE.with(|wt| {
8840 let mut wtile = wt.borrow_mut();
8841 wtile.resize(TR * cols, 0.0);
8842 let mut r0 = 0usize;
8843 while r0 < rows {
8844 let tr = TR.min(rows - r0);
8845 let wt_addr = SendMut(wtile.as_mut_ptr());
8846 let run = |start: usize, end: usize| {
8847 for r in start..end {
8848 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
8850 dequant_row(r0 + r, dst);
8851 }
8852 };
8853 dispatch_rows(pool, tr, &run);
8854 unsafe {
8855 accel_blas::cblas_sgemm(
8856 101, 111, 112, b as i32,
8860 tr as i32,
8861 cols as i32,
8862 1.0,
8863 xs_all.as_ptr(),
8864 cols as i32,
8865 wtile.as_ptr(),
8866 cols as i32,
8867 0.0,
8868 out.as_mut_ptr().add(r0),
8869 rows as i32,
8870 );
8871 }
8872 r0 += tr;
8873 }
8874 });
8875}
8876
8877fn q4t_matmat(
8878 bytes: &[u8],
8879 xs_all: &[f32],
8880 b: usize,
8881 rows: usize,
8882 cols: usize,
8883 out: &mut [f32],
8884 pool: Option<&Pool>,
8885) {
8886 debug_assert_eq!(out.len(), b * rows);
8887 let gpr = cols / GROUP_SIZE;
8888 #[cfg(target_os = "macos")]
8892 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8893 dequant_matmat_accel(
8894 &|r, dst| {
8895 for gi in 0..gpr {
8896 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
8897 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8898 for (k, &bb) in tile[2..].iter().enumerate() {
8899 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
8900 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
8901 }
8902 }
8903 },
8904 xs_all,
8905 b,
8906 rows,
8907 cols,
8908 out,
8909 pool,
8910 );
8911 return;
8912 }
8913 let out_addr = SendMut(out.as_mut_ptr());
8914 if a8w8_enabled() {
8915 let acts: Vec<SplitAct> = (0..b)
8916 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8917 .collect();
8918 let acts = &acts;
8919 #[cfg(target_arch = "x86_64")]
8920 let blocked_ok = avx2_enabled() && blocked_enabled();
8921 #[cfg(target_arch = "aarch64")]
8922 let blocked_ok = sdot_enabled() && blocked_enabled();
8923 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
8924 let blocked_ok = false;
8925 let run = move |start: usize, end: usize| {
8926 for r in start..end {
8927 let mut bi = 0usize;
8928 #[cfg(target_arch = "aarch64")]
8929 if blocked_ok {
8930 while bi + 4 <= acts.len() {
8931 let xs = [
8932 acts[bi].xq.as_slice(),
8933 acts[bi + 1].xq.as_slice(),
8934 acts[bi + 2].xq.as_slice(),
8935 acts[bi + 3].xq.as_slice(),
8936 ];
8937 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
8938 for k in 0..4 {
8939 let act = &acts[bi + k];
8940 let mut acc = d[k] * act.sx;
8941 for &(j, xv) in &act.outliers {
8942 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
8943 acc += w * sc * xv;
8944 }
8945 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8947 }
8948 bi += 4;
8949 }
8950 }
8951 #[cfg(target_arch = "x86_64")]
8952 if blocked_ok {
8953 while bi + 4 <= acts.len() {
8954 let xs = [
8955 acts[bi].xq.as_slice(),
8956 acts[bi + 1].xq.as_slice(),
8957 acts[bi + 2].xq.as_slice(),
8958 acts[bi + 3].xq.as_slice(),
8959 ];
8960 let d = unsafe {
8961 if vnni_tiles_enabled() {
8962 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
8963 } else {
8964 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
8965 }
8966 };
8967 for k in 0..4 {
8968 let act = &acts[bi + k];
8969 let mut acc = d[k] * act.sx;
8970 for &(j, xv) in &act.outliers {
8971 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
8972 acc += w * sc * xv;
8973 }
8974 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8976 }
8977 bi += 4;
8978 }
8979 }
8980 let _ = blocked_ok;
8981 while bi < acts.len() {
8982 let act = &acts[bi];
8983 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8984 for &(j, xv) in &act.outliers {
8985 let (w, s) = q4t_outlier(bytes, r, gpr, j);
8986 acc += w * s * xv;
8987 }
8988 unsafe { *out_addr.at(bi * rows + r) = acc };
8990 bi += 1;
8991 }
8992 }
8993 };
8994 dispatch_rows(pool, rows, &run);
8995 return;
8996 }
8997 let run = move |start: usize, end: usize| {
8998 for r in start..end {
8999 for bi in 0..b {
9000 let x = &xs_all[bi * cols..(bi + 1) * cols];
9001 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
9003 }
9004 }
9005 };
9006 dispatch_rows(pool, rows, &run);
9007}
9008
9009fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
9018 (0..gpr)
9019 .map(|gi| {
9020 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
9021 .iter()
9022 .map(|&v| v as i32)
9023 .sum()
9024 })
9025 .collect()
9026}
9027
9028#[inline]
9032#[allow(unreachable_code)]
9033#[cfg(target_arch = "x86_64")]
9038#[target_feature(enable = "avx2")]
9039unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9040 unsafe {
9042 use core::arch::x86_64::*;
9043 let expand = _mm256_setr_epi8(
9045 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,
9046 3, 3, 3,
9047 );
9048 let bitsel = _mm256_setr_epi8(
9049 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9050 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9051 );
9052 let ones8 = _mm256_set1_epi8(1);
9053 let ones16 = _mm256_set1_epi16(1);
9054 let mut acc = 0f32;
9055 for gi in 0..gpr {
9056 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9057 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9058 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9059 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9060 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9061 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9062 let sel = _mm256_and_si256(x, mask);
9063 let p16 = _mm256_maddubs_epi16(ones8, sel);
9065 let d32 = _mm256_madd_epi16(p16, ones16);
9066 let hi128 = _mm256_extracti128_si256::<1>(d32);
9067 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
9068 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9069 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9070 let msum = _mm_cvtsi128_si32(s32);
9071 let d = 2 * msum - gsum[gi];
9074 acc += d as f32 * s;
9075 }
9076 acc
9077 }
9078}
9079
9080#[cfg(target_arch = "x86_64")]
9083#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9084unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9085 unsafe {
9087 use core::arch::x86_64::*;
9088 let expand = _mm256_setr_epi8(
9089 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,
9090 3, 3, 3,
9091 );
9092 let bitsel = _mm256_setr_epi8(
9093 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9094 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9095 );
9096 let ones8 = _mm256_set1_epi8(1);
9097 let mut acc = 0f32;
9098 for gi in 0..gpr {
9099 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9100 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9101 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9102 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9103 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9104 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9105 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
9106 let d = 2 * msum - gsum[gi];
9107 acc += d as f32 * s;
9108 }
9109 acc
9110 }
9111}
9112
9113#[cfg(target_arch = "x86_64")]
9115#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9116unsafe fn dot_q1_row_1x4_vnni(
9117 bytes: &[u8],
9118 r: usize,
9119 gpr: usize,
9120 xs: [&[i8]; 4],
9121 gsums: [&[i32]; 4],
9122) -> [f32; 4] {
9123 unsafe {
9125 use core::arch::x86_64::*;
9126 let expand = _mm256_setr_epi8(
9127 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,
9128 3, 3, 3,
9129 );
9130 let bitsel = _mm256_setr_epi8(
9131 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9132 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9133 );
9134 let ones8 = _mm256_set1_epi8(1);
9135 let mut acc = [0f32; 4];
9136 for gi in 0..gpr {
9137 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9138 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9139 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9140 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9141 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9142 for (k, xq) in xs.iter().enumerate() {
9143 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9144 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
9145 let d = 2 * msum - gsums[k][gi];
9146 acc[k] += d as f32 * s;
9147 }
9148 }
9149 acc
9150 }
9151}
9152
9153#[cfg(target_arch = "x86_64")]
9156#[target_feature(enable = "avx2")]
9157unsafe fn dot_q1_row_1x4_avx2(
9158 bytes: &[u8],
9159 r: usize,
9160 gpr: usize,
9161 xs: [&[i8]; 4],
9162 gsums: [&[i32]; 4],
9163) -> [f32; 4] {
9164 unsafe {
9166 use core::arch::x86_64::*;
9167 let expand = _mm256_setr_epi8(
9168 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,
9169 3, 3, 3,
9170 );
9171 let bitsel = _mm256_setr_epi8(
9172 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9173 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9174 );
9175 let ones8 = _mm256_set1_epi8(1);
9176 let ones16 = _mm256_set1_epi16(1);
9177 let mut acc = [0f32; 4];
9178 for gi in 0..gpr {
9179 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9180 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9181 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9182 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9183 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9184 for (k, xq) in xs.iter().enumerate() {
9185 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9186 let sel = _mm256_and_si256(x, mask);
9187 let p16 = _mm256_maddubs_epi16(ones8, sel);
9188 let d32 = _mm256_madd_epi16(p16, ones16);
9189 let hi128 = _mm256_extracti128_si256::<1>(d32);
9190 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
9191 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9192 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9193 let msum = _mm_cvtsi128_si32(s32);
9194 let d = 2 * msum - gsums[k][gi];
9195 acc[k] += d as f32 * s;
9196 }
9197 }
9198 acc
9199 }
9200}
9201
9202#[allow(unreachable_code)]
9203fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9204 #[cfg(target_arch = "aarch64")]
9205 unsafe {
9206 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
9207 }
9208 #[cfg(target_arch = "x86_64")]
9209 if avx2_enabled() {
9210 unsafe {
9211 if vnni_tiles_enabled() {
9212 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
9213 }
9214 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
9215 }
9216 }
9217 let _ = gsum;
9218 let mut acc = 0f32;
9219 for gi in 0..gpr {
9220 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9221 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9222 let mut d = 0i32;
9223 for (j, &b) in tile[2..].iter().enumerate() {
9224 for k in 0..8 {
9225 let w = ((b >> k) & 1) as i32 * 2 - 1;
9226 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
9227 }
9228 }
9229 acc += d as f32 * s;
9230 }
9231 acc
9232}
9233
9234#[cfg(target_arch = "aarch64")]
9243#[target_feature(enable = "neon,dotprod")]
9244unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9245 unsafe {
9248 use core::arch::aarch64::*;
9249 use core::arch::asm;
9250 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9251 let m = vld1q_u8(MASKS.as_ptr());
9252 macro_rules! tile_dot {
9254 ($t:expr, $x:expr) => {{
9255 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
9256 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
9257 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9258 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9259 let x0 = vld1q_s8($x);
9260 let x1 = vld1q_s8($x.add(16));
9261 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9262 asm!(
9263 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9264 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9265 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9266 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9267 options(pure, nomem, nostack),
9268 );
9269 vaddq_s32(a0, a1)
9270 }};
9271 }
9272 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
9281 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
9282 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
9283 const IW11: [u8; 16] = [
9284 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9285 ];
9286 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
9287 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
9288 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
9289 let isc = vld1_u8(ISC.as_ptr());
9290 macro_rules! tile_dot_tbl {
9292 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
9293 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
9294 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
9295 let x0 = vld1q_s8($x);
9296 let x1 = vld1q_s8($x.add(16));
9297 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9298 asm!(
9299 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9300 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9301 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9302 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9303 options(pure, nomem, nostack),
9304 );
9305 vaddq_s32(a0, a1)
9306 }};
9307 }
9308 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
9309 let row_base = r * gpr * Q1_TILE;
9310 let abs_end = bytes.len();
9311 let xp = xq.as_ptr();
9312 let gp = gsum.as_ptr();
9313 let mut accv = vdupq_n_f32(0.0);
9314 let mut gi = 0;
9315 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
9318 let t0 = base.add(gi * Q1_TILE);
9319 let ld_a = vld1q_u8(t0);
9320 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
9321 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
9322 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
9323 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
9324 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
9325 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9327 let g = vld1q_s32(gp.add(gi));
9328 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9329 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
9330 let scf: float32x4_t;
9331 asm!(
9332 "fcvtl {o:v}.4s, {i:v}.4h",
9333 o = out(vreg) scf, i = in(vreg) sc16,
9334 options(pure, nomem, nostack),
9335 );
9336 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
9337 gi += 4;
9338 }
9339 let mut acc = vaddvq_f32(accv);
9340 while gi < gpr {
9341 let t = base.add(gi * Q1_TILE);
9342 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9343 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
9344 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
9345 gi += 1;
9346 }
9347 acc
9348 }
9349}
9350
9351#[cfg(target_arch = "aarch64")]
9356#[target_feature(enable = "neon,dotprod")]
9357unsafe fn dot_q1_row_1x4_sdot(
9358 bytes: &[u8],
9359 r: usize,
9360 gpr: usize,
9361 xs: [&[i8]; 4],
9362 gs: [&[i32]; 4],
9363) -> [f32; 4] {
9364 unsafe {
9366 use core::arch::aarch64::*;
9367 use core::arch::asm;
9368 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9369 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
9370 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
9371 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
9372 const IW11: [u8; 16] = [
9373 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9374 ];
9375 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
9376 let m = vld1q_u8(MASKS.as_ptr());
9377 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
9378 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
9379 let isc = vld1_u8(ISC.as_ptr());
9380 macro_rules! sdot2 {
9381 ($w0:expr, $w1:expr, $x:expr) => {{
9382 let x0 = vld1q_s8($x);
9383 let x1 = vld1q_s8($x.add(16));
9384 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9385 asm!(
9386 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9387 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9388 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9389 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9390 options(pure, nomem, nostack),
9391 );
9392 vaddq_s32(a0, a1)
9393 }};
9394 }
9395 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
9396 let row_base = r * gpr * Q1_TILE;
9397 let abs_end = bytes.len();
9398 let mut accv = [vdupq_n_f32(0.0); 4];
9399 let mut gi = 0;
9400 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
9401 let t0 = base.add(gi * Q1_TILE);
9402 let ld_a = vld1q_u8(t0);
9403 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
9404 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
9406 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
9407 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
9408 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
9409 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
9410 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
9411 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
9412 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
9413 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
9414 let scf: float32x4_t;
9415 asm!(
9416 "fcvtl {o:v}.4s, {i:v}.4h",
9417 o = out(vreg) scf, i = in(vreg) sc16,
9418 options(pure, nomem, nostack),
9419 );
9420 for k in 0..4 {
9421 let xp = xs[k].as_ptr();
9422 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
9423 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
9424 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
9425 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
9426 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9427 let g = vld1q_s32(gs[k].as_ptr().add(gi));
9428 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9429 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
9430 }
9431 gi += 4;
9432 }
9433 let mut acc = [
9434 vaddvq_f32(accv[0]),
9435 vaddvq_f32(accv[1]),
9436 vaddvq_f32(accv[2]),
9437 vaddvq_f32(accv[3]),
9438 ];
9439 while gi < gpr {
9440 let t = base.add(gi * Q1_TILE);
9441 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9442 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
9443 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
9444 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9445 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9446 for k in 0..4 {
9447 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
9448 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
9449 }
9450 gi += 1;
9451 }
9452 acc
9453 }
9454}
9455
9456#[inline]
9458fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
9459 let gi = j / GROUP_SIZE;
9460 let k = j % GROUP_SIZE;
9461 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9462 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9463 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
9464 ((bit as i32 * 2 - 1) as f32, s)
9465}
9466
9467#[inline]
9469fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
9470 let mut acc = 0f32;
9471 for gi in 0..gpr {
9472 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9473 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9474 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9475 let mut ga = 0f32;
9476 for (j, &b) in tile[2..].iter().enumerate() {
9477 for k in 0..8 {
9478 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
9479 }
9480 }
9481 acc += ga * s;
9482 }
9483 acc
9484}
9485
9486#[allow(clippy::too_many_arguments)]
9489fn bin_range_a8w8(
9490 w: BinW,
9491 gpr: usize,
9492 act: &SplitAct,
9493 gsum: &[i32],
9494 out: SendMut,
9495 start: usize,
9496 end: usize,
9497) {
9498 for r in start..end {
9499 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9500 for &(j, xv) in &act.outliers {
9501 let (wv, s) = w.outlier(r, gpr, j);
9502 acc += wv * s * xv;
9503 }
9504 unsafe { *out.at(r) = acc };
9506 }
9507}
9508
9509fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
9511 for r in start..end {
9512 unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
9514 }
9515}
9516
9517#[derive(Clone, Copy)]
9525pub(crate) enum BinW<'a> {
9526 Q1(&'a [u8]),
9527 Q1g {
9528 bytes: &'a [u8],
9529 sc_off: usize,
9531 rows: usize,
9534 },
9535}
9536
9537impl<'a> BinW<'a> {
9538 fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
9539 Self::Q1g {
9540 bytes,
9541 sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
9542 rows,
9543 }
9544 }
9545
9546 #[inline]
9547 fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9548 match self {
9549 Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
9550 Self::Q1g {
9551 bytes,
9552 sc_off,
9553 rows,
9554 } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
9555 }
9556 }
9557
9558 #[inline]
9559 fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
9560 match self {
9561 Self::Q1(b) => q1_row_exact(b, r, gpr, x),
9562 Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
9563 }
9564 }
9565
9566 #[inline]
9567 fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
9568 match self {
9569 Self::Q1(b) => q1_outlier(b, r, gpr, j),
9570 Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
9571 }
9572 }
9573
9574 #[cfg(target_arch = "aarch64")]
9579 #[inline]
9580 unsafe fn dot_1x4_sdot(
9581 self,
9582 r: usize,
9583 gpr: usize,
9584 xs: [&[i8]; 4],
9585 gs: [&[i32]; 4],
9586 ) -> [f32; 4] {
9587 unsafe {
9589 match self {
9590 Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
9591 Self::Q1g {
9592 bytes,
9593 sc_off,
9594 rows,
9595 } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
9596 }
9597 }
9598 }
9599
9600 #[cfg(target_arch = "x86_64")]
9605 #[inline]
9606 unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
9607 unsafe {
9609 match (self, vnni_tiles_enabled()) {
9610 (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
9611 (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
9612 (Self::Q1g { bytes, sc_off, .. }, true) => {
9613 dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
9614 }
9615 (Self::Q1g { bytes, sc_off, .. }, false) => {
9616 dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
9617 }
9618 }
9619 }
9620 }
9621}
9622
9623#[inline]
9633fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
9634 u32::from_le_bytes([
9635 bytes[t * 4],
9636 bytes[t * 4 + 1],
9637 bytes[t * 4 + 2],
9638 bytes[t * 4 + 3],
9639 ])
9640}
9641
9642#[inline]
9644fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
9645 let o = sc_off + (t / 4) * 2;
9646 u16::from_le_bytes([bytes[o], bytes[o + 1]])
9647}
9648
9649#[allow(unreachable_code)]
9651#[allow(clippy::too_many_arguments)]
9652fn dot_q1g_row_i8(
9653 bytes: &[u8],
9654 sc_off: usize,
9655 rows: usize,
9656 r: usize,
9657 gpr: usize,
9658 xq: &[i8],
9659 gsum: &[i32],
9660) -> f32 {
9661 #[cfg(target_arch = "aarch64")]
9662 unsafe {
9663 return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
9664 }
9665 #[cfg(target_arch = "x86_64")]
9666 if avx2_enabled() {
9667 unsafe {
9668 if vnni_tiles_enabled() {
9669 return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
9670 }
9671 return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
9672 }
9673 }
9674 let _ = (gsum, rows);
9675 let mut acc = 0f32;
9676 for gi in 0..gpr {
9677 let t = r * gpr + gi;
9678 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9679 let bits = q1g_bits(bytes, t);
9680 let mut d = 0i32;
9681 for k in 0..GROUP_SIZE {
9682 let w = ((bits >> k) & 1) as i32 * 2 - 1;
9683 d += w * xq[gi * GROUP_SIZE + k] as i32;
9684 }
9685 acc += d as f32 * s;
9686 }
9687 acc
9688}
9689
9690#[cfg(target_arch = "x86_64")]
9695#[target_feature(enable = "avx2")]
9696#[inline]
9697unsafe fn q1g_hsum4(
9698 d0: core::arch::x86_64::__m256i,
9699 d1: core::arch::x86_64::__m256i,
9700 d2: core::arch::x86_64::__m256i,
9701 d3: core::arch::x86_64::__m256i,
9702) -> core::arch::x86_64::__m128i {
9703 unsafe {
9705 use core::arch::x86_64::*;
9706 let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
9707 _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
9708 }
9709}
9710
9711#[cfg(target_arch = "x86_64")]
9715#[target_feature(enable = "avx2")]
9716#[inline]
9717unsafe fn q1g_group_acc(
9718 acc: f32,
9719 msum4: core::arch::x86_64::__m128i,
9720 gsum4: &[i32],
9721 s: f32,
9722) -> f32 {
9723 unsafe {
9725 use core::arch::x86_64::*;
9726 let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
9727 let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
9728 let mut f = [0f32; 4];
9729 _mm_storeu_ps(f.as_mut_ptr(), d);
9730 let mut acc = acc;
9731 for v in f {
9732 acc += v * s;
9733 }
9734 acc
9735 }
9736}
9737
9738#[cfg(target_arch = "x86_64")]
9743#[target_feature(enable = "avx2")]
9744unsafe fn dot_q1g_row_avx2(
9745 bytes: &[u8],
9746 sc_off: usize,
9747 r: usize,
9748 gpr: usize,
9749 xq: &[i8],
9750 gsum: &[i32],
9751) -> f32 {
9752 unsafe {
9755 use core::arch::x86_64::*;
9756 let expand = _mm256_setr_epi8(
9757 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,
9758 3, 3, 3,
9759 );
9760 let bitsel = _mm256_setr_epi8(
9761 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9762 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9763 );
9764 let ones8 = _mm256_set1_epi8(1);
9765 let ones16 = _mm256_set1_epi16(1);
9766 let bp = bytes.as_ptr().add(r * gpr * 4);
9767 let mut acc = 0f32;
9768 let mut gi = 0;
9769 while gi < gpr {
9770 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9771 let mut d = [_mm256_setzero_si256(); 4];
9772 for (k, dk) in d.iter_mut().enumerate() {
9773 let t = gi + k;
9774 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9775 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9776 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9777 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9778 let sel = _mm256_and_si256(x, mask);
9779 *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9780 }
9781 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9782 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9783 gi += 4;
9784 }
9785 acc
9786 }
9787}
9788
9789#[cfg(target_arch = "x86_64")]
9792#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9793unsafe fn dot_q1g_row_vnni(
9794 bytes: &[u8],
9795 sc_off: usize,
9796 r: usize,
9797 gpr: usize,
9798 xq: &[i8],
9799 gsum: &[i32],
9800) -> f32 {
9801 unsafe {
9803 use core::arch::x86_64::*;
9804 let expand = _mm256_setr_epi8(
9805 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,
9806 3, 3, 3,
9807 );
9808 let bitsel = _mm256_setr_epi8(
9809 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9810 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9811 );
9812 let ones8 = _mm256_set1_epi8(1);
9813 let bp = bytes.as_ptr().add(r * gpr * 4);
9814 let mut acc = 0f32;
9815 let mut gi = 0;
9816 while gi < gpr {
9817 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9818 let mut d = [_mm256_setzero_si256(); 4];
9819 for (k, dk) in d.iter_mut().enumerate() {
9820 let t = gi + k;
9821 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9822 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9823 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9824 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9825 *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
9826 }
9827 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9828 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9829 gi += 4;
9830 }
9831 acc
9832 }
9833}
9834
9835#[cfg(target_arch = "x86_64")]
9839#[target_feature(enable = "avx2")]
9840unsafe fn dot_q1g_row_1x4_avx2(
9841 bytes: &[u8],
9842 sc_off: usize,
9843 r: usize,
9844 gpr: usize,
9845 xs: [&[i8]; 4],
9846 gsums: [&[i32]; 4],
9847) -> [f32; 4] {
9848 unsafe {
9850 use core::arch::x86_64::*;
9851 let expand = _mm256_setr_epi8(
9852 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,
9853 3, 3, 3,
9854 );
9855 let bitsel = _mm256_setr_epi8(
9856 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9857 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9858 );
9859 let ones8 = _mm256_set1_epi8(1);
9860 let ones16 = _mm256_set1_epi16(1);
9861 let bp = bytes.as_ptr().add(r * gpr * 4);
9862 let mut acc = [0f32; 4];
9863 let mut gi = 0;
9864 while gi < gpr {
9865 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9866 let mut d = [[_mm256_setzero_si256(); 4]; 4];
9868 for t in 0..4 {
9869 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9870 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9871 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9872 for (k, xq) in xs.iter().enumerate() {
9873 let x = _mm256_loadu_si256(
9874 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9875 );
9876 let sel = _mm256_and_si256(x, mask);
9877 d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9878 }
9879 }
9880 for k in 0..4 {
9881 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9882 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9883 }
9884 gi += 4;
9885 }
9886 acc
9887 }
9888}
9889
9890#[cfg(target_arch = "x86_64")]
9892#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9893unsafe fn dot_q1g_row_1x4_vnni(
9894 bytes: &[u8],
9895 sc_off: usize,
9896 r: usize,
9897 gpr: usize,
9898 xs: [&[i8]; 4],
9899 gsums: [&[i32]; 4],
9900) -> [f32; 4] {
9901 unsafe {
9903 use core::arch::x86_64::*;
9904 let expand = _mm256_setr_epi8(
9905 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,
9906 3, 3, 3,
9907 );
9908 let bitsel = _mm256_setr_epi8(
9909 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9910 -128, 1, 2, 4, 8, 16, 32, 64, -128,
9911 );
9912 let ones8 = _mm256_set1_epi8(1);
9913 let bp = bytes.as_ptr().add(r * gpr * 4);
9914 let mut acc = [0f32; 4];
9915 let mut gi = 0;
9916 while gi < gpr {
9917 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9918 let mut d = [[_mm256_setzero_si256(); 4]; 4];
9919 for t in 0..4 {
9920 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9921 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9922 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9923 for (k, xq) in xs.iter().enumerate() {
9924 let x = _mm256_loadu_si256(
9925 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9926 );
9927 d[k][t] = _mm256_dpbusd_epi32(
9928 _mm256_setzero_si256(),
9929 ones8,
9930 _mm256_and_si256(x, mask),
9931 );
9932 }
9933 }
9934 for k in 0..4 {
9935 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9936 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9937 }
9938 gi += 4;
9939 }
9940 acc
9941 }
9942}
9943
9944#[cfg(target_arch = "aarch64")]
9955#[target_feature(enable = "neon,dotprod")]
9956#[allow(clippy::too_many_arguments)]
9957unsafe fn dot_q1g_row_sdot(
9958 bytes: &[u8],
9959 sc_off: usize,
9960 rows: usize,
9961 r: usize,
9962 gpr: usize,
9963 xq: &[i8],
9964 gsum: &[i32],
9965) -> f32 {
9966 unsafe {
9969 use core::arch::aarch64::*;
9970 use core::arch::asm;
9971 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9972 const IW: [[u8; 16]; 8] = [
9974 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
9975 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
9976 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
9977 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
9978 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
9979 [
9980 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9981 ],
9982 [
9983 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
9984 ],
9985 [
9986 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
9987 ],
9988 ];
9989 let m = vld1q_u8(MASKS.as_ptr());
9990 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
9991 macro_rules! tile_dot {
9992 ($w0:expr, $w1:expr, $x:expr) => {{
9993 let x0 = vld1q_s8($x);
9994 let x1 = vld1q_s8($x.add(16));
9995 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9996 asm!(
9997 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9998 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9999 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10000 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10001 options(pure, nomem, nostack),
10002 );
10003 vaddq_s32(a0, a1)
10004 }};
10005 }
10006 macro_rules! tbl_mask {
10007 ($ld:expr, $i:expr) => {
10008 vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
10009 };
10010 }
10011 let bp = bytes.as_ptr().add(r * gpr * 4);
10012 let xp = xq.as_ptr();
10013 let gp = gsum.as_ptr();
10014 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
10015 let mut accv = vdupq_n_f32(0.0);
10016 let mut gi = 0;
10017 while gi + 4 <= vec_end {
10018 let ld = vld1q_u8(bp.add(gi * 4));
10019 let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
10020 let d1 = tile_dot!(
10021 tbl_mask!(ld, 2),
10022 tbl_mask!(ld, 3),
10023 xp.add((gi + 1) * GROUP_SIZE)
10024 );
10025 let d2 = tile_dot!(
10026 tbl_mask!(ld, 4),
10027 tbl_mask!(ld, 5),
10028 xp.add((gi + 2) * GROUP_SIZE)
10029 );
10030 let d3 = tile_dot!(
10031 tbl_mask!(ld, 6),
10032 tbl_mask!(ld, 7),
10033 xp.add((gi + 3) * GROUP_SIZE)
10034 );
10035 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
10036 let g = vld1q_s32(gp.add(gi));
10037 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
10038 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10039 let scf: float32x4_t;
10040 asm!(
10041 "fcvtl {o:v}.4s, {i:v}.4h",
10042 o = out(vreg) scf, i = in(vreg) h,
10043 options(pure, nomem, nostack),
10044 );
10045 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
10046 gi += 4;
10047 }
10048 let mut acc = vaddvq_f32(accv);
10049 while gi < gpr {
10050 let t = bp.add(gi * 4);
10051 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10052 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
10053 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
10054 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
10055 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
10056 let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
10057 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
10058 gi += 1;
10059 }
10060 acc
10061 }
10062}
10063
10064#[cfg(target_arch = "aarch64")]
10068#[target_feature(enable = "neon,dotprod")]
10069#[allow(clippy::too_many_arguments)]
10070unsafe fn dot_q1g_row_1x4_sdot(
10071 bytes: &[u8],
10072 sc_off: usize,
10073 rows: usize,
10074 r: usize,
10075 gpr: usize,
10076 xs: [&[i8]; 4],
10077 gs: [&[i32]; 4],
10078) -> [f32; 4] {
10079 unsafe {
10081 use core::arch::aarch64::*;
10082 use core::arch::asm;
10083 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
10084 const IW: [[u8; 16]; 8] = [
10085 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
10086 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
10087 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
10088 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
10089 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
10090 [
10091 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
10092 ],
10093 [
10094 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
10095 ],
10096 [
10097 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
10098 ],
10099 ];
10100 let m = vld1q_u8(MASKS.as_ptr());
10101 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
10102 macro_rules! sdot2 {
10103 ($w0:expr, $w1:expr, $x:expr) => {{
10104 let x0 = vld1q_s8($x);
10105 let x1 = vld1q_s8($x.add(16));
10106 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10107 asm!(
10108 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10109 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10110 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10111 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10112 options(pure, nomem, nostack),
10113 );
10114 vaddq_s32(a0, a1)
10115 }};
10116 }
10117 let bp = bytes.as_ptr().add(r * gpr * 4);
10118 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
10119 let mut accv = [vdupq_n_f32(0.0); 4];
10120 let mut gi = 0;
10121 while gi + 4 <= vec_end {
10122 let ld = vld1q_u8(bp.add(gi * 4));
10123 let w: [int8x16_t; 8] =
10125 core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
10126 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10127 let scf: float32x4_t;
10128 asm!(
10129 "fcvtl {o:v}.4s, {i:v}.4h",
10130 o = out(vreg) scf, i = in(vreg) h,
10131 options(pure, nomem, nostack),
10132 );
10133 for k in 0..4 {
10134 let xp = xs[k].as_ptr();
10135 let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
10136 let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
10137 let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
10138 let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
10139 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
10140 let g = vld1q_s32(gs[k].as_ptr().add(gi));
10141 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
10142 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
10143 }
10144 gi += 4;
10145 }
10146 let mut acc = [
10147 vaddvq_f32(accv[0]),
10148 vaddvq_f32(accv[1]),
10149 vaddvq_f32(accv[2]),
10150 vaddvq_f32(accv[3]),
10151 ];
10152 while gi < gpr {
10153 let t = bp.add(gi * 4);
10154 let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10155 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
10156 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
10157 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
10158 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
10159 for k in 0..4 {
10160 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
10161 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
10162 }
10163 gi += 1;
10164 }
10165 acc
10166 }
10167}
10168
10169#[inline]
10171fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
10172 let t = r * gpr + j / GROUP_SIZE;
10173 let k = j % GROUP_SIZE;
10174 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
10175 let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
10176 ((bit as i32 * 2 - 1) as f32, s)
10177}
10178
10179#[inline]
10181fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
10182 let mut acc = 0f32;
10183 for gi in 0..gpr {
10184 let t = r * gpr + gi;
10185 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
10186 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10187 let mut ga = 0f32;
10188 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
10189 for k in 0..8 {
10190 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
10191 }
10192 }
10193 acc += ga * s;
10194 }
10195 acc
10196}
10197
10198fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
10203 let entries = base_len + (rows + 1) * 4;
10204 (base_len, entries, entries <= bytes.len())
10205}
10206
10207#[inline]
10209fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
10210 let o = rp_off + r * 4;
10211 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
10212}
10213
10214const SIGN5: [[f32; 5]; 256] = {
10218 let mut lut = [[0.0f32; 5]; 256];
10219 let pow3 = [1u16, 3, 9, 27, 81];
10220 let mut byte = 0usize;
10221 while byte < 256 {
10222 let mut i = 0usize;
10223 while i < 5 {
10224 let code = (byte as u16 / pow3[i]) % 3;
10225 lut[byte][i] = if code == 1 {
10226 1.0
10227 } else if code == 2 {
10228 -1.0
10229 } else {
10230 0.0
10231 };
10232 i += 1;
10233 }
10234 byte += 1;
10235 }
10236 lut
10237};
10238
10239const SIGN5_I8: [[i8; 5]; 256] = {
10241 let mut lut = [[0i8; 5]; 256];
10242 let pow3 = [1u16, 3, 9, 27, 81];
10243 let mut byte = 0usize;
10244 while byte < 256 {
10245 let mut i = 0usize;
10246 while i < 5 {
10247 let code = (byte as u16 / pow3[i]) % 3;
10248 lut[byte][i] = if code == 1 {
10249 1
10250 } else if code == 2 {
10251 -1
10252 } else {
10253 0
10254 };
10255 i += 1;
10256 }
10257 byte += 1;
10258 }
10259 lut
10260};
10261
10262const SIGN5_U64: [u64; 256] = {
10268 let mut lut = [0u64; 256];
10269 let pow3 = [1u16, 3, 9, 27, 81];
10270 let mut byte = 0usize;
10271 while byte < 256 {
10272 let mut v = 0u64;
10273 let mut i = 0usize;
10274 while i < 5 {
10275 let code = (byte as u16 / pow3[i]) % 3;
10276 let s: u8 = if code == 1 {
10277 1
10278 } else if code == 2 {
10279 0xFF
10280 } else {
10281 0
10282 };
10283 v |= (s as u64) << (i * 8);
10284 i += 1;
10285 }
10286 lut[byte] = v;
10287 byte += 1;
10288 }
10289 lut
10290};
10291
10292#[inline]
10297fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
10298 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10299 let off = (r * gpr + j / GROUP_SIZE) * TILE;
10300 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10301 let within = j % GROUP_SIZE;
10302 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
10303}
10304
10305#[cfg(target_arch = "aarch64")]
10308#[target_feature(enable = "neon,dotprod")]
10309#[inline]
10310unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
10311 unsafe {
10313 use core::arch::aarch64::*;
10314 use core::arch::asm;
10315 let w0 = vld1q_s8(w);
10316 let w1 = vld1q_s8(w.add(16));
10317 let x0 = vld1q_s8(x);
10318 let x1 = vld1q_s8(x.add(16));
10319 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10320 asm!(
10321 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10322 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10323 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10324 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10325 options(pure, nomem, nostack),
10326 );
10327 vaddvq_s32(vaddq_s32(a0, a1))
10328 }
10329}
10330
10331#[cfg(target_arch = "x86_64")]
10334#[target_feature(enable = "avx2")]
10335#[inline]
10336unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
10337 unsafe {
10339 use core::arch::x86_64::*;
10340 let wv = _mm256_loadu_si256(w as *const __m256i);
10341 let xv = _mm256_loadu_si256(x as *const __m256i);
10342 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10343 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
10344 let hi128 = _mm256_extracti128_si256::<1>(d);
10345 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10346 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10347 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10348 _mm_cvtsi128_si32(s32)
10349 }
10350}
10351
10352#[inline]
10357fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
10358 debug_assert!(dst.len() >= 40);
10359 unsafe {
10362 let p = dst.as_mut_ptr();
10363 for bi in 0..7 {
10364 core::ptr::write_unaligned(
10365 p.add(bi * 5) as *mut u64,
10366 SIGN5_U64[*codes.add(bi) as usize],
10367 );
10368 }
10369 }
10370}
10371
10372#[inline]
10377fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
10378 #[cfg(target_arch = "aarch64")]
10379 unsafe {
10380 return sdot32_i8(w, x);
10381 }
10382 #[cfg(target_arch = "x86_64")]
10383 unsafe {
10384 return i8dot32_avx2(w, x);
10385 }
10386 #[allow(unreachable_code)]
10387 unsafe {
10388 let mut s = 0i32;
10389 for k in 0..GROUP_SIZE {
10390 s += *w.add(k) as i32 * *x.add(k) as i32;
10391 }
10392 s
10393 }
10394}
10395
10396#[inline]
10397unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
10398 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
10399 (
10400 SIGN5_U64[*codes as usize],
10401 SIGN5_U64[*codes.add(1) as usize],
10402 SIGN5_U64[*codes.add(2) as usize],
10403 SIGN5_U64[*codes.add(3) as usize],
10404 SIGN5_U64[*codes.add(4) as usize],
10405 SIGN5_U64[*codes.add(5) as usize],
10406 SIGN5_U64[*codes.add(6) as usize],
10407 )
10408 };
10409
10410 let u0 = s0 | (s1 << 40);
10411 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
10412 let u2 = (s3 >> 8) | (s4 << 32);
10413 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
10414
10415 (u0, u1, u2, u3)
10416}
10417
10418#[cfg(target_arch = "aarch64")]
10422#[target_feature(enable = "neon,dotprod")]
10423unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10424 use core::arch::aarch64::*;
10425 use core::arch::asm;
10426 unsafe {
10427 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10428 let mut acc = 0f32;
10429 let bytes_ptr = bytes.as_ptr();
10430 let xq_ptr = xq.as_ptr();
10431 let row_off = r * gpr * TILE;
10432
10433 let gpr2 = gpr & !1;
10434 let mut gi = 0;
10435 while gi < gpr2 {
10436 let off0 = row_off + gi * TILE;
10437 let off1 = off0 + TILE;
10438 let s0 = f16_to_f32(u16::from_le_bytes([
10439 *bytes_ptr.add(off0),
10440 *bytes_ptr.add(off0 + 1),
10441 ]));
10442 let s1 = f16_to_f32(u16::from_le_bytes([
10443 *bytes_ptr.add(off1),
10444 *bytes_ptr.add(off1 + 1),
10445 ]));
10446
10447 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10448 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10449
10450 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10451 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10452 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10453 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10454
10455 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
10456 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
10457 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
10458 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
10459
10460 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
10461 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10462 asm!(
10463 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
10464 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
10465 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
10466 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
10467 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
10468 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
10469 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
10470 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
10471 options(pure, nomem, nostack),
10472 );
10473 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
10474 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
10475 acc += d0 as f32 * s0 + d1 as f32 * s1;
10476 gi += 2;
10477 }
10478
10479 if gi < gpr {
10480 let off = row_off + gi * TILE;
10481 let s = f16_to_f32(u16::from_le_bytes([
10482 *bytes_ptr.add(off),
10483 *bytes_ptr.add(off + 1),
10484 ]));
10485 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10486 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10487 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10488 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
10489 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
10490 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10491 asm!(
10492 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10493 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10494 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10495 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10496 options(pure, nomem, nostack),
10497 );
10498 let d = vaddvq_s32(vaddq_s32(a0, a1));
10499 acc += d as f32 * s;
10500 }
10501 acc
10502 }
10503}
10504
10505#[cfg(target_arch = "x86_64")]
10507#[target_feature(enable = "avx2")]
10508unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10509 use core::arch::x86_64::*;
10510 unsafe {
10511 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10512 let mut acc = 0f32;
10513 let bytes_ptr = bytes.as_ptr();
10514 let xq_ptr = xq.as_ptr();
10515 let row_off = r * gpr * TILE;
10516
10517 let ones = _mm256_set1_epi16(1);
10518 for gi in 0..gpr {
10519 let off = row_off + gi * TILE;
10520 let s = f16_to_f32(u16::from_le_bytes([
10521 *bytes_ptr.add(off),
10522 *bytes_ptr.add(off + 1),
10523 ]));
10524 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10525 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
10526 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
10527 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10528 let d256 = _mm256_madd_epi16(p16, ones);
10529 let d128 = _mm_add_epi32(
10530 _mm256_castsi256_si128(d256),
10531 _mm256_extracti128_si256(d256, 1),
10532 );
10533 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
10534 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
10535 acc += d32 as f32 * s;
10536 }
10537 acc
10538 }
10539}
10540
10541#[cfg(target_arch = "x86_64")]
10543#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10544unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10545 use core::arch::x86_64::*;
10546 unsafe {
10548 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10549 let mut acc = 0f32;
10550 let bytes_ptr = bytes.as_ptr();
10551 let xq_ptr = xq.as_ptr();
10552 let row_off = r * gpr * TILE;
10553 for gi in 0..gpr {
10554 let off = row_off + gi * TILE;
10555 let s = f16_to_f32(u16::from_le_bytes([
10556 *bytes_ptr.add(off),
10557 *bytes_ptr.add(off + 1),
10558 ]));
10559 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10560 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
10561 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
10562 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10563 acc += d as f32 * s;
10564 }
10565 acc
10566 }
10567}
10568
10569#[inline]
10573fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10574 #[cfg(target_arch = "aarch64")]
10575 unsafe {
10576 return q1t_dot_row_sdot(bytes, r, gpr, xq);
10577 }
10578 #[cfg(target_arch = "x86_64")]
10579 unsafe {
10580 if vnni_tiles_enabled() {
10581 return q1t_dot_row_vnni(bytes, r, gpr, xq);
10582 }
10583 return q1t_dot_row_avx2(bytes, r, gpr, xq);
10584 }
10585 #[allow(unreachable_code)]
10586 {
10587 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10588 let mut acc = 0f32;
10589 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
10591 let off = (r * gpr + gi) * TILE;
10592 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10593 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
10594 let mut d = 0i32;
10595 for k in 0..GROUP_SIZE {
10596 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
10597 }
10598 acc += d as f32 * s;
10599 }
10600 acc
10601 }
10602}
10603
10604fn q1t_row_outlier_correction(
10611 bytes: &[u8],
10612 r: usize,
10613 rp_off: usize,
10614 entries_off: usize,
10615 has_ov: bool,
10616 x: &[f32],
10617) -> f32 {
10618 if !has_ov {
10619 return 0.0;
10620 }
10621 let (c0, c1) = (
10622 q1t_rowptr(bytes, rp_off, r),
10623 q1t_rowptr(bytes, rp_off, r + 1),
10624 );
10625 let mut corr = 0f32;
10626 for p in c0..c1 {
10627 let e = entries_off + p * 4;
10628 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10629 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10630 corr += val * x[col];
10631 }
10632 corr
10633}
10634
10635fn q1t_dequant_row(
10639 bytes: &[u8],
10640 r: usize,
10641 gpr: usize,
10642 rp_off: usize,
10643 entries_off: usize,
10644 has_ov: bool,
10645 buf: &mut [f32],
10646) {
10647 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10648 for g in 0..gpr {
10649 let off = (r * gpr + g) * TILE;
10650 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10651 let codes = &bytes[off + 2..off + TILE];
10652 let bc = g * GROUP_SIZE;
10653 for bi in 0..6 {
10655 let lut = &SIGN5[codes[bi] as usize];
10656 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
10657 for i in 0..5 {
10658 d[i] = lut[i] * s;
10659 }
10660 }
10661 let lut = &SIGN5[codes[6] as usize];
10662 buf[bc + 30] = lut[0] * s;
10663 buf[bc + 31] = lut[1] * s;
10664 }
10665 if !has_ov {
10666 return;
10667 }
10668 let (c0, c1) = (
10669 q1t_rowptr(bytes, rp_off, r),
10670 q1t_rowptr(bytes, rp_off, r + 1),
10671 );
10672 for p in c0..c1 {
10673 let e = entries_off + p * 4;
10674 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10675 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10676 }
10677}
10678
10679fn q1t_add_overlay(
10683 bytes: &[u8],
10684 x: &[f32],
10685 rows: usize,
10686 cols: usize,
10687 out: &mut [f32],
10688 pool: Option<&Pool>,
10689) {
10690 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10691 let gpr = cols / GROUP_SIZE;
10692 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10693 if !has_ov {
10694 return;
10695 }
10696 let out_addr = SendMut(out.as_mut_ptr());
10697 let run = move |start: usize, end: usize| {
10698 for r in start..end {
10699 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10700 unsafe { *out_addr.at(r) += corr };
10702 }
10703 };
10704 dispatch_rows(pool, rows, &run);
10705}
10706
10707#[allow(clippy::too_many_arguments)]
10710fn q1t_range_a8w8(
10711 bytes: &[u8],
10712 gpr: usize,
10713 rp_off: usize,
10714 ent_off: usize,
10715 has_ov: bool,
10716 act: &SplitAct,
10717 x: &[f32],
10718 out: SendMut,
10719 start: usize,
10720 end: usize,
10721) {
10722 for r in start..end {
10723 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10724 for &(j, xv) in &act.outliers {
10725 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10726 }
10727 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10728 unsafe { *out.at(r) = acc };
10730 }
10731}
10732
10733#[allow(clippy::too_many_arguments)]
10736fn q1t_range_f32_batch(
10737 bytes: &[u8],
10738 gpr: usize,
10739 rp_off: usize,
10740 ent_off: usize,
10741 has_ov: bool,
10742 x: &[f32],
10743 out: SendMut,
10744 start: usize,
10745 end: usize,
10746) {
10747 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10748 let mut sg = [0f32; GROUP_SIZE];
10749 for r in start..end {
10750 let mut acc = 0f32;
10751 for g in 0..gpr {
10752 let off = (r * gpr + g) * TILE;
10753 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10754 let codes = &bytes[off + 2..off + TILE];
10755 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10756 for bi in 0..6 {
10757 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10758 }
10759 let lut = &SIGN5[codes[6] as usize];
10760 sg[30] = lut[0];
10761 sg[31] = lut[1];
10762 let mut gsum = 0f32;
10763 for k in 0..GROUP_SIZE {
10764 gsum += sg[k] * xg[k];
10765 }
10766 acc += s * gsum;
10767 }
10768 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10769 unsafe { *out.at(r) = acc };
10771 }
10772}
10773
10774fn q1t_matvec(
10778 bytes: &[u8],
10779 x: &[f32],
10780 rows: usize,
10781 cols: usize,
10782 out: &mut [f32],
10783 pool: Option<&Pool>,
10784) {
10785 debug_assert_eq!(out.len(), rows);
10786 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10787 let gpr = cols / GROUP_SIZE;
10788 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10789 let out_addr = SendMut(out.as_mut_ptr());
10790 if a8w8_enabled() {
10794 let act = split_act(x);
10795 let act = &act;
10796 let run = move |start: usize, end: usize| {
10797 for r in start..end {
10798 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10799 for &(j, xv) in &act.outliers {
10800 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10801 }
10802 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10803 unsafe { *out_addr.at(r) = acc };
10805 }
10806 };
10807 dispatch_rows(pool, rows, &run);
10808 return;
10809 }
10810 let run = move |start: usize, end: usize| {
10811 let mut sg = [0f32; GROUP_SIZE];
10815 for r in start..end {
10816 let mut acc = 0f32;
10817 for g in 0..gpr {
10818 let off = (r * gpr + g) * TILE;
10819 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10820 let codes = &bytes[off + 2..off + TILE];
10821 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10822 for bi in 0..6 {
10823 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10824 }
10825 let lut = &SIGN5[codes[6] as usize];
10826 sg[30] = lut[0];
10827 sg[31] = lut[1];
10828 let mut gsum = 0f32;
10829 for k in 0..GROUP_SIZE {
10830 gsum += sg[k] * xg[k];
10831 }
10832 acc += s * gsum;
10833 }
10834 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10835 unsafe { *out_addr.at(r) = acc };
10836 }
10837 };
10838 dispatch_rows(pool, rows, &run);
10839}
10840
10841#[cfg(target_arch = "aarch64")]
10847#[target_feature(enable = "neon,dotprod")]
10848unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
10849 use core::arch::aarch64::*;
10850 use core::arch::asm;
10851 unsafe {
10853 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10854 let bytes_ptr = bytes.as_ptr();
10855 let row_off = r * gpr * TILE;
10856 let xp = [xa.as_ptr(), xb.as_ptr()];
10857 let mut acc = [0f32; 2];
10858 macro_rules! sdot2 {
10859 ($w0:expr, $w1:expr, $x:expr) => {{
10860 let x0 = vld1q_s8($x);
10861 let x1 = vld1q_s8($x.add(16));
10862 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10863 asm!(
10864 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10865 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10866 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10867 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10868 options(pure, nomem, nostack),
10869 );
10870 vaddvq_s32(vaddq_s32(a0, a1))
10871 }};
10872 }
10873 let gpr2 = gpr & !1;
10874 let mut gi = 0;
10875 while gi < gpr2 {
10876 let off0 = row_off + gi * TILE;
10877 let off1 = off0 + TILE;
10878 let s0 = f16_to_f32(u16::from_le_bytes([
10879 *bytes_ptr.add(off0),
10880 *bytes_ptr.add(off0 + 1),
10881 ]));
10882 let s1 = f16_to_f32(u16::from_le_bytes([
10883 *bytes_ptr.add(off1),
10884 *bytes_ptr.add(off1 + 1),
10885 ]));
10886 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10887 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10888 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10889 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10890 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10891 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10892 for k in 0..2 {
10893 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
10894 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
10895 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
10896 }
10897 gi += 2;
10898 }
10899 if gi < gpr {
10900 let off = row_off + gi * TILE;
10901 let s = f16_to_f32(u16::from_le_bytes([
10902 *bytes_ptr.add(off),
10903 *bytes_ptr.add(off + 1),
10904 ]));
10905 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10906 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10907 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10908 for k in 0..2 {
10909 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
10910 acc[k] += d as f32 * s;
10911 }
10912 }
10913 acc
10914 }
10915}
10916
10917fn q1t_matvec2(
10923 bytes: &[u8],
10924 x1: &[f32],
10925 x2: &[f32],
10926 rows: usize,
10927 cols: usize,
10928 o1: &mut [f32],
10929 o2: &mut [f32],
10930 pool: Option<&Pool>,
10931) {
10932 debug_assert_eq!(o1.len(), rows);
10933 debug_assert_eq!(o2.len(), rows);
10934 const TILE: usize = cortiq_core::quant::Q1T_TILE;
10935 let gpr = cols / GROUP_SIZE;
10936 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10937 let out1 = SendMut(o1.as_mut_ptr());
10938 let out2 = SendMut(o2.as_mut_ptr());
10939 if a8w8_enabled() {
10940 let a1 = split_act(x1);
10941 let a2 = split_act(x2);
10942 let (a1, a2) = (&a1, &a2);
10943 let run = move |start: usize, end: usize| {
10944 for r in start..end {
10945 #[cfg(target_arch = "aarch64")]
10946 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
10949 #[cfg(not(target_arch = "aarch64"))]
10950 let ds = [
10951 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
10952 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
10953 ];
10954 let mut acc1 = ds[0] * a1.sx;
10955 for &(j, xv) in &a1.outliers {
10956 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
10957 }
10958 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
10959 let mut acc2 = ds[1] * a2.sx;
10960 for &(j, xv) in &a2.outliers {
10961 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
10962 }
10963 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
10964 unsafe {
10966 *out1.at(r) = acc1;
10967 *out2.at(r) = acc2;
10968 }
10969 }
10970 };
10971 dispatch_rows(pool, rows, &run);
10972 return;
10973 }
10974 let run = move |start: usize, end: usize| {
10975 let mut sg = [0f32; GROUP_SIZE];
10978 for r in start..end {
10979 let mut acc1 = 0f32;
10980 let mut acc2 = 0f32;
10981 for g in 0..gpr {
10982 let off = (r * gpr + g) * TILE;
10983 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10984 let codes = &bytes[off + 2..off + TILE];
10985 for bi in 0..6 {
10986 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10987 }
10988 let lut = &SIGN5[codes[6] as usize];
10989 sg[30] = lut[0];
10990 sg[31] = lut[1];
10991 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10992 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10993 let mut gsum1 = 0f32;
10994 for k in 0..GROUP_SIZE {
10995 gsum1 += sg[k] * xg1[k];
10996 }
10997 acc1 += s * gsum1;
10998 let mut gsum2 = 0f32;
10999 for k in 0..GROUP_SIZE {
11000 gsum2 += sg[k] * xg2[k];
11001 }
11002 acc2 += s * gsum2;
11003 }
11004 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
11005 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
11006 unsafe {
11008 *out1.at(r) = acc1;
11009 *out2.at(r) = acc2;
11010 }
11011 }
11012 };
11013 dispatch_rows(pool, rows, &run);
11014}
11015
11016fn q1t_matmat(
11019 bytes: &[u8],
11020 xs: &[f32],
11021 b: usize,
11022 rows: usize,
11023 cols: usize,
11024 out: &mut [f32],
11025 pool: Option<&Pool>,
11026) {
11027 debug_assert_eq!(out.len(), b * rows);
11028 const TILE: usize = cortiq_core::quant::Q1T_TILE;
11029 let gpr = cols / GROUP_SIZE;
11030 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
11031 let out_addr = SendMut(out.as_mut_ptr());
11032 if a8w8_enabled() {
11036 let acts: Vec<SplitAct> = (0..b)
11037 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
11038 .collect();
11039 let acts = &acts;
11040 let run = move |start: usize, end: usize| {
11041 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
11045 for g in 0..gpr {
11046 let off = (r * gpr + g) * TILE;
11047 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11048 q1t_unpack_group_i8(
11049 bytes.as_ptr().wrapping_add(off + 2),
11050 &mut sg[g * GROUP_SIZE..],
11051 );
11052 }
11053 for bi in 0..b {
11054 let act = &acts[bi];
11055 let mut isum = 0f32;
11056 for g in 0..gpr {
11057 let d = q1t_i8dot32(
11058 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
11059 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
11060 );
11061 isum += d as f32 * sc[g];
11062 }
11063 let mut acc = isum * act.sx;
11064 for &(j, xv) in &act.outliers {
11065 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
11066 }
11067 accs[bi] = acc;
11068 }
11069 if has_ov {
11073 let (c0, c1) = (
11074 q1t_rowptr(bytes, rp_off, r),
11075 q1t_rowptr(bytes, rp_off, r + 1),
11076 );
11077 for p in c0..c1 {
11078 let e = ent_off + p * 4;
11079 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
11080 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
11081 for bi in 0..b {
11082 accs[bi] += val * xs[bi * cols + col];
11083 }
11084 }
11085 }
11086 for bi in 0..b {
11087 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
11088 }
11089 }
11090 };
11091 dispatch_rows(pool, rows, &run);
11092 return;
11093 }
11094 let run = move |start: usize, end: usize| {
11095 let mut buf = vec![0f32; cols];
11096 for r in start..end {
11097 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
11098 for bi in 0..b {
11099 let xr = &xs[bi * cols..(bi + 1) * cols];
11100 let mut acc = 0f32;
11101 for j in 0..cols {
11102 acc += buf[j] * xr[j];
11103 }
11104 unsafe { *out_addr.at(bi * rows + r) = acc };
11105 }
11106 }
11107 };
11108 dispatch_rows(pool, rows, &run);
11109}
11110
11111#[cfg(test)]
11114fn q1_matvec(
11115 bytes: &[u8],
11116 x: &[f32],
11117 rows: usize,
11118 cols: usize,
11119 out: &mut [f32],
11120 pool: Option<&Pool>,
11121) {
11122 bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
11123}
11124
11125fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
11128 debug_assert_eq!(out.len(), rows);
11129 let gpr = cols / GROUP_SIZE;
11130 let out_addr = SendMut(out.as_mut_ptr());
11131 if a8w8_enabled() {
11132 let act = split_act(x);
11133 let gsum = q1_group_sums(&act.xq, gpr);
11134 let (act, gsum) = (&act, &gsum);
11135 let run =
11136 move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
11137 dispatch_rows(pool, rows, &run);
11138 return;
11139 }
11140 let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
11141 dispatch_rows(pool, rows, &run);
11142}
11143
11144#[cfg(test)]
11146#[allow(clippy::too_many_arguments)]
11147fn q1_matvec2(
11148 bytes: &[u8],
11149 x1: &[f32],
11150 x2: &[f32],
11151 rows: usize,
11152 cols: usize,
11153 o1: &mut [f32],
11154 o2: &mut [f32],
11155 pool: Option<&Pool>,
11156) {
11157 bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
11158}
11159
11160#[allow(clippy::too_many_arguments)]
11162fn bin_matvec2(
11163 w: BinW,
11164 x1: &[f32],
11165 x2: &[f32],
11166 rows: usize,
11167 cols: usize,
11168 o1: &mut [f32],
11169 o2: &mut [f32],
11170 pool: Option<&Pool>,
11171) {
11172 let gpr = cols / GROUP_SIZE;
11173 let p1 = SendMut(o1.as_mut_ptr());
11174 let p2 = SendMut(o2.as_mut_ptr());
11175 if a8w8_enabled() {
11176 let a1 = split_act(x1);
11177 let a2 = split_act(x2);
11178 let g1 = q1_group_sums(&a1.xq, gpr);
11179 let g2 = q1_group_sums(&a2.xq, gpr);
11180 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
11181 let run = move |start: usize, end: usize| {
11182 for r in start..end {
11183 let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
11184 let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
11185 for &(j, xv) in &a1.outliers {
11186 let (wv, s) = w.outlier(r, gpr, j);
11187 v1 += wv * s * xv;
11188 }
11189 for &(j, xv) in &a2.outliers {
11190 let (wv, s) = w.outlier(r, gpr, j);
11191 v2 += wv * s * xv;
11192 }
11193 unsafe {
11195 *p1.at(r) = v1;
11196 *p2.at(r) = v2;
11197 }
11198 }
11199 };
11200 dispatch_rows(pool, rows, &run);
11201 return;
11202 }
11203 let run = move |start: usize, end: usize| {
11204 for r in start..end {
11205 unsafe {
11207 *p1.at(r) = w.row_exact(r, gpr, x1);
11208 *p2.at(r) = w.row_exact(r, gpr, x2);
11209 }
11210 }
11211 };
11212 dispatch_rows(pool, rows, &run);
11213}
11214
11215#[cfg(test)]
11217#[allow(clippy::too_many_arguments)]
11218fn q1_matmat(
11219 bytes: &[u8],
11220 xs_all: &[f32],
11221 b: usize,
11222 rows: usize,
11223 cols: usize,
11224 out: &mut [f32],
11225 pool: Option<&Pool>,
11226) {
11227 bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
11228}
11229
11230#[allow(clippy::too_many_arguments)]
11232fn bin_matmat(
11233 w: BinW,
11234 xs_all: &[f32],
11235 b: usize,
11236 rows: usize,
11237 cols: usize,
11238 out: &mut [f32],
11239 pool: Option<&Pool>,
11240) {
11241 debug_assert_eq!(out.len(), b * rows);
11242 let gpr = cols / GROUP_SIZE;
11243 let out_addr = SendMut(out.as_mut_ptr());
11244 if a8w8_enabled() {
11245 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
11246 .map(|bi| {
11247 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
11248 let gsum = q1_group_sums(&act.xq, gpr);
11249 (act, gsum)
11250 })
11251 .collect();
11252 let acts = &acts;
11253 #[cfg(target_arch = "x86_64")]
11254 let blocked_ok = avx2_enabled() && blocked_enabled();
11255 #[cfg(target_arch = "aarch64")]
11256 let blocked_ok = sdot_enabled() && blocked_enabled();
11257 let run = move |start: usize, end: usize| {
11258 for r in start..end {
11259 let mut bi = 0usize;
11260 #[cfg(target_arch = "aarch64")]
11263 if blocked_ok {
11264 while bi + 4 <= acts.len() {
11265 let xs = [
11266 acts[bi].0.xq.as_slice(),
11267 acts[bi + 1].0.xq.as_slice(),
11268 acts[bi + 2].0.xq.as_slice(),
11269 acts[bi + 3].0.xq.as_slice(),
11270 ];
11271 let gs = [
11272 acts[bi].1.as_slice(),
11273 acts[bi + 1].1.as_slice(),
11274 acts[bi + 2].1.as_slice(),
11275 acts[bi + 3].1.as_slice(),
11276 ];
11277 let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
11279 for k in 0..4 {
11280 let (act, _) = &acts[bi + k];
11281 let mut acc = d[k] * act.sx;
11282 for &(j, xv) in &act.outliers {
11283 let (wv, sc) = w.outlier(r, gpr, j);
11284 acc += wv * sc * xv;
11285 }
11286 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11288 }
11289 bi += 4;
11290 }
11291 }
11292 #[cfg(target_arch = "x86_64")]
11293 if blocked_ok {
11294 while bi + 4 <= acts.len() {
11295 let xs = [
11296 acts[bi].0.xq.as_slice(),
11297 acts[bi + 1].0.xq.as_slice(),
11298 acts[bi + 2].0.xq.as_slice(),
11299 acts[bi + 3].0.xq.as_slice(),
11300 ];
11301 let gs = [
11302 acts[bi].1.as_slice(),
11303 acts[bi + 1].1.as_slice(),
11304 acts[bi + 2].1.as_slice(),
11305 acts[bi + 3].1.as_slice(),
11306 ];
11307 let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
11309 for k in 0..4 {
11310 let (act, _) = &acts[bi + k];
11311 let mut acc = d[k] * act.sx;
11312 for &(j, xv) in &act.outliers {
11313 let (wv, sc) = w.outlier(r, gpr, j);
11314 acc += wv * sc * xv;
11315 }
11316 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11318 }
11319 bi += 4;
11320 }
11321 }
11322 while bi < acts.len() {
11323 let (act, gsum) = &acts[bi];
11324 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
11325 for &(j, xv) in &act.outliers {
11326 let (wv, s) = w.outlier(r, gpr, j);
11327 acc += wv * s * xv;
11328 }
11329 unsafe { *out_addr.at(bi * rows + r) = acc };
11331 bi += 1;
11332 }
11333 }
11334 };
11335 dispatch_rows(pool, rows, &run);
11336 return;
11337 }
11338 let run = move |start: usize, end: usize| {
11339 for r in start..end {
11340 for bi in 0..b {
11341 let x = &xs_all[bi * cols..(bi + 1) * cols];
11342 unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
11344 }
11345 }
11346 };
11347 dispatch_rows(pool, rows, &run);
11348}
11349
11350fn q4matvec(
11356 bytes: &[u8],
11357 x: &[f32],
11358 rows: usize,
11359 cols: usize,
11360 out: &mut [f32],
11361 pool: Option<&Pool>,
11362) {
11363 debug_assert_eq!(out.len(), rows);
11364 let (packed, scales) = q4_split(bytes, rows, cols);
11365 let gpr = cols / GROUP_SIZE;
11366 let out_addr = SendMut(out.as_mut_ptr());
11367
11368 if a8w8_enabled() {
11369 let act = split_act(x);
11370 let run = move |start: usize, end: usize| {
11371 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
11372 };
11373 dispatch_rows(pool, rows, &run);
11374 return;
11375 }
11376
11377 let run =
11378 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
11379 dispatch_rows(pool, rows, &run);
11380}
11381
11382#[inline]
11385#[allow(unreachable_code)]
11386#[cfg(target_arch = "x86_64")]
11391#[target_feature(enable = "avx2")]
11392unsafe fn dot_q4b_row_1x4_avx2(
11393 buf: &[u8],
11394 scales: &[u8],
11395 g0: usize,
11396 gpr: usize,
11397 xs: [&[i8]; 4],
11398) -> [f32; 4] {
11399 unsafe {
11401 use core::arch::x86_64::*;
11402 let ones = _mm256_set1_epi16(1);
11403 let mut acc = [0f32; 4];
11404 for gi in 0..gpr {
11405 let s = f16_to_f32(u16::from_le_bytes([
11406 scales[(g0 + gi) * 2],
11407 scales[(g0 + gi) * 2 + 1],
11408 ]));
11409 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11410 let aw = _mm256_abs_epi8(w);
11411 for (k, xq) in xs.iter().enumerate() {
11412 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11413 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
11414 let d = _mm256_madd_epi16(p16, ones);
11415 let hi128 = _mm256_extracti128_si256::<1>(d);
11416 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11417 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11418 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11419 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
11420 }
11421 }
11422 acc
11423 }
11424}
11425
11426#[cfg(target_arch = "x86_64")]
11428#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11429unsafe fn dot_q4b_row_1x4_vnni(
11430 buf: &[u8],
11431 scales: &[u8],
11432 g0: usize,
11433 gpr: usize,
11434 xs: [&[i8]; 4],
11435) -> [f32; 4] {
11436 unsafe {
11438 use core::arch::x86_64::*;
11439 let mut acc = [0f32; 4];
11440 for gi in 0..gpr {
11441 let s = f16_to_f32(u16::from_le_bytes([
11442 scales[(g0 + gi) * 2],
11443 scales[(g0 + gi) * 2 + 1],
11444 ]));
11445 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11446 let aw = _mm256_abs_epi8(w);
11447 for (k, xq) in xs.iter().enumerate() {
11448 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11449 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
11450 acc[k] += d as f32 * s;
11451 }
11452 }
11453 acc
11454 }
11455}
11456
11457#[cfg(target_arch = "x86_64")]
11463#[target_feature(enable = "avx2")]
11464unsafe fn dot_q4b_row_1x4_sx_avx2(
11465 buf: &[u8],
11466 scales: &[u8],
11467 g0: usize,
11468 gpr: usize,
11469 xs: [&[i8]; 4],
11470 sxs: [f32; 4],
11471) -> [f32; 4] {
11472 unsafe {
11474 use core::arch::x86_64::*;
11475 let ones = _mm256_set1_epi16(1);
11476 let mut acc = [0f32; 4];
11477 for gi in 0..gpr {
11478 let s = f16_to_f32(u16::from_le_bytes([
11479 scales[(g0 + gi) * 2],
11480 scales[(g0 + gi) * 2 + 1],
11481 ]));
11482 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11483 let aw = _mm256_abs_epi8(w);
11484 for (k, xq) in xs.iter().enumerate() {
11485 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11486 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
11487 let d = _mm256_madd_epi16(p16, ones);
11488 let hi128 = _mm256_extracti128_si256::<1>(d);
11489 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11490 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11491 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11492 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
11493 }
11494 }
11495 acc
11496 }
11497}
11498
11499#[cfg(target_arch = "x86_64")]
11502#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11503unsafe fn dot_q4b_row_1x4_sx_vnni(
11504 buf: &[u8],
11505 scales: &[u8],
11506 g0: usize,
11507 gpr: usize,
11508 xs: [&[i8]; 4],
11509 sxs: [f32; 4],
11510) -> [f32; 4] {
11511 unsafe {
11513 use core::arch::x86_64::*;
11514 let mut acc = [0f32; 4];
11515 for gi in 0..gpr {
11516 let s = f16_to_f32(u16::from_le_bytes([
11517 scales[(g0 + gi) * 2],
11518 scales[(g0 + gi) * 2 + 1],
11519 ]));
11520 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11521 let aw = _mm256_abs_epi8(w);
11522 for (k, xq) in xs.iter().enumerate() {
11523 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11524 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
11525 acc[k] += (d as f32 * sxs[k]) * s;
11526 }
11527 }
11528 acc
11529 }
11530}
11531
11532#[allow(unreachable_code)]
11533fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11534 #[cfg(target_arch = "aarch64")]
11535 unsafe {
11536 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
11537 }
11538 #[cfg(target_arch = "x86_64")]
11539 unsafe {
11540 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
11541 }
11542 let mut acc = 0f32;
11543 for gi in 0..gpr {
11544 let g = g0 + gi;
11545 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11546 let mut d = 0i32;
11547 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11548 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
11549 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
11550 }
11551 acc += d as f32 * s;
11552 }
11553 acc
11554}
11555
11556#[inline]
11558#[allow(unreachable_code)]
11559fn dot_q4_row_i8_2(
11560 packed: &[u8],
11561 scales: &[u8],
11562 g0: usize,
11563 gpr: usize,
11564 xq1: &[i8],
11565 xq2: &[i8],
11566) -> (f32, f32) {
11567 #[cfg(target_arch = "aarch64")]
11568 unsafe {
11569 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
11570 }
11571 #[cfg(target_arch = "x86_64")]
11572 unsafe {
11573 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
11574 }
11575 (
11576 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
11577 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
11578 )
11579}
11580
11581#[allow(clippy::too_many_arguments)]
11584fn q4_range_a8w8(
11585 packed: &[u8],
11586 scales: &[u8],
11587 gpr: usize,
11588 cols: usize,
11589 act: &SplitAct,
11590 out: SendMut,
11591 start: usize,
11592 end: usize,
11593) {
11594 for r in start..end {
11595 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
11596 for &(j, xv) in &act.outliers {
11598 let flat = r * cols + j;
11599 let byte = packed[flat / 2];
11600 let nib = if flat & 1 == 0 {
11601 byte & 0x0F
11602 } else {
11603 byte >> 4
11604 };
11605 let s = f16_to_f32(u16::from_le_bytes([
11606 scales[(flat / GROUP_SIZE) * 2],
11607 scales[(flat / GROUP_SIZE) * 2 + 1],
11608 ]));
11609 acc += ((nib as i32 - 8) as f32) * s * xv;
11610 }
11611 unsafe { *out.at(r) = acc };
11613 }
11614}
11615
11616#[allow(clippy::too_many_arguments)]
11619fn q4_range2_a8w8(
11620 packed: &[u8],
11621 scales: &[u8],
11622 gpr: usize,
11623 cols: usize,
11624 a1: &SplitAct,
11625 a2: &SplitAct,
11626 p1: SendMut,
11627 p2: SendMut,
11628 start: usize,
11629 end: usize,
11630) {
11631 for r in start..end {
11632 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
11633 let mut acc1 = s1 * a1.sx;
11634 let mut acc2 = s2 * a2.sx;
11635 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
11637 for &(j, xv) in outliers {
11638 let flat = r * cols + j;
11639 let byte = packed[flat / 2];
11640 let nib = if flat & 1 == 0 {
11641 byte & 0x0F
11642 } else {
11643 byte >> 4
11644 };
11645 let s = f16_to_f32(u16::from_le_bytes([
11646 scales[(flat / GROUP_SIZE) * 2],
11647 scales[(flat / GROUP_SIZE) * 2 + 1],
11648 ]));
11649 *acc += ((nib as i32 - 8) as f32) * s * xv;
11650 }
11651 };
11652 fix(&a1.outliers, &mut acc1);
11653 fix(&a2.outliers, &mut acc2);
11654 unsafe {
11656 *p1.at(r) = acc1;
11657 *p2.at(r) = acc2;
11658 }
11659 }
11660}
11661
11662fn q4_range_f32(
11664 packed: &[u8],
11665 scales: &[u8],
11666 gpr: usize,
11667 x: &[f32],
11668 out: SendMut,
11669 start: usize,
11670 end: usize,
11671) {
11672 for r in start..end {
11673 let mut acc = 0f32;
11674 for gi in 0..gpr {
11675 let g = r * gpr + gi;
11676 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11677 let pk = &packed[g * 16..(g + 1) * 16];
11678 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11679 let mut ga = 0f32;
11680 for (k, &b) in pk.iter().enumerate() {
11681 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
11682 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
11683 }
11684 acc += ga * s;
11685 }
11686 unsafe { *out.at(r) = acc };
11688 }
11689}
11690
11691#[allow(clippy::too_many_arguments)]
11695fn q4matvec2(
11696 bytes: &[u8],
11697 x1: &[f32],
11698 x2: &[f32],
11699 rows: usize,
11700 cols: usize,
11701 o1: &mut [f32],
11702 o2: &mut [f32],
11703 pool: Option<&Pool>,
11704) {
11705 debug_assert_eq!(o1.len(), rows);
11706 debug_assert_eq!(o2.len(), rows);
11707 let (packed, scales) = q4_split(bytes, rows, cols);
11708 let gpr = cols / GROUP_SIZE;
11709
11710 if a8w8_enabled() {
11711 let a1 = split_act(x1);
11712 let a2 = split_act(x2);
11713 let p1 = SendMut(o1.as_mut_ptr());
11714 let p2 = SendMut(o2.as_mut_ptr());
11715 let run = move |start: usize, end: usize| {
11716 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
11717 };
11718 dispatch_rows(pool, rows, &run);
11719 return;
11720 }
11721
11722 let p1 = SendMut(o1.as_mut_ptr());
11723 let p2 = SendMut(o2.as_mut_ptr());
11724 let run = move |start: usize, end: usize| {
11725 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
11726 };
11727 dispatch_rows(pool, rows, &run);
11728}
11729
11730#[allow(clippy::too_many_arguments)]
11732fn q4_range2_f32(
11733 packed: &[u8],
11734 scales: &[u8],
11735 gpr: usize,
11736 x1: &[f32],
11737 x2: &[f32],
11738 p1: SendMut,
11739 p2: SendMut,
11740 start: usize,
11741 end: usize,
11742) {
11743 for r in start..end {
11744 let (mut acc1, mut acc2) = (0f32, 0f32);
11745 for gi in 0..gpr {
11746 let g = r * gpr + gi;
11747 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11748 let pk = &packed[g * 16..(g + 1) * 16];
11749 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11750 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11751 let (mut g1, mut g2) = (0f32, 0f32);
11752 for (k, &b) in pk.iter().enumerate() {
11753 let wl = (b & 0x0F) as f32 - 8.0;
11754 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
11755 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
11756 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
11757 }
11758 acc1 += g1 * s;
11759 acc2 += g2 * s;
11760 }
11761 unsafe {
11763 *p1.at(r) = acc1;
11764 *p2.at(r) = acc2;
11765 }
11766 }
11767}
11768
11769thread_local! {
11770 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
11773 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11774}
11775
11776#[allow(clippy::too_many_arguments)]
11782fn q4matmat(
11783 bytes: &[u8],
11784 xs_all: &[f32],
11785 b: usize,
11786 rows: usize,
11787 cols: usize,
11788 out: &mut [f32],
11789 pool: Option<&Pool>,
11790) {
11791 debug_assert_eq!(xs_all.len(), b * cols);
11792 debug_assert_eq!(out.len(), b * rows);
11793 let (packed, scales) = q4_split(bytes, rows, cols);
11794 let gpr = cols / GROUP_SIZE;
11795 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11796
11797 if a8w8_enabled() {
11798 let acts: Vec<SplitAct> = (0..b)
11799 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11800 .collect();
11801 let acts = &acts;
11802 let out_addr = SendMut(out.as_mut_ptr());
11803 let run = move |start: usize, end: usize| {
11804 ROW_I8.with(|rb| {
11805 let mut buf = rb.borrow_mut();
11806 buf.resize(cols, 0);
11807 for r in start..end {
11808 for gi in 0..gpr {
11812 let g = r * gpr + gi;
11813 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11814 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
11815 buf[gi * GROUP_SIZE + k * 2 + 1] =
11816 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
11817 }
11818 }
11819 let mut bi = 0usize;
11820 #[cfg(target_arch = "x86_64")]
11821 if avx2_enabled() && blocked_enabled() {
11822 while bi + 4 <= acts.len() {
11823 let xs = [
11824 acts[bi].xq.as_slice(),
11825 acts[bi + 1].xq.as_slice(),
11826 acts[bi + 2].xq.as_slice(),
11827 acts[bi + 3].xq.as_slice(),
11828 ];
11829 let d = unsafe {
11830 if vnni_tiles_enabled() {
11831 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
11832 } else {
11833 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
11834 }
11835 };
11836 for k in 0..4 {
11837 let act = &acts[bi + k];
11838 let mut acc = d[k] * act.sx;
11839 for &(j, xv) in &act.outliers {
11840 acc += (buf[j] as i8) as f32
11841 * gscale((r * cols + j) / GROUP_SIZE)
11842 * xv;
11843 }
11844 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11846 }
11847 bi += 4;
11848 }
11849 }
11850 while bi < acts.len() {
11851 let act = &acts[bi];
11852 let mut acc = 0f32;
11853 for gi in 0..gpr {
11854 let d = dot_i8_i8(
11855 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11856 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11857 );
11858 acc += d as f32 * gscale(r * gpr + gi);
11859 }
11860 acc *= act.sx;
11861 for &(j, xv) in &act.outliers {
11863 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
11864 }
11865 unsafe { *out_addr.at(bi * rows + r) = acc };
11867 bi += 1;
11868 }
11869 }
11870 })
11871 };
11872 dispatch_rows(pool, rows, &run);
11873 return;
11874 }
11875
11876 let out_addr = SendMut(out.as_mut_ptr());
11877 let run = move |start: usize, end: usize| {
11878 ROW_F32.with(|rb| {
11879 let mut buf = rb.borrow_mut();
11880 buf.resize(cols, 0.0);
11881 for r in start..end {
11882 for gi in 0..gpr {
11885 let g = r * gpr + gi;
11886 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11887 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
11888 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
11889 }
11890 }
11891 for bi in 0..b {
11892 let x = &xs_all[bi * cols..(bi + 1) * cols];
11893 let mut acc = 0f32;
11894 for gi in 0..gpr {
11895 let mut ga = 0f32;
11896 for k in 0..GROUP_SIZE / 2 {
11901 let e = gi * GROUP_SIZE + k * 2;
11902 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
11903 }
11904 acc += ga * gscale(r * gpr + gi);
11905 }
11906 unsafe { *out_addr.at(bi * rows + r) = acc };
11908 }
11909 }
11910 })
11911 };
11912 dispatch_rows(pool, rows, &run);
11913}
11914
11915#[allow(clippy::too_many_arguments)]
11920fn vbitmatmat(
11921 bytes: &[u8],
11922 offsets: &[usize],
11923 xs_all: &[f32],
11924 b: usize,
11925 rows: usize,
11926 cols: usize,
11927 out: &mut [f32],
11928 pool: Option<&Pool>,
11929) {
11930 debug_assert_eq!(xs_all.len(), b * cols);
11931 debug_assert_eq!(out.len(), b * rows);
11932 debug_assert_eq!(offsets.len(), rows + 1);
11933 let ng = cols / GROUP_SIZE;
11934 let bits = &bytes[..rows];
11935 let sc_off = rows;
11936 let gscale = |r: usize, g: usize| {
11937 let so = (r * ng + g) * 2;
11938 f16_to_f32(u16::from_le_bytes([
11939 bytes[sc_off + so],
11940 bytes[sc_off + so + 1],
11941 ]))
11942 };
11943
11944 let decode_f32 = |r: usize, dst: &mut [f32]| {
11946 let bw = bits[r] as usize;
11947 let l = ((1i32 << (bw - 1)) - 1) as f32;
11948 let data = &bytes[offsets[r]..offsets[r + 1]];
11949 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
11950 for d in dst.iter_mut() {
11951 while nbits < bw {
11952 acc = (acc << 8) | data[idx] as u64;
11953 idx += 1;
11954 nbits += 8;
11955 }
11956 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
11957 nbits -= bw;
11958 *d = u - l;
11959 }
11960 };
11961
11962 if a8w8_enabled() {
11963 let acts: Vec<SplitAct> = (0..b)
11964 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11965 .collect();
11966 let acts = &acts;
11967 let out_addr = SendMut(out.as_mut_ptr());
11968 let run = move |start: usize, end: usize| {
11969 for r in start..end {
11970 let bw = bits[r] as usize;
11971 if bw == 8 {
11972 ROW_F32.with(|rb| {
11975 let mut buf = rb.borrow_mut();
11976 buf.resize(cols, 0.0);
11977 decode_f32(r, &mut buf);
11978 for bi in 0..b {
11979 let x = &xs_all[bi * cols..(bi + 1) * cols];
11980 let mut dot = 0f32;
11981 for g in 0..ng {
11982 let mut gd = 0f32;
11983 for k in 0..GROUP_SIZE {
11984 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
11985 }
11986 dot += gd * gscale(r, g);
11987 }
11988 unsafe { *out_addr.at(bi * rows + r) = dot };
11990 }
11991 });
11992 continue;
11993 }
11994 let l = (1i32 << (bw - 1)) - 1;
11995 let data = &bytes[offsets[r]..offsets[r + 1]];
11996 ROW_I8.with(|rb| {
11997 let mut buf = rb.borrow_mut();
11998 buf.resize(cols, 0);
11999 #[inline(always)]
12000 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
12001 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
12002 let u = unpack8::<B>(&data[blk * B..]);
12003 for k in 0..8 {
12004 chunk[k] = (u[k] - l) as i8 as u8;
12005 }
12006 }
12007 }
12008 match bw {
12009 3 => fill::<3>(data, l, &mut buf),
12010 4 => vbit_fill4(data, &mut buf),
12011 5 => fill::<5>(data, l, &mut buf),
12012 6 => fill::<6>(data, l, &mut buf),
12013 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
12014 }
12015 let mut bi = 0usize;
12016 #[cfg(target_arch = "x86_64")]
12020 if avx2_enabled() && blocked_enabled() {
12021 while bi + 4 <= acts.len() {
12022 let xs = [
12023 acts[bi].xq.as_slice(),
12024 acts[bi + 1].xq.as_slice(),
12025 acts[bi + 2].xq.as_slice(),
12026 acts[bi + 3].xq.as_slice(),
12027 ];
12028 let sxs = [
12029 acts[bi].sx,
12030 acts[bi + 1].sx,
12031 acts[bi + 2].sx,
12032 acts[bi + 3].sx,
12033 ];
12034 let d = unsafe {
12035 if vnni_tiles_enabled() {
12036 dot_q4b_row_1x4_sx_vnni(
12037 &buf,
12038 &bytes[sc_off..],
12039 r * ng,
12040 ng,
12041 xs,
12042 sxs,
12043 )
12044 } else {
12045 dot_q4b_row_1x4_sx_avx2(
12046 &buf,
12047 &bytes[sc_off..],
12048 r * ng,
12049 ng,
12050 xs,
12051 sxs,
12052 )
12053 }
12054 };
12055 for k in 0..4 {
12056 let act = &acts[bi + k];
12057 let mut dot = d[k];
12058 for &(j, xv) in &act.outliers {
12059 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
12060 }
12061 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
12063 }
12064 bi += 4;
12065 }
12066 }
12067 while bi < acts.len() {
12068 let act = &acts[bi];
12069 let mut dot = 0f32;
12070 for g in 0..ng {
12071 let d = dot_i8_i8(
12072 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
12073 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
12074 ) as f32
12075 * act.sx;
12076 dot += d * gscale(r, g);
12077 }
12078 for &(j, xv) in &act.outliers {
12079 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
12080 }
12081 unsafe { *out_addr.at(bi * rows + r) = dot };
12083 bi += 1;
12084 }
12085 });
12086 }
12087 };
12088 dispatch_rows(pool, rows, &run);
12089 return;
12090 }
12091
12092 let out_addr = SendMut(out.as_mut_ptr());
12093 let run = move |start: usize, end: usize| {
12094 ROW_F32.with(|rb| {
12095 let mut buf = rb.borrow_mut();
12096 buf.resize(cols, 0.0);
12097 for r in start..end {
12098 decode_f32(r, &mut buf);
12099 for bi in 0..b {
12100 let x = &xs_all[bi * cols..(bi + 1) * cols];
12101 let mut dot = 0f32;
12102 for g in 0..ng {
12103 let mut gd = 0f32;
12104 for k in 0..GROUP_SIZE {
12105 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
12106 }
12107 dot += gd * gscale(r, g);
12108 }
12109 unsafe { *out_addr.at(bi * rows + r) = dot };
12111 }
12112 }
12113 })
12114 };
12115 dispatch_rows(pool, rows, &run);
12116}
12117
12118pub(crate) fn gpu_batch_job<'a>(
12122 t: &'a QTensor,
12123 x: &[f32],
12124) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
12125 match t {
12126 QTensor::Mapped {
12127 model,
12128 idx,
12129 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
12130 rows,
12131 cols,
12132 row_scale,
12133 col_field,
12134 ..
12135 } => Some((
12136 model.clone(),
12137 crate::gpu::BatchJob {
12138 idx: *idx,
12139 rows: *rows,
12140 cols: *cols,
12141 row_scale,
12142 xs: prescale(x, col_field, *dt).into_owned(),
12143 layout: crate::gpu::BatchLayout::Q8,
12144 },
12145 )),
12146 QTensor::Mapped {
12148 model,
12149 idx,
12150 dtype: TensorDtype::Q1,
12151 rows,
12152 cols,
12153 ..
12154 } => Some((
12155 model.clone(),
12156 crate::gpu::BatchJob {
12157 idx: *idx,
12158 rows: *rows,
12159 cols: *cols,
12160 row_scale: &[],
12161 xs: x.to_vec(),
12162 layout: crate::gpu::BatchLayout::Q1,
12163 },
12164 )),
12165 QTensor::Mapped {
12167 model,
12168 idx,
12169 dtype: TensorDtype::Q1G,
12170 rows,
12171 cols,
12172 ..
12173 } => Some((
12174 model.clone(),
12175 crate::gpu::BatchJob {
12176 idx: *idx,
12177 rows: *rows,
12178 cols: *cols,
12179 row_scale: &[],
12180 xs: x.to_vec(),
12181 layout: crate::gpu::BatchLayout::Q1g,
12182 },
12183 )),
12184 QTensor::Mapped {
12189 model,
12190 idx,
12191 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
12192 rows,
12193 cols,
12194 ..
12195 } => Some((
12196 model.clone(),
12197 crate::gpu::BatchJob {
12198 idx: *idx,
12199 rows: *rows,
12200 cols: *cols,
12201 row_scale: &[],
12202 xs: x.to_vec(),
12203 layout: if *dt == TensorDtype::Q4Tiled {
12204 crate::gpu::BatchLayout::Q4t
12205 } else {
12206 crate::gpu::BatchLayout::Q4tp
12207 },
12208 },
12209 )),
12210 _ => None,
12211 }
12212}
12213
12214thread_local! {
12215 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
12216 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
12217}
12218
12219pub(crate) fn prescale<'a>(
12220 x: &'a [f32],
12221 col_field: &[f32],
12222 dtype: TensorDtype,
12223) -> std::borrow::Cow<'a, [f32]> {
12224 if dtype == TensorDtype::Q8_2f {
12225 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
12226 } else {
12227 std::borrow::Cow::Borrowed(x)
12228 }
12229}
12230
12231pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
12234 x: &[f32],
12235 col_field: &[f32],
12236 dtype: TensorDtype,
12237 buf_id: u8,
12238 f: F,
12239) -> R {
12240 if dtype == TensorDtype::Q8_2f {
12241 if buf_id == 1 {
12242 PRESCALE_BUF1.with(|b| {
12243 let mut buf = b.borrow_mut();
12244 buf.clear();
12245 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
12246 f(&buf)
12247 })
12248 } else {
12249 PRESCALE_BUF2.with(|b| {
12250 let mut buf = b.borrow_mut();
12251 buf.clear();
12252 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
12253 f(&buf)
12254 })
12255 }
12256 } else {
12257 f(x)
12258 }
12259}
12260
12261#[cfg(target_arch = "x86_64")]
12266pub(crate) fn avx2_enabled() -> bool {
12267 use std::sync::OnceLock;
12268 static ON: OnceLock<bool> = OnceLock::new();
12269 *ON.get_or_init(|| {
12270 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
12271 && std::arch::is_x86_feature_detected!("avx2")
12272 && std::arch::is_x86_feature_detected!("fma")
12273 })
12274}
12275
12276#[cfg(target_arch = "x86_64")]
12281fn avx2_a8w8_enabled() -> bool {
12282 if FLOAT_ACTIVATIONS.get() {
12283 return false;
12284 }
12285 use std::sync::OnceLock;
12286 static ON: OnceLock<bool> = OnceLock::new();
12287 *ON.get_or_init(|| {
12288 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
12289 })
12290}
12291
12292thread_local! {
12293 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
12294}
12295
12296pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
12299 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
12300 impl Drop for Restore {
12301 fn drop(&mut self) {
12302 FULL_GPU_Q8.set(self.0);
12303 }
12304 }
12305 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
12306}
12307
12308thread_local! {
12312 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
12313}
12314
12315fn exact_cpu_default_enabled() -> bool {
12324 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12325 *ON.get_or_init(|| std::env::var_os("CMF_SDOT").is_none())
12326}
12327
12328impl QTensor {
12329 #[inline]
12330 fn exact_cpu_default(&self) -> bool {
12331 !FLOAT_ACTIVATIONS.get()
12332 && exact_cpu_default_enabled()
12333 && matches!(
12334 self,
12335 Self::Mapped { model, dtype, .. } if matches!(model.header.arch.arch_name.as_str(), "flint" | "cortiq_embryo")
12336 || (model.header.arch.arch_name == "dory" && *dtype == TensorDtype::Q8_2f)
12339 )
12340 }
12341
12342 #[inline]
12343 fn any_exact_cpu_default(ts: &[&QTensor]) -> bool {
12344 ts.iter().any(|t| t.exact_cpu_default())
12345 }
12346
12347 #[inline]
12348 fn any_exact_cpu_default_pairs(pairs: &[(&QTensor, &QTensor)]) -> bool {
12349 pairs
12350 .iter()
12351 .any(|(gate, up)| gate.exact_cpu_default() || up.exact_cpu_default())
12352 }
12353}
12354
12355pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
12360 float_activations_scope(f)
12361}
12362
12363#[inline]
12367pub(crate) fn float_activations_active() -> bool {
12368 FLOAT_ACTIVATIONS.get()
12369}
12370
12371pub(crate) struct FloatActivationsScope {
12375 previous: bool,
12376 _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
12377}
12378
12379impl Drop for FloatActivationsScope {
12380 fn drop(&mut self) {
12381 FLOAT_ACTIVATIONS.set(self.previous);
12382 }
12383}
12384
12385#[inline]
12386pub(crate) fn enter_float_activations_scope() -> FloatActivationsScope {
12387 FloatActivationsScope {
12388 previous: FLOAT_ACTIVATIONS.replace(true),
12389 _not_send: std::marker::PhantomData,
12390 }
12391}
12392
12393#[doc(hidden)]
12399pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
12400 let _scope = enter_float_activations_scope();
12401 f()
12402}
12403
12404static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
12416
12417pub(crate) fn row_exact() -> bool {
12418 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
12419}
12420
12421fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
12422 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
12423 impl Drop for Restore<'_> {
12424 fn drop(&mut self) {
12425 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
12426 }
12427 }
12428 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
12429 let _restore = Restore(active);
12430 f()
12431}
12432
12433pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
12435 counted_row_exact_scope(&ROW_EXACT, f)
12436}
12437
12438#[inline]
12442pub(crate) fn a8w8_enabled() -> bool {
12443 #[cfg(target_arch = "aarch64")]
12444 {
12445 sdot_enabled()
12446 }
12447 #[cfg(target_arch = "x86_64")]
12448 {
12449 avx2_a8w8_enabled()
12450 }
12451 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
12452 {
12453 false
12454 }
12455}
12456
12457#[inline]
12460#[allow(unreachable_code)]
12461fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
12462 #[cfg(target_arch = "aarch64")]
12463 unsafe {
12464 return dot_i8_sdot(w, xq);
12465 }
12466 #[cfg(target_arch = "x86_64")]
12467 unsafe {
12468 if avx512vnni_enabled() {
12469 return dot_i8_i8_vnni(w, xq);
12470 }
12471 return dot_i8_i8_avx2(w, xq);
12472 }
12473 w.iter()
12474 .zip(xq)
12475 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
12476 .sum()
12477}
12478
12479#[cfg(target_arch = "x86_64")]
12483fn avx512vnni_enabled() -> bool {
12484 use std::sync::OnceLock;
12485 static ON: OnceLock<bool> = OnceLock::new();
12486 *ON.get_or_init(|| {
12487 std::env::var("CMF_AVX512")
12488 .map(|v| v != "0")
12489 .unwrap_or(true)
12490 && std::arch::is_x86_feature_detected!("avx512f")
12491 && std::arch::is_x86_feature_detected!("avx512bw")
12492 && std::arch::is_x86_feature_detected!("avx512vl")
12493 && std::arch::is_x86_feature_detected!("avx512vnni")
12494 })
12495}
12496
12497#[cfg(target_arch = "x86_64")]
12505fn vnni_tiles_enabled() -> bool {
12506 use std::sync::OnceLock;
12507 static ON: OnceLock<bool> = OnceLock::new();
12508 *ON.get_or_init(|| {
12509 std::env::var("CMF_VNNI_TILES")
12510 .map(|v| v != "0")
12511 .unwrap_or(true)
12512 && avx512vnni_enabled()
12513 })
12514}
12515
12516#[cfg(target_arch = "x86_64")]
12521#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12522#[inline]
12523unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
12524 unsafe {
12526 use core::arch::x86_64::*;
12527 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
12528 let hi128 = _mm256_extracti128_si256::<1>(d);
12529 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12530 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12531 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12532 _mm_cvtsi128_si32(s32)
12533 }
12534}
12535
12536#[cfg(target_arch = "x86_64")]
12541#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12542unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12543 unsafe {
12545 use core::arch::x86_64::*;
12546 let n = w.len();
12547 let mut j = 0usize;
12548 let mut total: i32;
12549 {
12554 #[inline(always)]
12555 unsafe fn step(
12556 w: *const u8,
12557 x: *const i8,
12558 acc: core::arch::x86_64::__m512i,
12559 ) -> core::arch::x86_64::__m512i {
12560 unsafe {
12561 use core::arch::x86_64::*;
12562 let wv = _mm512_loadu_si512(w as *const _);
12563 let xv = _mm512_loadu_si512(x as *const _);
12564 let aw = _mm512_abs_epi8(wv);
12565 let neg = _mm512_movepi8_mask(wv);
12566 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
12567 _mm512_dpbusd_epi32(acc, aw, sx)
12568 }
12569 }
12570 let (mut a0, mut a1, mut a2, mut a3) = (
12571 _mm512_setzero_si512(),
12572 _mm512_setzero_si512(),
12573 _mm512_setzero_si512(),
12574 _mm512_setzero_si512(),
12575 );
12576 while j + 256 <= n {
12577 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
12578 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
12579 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
12580 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
12581 j += 256;
12582 }
12583 while j + 64 <= n {
12584 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
12585 j += 64;
12586 }
12587 let s01 = _mm512_add_epi32(a0, a1);
12588 let s23 = _mm512_add_epi32(a2, a3);
12589 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
12590 }
12591 if j + 32 <= n {
12593 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12594 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12595 let d = _mm256_dpbusd_epi32(
12596 _mm256_setzero_si256(),
12597 _mm256_abs_epi8(wv),
12598 _mm256_sign_epi8(xv, wv),
12599 );
12600 let hi128 = _mm256_extracti128_si256::<1>(d);
12601 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12602 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12603 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12604 total += _mm_cvtsi128_si32(s32);
12605 j += 32;
12606 }
12607 while j < n {
12608 total += (w[j] as i8) as i32 * xq[j] as i32;
12609 j += 1;
12610 }
12611 total
12612 }
12613}
12614
12615#[cfg(target_arch = "x86_64")]
12617#[target_feature(enable = "avx2,fma")]
12618unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
12619 unsafe {
12621 use core::arch::x86_64::*;
12622 let n = x.len();
12623 let wp = w.as_ptr();
12624 let xp = x.as_ptr();
12625 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
12626 let mut j = 0usize;
12627 while j + 16 <= n {
12628 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
12629 let lo = _mm256_cvtepi8_epi32(wb);
12630 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
12631 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
12632 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
12633 j += 16;
12634 }
12635 let acc = _mm256_add_ps(a0, a1);
12636 let hi128 = _mm256_extractf128_ps::<1>(acc);
12637 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
12638 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
12639 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
12640 let mut sum = _mm_cvtss_f32(s32);
12641 while j < n {
12642 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
12643 j += 1;
12644 }
12645 sum
12646 }
12647}
12648
12649#[cfg(target_arch = "x86_64")]
12654#[target_feature(enable = "avx2")]
12655unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
12656 unsafe {
12658 use core::arch::x86_64::*;
12659 let n = w.len();
12660 let ones = _mm256_set1_epi16(1);
12661 let mut acc = _mm256_setzero_si256();
12662 let mut j = 0usize;
12663 while j + 32 <= n {
12664 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12665 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12666 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
12667 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
12668 j += 32;
12669 }
12670 let hi128 = _mm256_extracti128_si256::<1>(acc);
12671 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
12672 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12673 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12674 let mut s = _mm_cvtsi128_si32(s32);
12675 while j < n {
12676 s += (w[j] as i8) as i32 * xq[j] as i32;
12677 j += 1;
12678 }
12679 s
12680 }
12681}
12682
12683#[cfg(target_arch = "aarch64")]
12687#[target_feature(enable = "neon,i8mm")]
12688unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12689 unsafe {
12691 use core::arch::aarch64::*;
12692 use core::arch::asm;
12693 let n = w0.len();
12694 let w0p = w0.as_ptr() as *const i8;
12695 let w1p = w1.as_ptr() as *const i8;
12696 let mut acc01 = vdupq_n_s32(0);
12699 let mut acc23 = vdupq_n_s32(0);
12700 let mut i = 0usize;
12701 while i + 8 <= n {
12702 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
12703 let xb01 = vcombine_s8(
12704 vld1_s8(xs[0].as_ptr().add(i)),
12705 vld1_s8(xs[1].as_ptr().add(i)),
12706 );
12707 let xb23 = vcombine_s8(
12708 vld1_s8(xs[2].as_ptr().add(i)),
12709 vld1_s8(xs[3].as_ptr().add(i)),
12710 );
12711 asm!(
12712 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
12713 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
12714 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
12715 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
12716 options(pure, nomem, nostack),
12717 );
12718 i += 8;
12719 }
12720 let mut out = [[0i32; 4]; 2];
12721 let a01: [i32; 4] = core::mem::transmute(acc01);
12722 let a23: [i32; 4] = core::mem::transmute(acc23);
12723 out[0][0] = a01[0];
12724 out[0][1] = a01[1];
12725 out[1][0] = a01[2];
12726 out[1][1] = a01[3];
12727 out[0][2] = a23[0];
12728 out[0][3] = a23[1];
12729 out[1][2] = a23[2];
12730 out[1][3] = a23[3];
12731 if i < n {
12732 for (k, x) in xs.iter().enumerate() {
12733 for j in i..n {
12734 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12735 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12736 }
12737 }
12738 }
12739 out
12740 }
12741}
12742
12743#[cfg(target_arch = "aarch64")]
12747#[target_feature(enable = "neon,dotprod")]
12748unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12749 unsafe {
12751 use core::arch::aarch64::*;
12752 use core::arch::asm;
12753 let n = w0.len();
12754 let w0p = w0.as_ptr() as *const i8;
12755 let w1p = w1.as_ptr() as *const i8;
12756 let mut acc = [[vdupq_n_s32(0); 4]; 2];
12757 let mut i = 0usize;
12758 while i + 16 <= n {
12759 let wv0 = vld1q_s8(w0p.add(i));
12760 let wv1 = vld1q_s8(w1p.add(i));
12761 for (k, x) in xs.iter().enumerate() {
12762 let xv = vld1q_s8(x.as_ptr().add(i));
12763 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
12764 asm!(
12765 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
12766 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
12767 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12768 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
12769 options(pure, nomem, nostack),
12770 );
12771 acc[0][k] = a0;
12772 acc[1][k] = a1;
12773 }
12774 i += 16;
12775 }
12776 let mut out = [[0i32; 4]; 2];
12777 for r in 0..2 {
12778 for k in 0..4 {
12779 out[r][k] = vaddvq_s32(acc[r][k]);
12780 }
12781 }
12782 if i < n {
12783 for (k, x) in xs.iter().enumerate() {
12784 for j in i..n {
12785 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12786 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12787 }
12788 }
12789 }
12790 out
12791 }
12792}
12793
12794#[cfg(target_arch = "x86_64")]
12800#[target_feature(enable = "avx2")]
12801unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12802 unsafe {
12804 use core::arch::x86_64::*;
12805 let n = w0.len();
12806 let ones = _mm256_set1_epi16(1);
12807 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
12808 let mut j = 0usize;
12809 while j + 32 <= n {
12810 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
12811 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
12812 let aw0 = _mm256_abs_epi8(wv0);
12813 let aw1 = _mm256_abs_epi8(wv1);
12814 for (k, x) in xs.iter().enumerate() {
12815 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
12816 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
12817 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
12818 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
12819 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
12820 }
12821 j += 32;
12822 }
12823 let mut out = [[0i32; 4]; 2];
12824 for r in 0..2 {
12825 for k in 0..4 {
12826 let a = acc[r][k];
12827 let hi128 = _mm256_extracti128_si256::<1>(a);
12828 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
12829 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12830 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12831 out[r][k] = _mm_cvtsi128_si32(s32);
12832 }
12833 }
12834 if j < n {
12835 for (k, x) in xs.iter().enumerate() {
12836 for i in j..n {
12837 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
12838 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
12839 }
12840 }
12841 }
12842 out
12843 }
12844}
12845
12846#[cfg(target_arch = "x86_64")]
12851#[inline]
12852fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
12853 let dot = if avx512vnni_enabled() && row.len() >= 64 {
12854 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
12855 } else {
12856 unsafe { dot_i8_i8_avx2(row, &act.xq) }
12857 };
12858 let mut acc = dot as f32 * act.sx;
12859 for &(j, xv) in &act.outliers {
12860 acc += (row[j] as i8) as f32 * xv;
12861 }
12862 acc
12863}
12864
12865#[cfg(target_arch = "x86_64")]
12869#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12870unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12871 unsafe {
12873 use core::arch::x86_64::*;
12874 let n = w.len();
12875 let flip = _mm512_set1_epi8(-128); #[inline(always)]
12877 unsafe fn step(
12878 w: *const u8,
12879 x: *const i8,
12880 flip: core::arch::x86_64::__m512i,
12881 acc: core::arch::x86_64::__m512i,
12882 ) -> core::arch::x86_64::__m512i {
12883 unsafe {
12884 use core::arch::x86_64::*;
12885 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
12886 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
12887 }
12888 }
12889 let (mut a0, mut a1, mut a2, mut a3) = (
12890 _mm512_setzero_si512(),
12891 _mm512_setzero_si512(),
12892 _mm512_setzero_si512(),
12893 _mm512_setzero_si512(),
12894 );
12895 let mut j = 0usize;
12896 while j + 256 <= n {
12897 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12898 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
12899 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
12900 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
12901 j += 256;
12902 }
12903 while j + 64 <= n {
12904 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12905 j += 64;
12906 }
12907 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
12908 _mm512_add_epi32(a0, a1),
12909 _mm512_add_epi32(a2, a3),
12910 ));
12911 while j < n {
12913 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
12914 j += 1;
12915 }
12916 total
12917 }
12918}
12919
12920#[cfg(target_arch = "x86_64")]
12926#[target_feature(enable = "avx2")]
12927unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12928 unsafe {
12931 use core::arch::x86_64::*;
12932 let lomask = _mm_set1_epi8(0x0F);
12933 let eight = _mm256_set1_epi8(8);
12934 let ones = _mm256_set1_epi16(1);
12935 let mut acc = 0f32;
12936 for gi in 0..gpr {
12937 let g = g0 + gi;
12938 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12939 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
12940 let lo = _mm_and_si128(b, lomask);
12941 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
12942 let w = _mm256_sub_epi8(
12943 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
12944 eight,
12945 );
12946 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
12947 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
12948 let d = _mm256_madd_epi16(p16, ones);
12949 let hi128 = _mm256_extracti128_si256::<1>(d);
12950 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12951 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12952 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12953 acc += _mm_cvtsi128_si32(s32) as f32 * s;
12954 }
12955 acc
12956 }
12957}
12958
12959#[cfg(target_arch = "x86_64")]
12962#[target_feature(enable = "avx2")]
12963unsafe fn dot_q4_row_avx2_2(
12964 packed: &[u8],
12965 scales: &[u8],
12966 g0: usize,
12967 gpr: usize,
12968 xq1: &[i8],
12969 xq2: &[i8],
12970) -> (f32, f32) {
12971 unsafe {
12973 use core::arch::x86_64::*;
12974 let lomask = _mm_set1_epi8(0x0F);
12975 let eight = _mm256_set1_epi8(8);
12976 let ones = _mm256_set1_epi16(1);
12977 let (mut acc1, mut acc2) = (0f32, 0f32);
12978 #[inline(always)]
12979 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
12980 unsafe {
12981 use core::arch::x86_64::*;
12982 let hi128 = _mm256_extracti128_si256::<1>(d);
12983 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12984 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12985 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12986 _mm_cvtsi128_si32(s32)
12987 }
12988 }
12989 for gi in 0..gpr {
12990 let g = g0 + gi;
12991 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12992 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
12993 let lo = _mm_and_si128(b, lomask);
12994 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
12995 let w = _mm256_sub_epi8(
12996 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
12997 eight,
12998 );
12999 let aw = _mm256_abs_epi8(w);
13000 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13001 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13002 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
13003 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
13004 acc1 += hsum(d1) as f32 * s;
13005 acc2 += hsum(d2) as f32 * s;
13006 }
13007 (acc1, acc2)
13008 }
13009}
13010
13011#[cfg(target_arch = "x86_64")]
13013fn q8_range_avx2(
13014 q: &[u8],
13015 row_scale: &[f32],
13016 act: &SplitAct,
13017 cols: usize,
13018 out_addr: SendMut,
13019 start: usize,
13020 end: usize,
13021) {
13022 for o in start..end {
13023 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13024 unsafe { *out_addr.at(o) = v };
13026 }
13027}
13028
13029#[cfg(target_arch = "x86_64")]
13031#[allow(clippy::too_many_arguments)]
13032fn q8_range2_avx2(
13033 q: &[u8],
13034 row_scale: &[f32],
13035 a1: &SplitAct,
13036 a2: &SplitAct,
13037 cols: usize,
13038 p1: SendMut,
13039 p2: SendMut,
13040 start: usize,
13041 end: usize,
13042) {
13043 for o in start..end {
13044 let row = &q[o * cols..(o + 1) * cols];
13045 unsafe {
13047 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
13048 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
13049 }
13050 }
13051}
13052
13053#[cfg(target_arch = "aarch64")]
13064fn i8mm_enabled() -> bool {
13065 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13066 *ON.get_or_init(|| {
13067 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
13068 && std::arch::is_aarch64_feature_detected!("i8mm")
13069 })
13070}
13071
13072#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
13076fn sdot_enabled() -> bool {
13077 if FLOAT_ACTIVATIONS.get() {
13078 return false;
13079 }
13080 use std::sync::OnceLock;
13081 static ON: OnceLock<bool> = OnceLock::new();
13082 *ON.get_or_init(|| {
13083 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
13084 if !want {
13085 return false;
13086 }
13087
13088 #[cfg(target_arch = "aarch64")]
13089 {
13090 if std::arch::is_aarch64_feature_detected!("dotprod") {
13091 return true;
13092 }
13093 #[cfg(target_os = "android")]
13094 {
13095 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
13096 if cpuinfo.lines().any(|l| {
13097 (l.starts_with("Features") || l.starts_with("features"))
13098 && l.contains("asimddp")
13099 }) {
13100 return true;
13101 }
13102 }
13103 }
13104 false
13105 }
13106 #[cfg(not(target_arch = "aarch64"))]
13107 {
13108 false
13109 }
13110 })
13111}
13112
13113struct SplitAct {
13118 xq: Vec<i8>,
13119 sx: f32,
13120 outliers: Vec<(usize, f32)>,
13121 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
13124 xsum: i32,
13125 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
13135 xg8: Vec<f32>,
13136}
13137
13138thread_local! {
13139 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
13142 const { std::cell::RefCell::new(Vec::new()) };
13143 static XG8_FREE: std::cell::RefCell<Vec<Vec<f32>>> =
13145 const { std::cell::RefCell::new(Vec::new()) };
13146}
13147
13148fn xg8_take() -> Vec<f32> {
13150 let mut v = XG8_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13151 v.clear();
13152 v
13153}
13154
13155impl Drop for SplitAct {
13156 fn drop(&mut self) {
13157 let buf = std::mem::take(&mut self.xq);
13158 if buf.capacity() > 0 {
13159 XQ_FREE.with(|f| {
13160 let mut f = f.borrow_mut();
13161 if f.len() < 16 {
13162 f.push(buf);
13163 }
13164 });
13165 }
13166 let buf = std::mem::take(&mut self.xg8);
13167 if buf.capacity() > 0 {
13168 XG8_FREE.with(|f| {
13169 let mut f = f.borrow_mut();
13170 if f.len() < 16 {
13171 f.push(buf);
13172 }
13173 });
13174 }
13175 }
13176}
13177
13178fn xg8_fill(xq: &[i8], out: &mut Vec<f32>) {
13182 out.clear();
13183 out.extend(xq.chunks_exact(GROUP_SIZE).map(|c| {
13184 let mut s = 0i32;
13185 for &v in c {
13186 s += v as i32;
13187 }
13188 s as f32
13189 }));
13190}
13191
13192thread_local! {
13193 static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
13200 static KROWS: UnsafeCell<[Vec<f32>; 2]> =
13205 const { UnsafeCell::new([Vec::new(), Vec::new()]) };
13206}
13207
13208#[inline]
13211fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
13212 KROW.with(|s| {
13213 let b = unsafe { &mut *s.get() };
13216 if b.len() < n {
13217 b.resize(n, 0.0);
13218 }
13219 f(&mut b[..n])
13220 })
13221}
13222
13223#[inline(always)]
13234fn q8_round(t: f32) -> i8 {
13235 let t = t.clamp(-127.0, 127.0);
13236 let i = t as i32;
13237 let f = t - i as f32;
13238 let r = if f >= 0.5 {
13239 i + 1
13240 } else if f <= -0.5 {
13241 i - 1
13242 } else {
13243 i
13244 };
13245 r as i8
13246}
13247
13248#[inline]
13249fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
13250 KROWS.with(|s| {
13251 let b = unsafe { &mut *s.get() };
13254 for row in b.iter_mut() {
13255 if row.len() < n {
13256 row.resize(n, 0.0);
13257 }
13258 }
13259 let (a, b) = b.split_at_mut(1);
13260 f(&mut a[0][..n], &mut b[0][..n])
13261 })
13262}
13263
13264#[inline]
13265fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
13266 if limit > 0.0 {
13267 up = up.clamp(-limit, limit);
13268 gate = gate.min(limit);
13269 }
13270 gate / (1.0 + (-gate).exp()) * up
13271}
13272
13273fn split_acts_pooled(xs_all: &[f32], b: usize, cols: usize, pool: Option<&Pool>) -> Vec<SplitAct> {
13278 match pool {
13279 Some(p) if b >= 16 => {
13280 let slots: Vec<std::sync::Mutex<Option<SplitAct>>> =
13281 (0..b).map(|_| std::sync::Mutex::new(None)).collect();
13282 p.run_rows(b, &|start, end| {
13283 for bi in start..end {
13284 let a = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
13285 *slots[bi].lock().unwrap_or_else(|e| e.into_inner()) = Some(a);
13286 }
13287 });
13288 slots
13289 .into_iter()
13290 .map(|m| {
13291 m.into_inner()
13292 .unwrap_or_else(|e| e.into_inner())
13293 .expect("every column split")
13294 })
13295 .collect()
13296 }
13297 _ => (0..b)
13298 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
13299 .collect(),
13300 }
13301}
13302
13303fn split_act(x: &[f32]) -> SplitAct {
13304 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
13305 let n = x.len();
13306 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
13307 let thr = 8.0 * rms;
13308 let mut outliers: Vec<(usize, f32)> = Vec::new();
13312 let mut amax = 0f32;
13313 for (j, &v) in x.iter().enumerate() {
13314 let a = v.abs();
13315 if a > thr {
13316 outliers.push((j, v));
13317 } else if a > amax {
13318 amax = a;
13319 }
13320 }
13321 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
13322 let inv = 1.0 / sx;
13323 crate::cpuprof::count_outliers(outliers.len() as u64);
13324 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13325 xq.clear();
13326 xq.reserve(n);
13327 if outliers.is_empty() {
13328 xq.extend(
13329 x.iter()
13330 .map(|&v| q8_round(v * inv)),
13331 );
13332 } else {
13333 xq.extend(x.iter().map(|&v| {
13335 if v.abs() > thr {
13336 0
13337 } else {
13338 q8_round(v * inv)
13339 }
13340 }));
13341 }
13342 let xsum = xq.iter().map(|&v| v as i32).sum();
13343 let mut xg8 = xg8_take();
13344 xg8_fill(&xq, &mut xg8);
13345 SplitAct {
13346 xq,
13347 sx,
13348 outliers,
13349 xsum,
13350 xg8,
13351 }
13352}
13353
13354fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
13355 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
13356 let n = x.len();
13357 let rms = (x
13358 .iter()
13359 .zip(col)
13360 .map(|(&a, &c)| {
13361 let v = a * c;
13362 (v * v) as f64
13363 })
13364 .sum::<f64>()
13365 / n.max(1) as f64)
13366 .sqrt() as f32;
13367 let thr = 8.0 * rms;
13368
13369 let mut outliers = Vec::new();
13370 let mut amax = 0f32;
13371 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
13372 let v = a * c;
13373 let s = v.abs();
13374 if s > thr {
13375 outliers.push((j, v));
13376 } else if s > amax {
13377 amax = s;
13378 }
13379 }
13380
13381 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
13382 let inv = 1.0 / sx;
13383 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13384 xq.clear();
13385 xq.reserve(n);
13386 if outliers.is_empty() {
13387 xq.extend(
13388 x.iter()
13389 .zip(col)
13390 .map(|(&a, &c)| q8_round((a * c) * inv)),
13391 );
13392 } else {
13393 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
13394 let v = a * c;
13395 if v.abs() > thr {
13396 0
13397 } else {
13398 q8_round(v * inv)
13399 }
13400 }));
13401 }
13402 let xsum = xq.iter().map(|&v| v as i32).sum();
13403 let mut xg8 = xg8_take();
13404 xg8_fill(&xq, &mut xg8);
13405 SplitAct {
13406 xq,
13407 sx,
13408 outliers,
13409 xsum,
13410 xg8,
13411 }
13412}
13413
13414#[cfg(target_arch = "aarch64")]
13417#[target_feature(enable = "neon,dotprod")]
13418unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
13419 unsafe {
13421 use core::arch::aarch64::*;
13422 use core::arch::asm;
13423 let wp = w.as_ptr() as *const i8;
13424 let n = w.len();
13425 let (mut a0, mut a1, mut a2, mut a3) = (
13426 vdupq_n_s32(0),
13427 vdupq_n_s32(0),
13428 vdupq_n_s32(0),
13429 vdupq_n_s32(0),
13430 );
13431 let mut i = 0;
13432 while i + 64 <= n {
13433 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
13434 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
13435 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
13436 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
13437 asm!(
13438 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
13439 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
13440 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
13441 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
13442 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13443 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
13444 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
13445 options(pure, nomem, nostack),
13446 );
13447 i += 64;
13448 }
13449 while i + 16 <= n {
13450 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
13451 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
13452 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
13453 i += 16;
13454 }
13455 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
13456 while i < n {
13457 s += (*wp.add(i)) as i32 * xq[i] as i32;
13458 i += 1;
13459 }
13460 s
13461 }
13462}
13463
13464#[cfg(target_arch = "aarch64")]
13468#[target_feature(enable = "neon,dotprod")]
13469unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
13470 unsafe {
13472 use core::arch::aarch64::*;
13473 use core::arch::asm;
13474 let n = xq.len();
13475 let px = xq.as_ptr();
13476 let (p0, p1, p2, p3) = (
13477 w0.as_ptr() as *const i8,
13478 w1.as_ptr() as *const i8,
13479 w2.as_ptr() as *const i8,
13480 w3.as_ptr() as *const i8,
13481 );
13482 let (mut a0, mut a1, mut a2, mut a3) = (
13483 vdupq_n_s32(0),
13484 vdupq_n_s32(0),
13485 vdupq_n_s32(0),
13486 vdupq_n_s32(0),
13487 );
13488 let mut i = 0;
13489 while i + 16 <= n {
13490 let x = vld1q_s8(px.add(i));
13491 let v0 = vld1q_s8(p0.add(i));
13492 let v1 = vld1q_s8(p1.add(i));
13493 let v2 = vld1q_s8(p2.add(i));
13494 let v3 = vld1q_s8(p3.add(i));
13495 asm!(
13496 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
13497 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
13498 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
13499 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
13500 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13501 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
13502 options(pure, nomem, nostack),
13503 );
13504 i += 16;
13505 }
13506 let mut r = [
13507 vaddvq_s32(a0),
13508 vaddvq_s32(a1),
13509 vaddvq_s32(a2),
13510 vaddvq_s32(a3),
13511 ];
13512 while i < n {
13513 let xi = *px.add(i) as i32;
13514 r[0] += (*p0.add(i)) as i32 * xi;
13515 r[1] += (*p1.add(i)) as i32 * xi;
13516 r[2] += (*p2.add(i)) as i32 * xi;
13517 r[3] += (*p3.add(i)) as i32 * xi;
13518 i += 1;
13519 }
13520 r
13521 }
13522}
13523
13524#[cfg(target_arch = "aarch64")]
13531#[target_feature(enable = "neon,dotprod")]
13532unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
13533 unsafe {
13536 use core::arch::aarch64::*;
13537 use core::arch::asm;
13538 let n = xq.len();
13539 let px = xq.as_ptr();
13540 let pg = g.as_ptr() as *const i8;
13541 let (mut a0, mut a1, mut a2, mut a3) = (
13542 vdupq_n_s32(0),
13543 vdupq_n_s32(0),
13544 vdupq_n_s32(0),
13545 vdupq_n_s32(0),
13546 );
13547 let mut i = 0;
13548 while i + 16 <= n {
13549 let x = vld1q_s8(px.add(i));
13550 let base = pg.add(4 * i);
13551 let v0 = vld1q_s8(base);
13552 let v1 = vld1q_s8(base.add(16));
13553 let v2 = vld1q_s8(base.add(32));
13554 let v3 = vld1q_s8(base.add(48));
13555 asm!(
13556 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
13557 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
13558 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
13559 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
13560 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13561 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
13562 options(pure, nomem, nostack),
13563 );
13564 i += 16;
13565 }
13566 [
13567 vaddvq_s32(a0),
13568 vaddvq_s32(a1),
13569 vaddvq_s32(a2),
13570 vaddvq_s32(a3),
13571 ]
13572 }
13573}
13574
13575#[cfg(target_arch = "aarch64")]
13581fn q8_range_sdot(
13582 q: &[u8],
13583 rep: &[u8],
13584 row_scale: &[f32],
13585 act: &SplitAct,
13586 cols: usize,
13587 out_addr: SendMut,
13588 start: usize,
13589 end: usize,
13590) {
13591 let mut o = start;
13592 if !rep.is_empty() {
13595 while o < end && o % 4 != 0 {
13596 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13597 unsafe { *out_addr.at(o) = v };
13598 o += 1;
13599 }
13600 }
13601 while o + 4 <= end {
13602 let r = if rep.is_empty() {
13603 unsafe {
13604 dot_i8_sdot_4rows(
13605 &q[o * cols..(o + 1) * cols],
13606 &q[(o + 1) * cols..(o + 2) * cols],
13607 &q[(o + 2) * cols..(o + 3) * cols],
13608 &q[(o + 3) * cols..(o + 4) * cols],
13609 &act.xq,
13610 )
13611 }
13612 } else {
13613 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
13614 };
13615 for k in 0..4 {
13616 let mut acc = r[k] as f32 * act.sx;
13617 for &(j, xv) in &act.outliers {
13618 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
13619 }
13620 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
13622 }
13623 o += 4;
13624 }
13625 while o < end {
13626 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13627 unsafe { *out_addr.at(o) = v };
13628 o += 1;
13629 }
13630}
13631
13632#[cfg(target_arch = "aarch64")]
13635#[allow(clippy::too_many_arguments)]
13636fn q8_range2_sdot(
13637 q: &[u8],
13638 row_scale: &[f32],
13639 a1: &SplitAct,
13640 a2: &SplitAct,
13641 cols: usize,
13642 p1: SendMut,
13643 p2: SendMut,
13644 start: usize,
13645 end: usize,
13646) {
13647 for o in start..end {
13648 let row = &q[o * cols..(o + 1) * cols];
13649 unsafe {
13651 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
13652 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
13653 }
13654 }
13655}
13656
13657#[allow(clippy::too_many_arguments)]
13659fn q8_range2_f32(
13660 q: &[u8],
13661 row_scale: &[f32],
13662 x1: &[f32],
13663 x2: &[f32],
13664 cols: usize,
13665 p1: SendMut,
13666 p2: SendMut,
13667 start: usize,
13668 end: usize,
13669) {
13670 for o in start..end {
13671 let row = &q[o * cols..(o + 1) * cols];
13672 unsafe {
13674 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
13675 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
13676 }
13677 }
13678}
13679
13680fn q8_range_f32(
13682 q: &[u8],
13683 row_scale: &[f32],
13684 xs: &[f32],
13685 cols: usize,
13686 out_addr: SendMut,
13687 start: usize,
13688 end: usize,
13689) {
13690 for o in start..end {
13691 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
13692 unsafe { *out_addr.at(o) = v };
13694 }
13695}
13696
13697#[inline]
13701fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
13702 #[cfg(target_arch = "aarch64")]
13703 return row_dot_sdot(row, act);
13704 #[cfg(target_arch = "x86_64")]
13705 return row_dot_avx2(row, act);
13706 #[allow(unreachable_code)]
13707 q8_row_dot_scalar(row, act)
13708}
13709
13710#[allow(dead_code)]
13711fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
13712 let mut acc = 0i32;
13713 for (k, &b) in row.iter().enumerate() {
13714 acc += (b as i8) as i32 * act.xq[k] as i32;
13715 }
13716 let mut acc = acc as f32 * act.sx;
13717 for &(j, xv) in &act.outliers {
13718 acc += (row[j] as i8) as f32 * xv;
13719 }
13720 acc
13721}
13722
13723#[cfg(target_arch = "aarch64")]
13726#[inline]
13727fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
13728 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
13729 for &(j, xv) in &act.outliers {
13730 acc += (row[j] as i8) as f32 * xv;
13731 }
13732 acc
13733}
13734
13735#[cfg(target_arch = "aarch64")]
13743#[target_feature(enable = "neon,dotprod")]
13744unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
13745 unsafe {
13748 use core::arch::aarch64::*;
13749 use core::arch::asm;
13750 let lomask = vdupq_n_u8(0x0F);
13751 let eight = vdupq_n_s8(8);
13752 let mut acc = 0f32;
13753 for gi in 0..gpr {
13754 let g = g0 + gi;
13755 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13756 let b = vld1q_u8(packed.as_ptr().add(g * 16));
13757 let lo = vandq_u8(b, lomask);
13758 let hi = vshrq_n_u8::<4>(b);
13759 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13760 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13761 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
13762 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
13763 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
13764 asm!(
13765 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
13766 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
13767 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13768 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
13769 options(pure, nomem, nostack),
13770 );
13771 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13772 }
13773 acc
13774 }
13775}
13776
13777#[cfg(target_arch = "aarch64")]
13782#[target_feature(enable = "neon,dotprod")]
13783unsafe fn dot_q4_row_sdot2(
13784 packed: &[u8],
13785 scales: &[u8],
13786 g0: usize,
13787 gpr: usize,
13788 xq1: &[i8],
13789 xq2: &[i8],
13790) -> (f32, f32) {
13791 unsafe {
13794 use core::arch::aarch64::*;
13795 use core::arch::asm;
13796 let lomask = vdupq_n_u8(0x0F);
13797 let eight = vdupq_n_s8(8);
13798 let (mut acc1, mut acc2) = (0f32, 0f32);
13799 for gi in 0..gpr {
13800 let g = g0 + gi;
13801 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13802 let b = vld1q_u8(packed.as_ptr().add(g * 16));
13803 let lo = vandq_u8(b, lomask);
13804 let hi = vshrq_n_u8::<4>(b);
13805 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13806 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13807 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
13808 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
13809 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
13810 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
13811 let (mut a0, mut a1, mut b0, mut b1) = (
13812 vdupq_n_s32(0),
13813 vdupq_n_s32(0),
13814 vdupq_n_s32(0),
13815 vdupq_n_s32(0),
13816 );
13817 asm!(
13818 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
13819 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
13820 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
13821 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
13822 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13823 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
13824 e0 = in(vreg) e0, e1 = in(vreg) e1,
13825 x10 = in(vreg) x10, x11 = in(vreg) x11,
13826 x20 = in(vreg) x20, x21 = in(vreg) x21,
13827 options(pure, nomem, nostack),
13828 );
13829 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13830 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
13831 }
13832 (acc1, acc2)
13833 }
13834}
13835
13836#[inline]
13841pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
13842 #[cfg(target_arch = "aarch64")]
13843 unsafe {
13844 return axpy_i8_f32_neon(acc, row, w);
13845 }
13846 #[cfg(target_arch = "x86_64")]
13847 if avx2_enabled() {
13848 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
13849 }
13850 #[allow(unreachable_code)]
13851 {
13852 for (a, &b) in acc.iter_mut().zip(row) {
13853 *a += w * b as f32;
13854 }
13855 }
13856}
13857
13858#[cfg(target_arch = "x86_64")]
13860#[target_feature(enable = "avx2,fma")]
13861unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
13862 unsafe {
13864 use core::arch::x86_64::*;
13865 let n = acc.len().min(row.len());
13866 let ap = acc.as_mut_ptr();
13867 let rp = row.as_ptr();
13868 let wv = _mm256_set1_ps(w);
13869 let mut j = 0usize;
13870 while j + 16 <= n {
13871 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
13872 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
13873 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
13874 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
13875 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
13876 _mm256_storeu_ps(ap.add(j), v0);
13877 _mm256_storeu_ps(ap.add(j + 8), v1);
13878 j += 16;
13879 }
13880 while j < n {
13881 *ap.add(j) += w * (*rp.add(j)) as f32;
13882 j += 1;
13883 }
13884 }
13885}
13886
13887#[cfg(target_arch = "aarch64")]
13888#[target_feature(enable = "neon")]
13889unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
13890 unsafe {
13892 use core::arch::aarch64::*;
13893 let n = acc.len().min(row.len());
13894 let ap = acc.as_mut_ptr();
13895 let rp = row.as_ptr();
13896 let wv = vdupq_n_f32(w);
13897 let mut j = 0usize;
13898 while j + 16 <= n {
13899 let rb = vld1q_s8(rp.add(j));
13900 let lo = vmovl_s8(vget_low_s8(rb));
13901 let hi = vmovl_s8(vget_high_s8(rb));
13902 for (off, half) in [(0, lo), (8, hi)] {
13903 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
13904 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
13905 let o = j + off;
13906 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
13907 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
13908 }
13909 j += 16;
13910 }
13911 while j < n {
13912 *ap.add(j) += w * (*rp.add(j)) as f32;
13913 j += 1;
13914 }
13915 }
13916}
13917
13918#[inline]
13921pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
13922 #[cfg(target_arch = "aarch64")]
13923 unsafe {
13924 return dot_i8_f32_neon(w, x);
13925 }
13926 #[cfg(target_arch = "x86_64")]
13927 if avx2_enabled() {
13928 return unsafe { dot_i8_f32_avx2(w, x) };
13929 }
13930 #[allow(unreachable_code)]
13931 {
13932 let mut sum = 0.0f32;
13933 for (j, &b) in w.iter().enumerate() {
13934 sum += (b as i8) as f32 * x[j];
13935 }
13936 sum
13937 }
13938}
13939
13940#[inline]
13944fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
13945 #[cfg(target_arch = "aarch64")]
13946 unsafe {
13947 return dot_i8_col_f32_neon(w, x, col);
13948 }
13949 #[allow(unreachable_code)]
13950 {
13951 let mut sum = 0.0f32;
13952 for (j, &b) in w.iter().enumerate() {
13953 sum += (b as i8) as f32 * x[j] * col[j];
13954 }
13955 sum
13956 }
13957}
13958
13959#[cfg(target_arch = "aarch64")]
13960#[target_feature(enable = "neon")]
13961unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
13962 unsafe {
13964 use core::arch::aarch64::*;
13965 let n = x.len();
13966 let wp = w.as_ptr() as *const i8;
13967 let xp = x.as_ptr();
13968 let cp = col.as_ptr();
13969 let (mut a0, mut a1, mut a2, mut a3) = (
13970 vdupq_n_f32(0.0),
13971 vdupq_n_f32(0.0),
13972 vdupq_n_f32(0.0),
13973 vdupq_n_f32(0.0),
13974 );
13975 let mut j = 0usize;
13976 while j + 16 <= n {
13977 let wb = vld1q_s8(wp.add(j));
13978 let lo = vmovl_s8(vget_low_s8(wb));
13979 let hi = vmovl_s8(vget_high_s8(wb));
13980 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
13981 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
13982 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
13983 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
13984 a0 = vfmaq_f32(
13985 a0,
13986 w0,
13987 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
13988 );
13989 a1 = vfmaq_f32(
13990 a1,
13991 w1,
13992 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
13993 );
13994 a2 = vfmaq_f32(
13995 a2,
13996 w2,
13997 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
13998 );
13999 a3 = vfmaq_f32(
14000 a3,
14001 w3,
14002 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
14003 );
14004 j += 16;
14005 }
14006 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
14007 while j < n {
14008 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
14009 j += 1;
14010 }
14011 sum
14012 }
14013}
14014
14015#[cfg(target_arch = "aarch64")]
14016#[target_feature(enable = "neon")]
14017unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
14018 unsafe {
14020 use core::arch::aarch64::*;
14021 let n = x.len();
14022 let wp = w.as_ptr() as *const i8;
14023 let xp = x.as_ptr();
14024 let (mut a0, mut a1, mut a2, mut a3) = (
14025 vdupq_n_f32(0.0),
14026 vdupq_n_f32(0.0),
14027 vdupq_n_f32(0.0),
14028 vdupq_n_f32(0.0),
14029 );
14030 let mut j = 0usize;
14031 while j + 16 <= n {
14032 let wb = vld1q_s8(wp.add(j));
14033 let lo = vmovl_s8(vget_low_s8(wb));
14034 let hi = vmovl_s8(vget_high_s8(wb));
14035 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
14036 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
14037 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
14038 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
14039 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
14040 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
14041 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
14042 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
14043 j += 16;
14044 }
14045 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
14046 while j < n {
14047 sum += (*wp.add(j)) as f32 * *xp.add(j);
14048 j += 1;
14049 }
14050 sum
14051 }
14052}
14053
14054#[allow(clippy::too_many_arguments)]
14055fn qmatvec(
14056 q: &[u8],
14057 rep: &[u8],
14058 row_scale: &[f32],
14059 x: &[f32],
14060 col_field: &[f32],
14061 dtype: TensorDtype,
14062 rows: usize,
14063 cols: usize,
14064 out: &mut [f32],
14065 pool: Option<&Pool>,
14066) {
14067 debug_assert_eq!(out.len(), rows);
14068 #[cfg(not(target_arch = "aarch64"))]
14069 let _ = rep;
14070
14071 #[cfg(target_arch = "aarch64")]
14072 if sdot_enabled() {
14073 let act = if dtype == TensorDtype::Q8_2f {
14074 split_act_q8_2f(x, col_field)
14075 } else {
14076 split_act(x)
14077 };
14078 let out_addr = SendMut(out.as_mut_ptr());
14079 let run_range = |start: usize, end: usize| {
14080 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
14081 };
14082 match pool {
14083 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14084 _ => run_range(0, rows),
14085 }
14086 return;
14087 }
14088 #[cfg(target_arch = "x86_64")]
14091 if avx2_a8w8_enabled() {
14092 let act = if dtype == TensorDtype::Q8_2f {
14093 split_act_q8_2f(x, col_field)
14094 } else {
14095 split_act(x)
14096 };
14097 let out_addr = SendMut(out.as_mut_ptr());
14098 let run_range = |start: usize, end: usize| {
14099 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
14100 };
14101 match pool {
14102 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14103 _ => run_range(0, rows),
14104 }
14105 return;
14106 }
14107
14108 prescale_with(x, col_field, dtype, 1, |xs| {
14109 let out_addr = SendMut(out.as_mut_ptr());
14110 let run_range = move |start: usize, end: usize| {
14111 for o in start..end {
14112 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
14113 unsafe { *out_addr.at(o) = v };
14115 }
14116 };
14117 match pool {
14118 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14119 _ => run_range(0, rows),
14120 }
14121 });
14122}
14123
14124#[allow(clippy::too_many_arguments)]
14125fn qmatvec2(
14126 q: &[u8],
14127 row_scale: &[f32],
14128 x1: &[f32],
14129 x2: &[f32],
14130 col_field: &[f32],
14131 dtype: TensorDtype,
14132 rows: usize,
14133 cols: usize,
14134 o1: &mut [f32],
14135 o2: &mut [f32],
14136 pool: Option<&Pool>,
14137) {
14138 #[cfg(target_arch = "aarch64")]
14139 if sdot_enabled() {
14140 let a1s = if dtype == TensorDtype::Q8_2f {
14141 split_act_q8_2f(x1, col_field)
14142 } else {
14143 split_act(x1)
14144 };
14145 let a2s = if dtype == TensorDtype::Q8_2f {
14146 split_act_q8_2f(x2, col_field)
14147 } else {
14148 split_act(x2)
14149 };
14150 let p1 = SendMut(o1.as_mut_ptr());
14151 let p2 = SendMut(o2.as_mut_ptr());
14152 let run_range = |start: usize, end: usize| {
14153 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
14154 };
14155 match pool {
14156 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14157 _ => run_range(0, rows),
14158 }
14159 return;
14160 }
14161 #[cfg(target_arch = "x86_64")]
14162 if avx2_a8w8_enabled() {
14163 let a1s = if dtype == TensorDtype::Q8_2f {
14164 split_act_q8_2f(x1, col_field)
14165 } else {
14166 split_act(x1)
14167 };
14168 let a2s = if dtype == TensorDtype::Q8_2f {
14169 split_act_q8_2f(x2, col_field)
14170 } else {
14171 split_act(x2)
14172 };
14173 let p1 = SendMut(o1.as_mut_ptr());
14174 let p2 = SendMut(o2.as_mut_ptr());
14175 let run_range = |start: usize, end: usize| {
14176 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
14177 };
14178 match pool {
14179 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14180 _ => run_range(0, rows),
14181 }
14182 return;
14183 }
14184
14185 prescale_with(x1, col_field, dtype, 1, |x1s| {
14186 prescale_with(x2, col_field, dtype, 2, |x2s| {
14187 let p1 = SendMut(o1.as_mut_ptr());
14188 let p2 = SendMut(o2.as_mut_ptr());
14189 let run_range = move |start: usize, end: usize| {
14190 for o in start..end {
14191 let row = &q[o * cols..(o + 1) * cols];
14192 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
14193 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
14194 unsafe {
14196 *p1.at(o) = s1;
14197 *p2.at(o) = s2;
14198 }
14199 }
14200 };
14201 match pool {
14202 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14203 _ => run_range(0, rows),
14204 }
14205 });
14206 });
14207}
14208
14209#[derive(Clone, Copy)]
14210struct SendMut(*mut f32);
14211unsafe impl Send for SendMut {}
14212unsafe impl Sync for SendMut {}
14213
14214impl SendMut {
14215 #[inline]
14216 fn at(self, i: usize) -> *mut f32 {
14217 unsafe { self.0.add(i) }
14218 }
14219}
14220
14221#[cfg(test)]
14222mod tests {
14223
14224 #[test]
14228 fn q8_round_is_round_clamp() {
14229 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
14230 let mut probe = vec![
14231 0.0f32,
14232 -0.0,
14233 f32::NAN,
14234 f32::INFINITY,
14235 f32::NEG_INFINITY,
14236 f32::MAX,
14237 f32::MIN,
14238 1e30,
14239 -1e30,
14240 f32::MIN_POSITIVE,
14241 -f32::MIN_POSITIVE,
14242 ];
14243 for k in -300i32..=300 {
14244 let h = k as f32 * 0.5;
14245 let up = f32::from_bits(h.to_bits() + 1);
14246 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
14247 for t in [h, up, down] {
14248 probe.push(t);
14249 probe.push(-t);
14250 }
14251 }
14252 let mut t = -140.0f32;
14253 while t < 140.0 {
14254 probe.push(t);
14255 t += 0.000_731;
14256 }
14257 for t in probe {
14258 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
14259 }
14260 }
14261
14262 use super::*;
14263
14264 #[test]
14268 fn f32_matmat_equals_matvec_bitwise() {
14269 let (rows, cols, b) = (5usize, 37usize, 7usize);
14270 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
14271 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
14272 let t = QTensor::from_f32(w, rows, cols);
14273 let mut out = vec![0f32; b * rows];
14274 t.matmat(&xs, b, &mut out, None);
14275 for bi in 0..b {
14276 let mut one = vec![0f32; rows];
14277 t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
14278 for o in 0..rows {
14279 assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
14280 }
14281 }
14282 }
14283
14284 #[test]
14287 fn f32_matmat_row_split_on_pool_is_bitwise() {
14288 let (rows, cols, b) = (300usize, 37usize, 7usize);
14289 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
14290 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
14291 let t = QTensor::from_f32(w, rows, cols);
14292 let mut alone = vec![0f32; b * rows];
14293 t.matmat(&xs, b, &mut alone, None);
14294 let pool = Pool::new(3);
14295 let mut split = vec![0f32; b * rows];
14296 t.matmat(&xs, b, &mut split, Some(&pool));
14297 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14298 assert_eq!(bits(&alone), bits(&split));
14299 }
14300
14301 #[test]
14302 fn q2tp_i8_dot_matches_exact_on_grid() {
14303 let (rows, cols) = (5, 64);
14307 let gpr = cols / GROUP_SIZE;
14308 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
14312 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
14313 .collect();
14314 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
14315 let x: Vec<f32> = (0..cols)
14316 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
14317 .collect();
14318 let act = split_act(&x);
14319 assert!(
14320 act.outliers.is_empty(),
14321 "on-grid input must have no outliers"
14322 );
14323 let gsum = q1_group_sums(&act.xq, gpr);
14324 for r in 0..rows {
14325 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
14326 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
14327 assert!(
14328 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
14329 "row {r}: exact {exact} vs i8 {fast}"
14330 );
14331 }
14332 }
14333
14334 #[cfg(target_arch = "x86_64")]
14335 #[test]
14336 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
14337 if !std::arch::is_x86_feature_detected!("avx2") {
14342 return;
14343 }
14344 let mut seed = 0x9e3779b9u32;
14345 let mut next = || {
14346 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
14347 seed
14348 };
14349 for _ in 0..20_000 {
14350 let mut ch = [0u8; Q2TP_CHUNK];
14351 let mut x = [0i8; GROUP_SIZE];
14352 for b in &mut ch {
14353 *b = next() as u8;
14354 }
14355 for v in &mut x {
14356 *v = (next() >> 24) as i8;
14357 }
14358 let mut reference = 0i32;
14359 for (k, &b) in ch.iter().enumerate() {
14360 reference += (b & 3) as i32 * x[k * 4] as i32;
14361 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
14362 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
14363 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
14364 }
14365 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
14368 assert_eq!(got, reference, "packed q2 lane mismatch");
14369 }
14370 }
14371
14372 #[test]
14373 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
14374 let (rows, cols) = (1usize, GROUP_SIZE);
14375 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
14376 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
14380 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
14381 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
14383 let mut raw = vec![0.0f32; rows];
14384 let mut affine = vec![0.0f32; rows];
14385 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
14386 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
14387 assert_eq!(raw, vec![-24.0]);
14388 assert_eq!(affine, vec![-8.0]);
14389 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
14390 }
14391
14392 #[test]
14393 fn q8_row_dot_fast_matches_scalar() {
14394 let cols = 96;
14397 let row: Vec<u8> = (0..cols)
14398 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
14399 .collect();
14400 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
14401 let act = split_act(&x);
14402 let fast = q8_row_dot(&row, &act);
14403 let scalar = q8_row_dot_scalar(&row, &act);
14404 assert!(
14405 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
14406 "fast {fast} vs scalar {scalar}"
14407 );
14408 }
14409
14410 #[test]
14411 fn f32_matvec_matches_matvec_rows_bitexact() {
14412 let (rows, cols) = (300, 40);
14413 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
14414 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
14415 let qt = QTensor::from_f32(w.clone(), rows, cols);
14416
14417 let mut a = vec![0.0f32; rows];
14418 matvec_rows(None, &w, &x, &mut a);
14419 let mut b = vec![0.0f32; rows];
14420 qt.matvec(&x, &mut b, None);
14421 assert_eq!(a, b);
14422 }
14423
14424 #[test]
14425 fn sdot_kernel_exact_on_grid() {
14426 eprintln!("sdot_enabled = {}", sdot_enabled());
14431 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
14433 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
14434 .collect();
14435 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
14436 let x: Vec<f32> = (0..cols)
14437 .map(|i| match i % 3 {
14438 0 => 1.0,
14439 1 => -1.0,
14440 _ => 0.0,
14441 })
14442 .collect();
14443 let mut a = vec![0.0f32; rows];
14444 qmatvec(
14445 &w,
14446 &[],
14447 &scales,
14448 &x,
14449 &[],
14450 TensorDtype::Q8Row,
14451 rows,
14452 cols,
14453 &mut a,
14454 None,
14455 );
14456 for o in 0..rows {
14457 let mut acc = 0.0f32;
14458 for j in 0..cols {
14459 acc += (w[o * cols + j] as i8) as f32 * x[j];
14460 }
14461 let expect = acc * scales[o];
14462 assert!(
14463 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
14464 "row {o}: {} vs {expect}",
14465 a[o]
14466 );
14467 }
14468 }
14469
14470 #[test]
14471 fn q1_tbl_fast_path_matches_reference() {
14472 let (rows, cols) = (5, 256);
14477 let gpr = cols / GROUP_SIZE;
14478 let mut bytes = Vec::new();
14479 for t in 0..rows * gpr {
14480 let s = 0.007 + (t % 11) as f32 * 0.004;
14481 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14482 for j in 0..4 {
14483 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
14484 }
14485 }
14486 let x: Vec<f32> = (0..cols)
14487 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
14488 .collect();
14489 let mut w = vec![0.0f32; rows * cols];
14490 cortiq_core::quant::dequant_q1(&bytes, &mut w);
14491 let mut got = vec![0.0f32; rows];
14492 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
14493 for o in 0..rows {
14494 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
14495 assert!(
14496 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
14497 "row {o}: {} vs {expect}",
14498 got[o]
14499 );
14500 }
14501 let b = 5usize;
14504 let mut xs_all = Vec::new();
14505 for bi in 0..b {
14506 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
14507 }
14508 let mut mm = vec![0.0f32; b * rows];
14509 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
14510 for bi in 0..b {
14511 let mut single = vec![0.0f32; rows];
14512 q1_matvec(
14513 &bytes,
14514 &xs_all[bi * cols..(bi + 1) * cols],
14515 rows,
14516 cols,
14517 &mut single,
14518 None,
14519 );
14520 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
14521 }
14522 }
14523
14524 #[test]
14525 fn q1_kernels_match_exact_reference() {
14526 let (rows, cols) = (7, 96);
14528 let gpr = cols / GROUP_SIZE;
14529 let mut bytes = Vec::new();
14530 for t in 0..rows * gpr {
14531 let s = 0.01 + (t % 13) as f32 * 0.003;
14532 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14533 for j in 0..4 {
14534 bytes.push(((t * 31 + j * 97) % 251) as u8);
14535 }
14536 }
14537 let x: Vec<f32> = (0..cols)
14539 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
14540 .collect();
14541 let mut w = vec![0.0f32; rows * cols];
14543 cortiq_core::quant::dequant_q1(&bytes, &mut w);
14544 let mut expect = vec![0.0f32; rows];
14545 for o in 0..rows {
14546 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
14547 }
14548 let mut got = vec![0.0f32; rows];
14549 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
14550 for o in 0..rows {
14551 assert!(
14552 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
14553 "row {o}: {} vs {}",
14554 got[o],
14555 expect[o]
14556 );
14557 }
14558 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
14560 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
14561 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14562 assert_eq!(a1, got);
14563 let mut xs = x.clone();
14564 xs.extend_from_slice(&x2);
14565 let mut mm = vec![0.0f32; 2 * rows];
14566 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
14567 assert_eq!(&mm[..rows], got.as_slice());
14568 assert_eq!(&mm[rows..], a2.as_slice());
14569 }
14570
14571 fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
14574 let tiles = rows * cols / GROUP_SIZE;
14575 let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
14576 for t in 0..tiles {
14577 let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
14578 q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14579 let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
14580 q1.extend_from_slice(&bits.to_le_bytes());
14581 }
14582 let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
14583 assert_eq!(rp.differing, 0, "uniform scales repack exactly");
14584 (q1, rp.data)
14585 }
14586
14587 fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
14590 let mut x: Vec<f32> = (0..cols)
14591 .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
14592 .collect();
14593 x[3 % cols] = 41.0;
14594 x[(cols / 2 + 5) % cols] = -37.5;
14595 x
14596 }
14597
14598 #[test]
14605 fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
14606 for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
14607 let gpr = cols / GROUP_SIZE;
14608 let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
14609 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14610 let x = q1g_test_x(cols, rows);
14611 let act = split_act(&x);
14612 assert!(!act.outliers.is_empty(), "test wants the outlier path");
14613 let gsum = q1_group_sums(&act.xq, gpr);
14614 for r in 0..rows {
14615 let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
14616 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
14617 let (a, b) = (
14618 wa.dot_i8(r, gpr, &act.xq, &gsum),
14619 wb.dot_i8(r, gpr, &act.xq, &gsum),
14620 );
14621 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
14622 for j in [0, 31, 32, 127, cols - 1] {
14623 let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
14624 assert_eq!(
14625 (a.0.to_bits(), a.1.to_bits()),
14626 (b.0.to_bits(), b.1.to_bits())
14627 );
14628 }
14629 #[cfg(target_arch = "x86_64")]
14630 if avx2_enabled() {
14631 let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
14632 unsafe {
14634 let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
14635 let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
14636 assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
14637 }
14638 if std::is_x86_feature_detected!("avx512vnni")
14639 && std::is_x86_feature_detected!("avx512bw")
14640 && std::is_x86_feature_detected!("avx512vl")
14641 {
14642 unsafe {
14644 let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
14645 let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
14646 assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
14647 }
14648 }
14649 }
14650 }
14651 let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
14653 bin_matvec(wa, &x, rows, cols, &mut ga, None);
14654 bin_matvec(wb, &x, rows, cols, &mut gb, None);
14655 assert_eq!(ga, gb, "{rows}x{cols} matvec");
14656 let x2 = q1g_test_x(cols, rows + 1);
14657 let (mut a1, mut a2, mut b1, mut b2) = (
14658 vec![0f32; rows],
14659 vec![0f32; rows],
14660 vec![0f32; rows],
14661 vec![0f32; rows],
14662 );
14663 bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14664 bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
14665 assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
14666 let b = 5usize;
14667 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
14668 let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14669 bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
14670 bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
14671 assert_eq!(ma, mb, "{rows}x{cols} matmat");
14672 let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
14675 .map(|k| {
14676 let a = split_act(&xs[k * cols..(k + 1) * cols]);
14677 let g = q1_group_sums(&a.xq, gpr);
14678 (a, g)
14679 })
14680 .collect();
14681 let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
14682 let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
14683 for r in 0..rows {
14684 #[cfg(target_arch = "x86_64")]
14685 if avx2_enabled() {
14686 let (a, b) = unsafe {
14688 (
14689 wa.dot_1x4_x86(r, gpr, xq4, gs4),
14690 wb.dot_1x4_x86(r, gpr, xq4, gs4),
14691 )
14692 };
14693 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14694 }
14695 #[cfg(target_arch = "aarch64")]
14696 if sdot_enabled() {
14697 let (a, b) = unsafe {
14699 (
14700 wa.dot_1x4_sdot(r, gpr, xq4, gs4),
14701 wb.dot_1x4_sdot(r, gpr, xq4, gs4),
14702 )
14703 };
14704 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14705 }
14706 let _ = (r, &xq4, &gs4);
14707 }
14708 }
14709 }
14710
14711 #[test]
14717 #[ignore]
14718 fn q1g_cpu_matvec_speed() {
14719 use crate::pool::Pool;
14720 let threads = std::env::var("CMF_THREADS")
14721 .ok()
14722 .and_then(|v| v.parse().ok())
14723 .unwrap_or(8usize);
14724 let pool = Pool::new(threads);
14725 for &(name, rows, cols, reps) in &[
14726 ("gate/up_proj", 6144usize, 2048usize, 41usize),
14727 ("down_proj", 2048, 6144, 41),
14728 ("lm_head", 151_669, 2048, 9),
14729 ] {
14730 let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
14731 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14732 let x = q1g_test_x(cols, 1);
14733 let mut out = vec![0f32; rows];
14734 let mut time = |w: BinW| {
14735 let t0 = std::time::Instant::now();
14736 bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
14737 t0.elapsed().as_secs_f64()
14738 };
14739 let _ = (time(wa), time(wb));
14740 let (mut ta, mut tb) = (Vec::new(), Vec::new());
14741 for r in 0..reps {
14742 if r % 2 == 0 {
14743 ta.push(time(wa));
14744 tb.push(time(wb));
14745 } else {
14746 tb.push(time(wb));
14747 ta.push(time(wa));
14748 }
14749 }
14750 ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
14751 tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
14752 let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
14753 eprintln!(
14754 "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
14755 a8w8_enabled(),
14756 ma * 1e6,
14757 mb * 1e6,
14758 mb / ma
14759 );
14760 }
14761 }
14762
14763 #[test]
14768 fn q1g_mapped_tensor_paths_equal_q1() {
14769 use crate::pool::Pool;
14770 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14771 let (rows, cols) = (300usize, 256usize);
14772 let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
14773 let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
14774 let header: CmfHeader = serde_json::from_value(serde_json::json!({
14775 "format": "cmf", "version": cortiq_core::CMF_VERSION,
14776 "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
14777 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14778 "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
14779 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14780 "max_position_embeddings": 8},
14781 "quant_type": "VBIT"
14782 }))
14783 .unwrap();
14784 let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
14785 name: name.into(),
14786 dtype,
14787 shape: vec![rows, cols],
14788 data: data.clone(),
14789 };
14790 let specs = [
14791 spec("g1", TensorDtype::Q1, &gq1),
14792 spec("u1", TensorDtype::Q1, &uq1),
14793 spec("gg", TensorDtype::Q1G, &gq1g),
14794 spec("ug", TensorDtype::Q1G, &uq1g),
14795 ];
14796 let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
14797 std::fs::create_dir_all(&dir).unwrap();
14798 let path = dir.join("m.cmf");
14799 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14800 let model = Arc::new(CmfModel::open(&path).unwrap());
14801 assert!(model.verify().is_empty(), "{:?}", model.verify());
14802 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14803 let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
14804 assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
14805 assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
14806 let pool = Pool::new(3);
14807 let x = q1g_test_x(cols, 4);
14808 let run = |a: &QTensor| {
14809 let mut o = vec![0f32; rows];
14810 a.matvec(&x, &mut o, Some(&pool));
14811 o
14812 };
14813 assert_eq!(run(&g1), run(&gg), "matvec");
14814 let many = |a: &QTensor, b: &QTensor| {
14815 let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
14816 QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
14817 (oa, ob)
14818 };
14819 let want = many(&g1, &u1);
14820 assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
14821 assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
14822 assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
14823 let fused = |g: &QTensor, u: &QTensor| {
14824 let mut o = vec![0f32; rows];
14825 let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
14826 (ok, o)
14827 };
14828 let (ok_a, fa) = fused(&g1, &u1);
14829 let (ok_b, fb) = fused(&gg, &ug);
14830 assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
14831 assert_eq!(fa, fb, "fused gate/up");
14832 let b = 6usize;
14833 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
14834 let mm = |a: &QTensor| {
14835 let mut o = vec![0f32; b * rows];
14836 a.matmat(&xs, b, &mut o, Some(&pool));
14837 o
14838 };
14839 assert_eq!(mm(&g1), mm(&gg), "matmat");
14840 let x2 = q1g_test_x(cols, 77);
14841 let pair = |a: &QTensor| {
14842 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14843 a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14844 (o1, o2)
14845 };
14846 assert_eq!(pair(&g1), pair(&gg), "matvec2");
14847 for r in [0, 1, rows / 2, rows - 1] {
14848 let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
14849 g1.row_f32(r, &mut ra);
14850 gg.row_f32(r, &mut rb);
14851 assert_eq!(ra, rb, "embedding row {r}");
14852 }
14853 drop((g1, u1, gg, ug, model));
14854 let _ = std::fs::remove_dir_all(&dir);
14855 }
14856
14857 #[test]
14864 fn blr_mapped_tensor_paths_equal_dense() {
14865 use crate::pool::Pool;
14866 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14867 let (rows, cols, r) = (96usize, 128usize, 16usize);
14868 let mut rng = crate::blr::tests::Rng(41);
14869 let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
14870 let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
14871 let mut q8: Vec<u8> = (0..rows * cols)
14873 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14874 .collect();
14875 for o in 0..rows {
14876 q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
14877 }
14878 let header: CmfHeader = serde_json::from_value(serde_json::json!({
14879 "format": "cmf", "version": cortiq_core::CMF_VERSION,
14880 "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
14881 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14882 "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
14883 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14884 "max_position_embeddings": 8},
14885 "quant_type": "Q8_2F"
14886 }))
14887 .unwrap();
14888 let specs = [
14889 TensorSpec {
14890 name: "g".into(),
14891 dtype: TensorDtype::BinaryLowRank,
14892 shape: vec![rows, cols],
14893 data: gate.clone(),
14894 },
14895 TensorSpec {
14896 name: "u".into(),
14897 dtype: TensorDtype::BinaryLowRank,
14898 shape: vec![rows, cols],
14899 data: up.clone(),
14900 },
14901 TensorSpec {
14902 name: "q".into(),
14903 dtype: TensorDtype::Q8Row,
14904 shape: vec![rows, cols],
14905 data: q8,
14906 },
14907 ];
14908 let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
14909 std::fs::create_dir_all(&dir).unwrap();
14910 let path = dir.join("m.cmf");
14911 CmfModel::write(&path, &header, &specs, None, None).unwrap();
14912 let model = Arc::new(CmfModel::open(&path).unwrap());
14913 assert!(model.verify().is_empty(), "{:?}", model.verify());
14914 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14915 let (g, u, q) = (t("g"), t("u"), t("q"));
14916 assert!(g.is_blr() && u.is_blr() && !q.is_blr());
14917 assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
14918 assert!(g.as_f32().is_none() && g.model_idx().is_none());
14919 assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
14920 assert!(!g.sparse_col_ok());
14921 assert_eq!((g.rows(), g.cols()), (rows, cols));
14922 let deq = |bytes: &[u8]| {
14923 let mut w = vec![0f32; rows * cols];
14924 cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
14925 w
14926 };
14927 let (wg, wu) = (deq(&gate), deq(&up));
14928 let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
14929 (0..rows)
14930 .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
14931 .collect()
14932 };
14933 let close = |a: &[f32], b: &[f32], what: &str| {
14934 let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
14935 let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
14936 assert!(e / s < 1e-5, "{what}: rel {}", e / s);
14937 };
14938 let pool = Pool::new(3);
14939 let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
14940 let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
14941 let mut o = vec![0f32; rows];
14942 g.matvec(&x, &mut o, Some(&pool));
14943 close(&o, &mv(&wg, &x), "matvec");
14944 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14945 g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14946 close(&o1, &mv(&wg, &x), "matvec2 a");
14947 close(&o2, &mv(&wg, &x2), "matvec2 b");
14948 for b in [2usize, 5] {
14949 let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
14950 let mut om = vec![0f32; b * rows];
14951 g.matmat(&xs, b, &mut om, Some(&pool));
14952 for t in 0..b {
14953 close(
14954 &om[t * rows..(t + 1) * rows],
14955 &mv(&wg, &xs[t * cols..(t + 1) * cols]),
14956 "matmat",
14957 );
14958 }
14959 }
14960 let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
14963 QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
14964 close(&a, &mv(&wg, &x), "matvec_many blr gate");
14965 close(&c, &mv(&wu, &x), "matvec_many blr up");
14966 let mut qo = vec![0f32; rows];
14967 q.matvec(&x, &mut qo, Some(&pool));
14968 assert_eq!(bq, qo, "matvec_many q8_row sibling");
14969 let mut f = vec![0f32; rows];
14970 assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
14971 let mut row = vec![0f32; cols];
14972 g.row_f32(rows / 3, &mut row);
14973 close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
14974 let mut scratch = vec![0f32; cols];
14975 let d = g.row_dot(7, &x, &mut scratch);
14976 let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
14977 assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
14978 let mut col = vec![0f32; rows];
14979 g.add_col_scaled(cols / 2, 1.5, &mut col);
14980 let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
14981 close(&col, &want, "add_col_scaled");
14982 drop((g, u, q, model));
14983 let _ = std::fs::remove_dir_all(&dir);
14984 }
14985
14986 #[test]
14987 fn repack_is_bit_identical() {
14988 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
14994 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14995 .collect();
14996 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
14997 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
14998 let rep = q8_repack_layout(&w, rows, cols);
14999 for g in 0..rows / 4 {
15001 for c in 0..cols / 16 {
15002 for lane in 0..4 {
15003 assert_eq!(
15004 &rep[g * 4 * cols + c * 64 + lane * 16
15005 ..g * 4 * cols + c * 64 + lane * 16 + 16],
15006 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
15007 );
15008 }
15009 }
15010 }
15011 let mut a = vec![0.0f32; rows];
15012 qmatvec(
15013 &w,
15014 &[],
15015 &scales,
15016 &x,
15017 &[],
15018 TensorDtype::Q8Row,
15019 rows,
15020 cols,
15021 &mut a,
15022 None,
15023 );
15024 let mut b = vec![0.0f32; rows];
15025 qmatvec(
15026 &w,
15027 &rep,
15028 &scales,
15029 &x,
15030 &[],
15031 TensorDtype::Q8Row,
15032 rows,
15033 cols,
15034 &mut b,
15035 None,
15036 );
15037 assert_eq!(a, b, "full-range repack output diverged");
15038
15039 #[cfg(target_arch = "aarch64")]
15040 if sdot_enabled() {
15041 let act = split_act(&x);
15043 let mut c1 = vec![0.0f32; rows];
15044 let mut c2 = vec![0.0f32; rows];
15045 q8_range_sdot(
15046 &w,
15047 &[],
15048 &scales,
15049 &act,
15050 cols,
15051 SendMut(c1.as_mut_ptr()),
15052 3,
15053 rows - 2,
15054 );
15055 q8_range_sdot(
15056 &w,
15057 &rep,
15058 &scales,
15059 &act,
15060 cols,
15061 SendMut(c2.as_mut_ptr()),
15062 3,
15063 rows - 2,
15064 );
15065 assert_eq!(c1, c2, "unaligned-range repack output diverged");
15066 }
15067 }
15068
15069 #[test]
15070 fn sdot_a8w8_noise_is_bounded() {
15071 let (rows, cols) = (16, 512);
15075 let w: Vec<u8> = (0..rows * cols)
15076 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
15077 .collect();
15078 let scales = vec![0.01f32; rows];
15079 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
15080 let mut a = vec![0.0f32; rows];
15081 qmatvec(
15082 &w,
15083 &[],
15084 &scales,
15085 &x,
15086 &[],
15087 TensorDtype::Q8Row,
15088 rows,
15089 cols,
15090 &mut a,
15091 None,
15092 );
15093 let (mut num, mut den) = (0f64, 0f64);
15094 for o in 0..rows {
15095 let mut acc = 0.0f32;
15096 for j in 0..cols {
15097 acc += (w[o * cols + j] as i8) as f32 * x[j];
15098 }
15099 let expect = acc * scales[o];
15100 num += ((a[o] - expect) as f64).powi(2);
15101 den += (expect as f64).powi(2);
15102 }
15103 let rel = (num / den.max(1e-12)).sqrt();
15104 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
15105 }
15106
15107 #[test]
15108 fn i8_dot_neon_matches_scalar() {
15109 let n = 100;
15110 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
15111 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
15112 let mut scalar = 0.0f32;
15113 for j in 0..n {
15114 scalar += (w[j] as i8) as f32 * x[j];
15115 }
15116 let fast = dot_i8_f32(&w, &x);
15117 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
15118 }
15119
15120 #[test]
15122 fn vbitmatvec_matches_full_dequant() {
15123 let (rows, cols) = (6, 64);
15124 let ng = cols / GROUP_SIZE;
15125 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
15127 let mut bytes = bits.clone();
15128 for g in 0..rows * ng {
15129 let s = 0.02 + 0.001 * g as f32;
15130 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15131 }
15132 for r in 0..rows {
15133 let b = bits[r] as usize;
15134 let (mut acc, mut nb) = (0u64, 0usize);
15135 let mut rowbytes = Vec::new();
15136 for i in 0..cols {
15137 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
15138 acc = (acc << b) | v;
15139 nb += b;
15140 while nb >= 8 {
15141 nb -= 8;
15142 rowbytes.push(((acc >> nb) & 0xFF) as u8);
15143 }
15144 }
15145 if nb > 0 {
15146 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15147 }
15148 bytes.extend_from_slice(&rowbytes);
15149 }
15150 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
15151
15152 let mut reference = vec![0f32; rows * cols];
15153 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
15154 let mut expect = vec![0f32; rows];
15155 for r in 0..rows {
15156 expect[r] = reference[r * cols..(r + 1) * cols]
15157 .iter()
15158 .zip(&x)
15159 .map(|(w, xv)| w * xv)
15160 .sum();
15161 }
15162 let mut got = vec![0f32; rows];
15163 let offsets = vbit_row_offsets(&bytes, rows, cols);
15164 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
15165 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
15169 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
15170 for r in 0..rows {
15171 assert!(
15172 (got[r] - expect[r]).abs() < tol * scale,
15173 "row {r}: {} vs {}",
15174 got[r],
15175 expect[r]
15176 );
15177 }
15178 }
15179
15180 #[test]
15185 #[cfg(target_arch = "x86_64")]
15186 fn vbit_matmat_blocked_matches_per_row() {
15187 let (rows, cols, b) = (64usize, 128usize, 9usize);
15188 let ng = cols / GROUP_SIZE;
15189 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
15190 let mut bytes = bits.clone();
15191 for g in 0..rows * ng {
15192 let sc = 0.02 + 0.0005 * g as f32;
15193 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15194 }
15195 for r in 0..rows {
15196 let bw = bits[r] as usize;
15197 let (mut acc, mut nb) = (0u64, 0usize);
15198 let mut rowbytes = Vec::new();
15199 for i in 0..cols {
15200 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
15201 acc = (acc << bw) | v;
15202 nb += bw;
15203 while nb >= 8 {
15204 nb -= 8;
15205 rowbytes.push(((acc >> nb) & 0xFF) as u8);
15206 }
15207 }
15208 if nb > 0 {
15209 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15210 }
15211 bytes.extend_from_slice(&rowbytes);
15212 }
15213 let x: Vec<f32> = (0..b * cols)
15214 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15215 .collect();
15216 let offsets = vbit_row_offsets(&bytes, rows, cols);
15217 let mut y_a = vec![0f32; b * rows];
15218 let mut y_b = vec![0f32; b * rows];
15219 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
15220 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
15221 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
15222 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
15223 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
15224 let max_d = y_a
15225 .iter()
15226 .zip(&y_b)
15227 .map(|(p, q)| (p - q).abs())
15228 .fold(0.0f32, f32::max);
15229 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
15230 }
15231
15232 #[test]
15240 fn q4t_matmat_blocked_matches_per_row() {
15241 let (rows, cols, b) = (16usize, 64usize, 9usize);
15242 let gpr = cols / GROUP_SIZE;
15243 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
15244 for r in 0..rows {
15245 for g in 0..gpr {
15246 let t = (r * gpr + g) * Q4_TILE;
15247 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
15248 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15249 for k in 0..16 {
15250 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15251 }
15252 }
15253 }
15254 let x: Vec<f32> = (0..b * cols)
15255 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15256 .collect();
15257 let mut y_blk = vec![0f32; b * rows];
15258 let mut y_row = vec![0f32; b * rows];
15259 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
15260 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
15261 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
15262 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
15263 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
15264 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
15265 }
15266
15267 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
15274 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
15275 let gpr = cols / GROUP_SIZE;
15276 let stride = q4tp_code_stride(gpr);
15277 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
15278 let mut b = vec![0u8; codes_off + rows * stride];
15279 for r in 0..rows {
15280 for g in 0..gpr {
15281 let t = (r * gpr + g) * Q4TP_NIB;
15282 for k in 0..16 {
15283 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15284 }
15285 }
15286 let lo = -6.0 - 0.03 * (r % 17) as f32;
15287 let step = 0.01 + 0.004 * (r % 11) as f32;
15288 let p = params_off + r * 4;
15289 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
15290 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
15291 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
15292 for g in 0..gpr {
15293 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
15294 }
15295 }
15296 b
15297 }
15298
15299 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
15303 let gpr = cols / GROUP_SIZE;
15304 let v = Q4tpView::new(bytes, rows, cols);
15305 let mut out = vec![0u8; rows * gpr * Q4_TILE];
15306 let mut sc = vec![0f32; gpr];
15307 for r in 0..rows {
15308 v.scales_into(r, gpr, &mut sc);
15309 for g in 0..gpr {
15310 let t = (r * gpr + g) * Q4_TILE;
15311 let s = sc[g];
15312 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15313 let src = (r * gpr + g) * Q4TP_NIB;
15314 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
15315 }
15316 }
15317 out
15318 }
15319
15320 #[test]
15326 fn q4tp_exact_path_matches_dequant_reference() {
15327 let (rows, cols) = (256usize, 512usize);
15328 let gpr = cols / GROUP_SIZE;
15329 let bytes = synth_q4tp(rows, cols);
15330 let mut w = vec![0f32; rows * cols];
15331 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15332
15333 let x: Vec<f32> = (0..cols)
15334 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15335 .collect();
15336 let v = Q4tpView::new(&bytes, rows, cols);
15337 let mut sc = vec![0f32; gpr];
15338 for r in 0..rows {
15339 v.scales_into(r, gpr, &mut sc);
15340 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
15341 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
15342 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
15346 assert!(
15347 (got - want).abs() <= 1e-5 * mag,
15348 "row {r}: kernel {got} vs dequant {want}"
15349 );
15350 }
15351 }
15352
15353 #[test]
15359 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
15360 let (rows, cols) = (256usize, 512usize);
15361 let bytes = synth_q4tp(rows, cols);
15362 let twin = q4tp_as_q4t(&bytes, rows, cols);
15363 let x: Vec<f32> = (0..cols)
15364 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15365 .collect();
15366
15367 let mut got = vec![0f32; rows];
15368 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
15369 let mut want = vec![0f32; rows];
15370 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
15371
15372 let mut w = vec![0f32; rows * cols];
15375 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15376 for r in 0..rows {
15377 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
15378 assert!(
15379 (got[r] - want[r]).abs() <= 1e-3 * mag,
15380 "row {r}: q4tp {} vs q4t {}",
15381 got[r],
15382 want[r]
15383 );
15384 }
15385 }
15386
15387 #[test]
15392 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
15393 let (rows, cols, b) = (256usize, 512usize, 5usize);
15394 let bytes = synth_q4tp(rows, cols);
15395 let twin = q4tp_as_q4t(&bytes, rows, cols);
15396 let xs: Vec<f32> = (0..b * cols)
15397 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
15398 .collect();
15399
15400 let mut got = vec![0f32; b * rows];
15401 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
15402 let mut want = vec![0f32; b * rows];
15403 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
15404
15405 let mut w = vec![0f32; rows * cols];
15406 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15407 for t in 0..b {
15408 for r in 0..rows {
15409 let mag: f32 = (0..cols)
15410 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
15411 .sum();
15412 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
15413 assert!(
15414 (g - wa).abs() <= 1e-3 * mag,
15415 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
15416 );
15417 }
15418 }
15419 }
15420
15421 #[test]
15422 fn q4tp_matvec2_matches_the_single_stream_kernel() {
15423 let (rows, cols) = (128usize, 256usize);
15424 let gpr = cols / GROUP_SIZE;
15425 let bytes = synth_q4tp(rows, cols);
15426 let xs: Vec<f32> = (0..2 * cols)
15427 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
15428 .collect();
15429
15430 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
15431 q4tp_matvec2(
15432 &bytes,
15433 &xs[..cols],
15434 &xs[cols..],
15435 rows,
15436 cols,
15437 &mut o1,
15438 &mut o2,
15439 None,
15440 );
15441
15442 let v = Q4tpView::new(&bytes, rows, cols);
15445 let mut sc = vec![0f32; gpr];
15446 for r in 0..rows {
15447 v.scales_into(r, gpr, &mut sc);
15448 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
15449 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
15450 }
15451 }
15452
15453 #[test]
15460 fn q4tp_matvec_keeps_pace_with_q4t() {
15461 let (rows, cols) = (4096usize, 3072usize);
15462 let bytes = synth_q4tp(rows, cols);
15463 let twin = q4tp_as_q4t(&bytes, rows, cols);
15464 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
15465 let mut o = vec![0f32; rows];
15466 let n = 12;
15467 let mut best = (f64::MAX, f64::MAX);
15468 for _ in 0..3 {
15471 let t0 = std::time::Instant::now();
15472 for _ in 0..n {
15473 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
15474 }
15475 best.0 = best.0.min(t0.elapsed().as_secs_f64());
15476 let t0 = std::time::Instant::now();
15477 for _ in 0..n {
15478 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
15479 }
15480 best.1 = best.1.min(t0.elapsed().as_secs_f64());
15481 }
15482 let ratio = best.1 / best.0;
15483 println!(
15484 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
15485 best.0 * 1e3 / n as f64,
15486 best.1 * 1e3 / n as f64
15487 );
15488 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
15489 }
15490
15491 #[cfg(target_os = "macos")]
15492 #[test]
15493 fn q4t_matmat_accel_matches_dequant_reference() {
15494 if !accel_gemm_enabled() {
15495 return; }
15497 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
15499 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
15500 for r in 0..rows {
15501 for g in 0..gpr {
15502 let t = (r * gpr + g) * Q4_TILE;
15503 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
15504 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15505 for k in 0..16 {
15506 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15507 }
15508 }
15509 }
15510 let x: Vec<f32> = (0..b * cols)
15511 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15512 .collect();
15513 let mut got = vec![0f32; b * rows];
15514 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
15515 let mut w = vec![0f32; rows * cols];
15517 for r in 0..rows {
15518 for g in 0..gpr {
15519 let t = (r * gpr + g) * Q4_TILE;
15520 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
15521 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
15522 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
15523 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
15524 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
15525 }
15526 }
15527 }
15528 for bi in 0..b {
15529 for r in 0..rows {
15530 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
15531 let d = (got[bi * rows + r] - want).abs();
15532 assert!(
15533 d <= want.abs().max(1.0) * 1e-4,
15534 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
15535 got[bi * rows + r]
15536 );
15537 }
15538 }
15539 }
15540
15541 #[test]
15542 fn q4matvec_matches_full_dequant() {
15543 let (rows, cols) = (8, 64);
15544 let groups = rows * cols / GROUP_SIZE;
15545 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15547 for i in 0..groups * 16 {
15548 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15549 }
15550 for g in 0..groups {
15551 let s = 0.01 + 0.003 * g as f32;
15552 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15553 }
15554 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15555
15556 let mut reference = vec![0.0f32; rows * cols];
15557 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
15558 let mut expect = vec![0.0f32; rows];
15559 for r in 0..rows {
15560 expect[r] = reference[r * cols..(r + 1) * cols]
15561 .iter()
15562 .zip(&x)
15563 .map(|(w, xv)| w * xv)
15564 .sum();
15565 }
15566
15567 let mut got = vec![0.0f32; rows];
15568 q4matvec(&bytes, &x, rows, cols, &mut got, None);
15569 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
15573 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
15574 for r in 0..rows {
15575 assert!(
15576 (got[r] - expect[r]).abs() < tol * scale,
15577 "row {r}: {} vs {}",
15578 got[r],
15579 expect[r]
15580 );
15581 }
15582 }
15583
15584 #[test]
15587 fn vbitmatvec2_equals_two_singles() {
15588 let (rows, cols) = (6, 64);
15589 let ng = cols / GROUP_SIZE;
15590 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
15591 let mut bytes = bits.clone();
15592 for g in 0..rows * ng {
15593 let s = 0.02 + 0.001 * g as f32;
15594 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15595 }
15596 for r in 0..rows {
15597 let b = bits[r] as usize;
15598 let (mut acc, mut nb) = (0u64, 0usize);
15599 let mut rowbytes = Vec::new();
15600 for i in 0..cols {
15601 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
15602 acc = (acc << b) | v;
15603 nb += b;
15604 while nb >= 8 {
15605 nb -= 8;
15606 rowbytes.push(((acc >> nb) & 0xFF) as u8);
15607 }
15608 }
15609 if nb > 0 {
15610 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15611 }
15612 bytes.extend_from_slice(&rowbytes);
15613 }
15614 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
15615 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
15616 let offsets = vbit_row_offsets(&bytes, rows, cols);
15617
15618 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15619 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
15620 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
15621 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
15622 vbitmatvec2(
15623 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
15624 );
15625 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
15626 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
15627 }
15628
15629 #[test]
15631 fn q4matvec2_equals_two_singles() {
15632 let (rows, cols) = (8, 128);
15633 let groups = rows * cols / GROUP_SIZE;
15634 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15635 for i in 0..groups * 16 {
15636 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15637 }
15638 for g in 0..groups {
15639 let s = 0.01 + 0.003 * g as f32;
15640 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15641 }
15642 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15645 x1[9] = 250.0;
15646 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
15647
15648 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15649 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
15650 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
15651 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
15652 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
15653 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
15654 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
15655 }
15656
15657 #[test]
15660 fn matvec_many_equals_separate_matvecs() {
15661 use crate::pool::Pool;
15662 let (r1, r2, cols) = (300, 200, 64);
15663 let mk = |salt: usize, rows: usize| {
15664 QTensor::from_f32(
15665 (0..rows * cols)
15666 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
15667 .collect(),
15668 rows,
15669 cols,
15670 )
15671 };
15672 let (a, b) = (mk(1, r1), mk(5, r2));
15673 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
15674 let pool = Pool::new(3);
15675
15676 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
15677 a.matvec(&x, &mut ea, Some(&pool));
15678 b.matvec(&x, &mut eb, Some(&pool));
15679 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
15680 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15681 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
15682 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
15683 }
15684
15685 #[test]
15690 fn q4tp_matvec_many_equals_separate_matvecs() {
15691 use crate::pool::Pool;
15692 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15693
15694 let (r1, r2, cols) = (300usize, 200usize, 64usize);
15695 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15696 "arch_name": "tiny-q4tp",
15697 "hidden_size": cols,
15698 "intermediate_size": cols * 2,
15699 "num_layers": 1,
15700 "num_attention_heads": 2,
15701 "num_kv_heads": 1,
15702 "head_dim": 32,
15703 "vocab_size": r1,
15704 "layer_types": ["FullAttention"],
15705 "rms_norm_eps": 1e-6,
15706 "max_position_embeddings": 8,
15707 "linear_conv_kernel_dim": 0,
15708 "linear_num_key_heads": 0,
15709 "linear_num_value_heads": 0
15710 }))
15711 .unwrap();
15712 let header = CmfHeader {
15713 format: "cmf".into(),
15714 version: CMF_VERSION,
15715 arch,
15716 quant_type: QuantType::Q4Block,
15717 provenance: None,
15718 tokenizer_config: None,
15719 section_hashes: None,
15720 skills: Vec::new(),
15721 shard: None,
15722 calibration: None,
15723 routing: None,
15724 genome: None,
15725 lineage: Vec::new(),
15726 router: None,
15727 segments: Vec::new(),
15728 };
15729 let specs = [
15730 TensorSpec {
15731 name: "q".into(),
15732 dtype: TensorDtype::Q4TiledP,
15733 shape: vec![r1, cols],
15734 data: synth_q4tp(r1, cols),
15735 },
15736 TensorSpec {
15737 name: "kv".into(),
15738 dtype: TensorDtype::Q4TiledP,
15739 shape: vec![r2, cols],
15740 data: synth_q4tp(r2, cols),
15741 },
15742 ];
15743 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
15744 std::fs::create_dir_all(&dir).unwrap();
15745 let path = dir.join("m.cmf");
15746 CmfModel::write(&path, &header, &specs, None, None).unwrap();
15747 let model = Arc::new(CmfModel::open(&path).unwrap());
15748 let (a, b) = (
15749 QTensor::from_model(&model, "q").unwrap(),
15750 QTensor::from_model(&model, "kv").unwrap(),
15751 );
15752 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
15753 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
15754 let x: Vec<f32> = (0..cols)
15755 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
15756 .collect();
15757 let pool = Pool::new(3);
15758 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15759 a.matvec(&x, &mut ea, Some(&pool));
15760 b.matvec(&x, &mut eb, Some(&pool));
15761 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15762 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15763 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
15764 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
15765 let _ = std::fs::remove_dir_all(&dir);
15766 }
15767
15768 #[cfg(target_arch = "x86_64")]
15774 #[test]
15775 fn fast_expert_rows_match_the_row_kernels() {
15776 use crate::pool::Pool;
15777 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15778
15779 for (rows, cols) in [(32usize, 2560usize), (16, 640), (48, 192), (16, 64)] {
15781 let bytes = synth_q4tp(rows, cols);
15782 let v = Q4tpView::new(&bytes, rows, cols);
15783 let gpr = cols / GROUP_SIZE;
15784 for r0 in (0..rows).step_by(16) {
15785 let Some(got) = moe_fast::block_scales(&v, r0, gpr) else {
15786 eprintln!("fast_expert_rows: no AVX-512 VNNI here, skipped");
15787 return;
15788 };
15789 for i in 0..16 {
15790 let mut want = vec![0f32; gpr];
15791 v.scales_into(r0 + i, gpr, &mut want);
15792 assert_eq!(
15793 want.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15794 got[i * gpr..(i + 1) * gpr].iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15795 "scales of row {} ({rows}x{cols})",
15796 r0 + i
15797 );
15798 }
15799 }
15800 }
15801
15802 let (h, inter, ne) = (192usize, 640usize, 5usize);
15803 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15804 "arch_name": "tiny-q4tp-moe",
15805 "hidden_size": h,
15806 "intermediate_size": inter,
15807 "num_layers": 1,
15808 "num_attention_heads": h / 32,
15809 "num_kv_heads": 1,
15810 "head_dim": 32,
15811 "vocab_size": 8,
15812 "layer_types": ["FullAttention"],
15813 "rms_norm_eps": 1e-6,
15814 "max_position_embeddings": 8,
15815 "linear_conv_kernel_dim": 0,
15816 "linear_num_key_heads": 0,
15817 "linear_num_value_heads": 0
15818 }))
15819 .unwrap();
15820 let header = CmfHeader {
15821 format: "cmf".into(),
15822 version: CMF_VERSION,
15823 arch,
15824 quant_type: QuantType::Q4Block,
15825 provenance: None,
15826 tokenizer_config: None,
15827 section_hashes: None,
15828 skills: Vec::new(),
15829 shard: None,
15830 calibration: None,
15831 routing: None,
15832 genome: None,
15833 lineage: Vec::new(),
15834 router: None,
15835 segments: Vec::new(),
15836 };
15837 let mut specs = Vec::new();
15838 for e in 0..ne {
15839 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
15840 .into_iter()
15841 .enumerate()
15842 {
15843 let mut data = synth_q4tp(r, c);
15844 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
15845 .iter_mut()
15846 .enumerate()
15847 {
15848 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
15849 }
15850 specs.push(TensorSpec {
15851 name: format!("{n}{e}"),
15852 dtype: TensorDtype::Q4TiledP,
15853 shape: vec![r, c],
15854 data,
15855 });
15856 }
15857 }
15858 let dir = std::env::temp_dir().join(format!("cmf-moe-fast-{}", std::process::id()));
15859 std::fs::create_dir_all(&dir).unwrap();
15860 let path = dir.join("m.cmf");
15861 CmfModel::write(&path, &header, &specs, None, None).unwrap();
15862 let model = Arc::new(CmfModel::open(&path).unwrap());
15863 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
15864 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
15865 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
15866 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
15867 let b = 6usize;
15868 let mut xs: Vec<f32> = (0..b * h)
15869 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
15870 .collect();
15871 xs[5] = 9.0; let groups: Vec<Vec<usize>> = vec![vec![0, 1, 2, 3, 4], vec![5], vec![], vec![2, 0, 5], vec![1, 3]];
15874 let lens: Vec<usize> = groups.iter().map(|v| v.len()).collect();
15875 let n_pairs: usize = lens.iter().sum();
15876 let pool = Pool::new(3);
15877 let close = |want: &[f32], got: &[f32], rel: f32, what: &str| {
15878 assert_eq!(want.len(), got.len(), "{what}: length");
15879 let top = want.iter().fold(0f32, |m, v| m.max(v.abs()));
15880 assert!(top > 0.0, "{what}: an all-zero reference proves nothing");
15881 for (i, (w, v)) in want.iter().zip(got).enumerate() {
15882 assert!((w - v).abs() <= rel * top, "{what}: [{i}] {w} vs {v} (scale {top})");
15883 }
15884 };
15885
15886 let pairs: Vec<(&QTensor, &QTensor)> = (0..ne).map(|e| (&g[e], &u[e])).collect();
15888 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
15889 assert!(QTensor::moe_gate_up_rows(&pairs, &groups, &xs, &mut gs, Some(&pool)));
15890 let downs: Vec<&QTensor> = d.iter().collect();
15891 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15892 assert!(QTensor::moe_down_rows(&downs, &lens, &gs, &mut ds, Some(&pool)));
15893
15894 let triples: Vec<(&QTensor, &QTensor, &QTensor)> = (0..ne).map(|e| (&g[e], &u[e], &d[e])).collect();
15895 let mut out = vec![0f32; n_pairs * h];
15896 let mut gs_fast = Vec::new();
15897 if !moe_fast::run_stages(&triples, &groups, &xs, &mut out, Some(&pool), None, Some(&mut gs_fast)) {
15898 eprintln!("fast_expert_rows: the int8 activation kernels are off here, skipped");
15899 let _ = std::fs::remove_dir_all(&dir);
15900 return;
15901 }
15902 close(&gs.concat(), &gs_fast, 1e-4, "gate and up");
15903 close(&ds.concat(), &out, 1e-2, "whole experts");
15904
15905 let mut gs_out = gs.clone();
15908 let top = gs.iter().flatten().fold(0f32, |m, v| m.max(v.abs()));
15909 gs_out[0][3] = 40.0 * top;
15910 let mut ds_out: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15911 assert!(QTensor::moe_down_rows(&downs, &lens, &gs_out, &mut ds_out, Some(&pool)));
15912 let mut out2 = vec![0f32; n_pairs * h];
15913 assert!(moe_fast::run_stages(&triples, &groups, &xs, &mut out2, Some(&pool), Some(&gs_out.concat()), None));
15914 close(&ds_out.concat(), &out2, 1e-4, "down");
15915
15916 let mut out3 = vec![0f32; n_pairs * h];
15918 assert!(QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut out3, None));
15919 assert_eq!(
15920 out.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15921 out3.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15922 "pooled and serial runs must agree bit for bit"
15923 );
15924 let mut short = vec![0f32; n_pairs * h - 1];
15925 assert!(!QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut short, Some(&pool)));
15926 let _ = std::fs::remove_dir_all(&dir);
15927 }
15928
15929 #[test]
15935 fn multi_token_moe_rows_equal_single_token_decode() {
15936 use crate::pool::Pool;
15937 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15938
15939 let (h, inter, ne) = (64usize, 128usize, 3usize);
15940 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15941 "arch_name": "tiny-q4tp-moe",
15942 "hidden_size": h,
15943 "intermediate_size": inter,
15944 "num_layers": 1,
15945 "num_attention_heads": 2,
15946 "num_kv_heads": 1,
15947 "head_dim": 32,
15948 "vocab_size": 8,
15949 "layer_types": ["FullAttention"],
15950 "rms_norm_eps": 1e-6,
15951 "max_position_embeddings": 8,
15952 "linear_conv_kernel_dim": 0,
15953 "linear_num_key_heads": 0,
15954 "linear_num_value_heads": 0
15955 }))
15956 .unwrap();
15957 let header = CmfHeader {
15958 format: "cmf".into(),
15959 version: CMF_VERSION,
15960 arch,
15961 quant_type: QuantType::Q4Block,
15962 provenance: None,
15963 tokenizer_config: None,
15964 section_hashes: None,
15965 skills: Vec::new(),
15966 shard: None,
15967 calibration: None,
15968 routing: None,
15969 genome: None,
15970 lineage: Vec::new(),
15971 router: None,
15972 segments: Vec::new(),
15973 };
15974 let mut specs = Vec::new();
15975 for e in 0..ne {
15976 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
15977 .into_iter()
15978 .enumerate()
15979 {
15980 let mut data = synth_q4tp(r, c);
15983 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
15984 .iter_mut()
15985 .enumerate()
15986 {
15987 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
15988 }
15989 specs.push(TensorSpec {
15990 name: format!("{n}{e}"),
15991 dtype: TensorDtype::Q4TiledP,
15992 shape: vec![r, c],
15993 data,
15994 });
15995 }
15996 }
15997 let dir = std::env::temp_dir().join(format!(
15998 "cmf-moe-rows-{}-{}",
15999 std::process::id(),
16000 FLOAT_ACTIVATIONS.get()
16001 ));
16002 std::fs::create_dir_all(&dir).unwrap();
16003 let path = dir.join("m.cmf");
16004 CmfModel::write(&path, &header, &specs, None, None).unwrap();
16005 let model = Arc::new(CmfModel::open(&path).unwrap());
16006 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
16007 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
16008 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
16009 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
16010 let b = 4usize;
16011 let mut xs: Vec<f32> = (0..b * h)
16012 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
16013 .collect();
16014 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
16017 (vec![2, 0], vec![0.6, 0.4]),
16018 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
16019 (vec![1], vec![1.0]),
16020 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
16021 ];
16022 let pool = Pool::new(3);
16023 let mut want = vec![0f32; b * h];
16025 for (tk, (idx, w)) in routes.iter().enumerate() {
16026 let x = &xs[tk * h..(tk + 1) * h];
16027 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
16028 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
16029 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
16030 if FLOAT_ACTIVATIONS.get() {
16031 for (slot, &e) in idx.iter().enumerate() {
16032 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
16033 g[e].matvec(x, &mut gate, Some(&pool));
16034 u[e].matvec(x, &mut up, Some(&pool));
16035 for (v, u) in gate.iter_mut().zip(up) {
16036 *v = (*v / (1.0 + (-*v).exp())) * u;
16037 }
16038 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
16039 }
16040 }
16041 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
16042 assert!(QTensor::moe_down_many(
16043 &downs,
16044 &gs,
16045 w,
16046 &mut want[tk * h..(tk + 1) * h],
16047 Some(&pool)
16048 ));
16049 }
16050 if FLOAT_ACTIVATIONS.get() {
16051 for (tk, (idx, w)) in routes.iter().enumerate() {
16052 let mut scalar = vec![0.0; h];
16053 for (&e, &weight) in idx.iter().zip(w) {
16054 let (mut gate, mut up, mut down) =
16055 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
16056 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
16057 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
16058 for (v, u) in gate.iter_mut().zip(up) {
16059 *v = (*v / (1.0 + (-*v).exp())) * u;
16060 }
16061 d[e].matvec(&gate, &mut down, Some(&pool));
16062 for (v, d) in scalar.iter_mut().zip(down) {
16063 *v += weight * d;
16064 }
16065 }
16066 assert_eq!(
16067 &want[tk * h..(tk + 1) * h],
16068 scalar,
16069 "float many equals scalar experts"
16070 );
16071 }
16072 }
16073 let mut experts: Vec<usize> = Vec::new();
16075 let mut groups: Vec<Vec<usize>> = Vec::new();
16076 for (tk, (idx, _)) in routes.iter().enumerate() {
16077 for &e in idx {
16078 match experts.iter().position(|&x| x == e) {
16079 Some(k) => groups[k].push(tk),
16080 None => {
16081 experts.push(e);
16082 groups.push(vec![tk]);
16083 }
16084 }
16085 }
16086 }
16087 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
16088 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
16089 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
16090 assert!(QTensor::moe_gate_up_rows(
16091 &pairs,
16092 &groups,
16093 &xs,
16094 &mut gs,
16095 Some(&pool)
16096 ));
16097 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
16098 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
16099 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
16100 assert!(QTensor::moe_down_rows(
16101 &downs,
16102 &lens,
16103 &gs,
16104 &mut ds,
16105 Some(&pool)
16106 ));
16107 let slot = |tk: usize, e: usize| {
16108 let k = experts.iter().position(|&x| x == e).unwrap();
16109 groups[..k].iter().map(|g| g.len()).sum::<usize>()
16110 + groups[k].iter().position(|&x| x == tk).unwrap()
16111 };
16112 let mut got = vec![0f32; b * h];
16113 for (tk, (idx, w)) in routes.iter().enumerate() {
16114 for i in 0..h {
16115 let mut acc = 0f32;
16116 for (&e, &we) in idx.iter().zip(w) {
16117 acc += we * ds[slot(tk, e)][i];
16118 }
16119 got[tk * h + i] = acc;
16120 }
16121 }
16122 assert!(want.iter().any(|v| *v != 0.0));
16123 assert_eq!(
16124 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16125 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16126 "multi-token MoE must equal decode bit for bit"
16127 );
16128
16129 let b5 = 5usize;
16132 let x5: Vec<f32> = (0..b5 * h)
16133 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
16134 .collect();
16135 let mut mm = vec![0f32; b5 * inter];
16136 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
16137 for tk in 0..b5 {
16138 let mut mv = vec![0f32; inter];
16139 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
16140 assert_eq!(
16141 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16142 mm[tk * inter..(tk + 1) * inter]
16143 .iter()
16144 .map(|v| v.to_bits())
16145 .collect::<Vec<_>>(),
16146 "row-exact matmat token {tk}"
16147 );
16148 }
16149 let _ = std::fs::remove_dir_all(&dir);
16152 }
16153
16154 #[test]
16162 fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
16163 use crate::pool::Pool;
16164 use std::sync::atomic::Ordering::Relaxed;
16165 let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
16166 Q4TP_ALT.store(2, Relaxed);
16169 let pool = Pool::new(3);
16170 for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
16174 let bytes = synth_q4tp(rows, cols);
16175 let mut xs: Vec<f32> = (0..b * cols)
16176 .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
16177 .collect();
16178 xs[3] = 7.5; let run = |exact: bool| {
16180 let mut out = vec![0f32; b * rows];
16181 q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
16182 out
16183 };
16184 let (fast, exact) = (run(false), run(true));
16185 #[cfg(not(target_arch = "aarch64"))]
16186 let _ = fast;
16187 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
16188 let mut matvecs = vec![0f32; b * rows];
16189 for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
16190 q4tp_matvec(
16191 &bytes,
16192 &xs[bi * cols..(bi + 1) * cols],
16193 rows,
16194 cols,
16195 o,
16196 Some(&pool),
16197 );
16198 }
16199 assert!(matvecs.iter().any(|v| *v != 0.0));
16200 assert_eq!(
16201 bits(&exact),
16202 bits(&matvecs),
16203 "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
16204 );
16205 #[cfg(target_arch = "aarch64")]
16206 {
16207 let gpr = cols / GROUP_SIZE;
16209 let v = Q4tpView::new(&bytes, rows, cols);
16210 let mut old = vec![0f32; b * rows];
16211 let mut sc = vec![0f32; gpr];
16212 let a8w8 = a8w8_enabled();
16213 let blocked = sdot_enabled() && blocked_enabled();
16214 let acts: Vec<SplitAct> = (0..b)
16215 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
16216 .collect();
16217 for r in 0..rows {
16218 v.scales_into(r, gpr, &mut sc);
16219 if !a8w8 {
16220 for bi in 0..b {
16221 let x = &xs[bi * cols..(bi + 1) * cols];
16222 old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
16223 }
16224 continue;
16225 }
16226 let finish = |d: f32, act: &SplitAct| {
16227 let mut acc = d * act.sx;
16228 for &(j, xv) in &act.outliers {
16229 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
16230 acc += w * s * xv;
16231 }
16232 acc
16233 };
16234 let mut bi = 0usize;
16235 while blocked && bi + 4 <= b {
16236 let xs4 = [
16237 acts[bi].xq.as_slice(),
16238 acts[bi + 1].xq.as_slice(),
16239 acts[bi + 2].xq.as_slice(),
16240 acts[bi + 3].xq.as_slice(),
16241 ];
16242 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
16243 for k in 0..4 {
16244 old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
16245 }
16246 bi += 4;
16247 }
16248 for (bi, act) in acts.iter().enumerate().skip(bi) {
16249 let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
16250 old[bi * rows + r] = finish(d, act);
16251 }
16252 }
16253 assert_eq!(
16254 bits(&fast),
16255 bits(&old),
16256 "{rows}x{cols} b={b}: the fast path changed outside row_exact"
16257 );
16258 if blocked {
16261 assert_ne!(
16262 bits(&fast),
16263 bits(&matvecs),
16264 "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
16265 );
16266 }
16267 }
16268 }
16269 Q4TP_ALT.store(0, Relaxed);
16270 }
16271
16272 #[test]
16273 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
16274 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
16275 let active = AtomicUsize::new(0);
16278 counted_row_exact_scope(&active, || {
16279 assert_eq!(active.load(Ordering::Acquire), 1);
16280 counted_row_exact_scope(&active, || {
16281 assert_eq!(active.load(Ordering::Acquire), 2);
16282 });
16283 assert_eq!(active.load(Ordering::Acquire), 1);
16284 });
16285 assert_eq!(active.load(Ordering::Acquire), 0);
16286
16287 let both_entered = Barrier::new(2);
16288 let release_last = Barrier::new(2);
16289 std::thread::scope(|s| {
16290 let first = s.spawn(|| counted_row_exact_scope(&active, || {
16291 both_entered.wait();
16292 }));
16293 let last = s.spawn(|| counted_row_exact_scope(&active, || {
16294 both_entered.wait();
16295 release_last.wait();
16296 }));
16297 first.join().unwrap();
16298 let after_first = active.load(Ordering::Acquire);
16299 release_last.wait();
16300 last.join().unwrap();
16301 assert_eq!(after_first, 1, "second request must remain exact");
16302 });
16303 assert_eq!(active.load(Ordering::Acquire), 0);
16304 let panic = std::panic::catch_unwind(|| {
16305 counted_row_exact_scope(&active, || panic!("scope unwind"));
16306 });
16307 assert!(panic.is_err());
16308 assert_eq!(active.load(Ordering::Acquire), 0);
16309 }
16310
16311 #[test]
16312 #[cfg(target_arch = "x86_64")]
16313 fn q4tp_float_avx2_is_bitwise_scalar() {
16314 if !avx2_enabled() {
16315 return;
16316 }
16317 for cols in [32, 64, 96, 2048, 4096] {
16318 let rows = 9;
16319 let bytes = synth_q4tp(rows, cols);
16320 let v = Q4tpView::new(&bytes, rows, cols);
16321 let gpr = cols / GROUP_SIZE;
16322 let mut sc = vec![0.0; gpr];
16323 for seed in 1..=5 {
16324 let xs: Vec<f32> = (0..cols)
16325 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
16326 .collect();
16327 for r in 0..rows {
16328 v.scales_into(r, gpr, &mut sc);
16329 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
16330 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
16331 assert_eq!(
16332 scalar.to_bits(),
16333 vector.to_bits(),
16334 "cols={cols} row={r} seed={seed}"
16335 );
16336 }
16337 }
16338 }
16339 }
16340
16341 #[test]
16342 fn multi_token_moe_rows_float_equal_single_token_decode() {
16343 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
16344 }
16345
16346 #[test]
16347 fn full_gpu_q8_scope_is_nested_and_thread_local() {
16348 assert!(!FULL_GPU_Q8.get());
16349 let before = gpu_split_frac();
16350 {
16351 let _guard = enter_full_gpu_q8_scope();
16352 assert_eq!(gpu_split_frac(), 1.0);
16353 {
16354 let _nested = enter_full_gpu_q8_scope();
16355 }
16356 assert_eq!(gpu_split_frac(), 1.0);
16357 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
16358 }
16359 assert!(!FULL_GPU_Q8.get());
16360 assert_eq!(gpu_split_frac(), before);
16361 }
16362
16363 #[test]
16364 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
16365 assert!(!FLOAT_ACTIVATIONS.get());
16366 let before = a8w8_enabled();
16367 float_activations_scope(|| {
16368 assert!(!a8w8_enabled());
16369 float_activations_scope(|| assert!(!a8w8_enabled()));
16370 assert!(FLOAT_ACTIVATIONS.get());
16371 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
16372 .join()
16373 .unwrap();
16374 });
16375 assert!(!FLOAT_ACTIVATIONS.get());
16376 assert_eq!(a8w8_enabled(), before);
16377 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
16378 assert!(!FLOAT_ACTIVATIONS.get());
16379 }
16380
16381 #[test]
16384 fn batched_matmat_equals_per_position_matvec() {
16385 let (rows, cols, b) = (8, 64, 5);
16386 let groups = rows * cols / GROUP_SIZE;
16388 let mut q4 = Vec::new();
16389 for i in 0..groups * 16 {
16390 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
16391 }
16392 for g in 0..groups {
16393 q4.extend_from_slice(
16394 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
16395 );
16396 }
16397 let ng = cols / GROUP_SIZE;
16399 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
16400 let mut vb = bits.clone();
16401 for g in 0..rows * ng {
16402 vb.extend_from_slice(
16403 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
16404 );
16405 }
16406 for r in 0..rows {
16407 let bw = bits[r] as usize;
16408 let (mut acc, mut nb) = (0u64, 0usize);
16409 let mut rowbytes = Vec::new();
16410 for i in 0..cols {
16411 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
16412 acc = (acc << bw) | v;
16413 nb += bw;
16414 while nb >= 8 {
16415 nb -= 8;
16416 rowbytes.push(((acc >> nb) & 0xFF) as u8);
16417 }
16418 }
16419 if nb > 0 {
16420 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
16421 }
16422 vb.extend_from_slice(&rowbytes);
16423 }
16424 let offsets = vbit_row_offsets(&vb, rows, cols);
16425
16426 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
16427
16428 let mut got = vec![0f32; b * rows];
16430 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
16431 for bi in 0..b {
16432 let mut expect = vec![0f32; rows];
16433 q4matvec(
16434 &q4,
16435 &xs[bi * cols..(bi + 1) * cols],
16436 rows,
16437 cols,
16438 &mut expect,
16439 None,
16440 );
16441 assert_eq!(
16442 &got[bi * rows..(bi + 1) * rows],
16443 &expect[..],
16444 "q4 batch pos {bi}"
16445 );
16446 }
16447
16448 let mut got = vec![0f32; b * rows];
16450 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
16451 for bi in 0..b {
16452 let mut expect = vec![0f32; rows];
16453 vbitmatvec(
16454 &vb,
16455 &offsets,
16456 &xs[bi * cols..(bi + 1) * cols],
16457 rows,
16458 cols,
16459 &mut expect,
16460 None,
16461 );
16462 assert_eq!(
16463 &got[bi * rows..(bi + 1) * rows],
16464 &expect[..],
16465 "vbit batch pos {bi}"
16466 );
16467 }
16468 }
16469
16470 #[test]
16474 fn q4_tiled_matches_q4_block_bitexact() {
16475 let (rows, cols, b) = (8usize, 128usize, 3usize);
16476 let groups = rows * cols / GROUP_SIZE;
16477 let mut split = Vec::with_capacity(groups * 18);
16478 for i in 0..groups * 16 {
16479 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
16480 }
16481 for g in 0..groups {
16482 split.extend_from_slice(
16483 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
16484 );
16485 }
16486 let (packed, scales) = split.split_at(groups * 16);
16488 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
16489 for g in 0..groups {
16490 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
16491 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
16492 }
16493
16494 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
16495 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
16497
16498 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
16499 q4matvec(&split, &x1, rows, cols, &mut a, None);
16500 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
16501 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
16502
16503 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
16504 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
16505 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
16506 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
16507 assert_eq!(a1, t1);
16508 assert_eq!(a2, t2);
16509
16510 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
16511 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
16512 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
16513 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
16514 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
16515 }
16516
16517 #[test]
16524 fn q4matvec_sdot_outlier_exact() {
16525 let (rows, cols) = (4, 128);
16526 let groups = rows * cols / GROUP_SIZE;
16527 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
16528 for i in 0..groups * 16 {
16529 bytes.push(((i * 11 + 5) % 256) as u8);
16530 }
16531 for g in 0..groups {
16532 let s = 0.02 + 0.002 * g as f32;
16533 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
16534 }
16535 let mut x: Vec<f32> = (0..cols)
16536 .map(|i| match i % 3 {
16537 0 => 1.0,
16538 1 => -1.0,
16539 _ => 0.0,
16540 })
16541 .collect();
16542 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
16545 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
16546 let mut expect = vec![0.0f32; rows];
16547 for r in 0..rows {
16548 expect[r] = reference[r * cols..(r + 1) * cols]
16549 .iter()
16550 .zip(&x)
16551 .map(|(w, xv)| w * xv)
16552 .sum();
16553 }
16554 let mut got = vec![0.0f32; rows];
16555 q4matvec(&bytes, &x, rows, cols, &mut got, None);
16556 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
16557 for r in 0..rows {
16558 assert!(
16559 (got[r] - expect[r]).abs() < 2e-3 * scale,
16560 "row {r}: {} vs {} (outlier term must be exact)",
16561 got[r],
16562 expect[r]
16563 );
16564 }
16565 }
16566
16567 #[test]
16571 fn q1t_matvec_matches_reference() {
16572 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
16573 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
16575 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
16576 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
16578 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
16579 let mut bytes = Vec::new();
16580 for r in 0..rows {
16581 for g in 0..gpr {
16582 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
16583 let mut c = [0u8; 7];
16584 for k in 0..GROUP_SIZE {
16585 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
16587 0
16588 } else {
16589 ((k + r * 3 + g) % 3) as u8 };
16591 cortiq_core::quant::q1t_pack(&mut c, k, code);
16592 }
16593 bytes.extend_from_slice(&c);
16594 }
16595 }
16596 let mut row_ptr = vec![0u32; rows + 1];
16599 for &(idx, _) in &outliers {
16600 row_ptr[idx as usize / cols + 1] += 1;
16601 }
16602 for r in 0..rows {
16603 row_ptr[r + 1] += row_ptr[r];
16604 }
16605 for &p in &row_ptr {
16606 bytes.extend_from_slice(&p.to_le_bytes());
16607 }
16608 for &(idx, v) in &outliers {
16609 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
16610 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
16611 }
16612
16613 let mut refw = vec![0f32; rows * cols];
16614 dequant_q1t(&bytes, rows, cols, &mut refw);
16615 let x: Vec<f32> = (0..cols)
16618 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
16619 .collect();
16620 let mut expect = vec![0f32; rows];
16621 for r in 0..rows {
16622 let mut a = 0.0f32;
16623 for j in 0..cols {
16624 a += refw[r * cols + j] * x[j];
16625 }
16626 expect[r] = a;
16627 }
16628 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
16629 let mut got = vec![0f32; rows];
16630 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
16631 for r in 0..rows {
16632 assert!(
16633 (got[r] - expect[r]).abs() < tol(expect[r]),
16634 "row {r}: {} vs {}",
16635 got[r],
16636 expect[r]
16637 );
16638 }
16639 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
16641 let mut gm = vec![0f32; 2 * rows];
16642 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
16643 for r in 0..rows {
16644 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
16645 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
16646 }
16647 let xb: Vec<f32> = (0..cols)
16651 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
16652 .collect();
16653 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16654 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
16655 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
16656 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16657 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
16658 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
16659 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
16660 }
16661
16662 #[test]
16665 fn q1t_matvec2_odd_gpr_matches_singles() {
16666 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
16667 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
16669 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
16670 for r in 0..rows {
16671 for g in 0..gpr {
16672 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
16673 let mut c = [0u8; 7];
16674 for k in 0..GROUP_SIZE {
16675 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
16676 }
16677 bytes.extend_from_slice(&c);
16678 }
16679 }
16680 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
16681 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
16682 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16683 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16684 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
16685 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16686 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16687 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
16688 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
16689 }
16690
16691 #[test]
16695 #[ignore]
16696 fn q1t_matvec2_speed() {
16697 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
16698 use std::time::Instant;
16699 let (rows, cols) = (8192usize, 4096usize);
16700 let gpr = cols / GROUP_SIZE;
16701 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
16702 for r in 0..rows {
16703 for g in 0..gpr {
16704 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
16705 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
16706 let mut c = [0u8; 7];
16707 for k in 0..GROUP_SIZE {
16708 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
16709 }
16710 bytes.extend_from_slice(&c);
16711 }
16712 }
16713 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
16714 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
16715 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16716 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16717 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16719 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16720 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
16721 for _ in 0..8 {
16722 let t0 = Instant::now();
16723 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16724 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
16725 let t1 = Instant::now();
16726 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16727 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
16728 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
16729 }
16730 assert_eq!(p1, s1);
16731 assert_eq!(p2, s2);
16732 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
16733 }
16734
16735 #[test]
16739 #[ignore]
16740 fn q1t_matvec_speed() {
16741 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
16742 use std::time::Instant;
16743 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
16745 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
16746 for r in 0..rows {
16747 for g in 0..gpr {
16748 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
16749 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
16750 let mut c = [0u8; 7];
16751 for k in 0..GROUP_SIZE {
16752 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
16753 }
16754 bytes.extend_from_slice(&c);
16755 }
16756 }
16757 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
16759 let mut idx = 0usize;
16760 while idx < n {
16761 row_ptr[idx / cols + 1] += 1;
16762 idx += stride;
16763 }
16764 for r in 0..rows {
16765 row_ptr[r + 1] += row_ptr[r];
16766 }
16767 for &p in &row_ptr {
16768 bytes.extend_from_slice(&p.to_le_bytes());
16769 }
16770 let mut idx = 0usize;
16771 while idx < n {
16772 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
16773 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
16774 idx += stride;
16775 }
16776 let x: Vec<f32> = (0..cols)
16779 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
16780 .collect();
16781 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
16782
16783 let slow = |out: &mut [f32]| {
16785 let mut buf = vec![0f32; cols];
16786 for r in 0..rows {
16787 for g in 0..gpr {
16788 let off = (r * gpr + g) * Q1T_TILE;
16789 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
16790 let codes = &bytes[off + 2..off + Q1T_TILE];
16791 for k in 0..GROUP_SIZE {
16792 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
16793 1 => s,
16794 2 => -s,
16795 _ => 0.0,
16796 };
16797 }
16798 }
16799 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
16800 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
16801 }
16802 };
16803 let iters = 5;
16804 let mut a = vec![0f32; rows];
16805 slow(&mut a); let t = Instant::now();
16807 for _ in 0..iters {
16808 slow(&mut a);
16809 }
16810 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16811
16812 let mut b = vec![0f32; rows];
16813 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
16815 for _ in 0..iters {
16816 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
16817 }
16818 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16819
16820 for r in 0..rows {
16821 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
16822 }
16823 println!(
16824 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
16825 slow_ms / fast_ms
16826 );
16827 }
16828}
16829
16830#[cfg(test)]
16831mod row_span_tests {
16832 use cortiq_core::TensorDtype;
16833
16834 #[test]
16838 fn q4tp_row_spans_cover_the_row() {
16839 let (rows, cols) = (37usize, 160usize);
16840 let gpr = cols / 32;
16841 let total = cortiq_core::quant::expected_nbytes(TensorDtype::Q4TiledP, &[rows, cols]).unwrap();
16842 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16843 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
16844 let lo = cortiq_core::quant::f32_to_f16(-4.0);
16845 let step = cortiq_core::quant::f32_to_f16(0.1);
16846 for r in 0..rows {
16847 let o = params_off + r * 4;
16848 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16849 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16850 }
16851 for r in [0usize, 1, 17, rows - 1] {
16852 let spans = super::row_span_offsets(TensorDtype::Q4TiledP, rows, cols, r).unwrap();
16853 assert_eq!(spans.iter().map(|s| s.1).sum::<usize>(), total / rows);
16854 let mut keep = vec![false; total];
16855 for &(o, n) in &spans {
16856 keep[o..o + n].fill(true);
16857 }
16858 let other: Vec<u8> =
16859 bytes.iter().zip(&keep).map(|(&b, &k)| if k { b } else { 0x5A }).collect();
16860 let (a, b) = (super::Q4tpView::new(&bytes, rows, cols), super::Q4tpView::new(&other, rows, cols));
16861 let (mut sa, mut sb) = (vec![0f32; gpr], vec![0f32; gpr]);
16862 a.scales_into(r, gpr, &mut sa);
16863 b.scales_into(r, gpr, &mut sb);
16864 assert_eq!(sa, sb, "row {r} scales");
16865 let nib = r * gpr * super::Q4TP_NIB..(r + 1) * gpr * super::Q4TP_NIB;
16866 assert_eq!(a.nib[nib.clone()], b.nib[nib], "row {r} nibbles");
16867 }
16868 assert!(super::row_span_offsets(TensorDtype::Q4TiledP, rows, cols, rows).is_none());
16869 assert!(super::row_span_offsets(TensorDtype::Q8_2f, rows, cols, 0).is_none());
16870 let t = super::row_span_offsets(TensorDtype::Q4Tiled, rows, cols, 3).unwrap();
16871 assert_eq!(t[0], (3 * gpr * super::Q4_TILE, gpr * super::Q4_TILE));
16872 }
16873}
16874
16875#[cfg(test)]
16876mod gemm_bench {
16877 #[test]
16888 #[ignore]
16889 fn q4tp_matmat_throughput() {
16890 let _alt = super::Q4TP_ALT_TEST_LOCK
16891 .lock()
16892 .unwrap_or_else(|e| e.into_inner());
16893 let b: usize = std::env::var("CMF_BENCH_B")
16897 .ok()
16898 .and_then(|v| v.parse().ok())
16899 .unwrap_or(296);
16900 let (rows, cols) = (9216usize, 2304usize);
16901 let (_, _, _) = (rows, cols, b);
16902 let total =
16903 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
16904 .unwrap();
16905 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16910 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16911 let lo = cortiq_core::quant::f32_to_f16(-4.0);
16912 let step = cortiq_core::quant::f32_to_f16(0.1);
16913 for r in 0..rows {
16914 let o = params_off + r * 4;
16915 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16916 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16917 }
16918 let _ = codes_off;
16919 let xs: Vec<f32> = (0..b * cols)
16920 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
16921 .collect();
16922 let mut out = vec![0f32; b * rows];
16923 let pool = crate::pool::Pool::from_env();
16924 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16930 let reps: usize = std::env::var("CMF_BENCH_REPS")
16931 .ok()
16932 .and_then(|v| v.parse().ok())
16933 .unwrap_or(10);
16934 let mut best = [f64::MAX; 2];
16935 let mut sums = [0f32; 2];
16936 for _ in 0..reps {
16937 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
16938 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
16939 let t = std::time::Instant::now();
16940 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16941 best[k] = best[k].min(t.elapsed().as_secs_f64());
16942 sums[k] = out.iter().take(64).sum::<f32>();
16943 }
16944 }
16945 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
16946 for (k, name) in ["previous", "tuned "].iter().enumerate() {
16947 println!(
16948 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
16949 best[k] * 1e3,
16950 flops / best[k] / 1e9,
16951 sums[k]
16952 );
16953 }
16954 assert!(
16955 (sums[0] - sums[1]).abs() < 1e-2,
16956 "the tuned kernel changed the result: {} vs {}",
16957 sums[0],
16958 sums[1]
16959 );
16960 }
16961
16962 #[test]
16968 fn q4tp_matmat_blocked_matches_scalar() {
16969 use std::sync::atomic::Ordering::Relaxed;
16970 let _alt = super::Q4TP_ALT_TEST_LOCK
16971 .lock()
16972 .unwrap_or_else(|e| e.into_inner());
16973 for &(rows, cols, b) in &[
16981 (64usize, 128usize, 7usize),
16982 (33, 96, 4),
16983 (16, 256, 9),
16984 (192, 2304, 37),
16985 ] {
16986 let total = cortiq_core::quant::expected_nbytes(
16987 cortiq_core::TensorDtype::Q4TiledP,
16988 &[rows, cols],
16989 )
16990 .unwrap();
16991 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16992 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
16993 let lo = cortiq_core::quant::f32_to_f16(-4.0);
16994 let step = cortiq_core::quant::f32_to_f16(0.1);
16995 for r in 0..rows {
16996 let o = params_off + r * 4;
16997 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16998 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16999 }
17000 let xs: Vec<f32> = (0..b * cols)
17001 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
17002 .collect();
17003 let mut got = vec![0f32; b * rows];
17004 let mut want = vec![0f32; b * rows];
17005 let gpr = cols / 32;
17006 let view = super::Q4tpView::new(&bytes, rows, cols);
17007 let pool = crate::pool::Pool::from_env();
17008 super::Q4TP_ALT.store(2, Relaxed);
17009 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
17010 super::Q4TP_ALT.store(1, Relaxed);
17011 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
17012 super::Q4TP_ALT.store(0, Relaxed);
17013 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
17020 let (mut worst, mut at) = (0f32, 0usize);
17021 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
17022 if (g - w).abs() > worst {
17023 worst = (g - w).abs();
17024 at = i;
17025 }
17026 }
17027 assert!(
17028 worst <= 1e-4 * scale,
17029 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
17030 (scale {scale:.3e}) at cell {at}: {} vs {}",
17031 got[at],
17032 want[at]
17033 );
17034
17035 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
17042 for bi in 0..b {
17043 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
17044 for r in 0..rows {
17045 let mut sc = vec![0f32; gpr];
17046 view.scales_into(r, gpr, &mut sc);
17047 let mut exact = 0f64;
17048 for j in 0..cols {
17049 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
17050 exact += w as f64 * sq as f64 * act.xq[j] as f64;
17051 }
17052 exact *= act.sx as f64;
17053 for &(j, xv) in &act.outliers {
17054 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
17055 exact += w as f64 * sq as f64 * xv as f64;
17056 }
17057 let i = bi * rows + r;
17058 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
17059 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
17060 }
17061 }
17062 println!(
17063 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
17064 per-column {e_scalar:.3e}"
17065 );
17066 assert!(
17071 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
17072 "{rows}x{cols} b={b}: error against f64 too large — blocked \
17073 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
17074 );
17075 }
17076 }
17077
17078 #[test]
17084 #[ignore]
17085 fn q4tp_matmat_row_exact_cost() {
17086 let _alt = super::Q4TP_ALT_TEST_LOCK
17087 .lock()
17088 .unwrap_or_else(|e| e.into_inner());
17089 let pool = crate::pool::Pool::from_env();
17090 let reps: usize = std::env::var("CMF_BENCH_REPS")
17091 .ok()
17092 .and_then(|v| v.parse().ok())
17093 .unwrap_or(30);
17094 for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
17095 let total = cortiq_core::quant::expected_nbytes(
17096 cortiq_core::TensorDtype::Q4TiledP,
17097 &[rows, cols],
17098 )
17099 .unwrap();
17100 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17101 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
17102 let lo = cortiq_core::quant::f32_to_f16(-4.0);
17103 let step = cortiq_core::quant::f32_to_f16(0.1);
17104 for r in 0..rows {
17105 let o = params_off + r * 4;
17106 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17107 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17108 }
17109 for &b in &[2usize, 4, 5, 8] {
17110 let xs: Vec<f32> = (0..b * cols)
17111 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17112 .collect();
17113 let mut out = vec![0f32; b * rows];
17114 let mut best = [f64::MAX; 3];
17115 for _ in 0..reps {
17116 for (k, best_k) in best.iter_mut().enumerate() {
17117 let t = std::time::Instant::now();
17118 match k {
17119 0 | 1 => super::q4tp_matmat_with(
17120 &bytes,
17121 &xs,
17122 b,
17123 rows,
17124 cols,
17125 &mut out,
17126 pool.as_deref(),
17127 k == 1,
17128 ),
17129 _ => {
17130 for (bi, o) in out.chunks_mut(rows).enumerate() {
17131 super::q4tp_matvec(
17132 &bytes,
17133 &xs[bi * cols..(bi + 1) * cols],
17134 rows,
17135 cols,
17136 o,
17137 pool.as_deref(),
17138 );
17139 }
17140 }
17141 }
17142 *best_k = best_k.min(t.elapsed().as_secs_f64());
17143 }
17144 }
17145 println!(
17146 "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
17147 {b} matvecs {:.3} ms",
17148 best[0] * 1e3,
17149 best[1] * 1e3,
17150 best[2] * 1e3
17151 );
17152 }
17153 }
17154 }
17155
17156 #[test]
17160 #[ignore]
17161 fn q4t_matmat_throughput() {
17162 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
17163 let total =
17164 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
17165 .unwrap();
17166 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
17169 let sc = cortiq_core::quant::f32_to_f16(0.02);
17170 for t in bytes.chunks_mut(super::Q4_TILE) {
17171 t[..2].copy_from_slice(&sc.to_le_bytes());
17172 }
17173 let xs: Vec<f32> = (0..b * cols)
17174 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17175 .collect();
17176 let mut out = vec![0f32; b * rows];
17177 let pool = crate::pool::Pool::from_env();
17178 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17179 let reps: usize = std::env::var("CMF_BENCH_REPS")
17180 .ok()
17181 .and_then(|v| v.parse().ok())
17182 .unwrap_or(10);
17183 let mut best = f64::MAX;
17184 for _ in 0..reps {
17185 let t = std::time::Instant::now();
17186 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17187 best = best.min(t.elapsed().as_secs_f64());
17188 }
17189 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
17190 println!(
17191 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
17192 best * 1e3,
17193 flops / best / 1e9,
17194 out.iter().take(64).sum::<f32>()
17195 );
17196 }
17197}
17198
17199#[cfg(all(test, target_arch = "x86_64"))]
17200mod x86_gemm_tests {
17201 use super::x86_gemm_rm;
17202
17203 #[test]
17210 fn q4tp_fma_batch_matches_the_row_exact_cell() {
17211 if !super::q4tp_fma_x86() || !super::a8w8_enabled() {
17212 return;
17213 }
17214 let pool = crate::pool::Pool::from_env();
17215 for &(rows, cols, b) in &[(33usize, 96usize, 7usize), (64, 256, 13), (257, 2560, 37)] {
17216 let total = cortiq_core::quant::expected_nbytes(
17217 cortiq_core::TensorDtype::Q4TiledP,
17218 &[rows, cols],
17219 )
17220 .unwrap();
17221 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17222 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
17223 let lo = cortiq_core::quant::f32_to_f16(-4.0);
17224 let step = cortiq_core::quant::f32_to_f16(0.1);
17225 for r in 0..rows {
17226 let o = params_off + r * 4;
17227 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17228 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17229 }
17230 let mut xs: Vec<f32> = (0..b * cols)
17231 .map(|i| ((i * 37 % 89) as f32 - 44.0) / 44.0)
17232 .collect();
17233 for bi in (0..b).step_by(3) {
17234 xs[bi * cols + (bi * 7) % cols] = 40.0; }
17236 let mut got = vec![0f32; b * rows];
17237 let mut want = vec![0f32; b * rows];
17238 super::q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref(), false);
17239 super::q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref(), true);
17240 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
17241 let worst = got.iter().zip(&want).map(|(g, w)| (g - w).abs()).fold(0f32, f32::max);
17242 assert!(
17243 worst <= 1e-5 * scale,
17244 "{rows}x{cols} b={b}: FMA batch and row-exact cell disagree by {worst:.3e} (scale {scale:.3e})"
17245 );
17246 }
17247 }
17248
17249 #[allow(clippy::too_many_arguments)]
17250 fn naive(
17251 m: usize,
17252 n: usize,
17253 k: usize,
17254 alpha: f32,
17255 a: &[f32],
17256 lda: usize,
17257 b: &[f32],
17258 ldb: usize,
17259 b_rows_are_n: bool,
17260 ) -> Vec<f64> {
17261 let mut c = vec![0f64; m * n];
17262 for i in 0..m {
17263 for j in 0..n {
17264 let mut acc = 0f64;
17265 for p in 0..k {
17266 let bv = if b_rows_are_n { b[j * ldb + p] } else { b[p * ldb + j] };
17267 acc += a[i * lda + p] as f64 * bv as f64;
17268 }
17269 c[i * n + j] = acc * alpha as f64;
17270 }
17271 }
17272 c
17273 }
17274
17275 #[test]
17280 fn x86_gemm_matches_naive_on_odd_shapes() {
17281 for &(m, n, k) in &[
17282 (1usize, 1usize, 1usize),
17283 (4, 16, 256),
17284 (7, 37, 19),
17285 (9, 300, 257),
17286 (64, 1000, 256),
17287 (5, 256, 1001),
17288 (13, 129, 512),
17289 (8, 16, 0),
17290 ] {
17291 for &bt in &[false, true] {
17292 let (lda, ldb, ldc) = (k + 3, if bt { k + 2 } else { n + 5 }, n + 1);
17293 let a: Vec<f32> = (0..m * lda)
17294 .map(|i| ((i * 7919) % 1000) as f32 / 500.0 - 1.0)
17295 .collect();
17296 let brows = if bt { n } else { k.max(1) };
17297 let b: Vec<f32> = (0..brows * ldb)
17298 .map(|i| ((i * 104_729) % 997) as f32 / 498.5 - 1.0)
17299 .collect();
17300 let mut c = vec![f32::NAN; m * ldc];
17301 x86_gemm_rm(m, n, k, 0.37, &a, lda, &b, ldb, bt, &mut c, ldc);
17302 let r = naive(m, n, k, 0.37, &a, lda, &b, ldb, bt);
17303 for i in 0..m {
17304 for j in 0..n {
17305 let (x, y) = (c[i * ldc + j] as f64, r[i * n + j]);
17306 assert!(
17307 (x - y).abs() <= 1e-3 * (1.0 + y.abs()),
17308 "m{m} n{n} k{k} bt{bt} [{i},{j}]: {x} vs {y}"
17309 );
17310 }
17311 assert!(c[i * ldc + n].is_nan(), "m{m} n{n} k{k}: padding column written");
17312 }
17313 }
17314 }
17315 }
17316}