1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16 GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17 q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::cell::UnsafeCell;
21use std::sync::Arc;
22
23pub enum QTensor {
24 F32 {
25 data: Vec<f32>,
26 rows: usize,
27 cols: usize,
28 },
29 Mapped {
30 model: Arc<CmfModel>,
31 idx: usize,
33 dtype: TensorDtype,
34 rows: usize,
35 cols: usize,
36 row_scale: Vec<f32>,
38 col_field: Vec<f32>,
40 vbit_offsets: Vec<usize>,
44 repack: Vec<u8>,
52 },
53 Blr(Box<crate::blr::BlrTensor>),
59}
60
61fn repack_enabled() -> bool {
68 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
69 *ON.get_or_init(|| {
70 std::env::var("CMF_REPACK")
71 .map(|v| v == "1")
72 .unwrap_or(cfg!(target_os = "android"))
73 })
74}
75
76fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
80 #[cfg(target_arch = "aarch64")]
81 let arch_ok = sdot_enabled();
82 #[cfg(not(target_arch = "aarch64"))]
83 let arch_ok = false;
84 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
85 return Vec::new();
86 }
87 q8_repack_layout(bytes, rows, cols)
88}
89
90fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
93 let groups = rows / 4;
94 let mut rep = vec![0u8; groups * 4 * cols];
95 for g in 0..groups {
96 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
97 for c in 0..cols / 16 {
98 for lane in 0..4 {
99 let src = (g * 4 + lane) * cols + c * 16;
100 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
101 .copy_from_slice(&bytes[src..src + 16]);
102 }
103 }
104 }
105 rep
106}
107
108fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
111 let ng = cols / GROUP_SIZE;
112 let bits = &bytes[..rows];
113 let mut offsets = Vec::with_capacity(rows + 1);
114 let mut off = rows + rows * ng * 2;
115 for r in 0..rows {
116 offsets.push(off);
117 off += (cols * bits[r] as usize).div_ceil(8);
118 }
119 offsets.push(off);
120 offsets
121}
122
123fn blocked_enabled() -> bool {
129 use std::sync::atomic::Ordering::Relaxed;
130 match BLOCKED_OVERRIDE.load(Relaxed) {
131 1 => false,
132 2 => true,
133 _ => {
134 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
135 *ON.get_or_init(|| {
136 std::env::var("CMF_X86_BLOCKED")
137 .map(|v| v != "0")
138 .unwrap_or(true)
139 })
140 }
141 }
142}
143
144static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
145
146pub fn set_blocked_override(on: Option<bool>) {
153 let v = match on {
154 None => 0,
155 Some(false) => 1,
156 Some(true) => 2,
157 };
158 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
159}
160
161fn gpu_lmhead_enabled() -> bool {
162 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163 *ON.get_or_init(|| {
164 std::env::var("CMF_GPU_LMHEAD")
165 .map(|v| v != "0")
166 .unwrap_or(true)
167 })
168}
169
170fn gpu_split_frac() -> f32 {
171 if FULL_GPU_Q8.get() {
174 return 1.0;
175 }
176 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
177 *F.get_or_init(|| {
178 std::env::var("CMF_GPU_SPLIT")
179 .ok()
180 .and_then(|v| v.parse::<f32>().ok())
181 .unwrap_or(0.5)
182 .clamp(0.0, 1.0)
183 })
184}
185
186impl QTensor {
187 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
188 debug_assert_eq!(data.len(), rows * cols);
189 Self::F32 { data, rows, cols }
190 }
191
192 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
195 let idx = model
198 .tensor_index(name)
199 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
200 let entry = &model.tensors[idx];
201 if entry.shape.len() != 2 {
202 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
203 }
204 let (rows, cols) = (entry.shape[0], entry.shape[1]);
205 let bytes = model.entry_bytes(entry);
206
207 match entry.dtype {
208 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
209 let n = rows * cols;
210 let scales_off = n;
211 let row_scale: Vec<f32> = (0..rows)
212 .map(|o| {
213 f16_to_f32(u16::from_le_bytes([
214 bytes[scales_off + o * 2],
215 bytes[scales_off + o * 2 + 1],
216 ]))
217 })
218 .collect();
219 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
220 let col_off = n + rows * 2;
221 (0..cols)
222 .map(|i| {
223 f16_to_f32(u16::from_le_bytes([
224 bytes[col_off + i * 2],
225 bytes[col_off + i * 2 + 1],
226 ]))
227 })
228 .collect()
229 } else {
230 Vec::new()
231 };
232 Ok(Self::Mapped {
233 model: model.clone(),
234 idx,
235 dtype: entry.dtype,
236 rows,
237 cols,
238 row_scale,
239 col_field,
240 vbit_offsets: Vec::new(),
241 repack: q8_repack(bytes, rows, cols),
242 })
243 }
244 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
246 model: model.clone(),
247 idx,
248 dtype: entry.dtype,
249 rows,
250 cols,
251 row_scale: Vec::new(),
252 col_field: Vec::new(),
253 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
254 repack: Vec::new(),
255 }),
256 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
260 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
261 let offsets: Vec<usize> = (0..=rows)
262 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
263 .collect();
264 Ok(Self::Mapped {
265 model: model.clone(),
266 idx,
267 dtype: entry.dtype,
268 rows,
269 cols,
270 row_scale: Vec::new(),
271 col_field: Vec::new(),
272 vbit_offsets: offsets,
273 repack: Vec::new(),
274 })
275 }
276 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
282 model: model.clone(),
283 idx,
284 dtype: entry.dtype,
285 rows,
286 cols,
287 row_scale: Vec::new(),
288 col_field: Vec::new(),
289 vbit_offsets: Vec::new(),
290 repack: Vec::new(),
291 }),
292 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
295 model: model.clone(),
296 idx,
297 dtype: entry.dtype,
298 rows,
299 cols,
300 row_scale: Vec::new(),
301 col_field: Vec::new(),
302 vbit_offsets: Vec::new(),
303 repack: Vec::new(),
304 }),
305 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
307 model: model.clone(),
308 idx,
309 dtype: entry.dtype,
310 rows,
311 cols,
312 row_scale: Vec::new(),
313 col_field: Vec::new(),
314 vbit_offsets: Vec::new(),
315 repack: Vec::new(),
316 }),
317 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
318 model: model.clone(),
319 idx,
320 dtype: entry.dtype,
321 rows,
322 cols,
323 row_scale: Vec::new(),
324 col_field: Vec::new(),
325 vbit_offsets: Vec::new(),
326 repack: Vec::new(),
327 }),
328 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
331 model: model.clone(),
332 idx,
333 dtype: entry.dtype,
334 rows,
335 cols,
336 row_scale: Vec::new(),
337 col_field: Vec::new(),
338 vbit_offsets: Vec::new(),
339 repack: Vec::new(),
340 }),
341 TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
342 model: model.clone(),
343 idx,
344 dtype: entry.dtype,
345 rows,
346 cols,
347 row_scale: Vec::new(),
348 col_field: Vec::new(),
349 vbit_offsets: Vec::new(),
350 repack: Vec::new(),
351 }),
352 TensorDtype::Q1T if cols % GROUP_SIZE == 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::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
369 Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
370 model, idx,
371 )?)))
372 }
373 _ => {
375 let mut data = vec![0.0f32; rows * cols];
376 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
377 Ok(Self::from_f32(data, rows, cols))
378 }
379 }
380 }
381
382 pub fn is_blr(&self) -> bool {
385 matches!(self, Self::Blr(_))
386 }
387
388 pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
394 match self {
395 Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
396 _ => None,
397 }
398 }
399
400 pub(crate) fn graph_weight_or_blr(
404 &self,
405 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
406 self.graph_weight()
407 .or_else(|| self.graph_weight_descriptor())
408 .or_else(|| {
409 self.graph_blr()
410 .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
411 })
412 }
413
414 pub(crate) fn is_q1(&self) -> bool {
418 matches!(
419 self,
420 Self::Mapped {
421 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
422 ..
423 }
424 )
425 }
426
427 fn bin_w(&self) -> Option<BinW<'_>> {
430 match self {
431 Self::Mapped {
432 dtype: TensorDtype::Q1,
433 ..
434 } => Some(BinW::Q1(self.quant_bytes())),
435 Self::Mapped {
436 dtype: TensorDtype::Q1G,
437 rows,
438 cols,
439 ..
440 } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
441 _ => None,
442 }
443 }
444
445 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
448 match self {
449 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
450 _ => None,
451 }
452 }
453
454 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
461 if self.has_prism_contract() {
462 return None;
463 }
464 match self {
465 #[cfg(target_os = "macos")]
466 Self::Mapped {
467 dtype: TensorDtype::Q1T,
468 ..
469 } if !crate::gpu::metal_q1t_enabled() => None,
470 Self::Mapped {
471 idx,
472 dtype:
473 TensorDtype::Q1
474 | TensorDtype::Q1G
475 | TensorDtype::Q1T
476 | TensorDtype::Q4Block
477 | TensorDtype::Q4Tiled
478 | TensorDtype::Q4TiledP
482 | TensorDtype::Q8Row
483 | TensorDtype::Q8_2f,
484 rows,
485 cols,
486 ..
487 } => Some((*idx, *rows, *cols)),
488 _ => None,
489 }
490 }
491
492 #[cfg(target_os = "macos")]
499 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
500 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
501 let name = &model.tensors[idx].name;
502 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
503 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
504 if forward && affine {
505 let e = model.tensors.get(idx)?;
506 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
507 }
508 }
509 self.q1_parts()
510 }
511
512 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
518 if self.has_prism_contract() {
519 return None;
520 }
521 match self {
522 Self::Mapped {
523 idx,
524 dtype: TensorDtype::Q4Tiled,
525 rows,
526 cols,
527 ..
528 } => Some((*idx, *rows, *cols)),
529 _ => None,
530 }
531 }
532
533 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
537 if self.has_prism_contract() {
538 return None;
539 }
540 match self {
541 Self::Mapped {
542 idx,
543 dtype: TensorDtype::Q4TiledP,
544 rows,
545 cols,
546 ..
547 } => Some((*idx, *rows, *cols)),
548 _ => None,
549 }
550 }
551
552 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
557 if self.has_prism_contract() {
558 return None;
559 }
560 match self {
561 Self::Mapped {
562 idx,
563 dtype: TensorDtype::Q8Row,
564 rows,
565 cols,
566 row_scale,
567 col_field,
568 ..
569 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
570 _ => None,
571 }
572 }
573
574 pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
577 if self.has_prism_contract() {
578 return None;
579 }
580 match self {
581 Self::Mapped {
582 idx,
583 dtype: TensorDtype::Q8_2f,
584 rows,
585 cols,
586 row_scale,
587 col_field,
588 ..
589 } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
590 _ => None,
591 }
592 }
593
594 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
598 match self {
599 Self::Mapped { dtype, .. } => Some(*dtype),
600 _ => None,
601 }
602 }
603
604 pub fn model_idx(&self) -> Option<usize> {
609 match self {
610 Self::Mapped { idx, .. } => Some(*idx),
611 _ => None,
612 }
613 }
614
615 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
620 match self {
621 Self::Mapped { model, .. } => Some(model.clone()),
622 _ => None,
623 }
624 }
625
626 pub(crate) fn has_prism_contract(&self) -> bool {
631 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
632 }
633
634 pub fn rows(&self) -> usize {
635 match self {
636 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
637 Self::Blr(t) => t.rows,
638 }
639 }
640
641 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
644 if self.has_prism_contract() {
645 return None;
646 }
647 match self {
648 Self::Mapped {
649 model,
650 idx,
651 dtype: TensorDtype::Q4Tiled,
652 ..
653 } => Some((model, *idx)),
654 _ => None,
655 }
656 }
657
658 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
661 if self.has_prism_contract() {
662 return None;
663 }
664 match self {
665 Self::Mapped {
666 model,
667 idx,
668 dtype: TensorDtype::Q4TiledP,
669 ..
670 } => Some((model, *idx)),
671 _ => None,
672 }
673 }
674
675 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
683 if self.has_prism_contract() {
684 return None;
685 }
686 match self {
687 Self::Mapped {
688 model,
689 idx,
690 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
691 ..
692 } => Some((model, *idx)),
693 _ => None,
694 }
695 }
696
697 pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
700 if self.has_prism_contract() {
701 return None;
702 }
703 match self {
704 Self::Mapped {
705 model,
706 idx,
707 dtype: TensorDtype::Q8_2f,
708 ..
709 } => Some((model, *idx)),
710 _ => None,
711 }
712 }
713
714 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
717 if self.has_prism_contract() {
718 return None;
719 }
720 match self {
721 Self::Mapped {
722 model,
723 idx,
724 dtype: TensorDtype::Q2TiledP,
725 ..
726 } => Some((model, *idx)),
727 _ => None,
728 }
729 }
730
731 pub fn cols(&self) -> usize {
732 match self {
733 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
734 Self::Blr(t) => t.cols,
735 }
736 }
737
738 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
741 if self.has_prism_contract() {
742 return None;
743 }
744 match self {
745 Self::Mapped {
746 model,
747 idx,
748 dtype: TensorDtype::Q1,
749 ..
750 } => Some((model, *idx)),
751 _ => None,
752 }
753 }
754
755 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
766 if self.has_prism_contract() {
767 return None;
768 }
769 self.graph_weight_descriptor()
770 }
771
772 pub(crate) fn graph_weight_descriptor(
777 &self,
778 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
779 match self {
780 Self::Mapped {
781 model,
782 idx,
783 dtype: TensorDtype::Q8Row,
784 row_scale,
785 ..
786 } => Some((model, *idx, 0, row_scale.as_slice())),
787 Self::Mapped {
788 model,
789 idx,
790 dtype: TensorDtype::Q1,
791 ..
792 } => Some((model, *idx, 1, &[])),
793 Self::Mapped {
796 model,
797 idx,
798 dtype: TensorDtype::Q1G,
799 ..
800 } => Some((model, *idx, 10, &[])),
801 Self::Mapped {
806 model,
807 idx,
808 dtype: TensorDtype::Q4Tiled,
809 ..
810 } => Some((model, *idx, 5, &[])),
811 Self::Mapped {
815 model,
816 idx,
817 dtype: TensorDtype::Q4TiledP,
818 ..
819 } => Some((model, *idx, 6, &[])),
820 Self::Mapped {
821 model,
822 idx,
823 dtype: TensorDtype::Q4Block,
824 ..
825 } => Some((model, *idx, 2, &[])),
826 Self::Mapped {
831 model,
832 idx,
833 dtype: TensorDtype::Q8_2f,
834 ..
835 } => Some((model, *idx, 7, &[])),
836 Self::Mapped {
837 model,
838 idx,
839 dtype: TensorDtype::Q1T,
840 ..
841 } => Some((model, *idx, 3, &[])),
842 Self::Mapped {
846 model,
847 idx,
848 dtype: TensorDtype::Q2TiledP,
849 ..
850 } => Some((model, *idx, 9, &[])),
851 _ => None,
852 }
853 }
854
855 pub fn as_f32(&self) -> Option<&[f32]> {
858 match self {
859 Self::F32 { data, .. } => Some(data),
860 Self::Mapped { .. } | Self::Blr(_) => None,
861 }
862 }
863
864 fn quant_bytes(&self) -> &[u8] {
865 match self {
866 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
867 Self::Blr(t) => t.bytes(),
868 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
869 }
870 }
871
872 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
874 let cols = self.cols();
875 debug_assert_eq!(dst.len(), cols);
876 match self {
877 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
878 Self::Blr(t) => t.row_f32(r, dst),
879 Self::Mapped {
880 model,
881 idx,
882 dtype,
883 row_scale,
884 col_field,
885 vbit_offsets,
886 ..
887 } => {
888 if *dtype == TensorDtype::Q4Tiled {
889 let bytes = self.quant_bytes();
890 let gpr = cols / GROUP_SIZE;
891 for gi in 0..gpr {
892 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
893 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
894 for (k, &b) in tile[2..].iter().enumerate() {
895 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
896 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
897 }
898 }
899 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
900 crate::prism::inverse_embedding(model, dst);
901 }
902 return;
903 }
904 if *dtype == TensorDtype::Q4TiledP {
905 let bytes = self.quant_bytes();
906 let gpr = cols / GROUP_SIZE;
907 let v = Q4tpView::new(bytes, self.rows(), cols);
908 let mut sc = vec![0f32; gpr];
909 v.scales_into(r, gpr, &mut sc);
910 for gi in 0..gpr {
911 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
912 let s = sc[gi];
913 for (k, &b) in tile.iter().enumerate() {
914 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
915 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
916 }
917 }
918 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
919 crate::prism::inverse_embedding(model, dst);
920 }
921 return;
922 }
923 if *dtype == TensorDtype::Q2TiledP {
924 let bytes = self.quant_bytes();
925 let gpr = cols / GROUP_SIZE;
926 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
927 let mut sc = vec![0f32; gpr];
928 v.scales_into(r, gpr, &mut sc);
929 for gi in 0..gpr {
930 let ch =
931 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
932 let s = sc[gi];
933 for (k, &b) in ch.iter().enumerate() {
934 for j in 0..4 {
935 let center = if crate::prism::is_affine_target(
936 model,
937 &model.tensors[*idx].name,
938 ) {
939 1.0
940 } else {
941 1.5
942 };
943 dst[gi * GROUP_SIZE + k * 4 + j] =
944 (((b >> (2 * j)) & 3) as f32 - center) * s;
945 }
946 }
947 }
948 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
949 crate::prism::inverse_embedding(model, dst);
950 }
951 return;
952 }
953 if *dtype == TensorDtype::Q4Block {
954 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
955 let gpr = cols / GROUP_SIZE;
956 for gi in 0..gpr {
957 let g = r * gpr + gi;
958 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
959 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
960 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
961 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
962 }
963 }
964 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
965 crate::prism::inverse_embedding(model, dst);
966 }
967 return;
968 }
969 if *dtype == TensorDtype::Q1 {
970 let bytes = self.quant_bytes();
971 let gpr = cols / GROUP_SIZE;
972 for gi in 0..gpr {
973 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
974 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
975 for (j, &b) in tile[2..].iter().enumerate() {
976 for k in 0..8 {
977 dst[gi * GROUP_SIZE + j * 8 + k] =
978 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
979 }
980 }
981 }
982 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
983 crate::prism::inverse_embedding(model, dst);
984 }
985 return;
986 }
987 if *dtype == TensorDtype::Q1G {
988 let bytes = self.quant_bytes();
991 let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
992 let gpr = cols / GROUP_SIZE;
993 for gi in 0..gpr {
994 let t = r * gpr + gi;
995 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
996 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
997 for k in 0..8 {
998 dst[gi * GROUP_SIZE + j * 8 + k] =
999 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1000 }
1001 }
1002 }
1003 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1004 crate::prism::inverse_embedding(model, dst);
1005 }
1006 return;
1007 }
1008 if *dtype == TensorDtype::Q1T {
1009 let bytes = self.quant_bytes();
1010 let gpr = cols / GROUP_SIZE;
1011 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1012 for gi in 0..gpr {
1013 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1014 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1015 bytes[off],
1016 bytes[off + 1],
1017 ]));
1018 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1019 for k in 0..GROUP_SIZE {
1020 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1021 {
1022 1 => s,
1023 2 => -s,
1024 _ => 0.0,
1025 };
1026 }
1027 }
1028 let rows = self.rows();
1030 let entries = base_len + (rows + 1) * 4;
1031 if entries <= bytes.len() {
1032 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1033 let r0 = u32::from_le_bytes([
1034 ptrs[r * 4],
1035 ptrs[r * 4 + 1],
1036 ptrs[r * 4 + 2],
1037 ptrs[r * 4 + 3],
1038 ]) as usize;
1039 let r1 = u32::from_le_bytes([
1040 ptrs[(r + 1) * 4],
1041 ptrs[(r + 1) * 4 + 1],
1042 ptrs[(r + 1) * 4 + 2],
1043 ptrs[(r + 1) * 4 + 3],
1044 ]) as usize;
1045 let off = entries + r0 * 4;
1046 for i in 0..r1 - r0 {
1047 let item = &bytes[off + i * 4..off + i * 4 + 4];
1048 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1049 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1050 item[2], item[3],
1051 ]));
1052 if c < cols {
1053 dst[c] = v;
1054 }
1055 }
1056 }
1057 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1058 crate::prism::inverse_embedding(model, dst);
1059 }
1060 return;
1061 }
1062 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1063 let bytes = self.quant_bytes();
1064 let rows = self.rows();
1065 let ng = cols / GROUP_SIZE;
1066 let bits = &bytes[..rows];
1067 let sc_off = rows;
1068 let off = vbit_offsets[r];
1071 let b = bits[r] as usize;
1072 let l = ((1usize << (b - 1)) - 1) as f32;
1073 let data = &bytes[off..];
1074 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1075 for (i, d) in dst.iter_mut().enumerate() {
1076 while nbits < b {
1077 acc = (acc << 8) | data[byte_idx] as u64;
1078 byte_idx += 1;
1079 nbits += 8;
1080 }
1081 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1082 nbits -= b;
1083 let so = (r * ng + i / GROUP_SIZE) * 2;
1084 let sv = f16_to_f32(u16::from_le_bytes([
1085 bytes[sc_off + so],
1086 bytes[sc_off + so + 1],
1087 ]));
1088 *d = (u - l) * sv;
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 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1096 let s = row_scale[r];
1097 match dtype {
1098 TensorDtype::Q8Row => {
1099 for (d, &b) in dst.iter_mut().zip(q) {
1100 *d = (b as i8) as f32 * s;
1101 }
1102 }
1103 TensorDtype::Q8_2f => {
1104 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1105 *d = (b as i8) as f32 * s * col_field[i];
1106 }
1107 }
1108 _ => unreachable!(),
1109 }
1110 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1111 crate::prism::inverse_embedding(model, dst);
1112 }
1113 }
1114 }
1115 }
1116
1117 pub fn sparse_col_ok(&self) -> bool {
1122 match self {
1123 Self::F32 { .. } => true,
1124 Self::Blr(_) => false,
1125 Self::Mapped { dtype, .. } => {
1126 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1127 }
1128 }
1129 }
1130
1131 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1135 let inter = self.cols();
1136 let hidden = self.rows();
1137 debug_assert_eq!(out.len(), hidden);
1138 match self {
1139 Self::F32 { data, .. } => {
1140 for (k, o) in out.iter_mut().enumerate() {
1141 *o += w * data[k * inter + c];
1142 }
1143 }
1144 Self::Blr(t) => t.add_col_scaled(c, w, out),
1145 Self::Mapped {
1146 dtype,
1147 row_scale,
1148 col_field,
1149 ..
1150 } => {
1151 let q = self.quant_bytes();
1152 let colf = if *dtype == TensorDtype::Q8_2f {
1153 col_field[c]
1154 } else {
1155 1.0
1156 };
1157 let wc = w * colf;
1158 for (k, o) in out.iter_mut().enumerate() {
1159 let b = q[k * inter + c] as i8 as f32;
1160 *o += wc * b * row_scale[k];
1161 }
1162 }
1163 }
1164 }
1165
1166 #[inline]
1175 pub fn prefetch_row(&self, r: usize) {
1176 let Self::Mapped { dtype, .. } = self else {
1177 return;
1178 };
1179 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1180 return;
1181 }
1182 let cols = self.cols();
1183 let q = self.quant_bytes();
1184 let (a, b) = (r * cols, (r + 1) * cols);
1185 if b > q.len() {
1186 return;
1187 }
1188 let mut j = a;
1189 while j < b {
1190 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1191 j += 512;
1192 }
1193 }
1194
1195 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1204 let cols = self.cols();
1205 debug_assert_eq!(out.len(), cols);
1206 match self {
1207 Self::F32 { data, .. } => {
1208 let row = &data[r * cols..(r + 1) * cols];
1209 for (o, v) in out.iter_mut().zip(row) {
1210 *o += w * v;
1211 }
1212 }
1213 Self::Blr(t) => {
1214 let tmp = &mut scratch[..cols];
1215 t.row_f32(r, tmp);
1216 for (o, v) in out.iter_mut().zip(tmp.iter()) {
1217 *o += w * v;
1218 }
1219 }
1220 Self::Mapped {
1221 dtype,
1222 row_scale,
1223 col_field,
1224 ..
1225 } => match dtype {
1226 TensorDtype::Q8Row => {
1227 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1228 let ws = w * row_scale[r];
1229 let row: &[i8] =
1230 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1231 axpy_i8_f32(out, row, ws);
1232 }
1233 TensorDtype::Q8_2f => {
1234 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1235 let ws = w * row_scale[r];
1236 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1237 *o += ws * c * (*b as i8 as f32);
1238 }
1239 }
1240 _ => {
1241 self.row_f32(r, scratch);
1242 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1243 *o += w * v;
1244 }
1245 }
1246 },
1247 }
1248 }
1249
1250 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1254 let cols = self.cols();
1255 match self {
1256 Self::F32 { data, .. } => {
1257 let row = &data[r * cols..(r + 1) * cols];
1258 row.iter().zip(x).map(|(w, v)| w * v).sum()
1259 }
1260 Self::Blr(t) => {
1261 let tmp = &mut scratch[..cols];
1262 t.row_f32(r, tmp);
1263 tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1264 }
1265 Self::Mapped {
1266 model,
1267 idx,
1268 dtype,
1269 row_scale,
1270 col_field,
1271 ..
1272 } => {
1273 let prism_forward =
1274 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1275 if prism_forward {
1276 let transformed = crate::prism::forward(model, &x[..cols]);
1277 let gpr = cols / GROUP_SIZE;
1278 match dtype {
1279 TensorDtype::Q2TiledP => {
1280 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1281 let mut sc = vec![0f32; gpr];
1282 v.scales_into(r, gpr, &mut sc);
1283 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1284 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1285 }
1286 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1287 }
1288 _ => {
1289 self.row_f32(r, scratch);
1290 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1291 }
1292 }
1293 }
1294 match dtype {
1295 TensorDtype::Q8Row => {
1296 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1297 dot_i8_f32(q, x) * row_scale[r]
1298 }
1299 TensorDtype::Q8_2f => {
1300 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1301 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1302 }
1303 _ => {
1304 self.row_f32(r, scratch);
1305 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1306 }
1307 }
1308 }
1309 }
1310 }
1311
1312 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1315 match self {
1316 Self::F32 { data, .. } => {
1320 if !crate::f32_backend::matvec(data, x, out) {
1321 matvec_rows(pool, data, x, out);
1322 }
1323 }
1324 Self::Blr(t) => {
1325 if crate::gpu::enabled_here()
1326 && crate::blr::gpu_wanted()
1327 && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1328 {
1329 return;
1330 }
1331 t.matvec(x, out, pool)
1332 }
1333 Self::Mapped {
1334 model,
1335 idx,
1336 dtype,
1337 rows,
1338 cols,
1339 row_scale,
1340 col_field,
1341 vbit_offsets,
1342 repack,
1343 } => {
1344 let _ = (model, idx);
1345 assert!(
1355 out.len() >= *rows && x.len() >= *cols,
1356 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1357 out.len(),
1358 x.len(),
1359 );
1360 let prism_forward =
1361 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1362 if *dtype == TensorDtype::Q2TiledP
1363 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1364 {
1365 use std::sync::atomic::{AtomicUsize, Ordering};
1366 static N: AtomicUsize = AtomicUsize::new(0);
1367 let n = N.fetch_add(1, Ordering::Relaxed);
1368 if n < 128 {
1369 eprintln!(
1370 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1371 model.tensors[*idx].name,
1372 prism_forward,
1373 rows,
1374 cols,
1375 crate::gpu::enabled_here(),
1376 crate::gpu::q2tp_gpu_opt_in(),
1377 crate::gpu::cur_layer(),
1378 );
1379 }
1380 }
1381 if prism_forward {
1386 let transformed = crate::prism::forward(model, &x[..*cols]);
1387 match dtype {
1388 TensorDtype::Q4Block => {
1389 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1390 }
1391 TensorDtype::Q4Tiled => {
1392 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1393 }
1394 TensorDtype::Q4TiledP => {
1395 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1396 }
1397 TensorDtype::Q2TiledP => {
1398 let affine =
1399 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1400 if *rows * *cols >= 8_388_608
1401 && crate::gpu::enabled_here()
1402 && crate::gpu::q2tp_gpu_opt_in()
1403 {
1404 let gpu_ok = if affine {
1405 crate::gpu::q2tp_affine_matvec(
1406 model,
1407 *idx,
1408 &transformed,
1409 *rows,
1410 *cols,
1411 out,
1412 )
1413 } else {
1414 crate::gpu::q2tp_matvec(
1415 model,
1416 *idx,
1417 &transformed,
1418 *rows,
1419 *cols,
1420 out,
1421 )
1422 };
1423 if gpu_ok {
1424 return;
1425 }
1426 }
1427 if affine {
1428 q2tp_affine_matvec(
1429 self.quant_bytes(),
1430 &transformed,
1431 *rows,
1432 *cols,
1433 out,
1434 pool,
1435 )
1436 } else {
1437 q2tp_matvec(
1438 self.quant_bytes(),
1439 &transformed,
1440 *rows,
1441 *cols,
1442 out,
1443 pool,
1444 )
1445 }
1446 }
1447 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1448 self.bin_w().expect("binary dtype"),
1449 &transformed,
1450 *rows,
1451 *cols,
1452 out,
1453 pool,
1454 ),
1455 TensorDtype::Q1T => {
1456 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1457 }
1458 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1459 self.quant_bytes(),
1460 vbit_offsets,
1461 &transformed,
1462 *rows,
1463 *cols,
1464 out,
1465 pool,
1466 ),
1467 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1468 self.quant_bytes(),
1469 repack,
1470 row_scale,
1471 &transformed,
1472 col_field,
1473 *dtype,
1474 *rows,
1475 *cols,
1476 out,
1477 pool,
1478 ),
1479 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1480 }
1481 return;
1482 }
1483 if *dtype == TensorDtype::Q4Block {
1484 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1488 let t0 = std::time::Instant::now();
1489 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1490 crate::gpu::ProbeArm::Gpu => {
1491 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1492 crate::gpu::probe_record(
1493 crate::gpu::OpClass::Matvec,
1494 true,
1495 t0.elapsed(),
1496 );
1497 return;
1498 }
1499 }
1500 crate::gpu::ProbeArm::CpuTimed => {
1501 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1502 crate::gpu::probe_record(
1503 crate::gpu::OpClass::Matvec,
1504 false,
1505 t0.elapsed(),
1506 );
1507 return;
1508 }
1509 crate::gpu::ProbeArm::Cpu => {}
1510 }
1511 }
1512 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1513 return;
1514 }
1515 if *dtype == TensorDtype::Q4Tiled {
1516 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1521 let t0 = std::time::Instant::now();
1522 let cls = crate::gpu::matvec_class(*rows, *cols);
1523 match crate::gpu::probe_arm(cls) {
1524 crate::gpu::ProbeArm::Gpu => {
1525 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1526 crate::gpu::probe_record(cls, true, t0.elapsed());
1527 return;
1528 }
1529 }
1530 crate::gpu::ProbeArm::CpuTimed => {
1531 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1532 crate::gpu::probe_record(cls, false, t0.elapsed());
1533 return;
1534 }
1535 crate::gpu::ProbeArm::Cpu => {}
1536 }
1537 }
1538 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1539 return;
1540 }
1541 if *dtype == TensorDtype::Q4TiledP {
1542 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1548 let t0 = std::time::Instant::now();
1549 let cls = crate::gpu::matvec_class(*rows, *cols);
1550 match crate::gpu::probe_arm(cls) {
1551 crate::gpu::ProbeArm::Gpu => {
1552 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1553 crate::gpu::probe_record(cls, true, t0.elapsed());
1554 return;
1555 }
1556 }
1557 crate::gpu::ProbeArm::CpuTimed => {
1558 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1559 crate::gpu::probe_record(cls, false, t0.elapsed());
1560 return;
1561 }
1562 crate::gpu::ProbeArm::Cpu => {}
1563 }
1564 }
1565 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1566 return;
1567 }
1568 if *dtype == TensorDtype::Q2TiledP {
1569 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1570 return;
1571 }
1572 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1573 let w = self.bin_w().expect("binary dtype");
1578 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1579 let t0 = std::time::Instant::now();
1580 let arm = if crate::gpu::q1_force() {
1581 crate::gpu::ProbeArm::Gpu
1582 } else {
1583 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1584 };
1585 match arm {
1586 crate::gpu::ProbeArm::Gpu => {
1587 let ok = if *dtype == TensorDtype::Q1G {
1588 crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1589 } else {
1590 crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1591 };
1592 if ok {
1593 crate::gpu::probe_record(
1594 crate::gpu::OpClass::Matvec,
1595 true,
1596 t0.elapsed(),
1597 );
1598 return;
1599 }
1600 }
1601 crate::gpu::ProbeArm::CpuTimed => {
1602 bin_matvec(w, x, *rows, *cols, out, pool);
1603 crate::gpu::probe_record(
1604 crate::gpu::OpClass::Matvec,
1605 false,
1606 t0.elapsed(),
1607 );
1608 return;
1609 }
1610 crate::gpu::ProbeArm::Cpu => {}
1611 }
1612 }
1613 bin_matvec(w, x, *rows, *cols, out, pool);
1614 return;
1615 }
1616 if *dtype == TensorDtype::Q1T {
1617 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1621 let t0 = std::time::Instant::now();
1622 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1623 crate::gpu::ProbeArm::Gpu => {
1624 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1625 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1626 crate::gpu::probe_record(
1627 crate::gpu::OpClass::Matvec,
1628 true,
1629 t0.elapsed(),
1630 );
1631 return;
1632 }
1633 }
1634 crate::gpu::ProbeArm::CpuTimed => {
1635 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1636 crate::gpu::probe_record(
1637 crate::gpu::OpClass::Matvec,
1638 false,
1639 t0.elapsed(),
1640 );
1641 return;
1642 }
1643 crate::gpu::ProbeArm::Cpu => {}
1644 }
1645 }
1646 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1647 return;
1648 }
1649 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1650 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1651 return;
1652 }
1653 let xs = prescale(x, col_field, *dtype);
1654 if *rows >= crate::gpu::min_rows()
1659 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1660 && gpu_lmhead_enabled()
1661 && crate::gpu::enabled_here()
1662 {
1663 let t0 = std::time::Instant::now();
1666 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1667 crate::gpu::ProbeArm::Gpu => {}
1668 crate::gpu::ProbeArm::CpuTimed => {
1669 qmatvec(
1670 self.quant_bytes(),
1671 repack,
1672 row_scale,
1673 x,
1674 col_field,
1675 *dtype,
1676 *rows,
1677 *cols,
1678 out,
1679 pool,
1680 );
1681 crate::gpu::probe_record(
1682 crate::gpu::OpClass::Matvec,
1683 false,
1684 t0.elapsed(),
1685 );
1686 return;
1687 }
1688 crate::gpu::ProbeArm::Cpu => {
1689 qmatvec(
1690 self.quant_bytes(),
1691 repack,
1692 row_scale,
1693 x,
1694 col_field,
1695 *dtype,
1696 *rows,
1697 *cols,
1698 out,
1699 pool,
1700 );
1701 return;
1702 }
1703 }
1704 let frac = gpu_split_frac();
1705 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1706 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1707 let bytes = self.quant_bytes();
1708 let ok = std::thread::scope(|sc| {
1709 let g = sc.spawn(|| {
1710 crate::gpu::q8_matvec_range(
1711 model,
1712 *idx,
1713 cpu_rows,
1714 &row_scale[cpu_rows..],
1715 &xs,
1716 *rows - cpu_rows,
1717 *cols,
1718 out_gpu,
1719 )
1720 });
1721 if cpu_rows > 0 {
1722 let rep_cpu = if repack.is_empty() {
1725 &[][..]
1726 } else {
1727 &repack[..(cpu_rows / 4) * 4 * *cols]
1728 };
1729 qmatvec(
1730 &bytes[..cpu_rows * *cols],
1731 rep_cpu,
1732 &row_scale[..cpu_rows],
1733 x,
1734 col_field,
1735 *dtype,
1736 cpu_rows,
1737 *cols,
1738 out_cpu,
1739 pool,
1740 );
1741 }
1742 g.join().unwrap_or(false)
1743 });
1744 if ok {
1745 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1746 return;
1747 }
1748 qmatvec(
1751 &bytes[cpu_rows * *cols..(*rows) * *cols],
1752 &[],
1753 &row_scale[cpu_rows..],
1754 x,
1755 col_field,
1756 *dtype,
1757 *rows - cpu_rows,
1758 *cols,
1759 out_gpu,
1760 pool,
1761 );
1762 return;
1763 }
1764 qmatvec(
1765 self.quant_bytes(),
1766 repack,
1767 row_scale,
1768 x,
1769 col_field,
1770 *dtype,
1771 *rows,
1772 *cols,
1773 out,
1774 pool,
1775 );
1776 }
1777 }
1778 }
1779
1780 pub fn matvec2(
1782 &self,
1783 x1: &[f32],
1784 x2: &[f32],
1785 o1: &mut [f32],
1786 o2: &mut [f32],
1787 pool: Option<&Pool>,
1788 ) {
1789 match self {
1790 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1791 Self::Blr(_) => {
1792 self.matvec(x1, o1, pool);
1793 self.matvec(x2, o2, pool);
1794 }
1795 Self::Mapped {
1796 model,
1797 idx,
1798 dtype,
1799 rows,
1800 cols,
1801 row_scale,
1802 col_field,
1803 vbit_offsets,
1804 ..
1805 } => {
1806 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1807 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1808 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1809 match dtype {
1810 TensorDtype::Q4Block => {
1811 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1812 }
1813 TensorDtype::Q4Tiled => {
1814 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1815 }
1816 TensorDtype::Q4TiledP => {
1817 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1818 }
1819 TensorDtype::Q2TiledP => {
1820 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1821 q2tp_affine_matvec2(
1822 self.quant_bytes(),
1823 &tx1,
1824 &tx2,
1825 *rows,
1826 *cols,
1827 o1,
1828 o2,
1829 pool,
1830 )
1831 } else {
1832 q2tp_matvec2(
1833 self.quant_bytes(),
1834 &tx1,
1835 &tx2,
1836 *rows,
1837 *cols,
1838 o1,
1839 o2,
1840 pool,
1841 )
1842 }
1843 }
1844 TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1845 self.bin_w().expect("binary dtype"),
1846 &tx1,
1847 &tx2,
1848 *rows,
1849 *cols,
1850 o1,
1851 o2,
1852 pool,
1853 ),
1854 TensorDtype::Q1T => {
1855 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1856 }
1857 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1858 self.quant_bytes(),
1859 vbit_offsets,
1860 &tx1,
1861 &tx2,
1862 *rows,
1863 *cols,
1864 o1,
1865 o2,
1866 pool,
1867 ),
1868 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1869 self.quant_bytes(),
1870 row_scale,
1871 &tx1,
1872 &tx2,
1873 col_field,
1874 *dtype,
1875 *rows,
1876 *cols,
1877 o1,
1878 o2,
1879 pool,
1880 ),
1881 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1882 }
1883 return;
1884 }
1885 if *dtype == TensorDtype::Q4Block {
1886 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1887 return;
1888 }
1889 if *dtype == TensorDtype::Q4Tiled {
1890 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1891 return;
1892 }
1893 if *dtype == TensorDtype::Q4TiledP {
1894 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1895 return;
1896 }
1897 if *dtype == TensorDtype::Q2TiledP {
1898 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1899 return;
1900 }
1901 if let Some(w) = self.bin_w() {
1902 bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1903 return;
1904 }
1905 if *dtype == TensorDtype::Q1T {
1906 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1912 return;
1913 }
1914 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1915 vbitmatvec2(
1916 self.quant_bytes(),
1917 vbit_offsets,
1918 x1,
1919 x2,
1920 *rows,
1921 *cols,
1922 o1,
1923 o2,
1924 pool,
1925 );
1926 return;
1927 }
1928 qmatvec2(
1929 self.quant_bytes(),
1930 row_scale,
1931 x1,
1932 x2,
1933 col_field,
1934 *dtype,
1935 *rows,
1936 *cols,
1937 o1,
1938 o2,
1939 pool,
1940 );
1941 }
1942 }
1943 }
1944}
1945
1946impl QTensor {
1947 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1955 if self.has_prism_contract() {
1956 return None;
1957 }
1958 match self {
1959 Self::Mapped {
1960 model, idx, dtype, ..
1961 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1962 _ => None,
1963 }
1964 }
1965
1966 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1967 let cols = self.cols();
1968 let rows = self.rows();
1969 debug_assert_eq!(xs_all.len(), b * cols);
1970 debug_assert_eq!(out.len(), b * rows);
1971 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1972 if crate::gptq_capture::capturing() {
1976 if let Self::Mapped { model, idx, .. } = self {
1977 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1978 }
1979 if let Self::Blr(t) = self {
1980 crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
1981 }
1982 }
1983 match self {
1984 Self::Blr(t) => {
1985 if crate::gpu::enabled_here()
1986 && crate::blr::gpu_wanted()
1987 && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
1988 {
1989 return;
1990 }
1991 t.matmat(xs_all, b, out, pool)
1992 }
1993 Self::F32 { data, .. } => {
1994 if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
1995 return;
1996 }
1997 let out_addr = SendMut(out.as_mut_ptr());
1998 let run = |start: usize, end: usize| {
2005 let mut bi = start;
2006 while bi < end {
2007 let n = (end - bi).min(4);
2008 for o in 0..rows {
2009 let row = &data[o * cols..(o + 1) * cols];
2010 if n == 4 {
2011 let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2012 let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2013 let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2014 for j in 0..cols {
2015 let w = row[j];
2016 a0 += w * x0[j];
2017 a1 += w * x1[j];
2018 a2 += w * x2[j];
2019 a3 += w * x3[j];
2020 }
2021 for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2022 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2023 }
2024 } else {
2025 for k in 0..n {
2026 let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2027 let mut acc = 0f32;
2028 for j in 0..cols {
2029 acc += row[j] * x[j];
2030 }
2031 unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2032 }
2033 }
2034 }
2035 bi += n;
2036 }
2037 };
2038 dispatch_rows(pool, b, &run);
2039 }
2040 Self::Mapped {
2041 model,
2042 idx,
2043 dtype,
2044 row_scale,
2045 col_field,
2046 vbit_offsets,
2047 ..
2048 } => {
2049 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2050 let mut transformed = Vec::with_capacity(xs_all.len());
2051 for bi in 0..b {
2052 transformed.extend_from_slice(&crate::prism::forward(
2053 model,
2054 &xs_all[bi * cols..(bi + 1) * cols],
2055 ));
2056 }
2057 match dtype {
2058 TensorDtype::Q4Block => {
2059 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2060 }
2061 TensorDtype::Q4Tiled => {
2062 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2063 }
2064 TensorDtype::Q4TiledP => {
2065 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2066 }
2067 TensorDtype::Q2TiledP => {
2068 let affine =
2069 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2070 let gpu_batch_ok = if affine {
2076 b >= 2
2077 } else {
2078 b >= 32 && b * rows * cols >= 128_000_000
2079 };
2080 if gpu_batch_ok
2081 && cols % 32 == 0
2082 && crate::gpu::enabled_here()
2083 && crate::gpu::q2tp_gpu_opt_in()
2084 {
2085 let gpu_ok = if affine {
2086 crate::gpu::q2tp_affine_matmat(
2087 model,
2088 *idx,
2089 &transformed,
2090 b,
2091 rows,
2092 cols,
2093 out,
2094 )
2095 } else {
2096 crate::gpu::q2tp_matmat(
2097 model,
2098 *idx,
2099 &transformed,
2100 b,
2101 rows,
2102 cols,
2103 out,
2104 )
2105 };
2106 if gpu_ok {
2107 return;
2108 }
2109 }
2110 if affine
2114 && b == 1
2115 && cols % 32 == 0
2116 && crate::gpu::enabled_here()
2117 && crate::gpu::q2tp_gpu_opt_in()
2118 && crate::gpu::q2tp_affine_matvec(
2119 model,
2120 *idx,
2121 &transformed[..cols],
2122 rows,
2123 cols,
2124 &mut out[..rows],
2125 )
2126 {
2127 return;
2128 }
2129 if affine {
2130 q2tp_affine_matmat(
2131 self.quant_bytes(),
2132 &transformed,
2133 b,
2134 rows,
2135 cols,
2136 out,
2137 pool,
2138 )
2139 } else {
2140 q2tp_matmat(
2141 self.quant_bytes(),
2142 &transformed,
2143 b,
2144 rows,
2145 cols,
2146 out,
2147 pool,
2148 )
2149 }
2150 }
2151 TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2152 self.bin_w().expect("binary dtype"),
2153 &transformed,
2154 b,
2155 rows,
2156 cols,
2157 out,
2158 pool,
2159 ),
2160 TensorDtype::Q1T => {
2161 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2162 }
2163 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2164 self.quant_bytes(),
2165 vbit_offsets,
2166 &transformed,
2167 b,
2168 rows,
2169 cols,
2170 out,
2171 pool,
2172 ),
2173 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2174 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2175 .map(|bi| {
2176 prescale(
2177 &transformed[bi * cols..(bi + 1) * cols],
2178 col_field,
2179 *dtype,
2180 )
2181 })
2182 .collect();
2183 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2184 }
2185 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2186 }
2187 return;
2188 }
2189 if *dtype == TensorDtype::Q4Block {
2190 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2191 return;
2192 }
2193 if *dtype == TensorDtype::Q4TiledP {
2194 if b >= 32
2208 && b * rows * cols >= 128_000_000
2209 && cols % 32 == 0
2210 && !row_exact()
2211 && !crate::gpu::mm_killed()
2212 && crate::gpu::enabled_here()
2213 {
2214 let class = if b >= 128 {
2215 crate::gpu::OpClass::MatmatWide
2216 } else {
2217 crate::gpu::OpClass::Matmat
2218 };
2219 if let Self::Mapped { model, idx, .. } = self {
2220 if crate::mm_ab::on() {
2232 let mut g = vec![0f32; b * rows];
2233 let t = std::time::Instant::now();
2234 let took = crate::gpu::q4tp_matmat(
2235 model, *idx, xs_all, b, rows, cols, &mut g,
2236 );
2237 let dg = t.elapsed();
2238 let t = std::time::Instant::now();
2239 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2240 let dc = t.elapsed();
2241 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2242 return;
2243 }
2244 let t0 = std::time::Instant::now();
2245 let resident = crate::gpu::weight_is_resident(model, *idx);
2249 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2250 crate::gpu::ProbeArm::Gpu => {
2251 if crate::gpu::q4tp_matmat(
2252 model, *idx, xs_all, b, rows, cols, out,
2253 ) {
2254 let el = t0.elapsed();
2255 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2265 let budget = std::time::Duration::from_secs_f64(
2266 flops / 1.5e12 * 8.0 + 0.020,
2267 );
2268 crate::gpu::mm_budget_check(
2269 "q4tp matmat",
2270 el,
2271 budget,
2272 crate::gpu::probe_was_cold() || !resident,
2273 );
2274 crate::gpu::probe_record(class, true, el);
2275 return;
2276 }
2277 }
2278 crate::gpu::ProbeArm::CpuTimed => {
2279 q4tp_matmat(
2280 self.quant_bytes(),
2281 xs_all,
2282 b,
2283 rows,
2284 cols,
2285 out,
2286 pool,
2287 );
2288 crate::gpu::probe_record(class, false, t0.elapsed());
2289 return;
2290 }
2291 crate::gpu::ProbeArm::Cpu => {}
2292 }
2293 }
2294 }
2295 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2296 return;
2297 }
2298 if *dtype == TensorDtype::Q2TiledP {
2299 if b >= 32
2305 && b * rows * cols >= 128_000_000
2306 && cols % 32 == 0
2307 && !crate::gpu::mm_killed()
2308 && crate::gpu::enabled_here()
2309 {
2310 let class = if b >= 128 {
2311 crate::gpu::OpClass::MatmatWide
2312 } else {
2313 crate::gpu::OpClass::Matmat
2314 };
2315 if let Self::Mapped { model, idx, .. } = self {
2316 let t0 = std::time::Instant::now();
2317 match crate::gpu::probe_arm(class) {
2318 crate::gpu::ProbeArm::Gpu => {
2319 if crate::gpu::q2tp_matmat(
2320 model, *idx, xs_all, b, rows, cols, out,
2321 ) {
2322 crate::gpu::probe_record(class, true, t0.elapsed());
2323 return;
2324 }
2325 }
2326 crate::gpu::ProbeArm::CpuTimed => {
2327 q2tp_matmat(
2328 self.quant_bytes(),
2329 xs_all,
2330 b,
2331 rows,
2332 cols,
2333 out,
2334 pool,
2335 );
2336 crate::gpu::probe_record(class, false, t0.elapsed());
2337 return;
2338 }
2339 crate::gpu::ProbeArm::Cpu => {}
2340 }
2341 }
2342 }
2343 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2348 return;
2349 }
2350 if *dtype == TensorDtype::Q4Tiled {
2351 if b >= 32
2363 && b * rows * cols >= 128_000_000
2364 && cols % 32 == 0
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 let t0 = std::time::Instant::now();
2375 match crate::gpu::probe_arm(class) {
2376 crate::gpu::ProbeArm::Gpu => {
2377 if crate::gpu::q4t_matmat(
2378 model, *idx, xs_all, b, rows, cols, out,
2379 ) {
2380 let el = t0.elapsed();
2381 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2391 let budget = std::time::Duration::from_secs_f64(
2392 flops / 1.5e12 * 8.0 + 0.020,
2393 );
2394 crate::gpu::mm_budget_check(
2395 "q4t matmat",
2396 el,
2397 budget,
2398 crate::gpu::probe_was_cold(),
2399 );
2400 crate::gpu::probe_record(class, true, el);
2401 return;
2402 }
2403 }
2404 crate::gpu::ProbeArm::CpuTimed => {
2405 q4t_matmat(
2406 self.quant_bytes(),
2407 xs_all,
2408 b,
2409 rows,
2410 cols,
2411 out,
2412 pool,
2413 );
2414 crate::gpu::probe_record(class, false, t0.elapsed());
2415 return;
2416 }
2417 crate::gpu::ProbeArm::Cpu => {}
2418 }
2419 }
2420 }
2421 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2422 return;
2423 }
2424 if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2425 let w = self.bin_w().expect("binary dtype");
2429 if b >= 32
2430 && b * rows * cols >= 128_000_000
2431 && cols % 64 == 0
2432 && crate::gpu::enabled_here()
2433 {
2434 if let Self::Mapped { model, idx, .. } = self {
2435 let t0 = std::time::Instant::now();
2436 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2437 crate::gpu::ProbeArm::Gpu => {
2438 let ok = if *dtype == TensorDtype::Q1G {
2439 crate::gpu::q1g_matmat(
2440 model, *idx, xs_all, b, rows, cols, out,
2441 )
2442 } else {
2443 crate::gpu::q1_matmat(
2444 model, *idx, xs_all, b, rows, cols, out,
2445 )
2446 };
2447 if ok {
2448 crate::gpu::probe_record(
2449 crate::gpu::OpClass::Matmat,
2450 true,
2451 t0.elapsed(),
2452 );
2453 return;
2454 }
2455 }
2456 crate::gpu::ProbeArm::CpuTimed => {
2457 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2458 crate::gpu::probe_record(
2459 crate::gpu::OpClass::Matmat,
2460 false,
2461 t0.elapsed(),
2462 );
2463 return;
2464 }
2465 crate::gpu::ProbeArm::Cpu => {}
2466 }
2467 }
2468 }
2469 bin_matmat(w, xs_all, b, rows, cols, out, pool);
2470 return;
2471 }
2472 if *dtype == TensorDtype::Q1T {
2473 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2476 if let Self::Mapped { model, idx, .. } = self {
2477 let t0 = std::time::Instant::now();
2478 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2479 crate::gpu::ProbeArm::Gpu => {
2480 if crate::gpu::q1t_matmat(
2481 model, *idx, xs_all, b, rows, cols, out,
2482 ) {
2483 crate::gpu::probe_record(
2484 crate::gpu::OpClass::Matmat,
2485 true,
2486 t0.elapsed(),
2487 );
2488 return;
2489 }
2490 }
2491 crate::gpu::ProbeArm::CpuTimed => {
2492 q1t_matmat(
2493 self.quant_bytes(),
2494 xs_all,
2495 b,
2496 rows,
2497 cols,
2498 out,
2499 pool,
2500 );
2501 crate::gpu::probe_record(
2502 crate::gpu::OpClass::Matmat,
2503 false,
2504 t0.elapsed(),
2505 );
2506 return;
2507 }
2508 crate::gpu::ProbeArm::Cpu => {}
2509 }
2510 }
2511 }
2512 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2513 return;
2514 }
2515 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2516 vbitmatmat(
2517 self.quant_bytes(),
2518 vbit_offsets,
2519 xs_all,
2520 b,
2521 rows,
2522 cols,
2523 out,
2524 pool,
2525 );
2526 return;
2527 }
2528 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2529 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2530 .collect();
2531 if row_exact()
2537 && (1..=4).contains(&b)
2538 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2539 && crate::gpu::enabled_here()
2540 && crate::gpu::wgpu_active()
2541 {
2542 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2543 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2544 return;
2545 }
2546 }
2547 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2553 if let Self::Mapped { model, idx, .. } = self {
2554 let t0 = std::time::Instant::now();
2555 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2556 crate::gpu::ProbeArm::Gpu
2557 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2558 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2559 {
2560 let q = self.quant_bytes();
2564 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2565 return;
2566 }
2567 crate::gpu::ProbeArm::Gpu => {
2568 if *dtype == TensorDtype::Q8_2f
2576 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2577 && crate::gpu::q8_matmat_2f(
2578 model, *idx, row_scale, col_field, xs_all, b, rows,
2579 cols, out,
2580 )
2581 {
2582 crate::gpu::probe_record(
2583 crate::gpu::OpClass::Matmat,
2584 true,
2585 t0.elapsed(),
2586 );
2587 return;
2588 }
2589 let flat: Vec<f32> =
2590 pre.iter().flat_map(|v| v.iter().copied()).collect();
2591 if crate::gpu::q8_matmat(
2592 model, *idx, row_scale, &flat, b, rows, cols, out,
2593 ) {
2594 crate::gpu::probe_record(
2595 crate::gpu::OpClass::Matmat,
2596 true,
2597 t0.elapsed(),
2598 );
2599 return;
2600 }
2601 }
2602 crate::gpu::ProbeArm::CpuTimed => {
2603 let q = self.quant_bytes();
2604 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2605 crate::gpu::probe_record(
2606 crate::gpu::OpClass::Matmat,
2607 false,
2608 t0.elapsed(),
2609 );
2610 return;
2611 }
2612 crate::gpu::ProbeArm::Cpu => {}
2613 }
2614 }
2615 }
2616 let q = self.quant_bytes();
2617 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2618 }
2619 }
2620 }
2621}
2622
2623impl QTensor {
2624 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2634 let (rows, cols) = (self.rows(), self.cols());
2635 let Self::Mapped {
2636 model,
2637 idx,
2638 dtype,
2639 row_scale,
2640 col_field,
2641 ..
2642 } = self
2643 else {
2644 return false;
2645 };
2646 if crate::prism::has_contract(model) {
2647 return false;
2648 }
2649 match *dtype {
2650 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2651 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2655 if *dtype == TensorDtype::Q8_2f
2658 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2659 && crate::gpu::q8_matmat_2f(
2660 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2661 )
2662 {
2663 return true;
2664 }
2665 let flat: Vec<f32> = (0..b)
2666 .flat_map(|bi| {
2667 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2668 })
2669 .collect();
2670 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2671 }
2672 _ => false,
2673 }
2674 }
2675
2676 pub fn matvec_many<const N: usize>(
2683 ts: [&QTensor; N],
2684 x: &[f32],
2685 mut outs: [&mut [f32]; N],
2686 pool: Option<&Pool>,
2687 ) {
2688 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2689 if ts.iter().any(|t| t.has_prism_contract() || t.is_blr()) {
2690 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2695 t.matvec(x, o, pool);
2696 }
2697 return;
2698 }
2699 let uniform_q8 = ts.iter().all(|t| {
2700 matches!(
2701 t,
2702 Self::Mapped {
2703 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2704 ..
2705 }
2706 )
2707 });
2708 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2709 if uniform_f32 && crate::f32_backend::active() {
2710 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2711 t.matvec(x, o, pool);
2712 }
2713 return;
2714 }
2715 let uniform_q4 = ts.iter().all(|t| {
2716 matches!(
2717 t,
2718 Self::Mapped {
2719 dtype: TensorDtype::Q4Block,
2720 ..
2721 }
2722 )
2723 });
2724 let uniform_vbit = ts.iter().all(|t| {
2725 matches!(
2726 t,
2727 Self::Mapped {
2728 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2729 ..
2730 }
2731 )
2732 });
2733 let uniform_q1 =
2736 ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2737 let uniform_q1t = ts.iter().all(|t| {
2738 matches!(
2739 t,
2740 Self::Mapped {
2741 dtype: TensorDtype::Q1T,
2742 ..
2743 }
2744 )
2745 });
2746 let uniform_q4tp = ts.iter().all(|t| {
2751 matches!(
2752 t,
2753 Self::Mapped {
2754 dtype: TensorDtype::Q4TiledP,
2755 ..
2756 }
2757 )
2758 }) && ts
2759 .iter()
2760 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2761 let Some(pool) = pool else {
2762 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2763 t.matvec(x, o, None);
2764 }
2765 return;
2766 };
2767 if total_rows < 256
2768 || !(uniform_q8
2769 || uniform_f32
2770 || uniform_q4
2771 || uniform_vbit
2772 || uniform_q1
2773 || uniform_q1t
2774 || uniform_q4tp)
2775 {
2776 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2777 t.matvec(x, o, Some(pool));
2778 }
2779 return;
2780 }
2781
2782 if uniform_q4tp {
2783 let cols = ts[0].cols();
2789 let gpr = cols / GROUP_SIZE;
2790 let views: [Q4tpView; N] =
2791 std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2792 let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2793 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2794 let locate = |flat: usize| -> (usize, usize) {
2796 let mut acc = 0;
2797 for (i, &r) in rows_of.iter().enumerate() {
2798 if flat < acc + r {
2799 return (i, flat - acc);
2800 }
2801 acc += r;
2802 }
2803 (rows_of.len() - 1, 0)
2804 };
2805 let (views, outs_addr) = (&views, &outs_addr);
2806 if a8w8_enabled() {
2807 let act = split_act(x);
2808 let act = &act;
2809 let run = |start: usize, end: usize| {
2810 with_krow(gpr, |sc| {
2811 for flat in start..end {
2812 let (t, r) = locate(flat);
2813 let v = &views[t];
2814 v.scales_into(r, gpr, sc);
2815 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2816 for &(j, xv) in &act.outliers {
2817 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2818 acc += w * s * xv;
2819 }
2820 unsafe { *outs_addr[t].at(r) = acc };
2822 }
2823 });
2824 };
2825 pool.run_rows(total_rows, &run);
2826 } else {
2827 let run = |start: usize, end: usize| {
2828 with_krow(gpr, |sc| {
2829 for flat in start..end {
2830 let (t, r) = locate(flat);
2831 let v = &views[t];
2832 v.scales_into(r, gpr, sc);
2833 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2835 }
2836 });
2837 };
2838 pool.run_rows(total_rows, &run);
2839 }
2840 return;
2841 }
2842
2843 if uniform_q1 {
2844 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2847 if a8w8_enabled() {
2848 let act = split_act(x);
2849 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2850 let (act, gsum) = (&act, &gsum);
2851 let closures: [_; N] = std::array::from_fn(|i| {
2852 let (w, gpr, out) = (
2853 ts[i].bin_w().expect("uniform_q1 checked"),
2854 ts[i].cols() / GROUP_SIZE,
2855 outs_addr[i],
2856 );
2857 move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
2858 });
2859 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2860 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2861 pool.run_many(&parts);
2862 } else {
2863 let closures: [_; N] = std::array::from_fn(|i| {
2864 let (w, gpr, out) = (
2865 ts[i].bin_w().expect("uniform_q1 checked"),
2866 ts[i].cols() / GROUP_SIZE,
2867 outs_addr[i],
2868 );
2869 move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
2870 });
2871 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2872 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2873 pool.run_many(&parts);
2874 }
2875 return;
2876 }
2877
2878 if uniform_q1t {
2879 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2883 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2884 if a8w8_enabled() {
2885 let act = split_act(x);
2886 let act = &act;
2887 let x_ref = x;
2888 let closures: [_; N] = std::array::from_fn(|i| {
2889 let bytes = ts[i].quant_bytes();
2890 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2891 let gpr = cols / GROUP_SIZE;
2892 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2893 let out = outs_addr[i];
2894 move |s: usize, e: usize| {
2895 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2896 }
2897 });
2898 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2899 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2900 pool.run_many(&parts);
2901 } else {
2902 let x_ref = x;
2903 let closures: [_; N] = std::array::from_fn(|i| {
2904 let bytes = ts[i].quant_bytes();
2905 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2906 let gpr = cols / GROUP_SIZE;
2907 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2908 let out = outs_addr[i];
2909 move |s: usize, e: usize| {
2910 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2911 }
2912 });
2913 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2914 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2915 pool.run_many(&parts);
2916 }
2917 return;
2918 }
2919
2920 if uniform_q4 || uniform_vbit {
2921 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2922 if a8w8_enabled() {
2924 let act = split_act(x);
2925 let act = &act;
2926 if uniform_q4 {
2927 let closures: [_; N] = std::array::from_fn(|i| {
2928 let (packed, scales) =
2929 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2930 let (gpr, cols, out) =
2931 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2932 move |s: usize, e: usize| {
2933 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2934 }
2935 });
2936 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2937 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2938 pool.run_many(&parts);
2939 } else {
2940 let closures: [_; N] = std::array::from_fn(|i| {
2941 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2942 unreachable!()
2943 };
2944 let (bytes, rows, cols, out) = (
2945 ts[i].quant_bytes(),
2946 ts[i].rows(),
2947 ts[i].cols(),
2948 outs_addr[i],
2949 );
2950 move |s: usize, e: usize| {
2951 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2952 }
2953 });
2954 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2955 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2956 pool.run_many(&parts);
2957 }
2958 return;
2959 }
2960 if uniform_q4 {
2961 let closures: [_; N] = std::array::from_fn(|i| {
2962 let (packed, scales) =
2963 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2964 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2965 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2966 });
2967 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2968 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2969 pool.run_many(&parts);
2970 } else {
2971 let closures: [_; N] = std::array::from_fn(|i| {
2972 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2973 unreachable!()
2974 };
2975 let (bytes, rows, cols, out) = (
2976 ts[i].quant_bytes(),
2977 ts[i].rows(),
2978 ts[i].cols(),
2979 outs_addr[i],
2980 );
2981 move |s: usize, e: usize| {
2982 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2983 }
2984 });
2985 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2986 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2987 pool.run_many(&parts);
2988 }
2989 return;
2990 }
2991
2992 if uniform_f32 {
2993 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2994 let closures: [_; N] = std::array::from_fn(|i| {
2995 let Self::F32 { data, cols, .. } = ts[i] else {
2996 unreachable!()
2997 };
2998 let out = outs_addr[i];
2999 move |start: usize, end: usize| {
3000 for o in start..end {
3001 let row = &data[o * cols..(o + 1) * cols];
3002 let mut sum = 0.0f32;
3003 for j in 0..*cols {
3004 sum += row[j] * x[j];
3005 }
3006 unsafe { *out.at(o) = sum };
3008 }
3009 }
3010 });
3011 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3012 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3013 pool.run_many(&parts);
3014 return;
3015 }
3016
3017 struct Ctx<'a> {
3020 bytes: &'a [u8],
3021 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3022 rep: &'a [u8],
3023 row_scale: &'a [f32],
3024 cols: usize,
3025 xs: std::borrow::Cow<'a, [f32]>,
3026 }
3027 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3028 let Self::Mapped {
3029 dtype,
3030 cols,
3031 row_scale,
3032 col_field,
3033 repack,
3034 ..
3035 } = ts[i]
3036 else {
3037 unreachable!()
3038 };
3039 Ctx {
3040 bytes: ts[i].quant_bytes(),
3041 rep: repack,
3042 row_scale,
3043 cols: *cols,
3044 xs: prescale(x, col_field, *dtype),
3045 }
3046 });
3047 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3048 #[cfg(target_arch = "aarch64")]
3049 if sdot_enabled() {
3050 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3051 let closures: [_; N] = std::array::from_fn(|i| {
3052 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3053 move |start: usize, end: usize| {
3054 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3055 }
3056 });
3057 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3058 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3059 pool.run_many(&parts);
3060 return;
3061 }
3062 #[cfg(target_arch = "x86_64")]
3063 if avx2_a8w8_enabled() {
3064 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3065 let closures: [_; N] = std::array::from_fn(|i| {
3066 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3067 move |start: usize, end: usize| {
3068 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3069 }
3070 });
3071 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3072 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3073 pool.run_many(&parts);
3074 return;
3075 }
3076 let closures: [_; N] = std::array::from_fn(|i| {
3077 let (c, out) = (&ctxs[i], outs_addr[i]);
3078 move |start: usize, end: usize| {
3079 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3080 }
3081 });
3082 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3083 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3084 pool.run_many(&parts);
3085 }
3086}
3087
3088impl QTensor {
3089 #[allow(clippy::needless_range_loop)]
3094 pub fn matvec2_many<const N: usize>(
3095 ts: [&QTensor; N],
3096 x1: &[f32],
3097 x2: &[f32],
3098 mut o1s: [&mut [f32]; N],
3099 mut o2s: [&mut [f32]; N],
3100 pool: Option<&Pool>,
3101 ) {
3102 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3103 if ts.iter().any(|t| t.has_prism_contract() || t.is_blr()) {
3104 for i in 0..N {
3105 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3106 }
3107 return;
3108 }
3109 let uniform_q8 = ts.iter().all(|t| {
3110 matches!(
3111 t,
3112 Self::Mapped {
3113 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3114 ..
3115 }
3116 )
3117 });
3118 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3119 let uniform_q4 = ts.iter().all(|t| {
3120 matches!(
3121 t,
3122 Self::Mapped {
3123 dtype: TensorDtype::Q4Block,
3124 ..
3125 }
3126 )
3127 });
3128 let uniform_vbit = ts.iter().all(|t| {
3129 matches!(
3130 t,
3131 Self::Mapped {
3132 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3133 ..
3134 }
3135 )
3136 });
3137 let fusable = pool.is_some()
3138 && total_rows >= 256
3139 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3140 if !fusable {
3141 for i in 0..N {
3142 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3143 }
3144 return;
3145 }
3146 let pool = pool.unwrap();
3147
3148 if uniform_q4 || uniform_vbit {
3149 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3150 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3151 if a8w8_enabled() {
3153 let a1 = split_act(x1);
3154 let a2 = split_act(x2);
3155 let (a1, a2) = (&a1, &a2);
3156 if uniform_q4 {
3157 let closures: [_; N] = std::array::from_fn(|i| {
3158 let (packed, scales) =
3159 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3160 let (gpr, cols, o1, o2) =
3161 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3162 move |s: usize, e: usize| {
3163 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3164 }
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 } else {
3170 let closures: [_; N] = std::array::from_fn(|i| {
3171 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3172 unreachable!()
3173 };
3174 let (bytes, rows, cols, o1, o2) = (
3175 ts[i].quant_bytes(),
3176 ts[i].rows(),
3177 ts[i].cols(),
3178 p1[i],
3179 p2[i],
3180 );
3181 move |s: usize, e: usize| {
3182 vbit_range2_a8w8(
3183 bytes,
3184 vbit_offsets,
3185 x1,
3186 x2,
3187 a1,
3188 a2,
3189 rows,
3190 cols,
3191 o1,
3192 o2,
3193 s,
3194 e,
3195 )
3196 }
3197 });
3198 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3199 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3200 pool.run_many(&parts);
3201 }
3202 return;
3203 }
3204 if uniform_q4 {
3205 let closures: [_; N] = std::array::from_fn(|i| {
3206 let (packed, scales) =
3207 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3208 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3209 move |s: usize, e: usize| {
3210 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3211 }
3212 });
3213 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3214 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3215 pool.run_many(&parts);
3216 } else {
3217 let closures: [_; N] = std::array::from_fn(|i| {
3218 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3219 unreachable!()
3220 };
3221 let (bytes, rows, cols, o1, o2) = (
3222 ts[i].quant_bytes(),
3223 ts[i].rows(),
3224 ts[i].cols(),
3225 p1[i],
3226 p2[i],
3227 );
3228 move |s: usize, e: usize| {
3229 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3230 }
3231 });
3232 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3233 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3234 pool.run_many(&parts);
3235 }
3236 return;
3237 }
3238
3239 if uniform_f32 {
3240 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3241 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3242 let closures: [_; N] = std::array::from_fn(|i| {
3243 let Self::F32 { data, cols, .. } = ts[i] else {
3244 unreachable!()
3245 };
3246 let (o1, o2) = (p1[i], p2[i]);
3247 move |start: usize, end: usize| {
3248 for o in start..end {
3249 let row = &data[o * cols..(o + 1) * cols];
3250 let (mut s1, mut s2) = (0.0f32, 0.0f32);
3251 for j in 0..*cols {
3252 s1 += row[j] * x1[j];
3253 s2 += row[j] * x2[j];
3254 }
3255 unsafe {
3257 *o1.at(o) = s1;
3258 *o2.at(o) = s2;
3259 }
3260 }
3261 }
3262 });
3263 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3264 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3265 pool.run_many(&parts);
3266 return;
3267 }
3268
3269 struct Ctx<'a> {
3270 bytes: &'a [u8],
3271 row_scale: &'a [f32],
3272 cols: usize,
3273 xs1: std::borrow::Cow<'a, [f32]>,
3274 xs2: std::borrow::Cow<'a, [f32]>,
3275 }
3276 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3277 let Self::Mapped {
3278 dtype,
3279 cols,
3280 row_scale,
3281 col_field,
3282 ..
3283 } = ts[i]
3284 else {
3285 unreachable!()
3286 };
3287 Ctx {
3288 bytes: ts[i].quant_bytes(),
3289 row_scale,
3290 cols: *cols,
3291 xs1: prescale(x1, col_field, *dtype),
3292 xs2: prescale(x2, col_field, *dtype),
3293 }
3294 });
3295 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3296 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3297 #[cfg(target_arch = "aarch64")]
3298 if sdot_enabled() {
3299 let acts: [(SplitAct, SplitAct); N] =
3300 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3301 let closures: [_; N] = std::array::from_fn(|i| {
3302 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3303 move |start: usize, end: usize| {
3304 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3305 }
3306 });
3307 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3308 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3309 pool.run_many(&parts);
3310 return;
3311 }
3312 #[cfg(target_arch = "x86_64")]
3313 if avx2_a8w8_enabled() {
3314 let acts: [(SplitAct, SplitAct); N] =
3315 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3316 let closures: [_; N] = std::array::from_fn(|i| {
3317 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3318 move |start: usize, end: usize| {
3319 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3320 }
3321 });
3322 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3323 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3324 pool.run_many(&parts);
3325 return;
3326 }
3327 let closures: [_; N] = std::array::from_fn(|i| {
3328 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3329 move |start: usize, end: usize| {
3330 q8_range2_f32(
3331 c.bytes,
3332 c.row_scale,
3333 &c.xs1,
3334 &c.xs2,
3335 c.cols,
3336 o1,
3337 o2,
3338 start,
3339 end,
3340 )
3341 }
3342 });
3343 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3344 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3345 pool.run_many(&parts);
3346 }
3347
3348 pub fn matvec_silu_mul(
3353 gate: &QTensor,
3354 up: &QTensor,
3355 x: &[f32],
3356 out: &mut [f32],
3357 pool: Option<&Pool>,
3358 ) -> bool {
3359 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3360 }
3361
3362 pub fn matvec_silu_mul_limited(
3368 gate: &QTensor,
3369 up: &QTensor,
3370 x: &[f32],
3371 out: &mut [f32],
3372 limit: f32,
3373 pool: Option<&Pool>,
3374 ) -> bool {
3375 if gate.has_prism_contract() || up.has_prism_contract() || gate.is_blr() || up.is_blr() {
3376 return false;
3380 }
3381 let inter = gate.rows();
3382 debug_assert_eq!(up.rows(), inter);
3383 debug_assert_eq!(out.len(), inter);
3384 debug_assert_eq!(gate.cols(), up.cols());
3385 if !a8w8_enabled() {
3386 return false;
3387 }
3388 let act = split_act(x);
3389 let act = &act;
3390 let x_ref = x;
3391 let out_addr = SendMut(out.as_mut_ptr());
3392
3393 match (gate, up) {
3394 (
3396 Self::Mapped {
3397 dtype: TensorDtype::Q4Block,
3398 ..
3399 },
3400 Self::Mapped {
3401 dtype: TensorDtype::Q4Block,
3402 ..
3403 },
3404 ) => {
3405 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3406 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3407 let gpr = gate.cols() / GROUP_SIZE;
3408 let cols = gate.cols();
3409 let run = move |start: usize, end: usize| {
3410 for r in start..end {
3411 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3412 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3413 for &(j, xv) in &act.outliers {
3414 let flat = r * cols + j;
3415 let gb = gp[flat / 2];
3416 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3417 let gsc = f16_to_f32(u16::from_le_bytes([
3418 gs[(flat / GROUP_SIZE) * 2],
3419 gs[(flat / GROUP_SIZE) * 2 + 1],
3420 ]));
3421 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3422 let ub = up_p[flat / 2];
3423 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3424 let usc = f16_to_f32(u16::from_le_bytes([
3425 up_s[(flat / GROUP_SIZE) * 2],
3426 up_s[(flat / GROUP_SIZE) * 2 + 1],
3427 ]));
3428 uv += ((un as i32 - 8) as f32) * usc * xv;
3429 }
3430 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3432 }
3433 };
3434 dispatch_rows(pool, inter, &run);
3435 true
3436 }
3437 (
3441 Self::Mapped {
3442 dtype: TensorDtype::Q4Tiled,
3443 ..
3444 },
3445 Self::Mapped {
3446 dtype: TensorDtype::Q4Tiled,
3447 ..
3448 },
3449 ) => {
3450 let g_bytes = gate.quant_bytes();
3451 let u_bytes = up.quant_bytes();
3452 let gpr = gate.cols() / GROUP_SIZE;
3453 let run = move |start: usize, end: usize| {
3454 for r in start..end {
3455 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3456 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3457 for &(j, xv) in &act.outliers {
3458 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3459 gv += w * s * xv;
3460 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3461 uv += w * s * xv;
3462 }
3463 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3465 }
3466 };
3467 dispatch_rows(pool, inter, &run);
3468 true
3469 }
3470 (
3473 Self::Mapped {
3474 dtype: TensorDtype::Q4TiledP,
3475 ..
3476 },
3477 Self::Mapped {
3478 dtype: TensorDtype::Q4TiledP,
3479 ..
3480 },
3481 ) => {
3482 let cols = gate.cols();
3483 let gpr = cols / GROUP_SIZE;
3484 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3485 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3486 let run = |start: usize, end: usize| {
3487 with_krows(gpr, |gsc, usc| {
3488 for r in start..end {
3489 gv_view.scales_into(r, gpr, gsc);
3490 uv_view.scales_into(r, gpr, usc);
3491 let mut gv =
3492 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3493 let mut uv =
3494 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3495 for &(j, xv) in &act.outliers {
3496 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3497 gv += w * s * xv;
3498 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3499 uv += w * s * xv;
3500 }
3501 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3503 }
3504 });
3505 };
3506 dispatch_rows(pool, inter, &run);
3507 true
3508 }
3509 (
3516 Self::Mapped {
3517 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3518 ..
3519 },
3520 Self::Mapped {
3521 dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3522 ..
3523 },
3524 ) => {
3525 let gw = gate.bin_w().expect("binary dtype matched");
3526 let uw = up.bin_w().expect("binary dtype matched");
3527 let gpr = gate.cols() / GROUP_SIZE;
3528 let gsum = q1_group_sums(&act.xq, gpr);
3529 let gsum = &gsum;
3530 let run = move |start: usize, end: usize| {
3531 for r in start..end {
3532 let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3533 let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3534 for &(j, xv) in &act.outliers {
3535 let (w, s) = gw.outlier(r, gpr, j);
3536 gv += w * s * xv;
3537 let (w, s) = uw.outlier(r, gpr, j);
3538 uv += w * s * xv;
3539 }
3540 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3542 }
3543 };
3544 dispatch_rows(pool, inter, &run);
3545 true
3546 }
3547 (
3551 Self::Mapped {
3552 dtype: TensorDtype::Q2TiledP,
3553 ..
3554 },
3555 Self::Mapped {
3556 dtype: TensorDtype::Q2TiledP,
3557 ..
3558 },
3559 ) => {
3560 let cols = gate.cols();
3561 let gpr = cols / GROUP_SIZE;
3562 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3563 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3564 let gsum = q1_group_sums(&act.xq, gpr);
3565 let gsum = &gsum;
3566 let run = move |start: usize, end: usize| {
3567 with_krows(gpr, |gsc, usc| {
3568 for r in start..end {
3569 gv_view.scales_into(r, gpr, gsc);
3570 uv_view.scales_into(r, gpr, usc);
3571 let mut gv =
3572 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3573 let mut uv =
3574 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3575 for &(j, xv) in &act.outliers {
3576 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3577 gv += w * s * xv;
3578 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3579 uv += w * s * xv;
3580 }
3581 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3583 }
3584 });
3585 };
3586 dispatch_rows(pool, inter, &run);
3587 true
3588 }
3589 (
3594 Self::Mapped {
3595 dtype: TensorDtype::Q8Row,
3596 row_scale: g_rs,
3597 ..
3598 },
3599 Self::Mapped {
3600 dtype: TensorDtype::Q8Row,
3601 row_scale: u_rs,
3602 ..
3603 },
3604 ) => {
3605 let g_bytes = gate.quant_bytes();
3606 let u_bytes = up.quant_bytes();
3607 let cols = gate.cols();
3608 let run = move |start: usize, end: usize| {
3609 for r in start..end {
3610 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3611 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3612 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3614 }
3615 };
3616 dispatch_rows(pool, inter, &run);
3617 true
3618 }
3619 (
3621 Self::Mapped {
3622 dtype: TensorDtype::Q1T,
3623 ..
3624 },
3625 Self::Mapped {
3626 dtype: TensorDtype::Q1T,
3627 ..
3628 },
3629 ) => {
3630 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3631 let g_bytes = gate.quant_bytes();
3632 let u_bytes = up.quant_bytes();
3633 let gpr = gate.cols() / GROUP_SIZE;
3634 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3635 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3636 let run = move |start: usize, end: usize| {
3637 for r in start..end {
3638 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3639 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3640 for &(j, xv) in &act.outliers {
3641 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3642 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3643 }
3644 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3645 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3646 unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3648 }
3649 };
3650 dispatch_rows(pool, inter, &run);
3651 true
3652 }
3653 _ => false,
3654 }
3655 }
3656
3657 pub fn moe_gate_up_many(
3672 pairs: &[(&QTensor, &QTensor)],
3673 x: &[f32],
3674 outs: &mut [Vec<f32>],
3675 pool: Option<&Pool>,
3676 ) -> bool {
3677 Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3678 }
3679
3680 pub fn moe_gate_up_many_limited(
3684 pairs: &[(&QTensor, &QTensor)],
3685 x: &[f32],
3686 outs: &mut [Vec<f32>],
3687 limit: f32,
3688 pool: Option<&Pool>,
3689 ) -> bool {
3690 if pairs.is_empty() || pairs.len() != outs.len() {
3691 return false;
3692 }
3693 if pairs.iter().any(|(g, u)| g.is_blr() || u.is_blr()) {
3694 return false;
3695 }
3696 if !a8w8_enabled() {
3697 if limit > 0.0 {
3698 return false;
3701 }
3702 let groups = vec![vec![0]; pairs.len()];
3703 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3704 }
3705 let inter = pairs[0].0.rows();
3706 let cols = pairs[0].0.cols();
3707 if cols % GROUP_SIZE != 0 {
3708 return false;
3709 }
3710 let gpr = cols / GROUP_SIZE;
3711 let q2 = matches!(
3714 pairs[0].0,
3715 Self::Mapped {
3716 dtype: TensorDtype::Q2TiledP,
3717 ..
3718 }
3719 );
3720 let want = if q2 {
3721 TensorDtype::Q2TiledP
3722 } else {
3723 TensorDtype::Q4TiledP
3724 };
3725 let mut views = Vec::with_capacity(pairs.len() * 2);
3726 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3727 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3728 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3729 if !both
3730 || g.rows() != inter
3731 || u.rows() != inter
3732 || g.cols() != cols
3733 || u.cols() != cols
3734 || o.len() != inter
3735 {
3736 return false;
3737 }
3738 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3739 views.push(mk(g.quant_bytes(), inter, cols));
3740 views.push(mk(u.quant_bytes(), inter, cols));
3741 }
3742 let act = split_act(x);
3743 let gsum = if q2 {
3744 q1_group_sums(&act.xq, gpr)
3745 } else {
3746 Vec::new()
3747 };
3748 let (act, gsum) = (&act, &gsum);
3749 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3750 let (views, ptrs) = (&views, &ptrs);
3751 let run = |start: usize, end: usize| {
3752 with_krows(gpr, |gsc, usc| {
3753 for flat in start..end {
3754 let (e, r) = (flat / inter, flat % inter);
3755 let gv_view = &views[e * 2];
3756 let uv_view = &views[e * 2 + 1];
3757 gv_view.scales_into(r, gpr, gsc);
3758 uv_view.scales_into(r, gpr, usc);
3759 let (mut gv, mut uv) = if q2 {
3760 (
3761 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3762 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3763 )
3764 } else {
3765 (
3766 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3767 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3768 )
3769 };
3770 for &(j, xv) in &act.outliers {
3771 let (og, ou) = if q2 {
3772 (
3773 q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3774 q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3775 )
3776 } else {
3777 (
3778 q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3779 q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3780 )
3781 };
3782 gv += og.0 * og.1 * xv;
3783 uv += ou.0 * ou.1 * xv;
3784 }
3785 unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3787 }
3788 });
3789 };
3790 dispatch_rows(pool, pairs.len() * inter, &run);
3791 true
3792 }
3793
3794 pub fn moe_down_many(
3803 downs: &[&QTensor],
3804 gs: &[Vec<f32>],
3805 weights: &[f32],
3806 out: &mut [f32],
3807 pool: Option<&Pool>,
3808 ) -> bool {
3809 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3810 return false;
3811 }
3812 if downs.iter().any(|d| d.is_blr()) {
3813 return false;
3814 }
3815 if !a8w8_enabled() {
3816 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3817 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3818 return false;
3819 }
3820 out.fill(0.0);
3821 for (row, &w) in terms.iter().zip(weights) {
3822 for (o, &v) in out.iter_mut().zip(row) {
3823 *o += w * v;
3824 }
3825 }
3826 return true;
3827 }
3828 let rows = out.len();
3829 let cols = downs[0].cols();
3830 if cols % GROUP_SIZE != 0 {
3831 return false;
3832 }
3833 let gpr = cols / GROUP_SIZE;
3834 let mut views = Vec::with_capacity(downs.len());
3835 for (d, g) in downs.iter().zip(gs.iter()) {
3836 if !matches!(
3837 d,
3838 Self::Mapped {
3839 dtype: TensorDtype::Q4TiledP,
3840 ..
3841 }
3842 ) || d.rows() != rows
3843 || d.cols() != cols
3844 || g.len() != cols
3845 {
3846 return false;
3847 }
3848 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3849 }
3850 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3852 let out_addr = SendMut(out.as_mut_ptr());
3860 let (views, acts, weights) = (&views, &acts, &weights);
3861 let run = |start: usize, end: usize| {
3862 with_krow(gpr, |sc| {
3863 for r in start..end {
3864 let mut acc = 0f32;
3865 for (e, v) in views.iter().enumerate() {
3866 v.scales_into(r, gpr, sc);
3867 let a = &acts[e];
3868 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
3869 for &(j, xv) in &a.outliers {
3870 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3871 d += w * s * xv;
3872 }
3873 acc += weights[e] * d;
3874 }
3875 unsafe { *out_addr.at(r) = acc };
3877 }
3878 });
3879 };
3880 dispatch_rows(pool, rows, &run);
3881 true
3882 }
3883
3884 pub fn moe_gate_up_rows(
3896 pairs: &[(&QTensor, &QTensor)],
3897 groups: &[Vec<usize>],
3898 xs: &[f32],
3899 outs: &mut [Vec<f32>],
3900 pool: Option<&Pool>,
3901 ) -> bool {
3902 if pairs.is_empty() || pairs.len() != groups.len() {
3903 return false;
3904 }
3905 if pairs.iter().any(|(g, u)| g.is_blr() || u.is_blr()) {
3906 return false;
3907 }
3908 let inter = pairs[0].0.rows();
3909 let cols = pairs[0].0.cols();
3910 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
3911 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
3912 return false;
3913 }
3914 let b = xs.len() / cols;
3915 let gpr = cols / GROUP_SIZE;
3916 let mut views = Vec::with_capacity(pairs.len() * 2);
3917 for (g, u) in pairs {
3918 let q4tp = |t: &QTensor| {
3919 matches!(
3920 t,
3921 Self::Mapped {
3922 dtype: TensorDtype::Q4TiledP,
3923 ..
3924 }
3925 )
3926 };
3927 if g.has_prism_contract()
3928 || u.has_prism_contract()
3929 || !q4tp(g)
3930 || !q4tp(u)
3931 || g.rows() != inter
3932 || u.rows() != inter
3933 || g.cols() != cols
3934 || u.cols() != cols
3935 {
3936 return false;
3937 }
3938 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
3939 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
3940 }
3941 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
3942 return false;
3943 }
3944 let quantized = a8w8_enabled();
3945 let acts: Vec<SplitAct> = if quantized {
3946 (0..b)
3947 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
3948 .collect()
3949 } else {
3950 Vec::new()
3951 };
3952 let mut offs = Vec::with_capacity(groups.len());
3953 let mut o = 0usize;
3954 for g in groups {
3955 offs.push(o);
3956 o += g.len();
3957 }
3958 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3959 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3960 let run = |start: usize, end: usize| {
3961 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3962 for flat in start..end {
3963 let (e, r) = (flat / inter, flat % inter);
3964 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
3965 gv_view.scales_into(r, gpr, &mut gsc);
3966 uv_view.scales_into(r, gpr, &mut usc);
3967 for (k, &t) in groups[e].iter().enumerate() {
3968 if !quantized {
3969 let x = &xs[t * cols..(t + 1) * cols];
3970 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
3971 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
3972 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
3973 continue;
3974 }
3975 let act = &acts[t];
3976 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3977 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3978 for &(j, xv) in &act.outliers {
3979 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3980 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3981 gv += og.0 * og.1 * xv;
3982 uv += ou.0 * ou.1 * xv;
3983 }
3984 let silu_g = gv / (1.0 + (-gv).exp());
3985 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
3988 }
3989 }
3990 };
3991 dispatch_rows(pool, pairs.len() * inter, &run);
3992 true
3993 }
3994
3995 pub fn moe_down_rows(
4002 downs: &[&QTensor],
4003 group_lens: &[usize],
4004 gs: &[Vec<f32>],
4005 outs: &mut [Vec<f32>],
4006 pool: Option<&Pool>,
4007 ) -> bool {
4008 if downs.is_empty() || downs.len() != group_lens.len() {
4009 return false;
4010 }
4011 if downs.iter().any(|d| d.is_blr()) {
4012 return false;
4013 }
4014 let rows = downs[0].rows();
4015 let cols = downs[0].cols();
4016 let n_pairs: usize = group_lens.iter().sum();
4017 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4018 return false;
4019 }
4020 let gpr = cols / GROUP_SIZE;
4021 let mut views = Vec::with_capacity(downs.len());
4022 for d in downs {
4023 if d.has_prism_contract()
4024 || !matches!(
4025 d,
4026 Self::Mapped {
4027 dtype: TensorDtype::Q4TiledP,
4028 ..
4029 }
4030 )
4031 || d.rows() != rows
4032 || d.cols() != cols
4033 {
4034 return false;
4035 }
4036 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4037 }
4038 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4039 return false;
4040 }
4041 let quantized = a8w8_enabled();
4042 let acts: Vec<SplitAct> = if quantized {
4043 gs.iter().map(|g| split_act(g)).collect()
4044 } else {
4045 Vec::new()
4046 };
4047 let mut offs = Vec::with_capacity(group_lens.len());
4048 let mut o = 0usize;
4049 for &l in group_lens {
4050 offs.push(o);
4051 o += l;
4052 }
4053 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4054 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4055 let run = |start: usize, end: usize| {
4056 let mut sc = vec![0f32; gpr];
4057 for flat in start..end {
4058 let (e, r) = (flat / rows, flat % rows);
4059 let v = &views[e];
4060 v.scales_into(r, gpr, &mut sc);
4061 for k in 0..group_lens[e] {
4062 if !quantized {
4063 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4064 unsafe { *ptrs[offs[e] + k].at(r) = d };
4065 continue;
4066 }
4067 let a = &acts[offs[e] + k];
4068 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4069 for &(j, xv) in &a.outliers {
4070 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4071 d += w * s * xv;
4072 }
4073 unsafe { *ptrs[offs[e] + k].at(r) = d };
4075 }
4076 }
4077 };
4078 dispatch_rows(pool, downs.len() * rows, &run);
4079 true
4080 }
4081}
4082
4083#[cfg(target_os = "macos")]
4088mod accel_blas {
4089 #[link(name = "Accelerate", kind = "framework")]
4090 unsafe extern "C" {
4091 pub fn cblas_sgemm(
4092 order: i32,
4093 trans_a: i32,
4094 trans_b: i32,
4095 m: i32,
4096 n: i32,
4097 k: i32,
4098 alpha: f32,
4099 a: *const f32,
4100 lda: i32,
4101 b: *const f32,
4102 ldb: i32,
4103 beta: f32,
4104 c: *mut f32,
4105 ldc: i32,
4106 );
4107 }
4108}
4109
4110#[cfg(target_os = "macos")]
4111pub(crate) fn accel_gemm_enabled() -> bool {
4112 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4113 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4114}
4115
4116#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4119pub(crate) fn accel_gemm_enabled() -> bool {
4120 true
4121}
4122
4123#[cfg(target_arch = "aarch64")]
4130#[allow(clippy::too_many_arguments)]
4131pub(crate) fn neon_gemm_rm(
4132 m: usize,
4133 n: usize,
4134 k: usize,
4135 alpha: f32,
4136 a: &[f32],
4137 lda: usize,
4138 b_mat: &[f32],
4139 ldb: usize,
4140 b_rows_are_n: bool,
4141 c: &mut [f32],
4142 ldc: usize,
4143) {
4144 debug_assert!(a.len() >= (m - 1) * lda + k);
4145 debug_assert!(c.len() >= (m - 1) * ldc + n);
4146 unsafe {
4148 use core::arch::aarch64::*;
4149 let mut i = 0usize;
4150 while i < m {
4151 let mi = (m - i).min(4);
4152 let mut j = 0usize;
4153 while j < n {
4154 let nj = (n - j).min(8);
4155 if mi == 4 && nj == 8 {
4156 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4157 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4158 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4159 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4160 for p in 0..k {
4161 let (b0, b1) = if b_rows_are_n {
4162 let base = b_mat.as_ptr().add(j * ldb + p);
4165 let g = |o: usize| *base.add(o * ldb);
4166 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4167 } else {
4168 let base = b_mat.as_ptr().add(p * ldb + j);
4169 (
4170 [*base, *base.add(1), *base.add(2), *base.add(3)],
4171 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4172 )
4173 };
4174 let bv0 = vld1q_f32(b0.as_ptr());
4175 let bv1 = vld1q_f32(b1.as_ptr());
4176 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4177 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4178 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4179 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4180 c0a = vfmaq_f32(c0a, a0, bv0);
4181 c0b = vfmaq_f32(c0b, a0, bv1);
4182 c1a = vfmaq_f32(c1a, a1, bv0);
4183 c1b = vfmaq_f32(c1b, a1, bv1);
4184 c2a = vfmaq_f32(c2a, a2, bv0);
4185 c2b = vfmaq_f32(c2b, a2, bv1);
4186 c3a = vfmaq_f32(c3a, a3, bv0);
4187 c3b = vfmaq_f32(c3b, a3, bv1);
4188 }
4189 let al = vdupq_n_f32(alpha);
4190 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4191 .iter()
4192 .enumerate()
4193 {
4194 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4195 vst1q_f32(dst, vmulq_f32(*ca, al));
4196 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4197 }
4198 } else {
4199 for r in 0..mi {
4200 for q in 0..nj {
4201 let mut acc = 0f32;
4202 for p in 0..k {
4203 let bv = if b_rows_are_n {
4204 b_mat[(j + q) * ldb + p]
4205 } else {
4206 b_mat[p * ldb + j + q]
4207 };
4208 acc += a[(i + r) * lda + p] * bv;
4209 }
4210 c[(i + r) * ldc + j + q] = acc * alpha;
4211 }
4212 }
4213 }
4214 j += nj;
4215 }
4216 i += mi;
4217 }
4218 }
4219}
4220
4221#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4223#[allow(clippy::too_many_arguments)]
4224pub(crate) fn sgemm_rm(
4225 m: usize,
4226 n: usize,
4227 k: usize,
4228 alpha: f32,
4229 a: &[f32],
4230 lda: usize,
4231 b_mat: &[f32],
4232 ldb: usize,
4233 b_rows_are_n: bool,
4234 c: &mut [f32],
4235 ldc: usize,
4236) {
4237 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4238}
4239
4240#[allow(clippy::too_many_arguments)]
4244pub fn sgemm_public(
4245 m: usize,
4246 n: usize,
4247 k: usize,
4248 alpha: f32,
4249 a: &[f32],
4250 lda: usize,
4251 b_mat: &[f32],
4252 ldb: usize,
4253 b_rows_are_n: bool,
4254 c: &mut [f32],
4255 ldc: usize,
4256) {
4257 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4258 {
4259 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4260 }
4261 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4266 {
4267 if m == 0 || n == 0 {
4268 return;
4269 }
4270 assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4271 let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4272 assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4273 unsafe {
4275 matrixmultiply::sgemm(
4276 m,
4277 k,
4278 n,
4279 alpha,
4280 a.as_ptr(),
4281 lda as isize,
4282 1,
4283 b_mat.as_ptr(),
4284 rsb,
4285 csb,
4286 0.0,
4287 c.as_mut_ptr(),
4288 ldc as isize,
4289 1,
4290 );
4291 }
4292 }
4293}
4294
4295#[cfg(target_os = "macos")]
4298#[allow(clippy::too_many_arguments)]
4299pub(crate) fn sgemm_rm(
4300 m: usize,
4301 n: usize,
4302 k: usize,
4303 alpha: f32,
4304 a: &[f32],
4305 lda: usize,
4306 b_mat: &[f32],
4307 ldb: usize,
4308 b_rows_are_n: bool,
4309 c: &mut [f32],
4310 ldc: usize,
4311) {
4312 debug_assert!(a.len() >= (m - 1) * lda + k);
4313 debug_assert!(c.len() >= (m - 1) * ldc + n);
4314 #[cfg(target_arch = "aarch64")]
4319 if std::env::var("CMF_FORCE_NEON_GEMM")
4320 .map(|v| v == "1")
4321 .unwrap_or(false)
4322 {
4323 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4324 }
4325 unsafe {
4326 accel_blas::cblas_sgemm(
4327 101, 111, if b_rows_are_n { 112 } else { 111 },
4330 m as i32,
4331 n as i32,
4332 k as i32,
4333 alpha,
4334 a.as_ptr(),
4335 lda as i32,
4336 b_mat.as_ptr(),
4337 ldb as i32,
4338 0.0,
4339 c.as_mut_ptr(),
4340 ldc as i32,
4341 );
4342 }
4343}
4344
4345#[cfg(target_os = "macos")]
4352fn qmatmat_accel(
4353 q: &[u8],
4354 row_scale: &[f32],
4355 pre: &[std::borrow::Cow<'_, [f32]>],
4356 rows: usize,
4357 cols: usize,
4358 out: &mut [f32],
4359 pool: Option<&Pool>,
4360) {
4361 const TR: usize = 2048;
4366 let b = pre.len();
4367 thread_local! {
4368 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4369 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4370 }
4371 XPANEL.with(|xp| {
4372 WTILE.with(|wt| {
4373 let mut xpanel = xp.borrow_mut();
4374 xpanel.clear();
4375 for x in pre {
4376 xpanel.extend_from_slice(x);
4377 }
4378 let mut wtile = wt.borrow_mut();
4379 wtile.resize(TR * cols, 0.0);
4380 let mut r0 = 0usize;
4381 while r0 < rows {
4382 let tr = TR.min(rows - r0);
4383 let wt_addr = SendMut(wtile.as_mut_ptr());
4385 let run = |start: usize, end: usize| {
4386 for r in start..end {
4387 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4388 let s = row_scale[r0 + r];
4389 let dst =
4391 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4392 for (d, &v) in dst.iter_mut().zip(row) {
4393 *d = (v as i8) as f32 * s;
4394 }
4395 }
4396 };
4397 dispatch_rows(pool, tr, &run);
4398 unsafe {
4400 accel_blas::cblas_sgemm(
4401 101, 111, 112, b as i32,
4405 tr as i32,
4406 cols as i32,
4407 1.0,
4408 xpanel.as_ptr(),
4409 cols as i32,
4410 wtile.as_ptr(),
4411 cols as i32,
4412 0.0,
4413 out.as_mut_ptr().add(r0),
4414 rows as i32,
4415 );
4416 }
4417 r0 += tr;
4418 }
4419 })
4420 });
4421}
4422
4423fn qmatmat(
4424 q: &[u8],
4425 row_scale: &[f32],
4426 pre: &[std::borrow::Cow<'_, [f32]>],
4427 rows: usize,
4428 cols: usize,
4429 out: &mut [f32],
4430 pool: Option<&Pool>,
4431) {
4432 let b = pre.len();
4433 debug_assert_eq!(out.len(), b * rows);
4434 #[cfg(target_os = "macos")]
4439 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4440 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4441 return;
4442 }
4443 #[cfg(target_arch = "aarch64")]
4444 if sdot_enabled() {
4445 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4446 let out_addr = SendMut(out.as_mut_ptr());
4447 let blocked_ok = blocked_enabled();
4450 let use_i8mm = i8mm_enabled();
4451 if blocked_ok {
4452 let run = |start: usize, end: usize| {
4453 let mut o = start;
4454 while o < end {
4455 if o + 2 <= end {
4456 let r0 = &q[o * cols..(o + 1) * cols];
4457 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4458 let mut bi = 0usize;
4459 while bi + 4 <= acts.len() {
4460 let xs = [
4461 acts[bi].xq.as_slice(),
4462 acts[bi + 1].xq.as_slice(),
4463 acts[bi + 2].xq.as_slice(),
4464 acts[bi + 3].xq.as_slice(),
4465 ];
4466 let d = if use_i8mm {
4467 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4468 } else {
4469 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4470 };
4471 for (r, row) in [r0, r1].into_iter().enumerate() {
4472 for k in 0..4 {
4473 let act = &acts[bi + k];
4474 let mut v = d[r][k] as f32 * act.sx;
4475 for &(j, xv) in &act.outliers {
4476 v += (row[j] as i8) as f32 * xv;
4477 }
4478 unsafe {
4479 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4480 };
4481 }
4482 }
4483 bi += 4;
4484 }
4485 while bi < acts.len() {
4486 for (r, row) in [r0, r1].into_iter().enumerate() {
4487 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4488 unsafe { *out_addr.at(bi * rows + o + r) = v };
4489 }
4490 bi += 1;
4491 }
4492 o += 2;
4493 } else {
4494 let row = &q[o * cols..(o + 1) * cols];
4495 for (bi, act) in acts.iter().enumerate() {
4496 let v = row_dot_sdot(row, act) * row_scale[o];
4497 unsafe { *out_addr.at(bi * rows + o) = v };
4498 }
4499 o += 1;
4500 }
4501 }
4502 };
4503 dispatch_rows(pool, rows, &run);
4504 return;
4505 }
4506 let run = |start: usize, end: usize| {
4507 for o in start..end {
4508 let row = &q[o * cols..(o + 1) * cols];
4509 for (bi, act) in acts.iter().enumerate() {
4510 let v = row_dot_sdot(row, act) * row_scale[o];
4511 unsafe { *out_addr.at(bi * rows + o) = v };
4512 }
4513 }
4514 };
4515 dispatch_rows(pool, rows, &run);
4516 return;
4517 }
4518 #[cfg(target_arch = "x86_64")]
4523 if avx2_a8w8_enabled() {
4524 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4525 let out_addr = SendMut(out.as_mut_ptr());
4526 let blocked_ok = blocked_enabled();
4529 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4530 let run = |start: usize, end: usize| {
4531 let mut o = start;
4532 while o < end {
4533 if o + 2 <= end {
4534 let r0 = &q[o * cols..(o + 1) * cols];
4535 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4536 let mut bi = 0usize;
4537 while bi + 4 <= acts.len() {
4538 let xs = [
4539 acts[bi].xq.as_slice(),
4540 acts[bi + 1].xq.as_slice(),
4541 acts[bi + 2].xq.as_slice(),
4542 acts[bi + 3].xq.as_slice(),
4543 ];
4544 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4545 for (r, row) in [r0, r1].into_iter().enumerate() {
4546 for k in 0..4 {
4547 let act = &acts[bi + k];
4548 let mut v = d[r][k] as f32 * act.sx;
4549 for &(j, xv) in &act.outliers {
4550 v += (row[j] as i8) as f32 * xv;
4551 }
4552 unsafe {
4553 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4554 };
4555 }
4556 }
4557 bi += 4;
4558 }
4559 while bi < acts.len() {
4560 for (r, row) in [r0, r1].into_iter().enumerate() {
4561 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4562 unsafe { *out_addr.at(bi * rows + o + r) = v };
4563 }
4564 bi += 1;
4565 }
4566 o += 2;
4567 } else {
4568 let row = &q[o * cols..(o + 1) * cols];
4569 for (bi, act) in acts.iter().enumerate() {
4570 let v = row_dot_avx2(row, act) * row_scale[o];
4571 unsafe { *out_addr.at(bi * rows + o) = v };
4572 }
4573 o += 1;
4574 }
4575 }
4576 };
4577 dispatch_rows(pool, rows, &run);
4578 return;
4579 }
4580 let run = |start: usize, end: usize| {
4581 for o in start..end {
4582 let row = &q[o * cols..(o + 1) * cols];
4583 for (bi, act) in acts.iter().enumerate() {
4584 let v = row_dot_avx2(row, act) * row_scale[o];
4585 unsafe { *out_addr.at(bi * rows + o) = v };
4586 }
4587 }
4588 };
4589 dispatch_rows(pool, rows, &run);
4590 return;
4591 }
4592 let out_addr = SendMut(out.as_mut_ptr());
4593 let run = |start: usize, end: usize| {
4594 for o in start..end {
4595 let row = &q[o * cols..(o + 1) * cols];
4596 for (bi, x) in pre.iter().enumerate() {
4597 let mut acc = 0f32;
4598 for j in 0..cols {
4599 acc += (row[j] as i8) as f32 * x[j];
4600 }
4601 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4602 }
4603 }
4604 };
4605 dispatch_rows(pool, rows, &run);
4606}
4607
4608fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4611 match pool {
4612 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4613 _ => run(0, rows),
4614 }
4615}
4616
4617fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4619 let groups = rows * cols / GROUP_SIZE;
4620 bytes.split_at(groups * 16)
4621}
4622
4623#[inline]
4628fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4629 #[cfg(target_arch = "aarch64")]
4630 unsafe {
4631 return vbit_fill4_neon(data, buf);
4632 }
4633 #[cfg(target_arch = "x86_64")]
4634 if avx2_enabled() {
4635 return unsafe { vbit_fill4_avx2(data, buf) };
4636 }
4637 #[allow(unreachable_code)]
4638 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4639 let u = unpack8::<4>(&data[blk * 4..]);
4640 for k in 0..8 {
4641 chunk[k] = (u[k] - 7) as i8 as u8;
4642 }
4643 }
4644}
4645
4646#[cfg(target_arch = "aarch64")]
4647#[target_feature(enable = "neon")]
4648unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4649 unsafe {
4652 use core::arch::aarch64::*;
4653 let n = buf.len();
4654 let mask = vdupq_n_u8(0x0F);
4655 let seven = vdupq_n_s8(7);
4656 let mut g = 0usize;
4657 while g * 32 + 32 <= n {
4658 let b = vld1q_u8(data.as_ptr().add(g * 16));
4659 let hi = vshrq_n_u8::<4>(b);
4660 let lo = vandq_u8(b, mask);
4661 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4662 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4663 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4664 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4665 g += 1;
4666 }
4667 }
4668}
4669
4670#[cfg(target_arch = "x86_64")]
4671#[target_feature(enable = "avx2")]
4672unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4673 unsafe {
4675 use core::arch::x86_64::*;
4676 let n = buf.len();
4677 let mask = _mm_set1_epi8(0x0F);
4678 let seven = _mm256_set1_epi8(7);
4679 let mut g = 0usize;
4680 while g * 32 + 32 <= n {
4681 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4682 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4683 let lo = _mm_and_si128(b, mask);
4684 let z = _mm256_sub_epi8(
4685 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4686 seven,
4687 );
4688 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4689 g += 1;
4690 }
4691 }
4692}
4693
4694#[inline(always)]
4699fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4700 let mut acc = 0u64;
4701 for i in 0..B {
4702 acc = (acc << 8) | data[i] as u64;
4703 }
4704 let mask = (1u64 << B) - 1;
4705 let mut out = [0i32; 8];
4706 for (k, o) in out.iter_mut().enumerate() {
4707 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4708 }
4709 out
4710}
4711
4712#[allow(clippy::too_many_arguments)]
4718fn vbitmatvec(
4719 bytes: &[u8],
4720 offsets: &[usize],
4721 x: &[f32],
4722 rows: usize,
4723 cols: usize,
4724 out: &mut [f32],
4725 pool: Option<&Pool>,
4726) {
4727 debug_assert_eq!(out.len(), rows);
4728 debug_assert_eq!(offsets.len(), rows + 1);
4729
4730 if a8w8_enabled() {
4734 let act = split_act(x);
4735 let out_addr = SendMut(out.as_mut_ptr());
4736 let run = move |start: usize, end: usize| {
4737 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4738 };
4739 dispatch_rows(pool, rows, &run);
4740 return;
4741 }
4742
4743 let out_addr = SendMut(out.as_mut_ptr());
4744 let run = move |start: usize, end: usize| {
4745 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4746 };
4747 dispatch_rows(pool, rows, &run);
4748}
4749
4750#[allow(clippy::too_many_arguments)]
4754fn vbit_range_a8w8(
4755 bytes: &[u8],
4756 offsets: &[usize],
4757 x: &[f32],
4758 act: &SplitAct,
4759 rows: usize,
4760 cols: usize,
4761 out: SendMut,
4762 start: usize,
4763 end: usize,
4764) {
4765 let ng = cols / GROUP_SIZE;
4766 let bits = &bytes[..rows];
4767 let sc_off = rows;
4768 let row_dot = |r: usize| -> f32 {
4769 let b = bits[r] as usize;
4770 let l = (1i32 << (b - 1)) - 1;
4771 let mask = (1u64 << b) - 1;
4772 let data = &bytes[offsets[r]..offsets[r + 1]];
4773 if b == 8 {
4774 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4776 let mut dot = 0f32;
4777 for g in 0..ng {
4778 let so = (r * ng + g) * 2;
4779 let sgf = f16_to_f32(u16::from_le_bytes([
4780 bytes[sc_off + so],
4781 bytes[sc_off + so + 1],
4782 ]));
4783 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4784 let mut gd = 0f32;
4785 for &xv in xg.iter() {
4786 if nbits < 8 {
4787 acc = (acc << 8) | data[idx] as u64;
4788 idx += 1;
4789 nbits += 8;
4790 }
4791 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4792 nbits -= 8;
4793 gd += (u - l) as f32 * xv;
4794 }
4795 dot += gd * sgf;
4796 }
4797 return dot;
4798 }
4799 thread_local! {
4803 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4804 const { std::cell::RefCell::new(Vec::new()) };
4805 }
4806 #[inline(always)]
4807 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4808 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4809 let u = unpack8::<B>(&data[blk * B..]);
4810 for k in 0..8 {
4811 chunk[k] = (u[k] - l) as i8 as u8;
4812 }
4813 }
4814 }
4815 let _ = mask;
4816 VBIT_SCRATCH.with(|scratch| {
4817 let mut buf = scratch.borrow_mut();
4818 buf.resize(cols, 0);
4819 match b {
4820 3 => fill::<3>(data, l, &mut buf),
4821 4 => vbit_fill4(data, &mut buf),
4822 5 => fill::<5>(data, l, &mut buf),
4823 6 => fill::<6>(data, l, &mut buf),
4824 _ => unreachable!(),
4825 }
4826 let mut dot = 0f32;
4827 for g in 0..ng {
4828 let so = (r * ng + g) * 2;
4829 let s = f16_to_f32(u16::from_le_bytes([
4830 bytes[sc_off + so],
4831 bytes[sc_off + so + 1],
4832 ]));
4833 let d = dot_i8_i8(
4834 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4835 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4836 ) as f32
4837 * act.sx;
4838 dot += d * s;
4839 }
4840 for &(j, xv) in &act.outliers {
4841 let so = (r * ng + j / GROUP_SIZE) * 2;
4842 let s = f16_to_f32(u16::from_le_bytes([
4843 bytes[sc_off + so],
4844 bytes[sc_off + so + 1],
4845 ]));
4846 dot += (buf[j] as i8) as f32 * s * xv;
4848 }
4849 dot
4850 })
4851 };
4852 for r in start..end {
4853 unsafe { *out.at(r) = row_dot(r) };
4855 }
4856}
4857
4858#[allow(clippy::too_many_arguments)]
4860fn vbit_range_f32(
4861 bytes: &[u8],
4862 offsets: &[usize],
4863 x: &[f32],
4864 rows: usize,
4865 cols: usize,
4866 out: SendMut,
4867 start: usize,
4868 end: usize,
4869) {
4870 let ng = cols / GROUP_SIZE;
4871 let bits = &bytes[..rows];
4872 let sc_off = rows;
4873 #[inline(always)]
4877 fn dot_row<const B: usize>(
4878 data: &[u8],
4879 bytes: &[u8],
4880 sc_off: usize,
4881 r: usize,
4882 ng: usize,
4883 x: &[f32],
4884 ) -> f32 {
4885 let l = ((1i32 << (B - 1)) - 1) as f32;
4886 let gbytes = GROUP_SIZE * B / 8;
4887 let mut dot = 0f32;
4888 for g in 0..ng {
4889 let so = (r * ng + g) * 2;
4890 let s = f16_to_f32(u16::from_le_bytes([
4891 bytes[sc_off + so],
4892 bytes[sc_off + so + 1],
4893 ]));
4894 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4895 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4896 let mut gd = 0f32;
4897 for blk in 0..GROUP_SIZE / 8 {
4898 let u = unpack8::<B>(&gd0[blk * B..]);
4899 let xb = &xg[blk * 8..blk * 8 + 8];
4900 for k in 0..8 {
4901 gd += (u[k] as f32 - l) * xb[k];
4902 }
4903 }
4904 dot += gd * s;
4905 }
4906 dot
4907 }
4908 for r in start..end {
4909 let data = &bytes[offsets[r]..offsets[r + 1]];
4910 let v = match bits[r] {
4911 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4912 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4913 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4914 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4915 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4916 b => unreachable!("vbit bit-width {b} (validated at load)"),
4917 };
4918 unsafe { *out.at(r) = v };
4920 }
4921}
4922
4923#[allow(clippy::too_many_arguments)]
4928fn vbitmatvec2(
4929 bytes: &[u8],
4930 offsets: &[usize],
4931 x1: &[f32],
4932 x2: &[f32],
4933 rows: usize,
4934 cols: usize,
4935 o1: &mut [f32],
4936 o2: &mut [f32],
4937 pool: Option<&Pool>,
4938) {
4939 debug_assert_eq!(o1.len(), rows);
4940 debug_assert_eq!(o2.len(), rows);
4941
4942 if a8w8_enabled() {
4943 let a1 = split_act(x1);
4944 let a2 = split_act(x2);
4945 let p1 = SendMut(o1.as_mut_ptr());
4946 let p2 = SendMut(o2.as_mut_ptr());
4947 let run = move |start: usize, end: usize| {
4948 vbit_range2_a8w8(
4949 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4950 )
4951 };
4952 dispatch_rows(pool, rows, &run);
4953 return;
4954 }
4955
4956 let p1 = SendMut(o1.as_mut_ptr());
4957 let p2 = SendMut(o2.as_mut_ptr());
4958 let run = move |start: usize, end: usize| {
4959 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4960 };
4961 dispatch_rows(pool, rows, &run);
4962}
4963
4964#[allow(clippy::too_many_arguments)]
4968fn vbit_range2_a8w8(
4969 bytes: &[u8],
4970 offsets: &[usize],
4971 x1: &[f32],
4972 x2: &[f32],
4973 a1: &SplitAct,
4974 a2: &SplitAct,
4975 rows: usize,
4976 cols: usize,
4977 p1: SendMut,
4978 p2: SendMut,
4979 start: usize,
4980 end: usize,
4981) {
4982 let ng = cols / GROUP_SIZE;
4983 let bits = &bytes[..rows];
4984 let sc_off = rows;
4985 let row_dots = |r: usize| -> (f32, f32) {
4986 let b = bits[r] as usize;
4987 let l = (1i32 << (b - 1)) - 1;
4988 let data = &bytes[offsets[r]..offsets[r + 1]];
4989 if b == 8 {
4990 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4993 let (mut d1, mut d2) = (0f32, 0f32);
4994 for g in 0..ng {
4995 let so = (r * ng + g) * 2;
4996 let sgf = f16_to_f32(u16::from_le_bytes([
4997 bytes[sc_off + so],
4998 bytes[sc_off + so + 1],
4999 ]));
5000 let (mut g1, mut g2) = (0f32, 0f32);
5001 for k in 0..GROUP_SIZE {
5002 if nbits < 8 {
5003 acc = (acc << 8) | data[idx] as u64;
5004 idx += 1;
5005 nbits += 8;
5006 }
5007 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5008 nbits -= 8;
5009 let w = (u - l) as f32;
5010 g1 += w * x1[g * GROUP_SIZE + k];
5011 g2 += w * x2[g * GROUP_SIZE + k];
5012 }
5013 d1 += g1 * sgf;
5014 d2 += g2 * sgf;
5015 }
5016 return (d1, d2);
5017 }
5018 thread_local! {
5019 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5020 const { std::cell::RefCell::new(Vec::new()) };
5021 }
5022 #[inline(always)]
5023 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5024 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5025 let u = unpack8::<B>(&data[blk * B..]);
5026 for k in 0..8 {
5027 chunk[k] = (u[k] - l) as i8 as u8;
5028 }
5029 }
5030 }
5031 VBIT_SCRATCH2.with(|scratch| {
5032 let mut buf = scratch.borrow_mut();
5033 buf.resize(cols, 0);
5034 match b {
5035 3 => fill::<3>(data, l, &mut buf),
5036 4 => vbit_fill4(data, &mut buf),
5037 5 => fill::<5>(data, l, &mut buf),
5038 6 => fill::<6>(data, l, &mut buf),
5039 _ => unreachable!(),
5040 }
5041 let (mut d1, mut d2) = (0f32, 0f32);
5042 for g in 0..ng {
5043 let so = (r * ng + g) * 2;
5044 let s = f16_to_f32(u16::from_le_bytes([
5045 bytes[sc_off + so],
5046 bytes[sc_off + so + 1],
5047 ]));
5048 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5049 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5050 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5051 d1 += v1 * s;
5052 d2 += v2 * s;
5053 }
5054 for &(j, xv) in &a1.outliers {
5055 let so = (r * ng + j / GROUP_SIZE) * 2;
5056 let s = f16_to_f32(u16::from_le_bytes([
5057 bytes[sc_off + so],
5058 bytes[sc_off + so + 1],
5059 ]));
5060 d1 += (buf[j] as i8) as f32 * s * xv;
5061 }
5062 for &(j, xv) in &a2.outliers {
5063 let so = (r * ng + j / GROUP_SIZE) * 2;
5064 let s = f16_to_f32(u16::from_le_bytes([
5065 bytes[sc_off + so],
5066 bytes[sc_off + so + 1],
5067 ]));
5068 d2 += (buf[j] as i8) as f32 * s * xv;
5069 }
5070 (d1, d2)
5071 })
5072 };
5073 for r in start..end {
5074 let (v1, v2) = row_dots(r);
5075 unsafe {
5077 *p1.at(r) = v1;
5078 *p2.at(r) = v2;
5079 }
5080 }
5081}
5082
5083#[allow(clippy::too_many_arguments)]
5087fn vbit_range2_f32(
5088 bytes: &[u8],
5089 offsets: &[usize],
5090 x1: &[f32],
5091 x2: &[f32],
5092 rows: usize,
5093 cols: usize,
5094 p1: SendMut,
5095 p2: SendMut,
5096 start: usize,
5097 end: usize,
5098) {
5099 let ng = cols / GROUP_SIZE;
5100 let bits = &bytes[..rows];
5101 let sc_off = rows;
5102 #[inline(always)]
5103 #[allow(clippy::too_many_arguments)]
5104 fn dot_row2<const B: usize>(
5105 data: &[u8],
5106 bytes: &[u8],
5107 sc_off: usize,
5108 r: usize,
5109 ng: usize,
5110 x1: &[f32],
5111 x2: &[f32],
5112 ) -> (f32, f32) {
5113 let l = ((1i32 << (B - 1)) - 1) as f32;
5114 let gbytes = GROUP_SIZE * B / 8;
5115 let (mut d1, mut d2) = (0f32, 0f32);
5116 for g in 0..ng {
5117 let so = (r * ng + g) * 2;
5118 let s = f16_to_f32(u16::from_le_bytes([
5119 bytes[sc_off + so],
5120 bytes[sc_off + so + 1],
5121 ]));
5122 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5123 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5124 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5125 let (mut g1, mut g2) = (0f32, 0f32);
5126 for blk in 0..GROUP_SIZE / 8 {
5127 let u = unpack8::<B>(&gd0[blk * B..]);
5128 for k in 0..8 {
5129 let w = u[k] as f32 - l;
5130 g1 += w * x1g[blk * 8 + k];
5131 g2 += w * x2g[blk * 8 + k];
5132 }
5133 }
5134 d1 += g1 * s;
5135 d2 += g2 * s;
5136 }
5137 (d1, d2)
5138 }
5139 for r in start..end {
5140 let data = &bytes[offsets[r]..offsets[r + 1]];
5141 let (v1, v2) = match bits[r] {
5142 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5143 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5144 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5145 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5146 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5147 b => unreachable!("vbit bit-width {b} (validated at load)"),
5148 };
5149 unsafe {
5151 *p1.at(r) = v1;
5152 *p2.at(r) = v2;
5153 }
5154 }
5155}
5156
5157#[inline]
5164#[allow(unreachable_code)]
5165fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5166 #[cfg(target_arch = "aarch64")]
5167 unsafe {
5168 return dot_q4t_row_sdot(bytes, r, gpr, xq);
5169 }
5170 #[cfg(target_arch = "x86_64")]
5171 unsafe {
5172 if vnni_tiles_enabled() {
5173 return dot_q4t_row_vnni(bytes, r, gpr, xq);
5174 }
5175 return dot_q4t_row_avx2(bytes, r, gpr, xq);
5176 }
5177 let mut acc = 0f32;
5178 for gi in 0..gpr {
5179 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5180 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5181 let mut d = 0i32;
5182 for (k, &b) in tile[2..].iter().enumerate() {
5183 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5184 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5185 }
5186 acc += d as f32 * s;
5187 }
5188 acc
5189}
5190
5191#[cfg(target_arch = "aarch64")]
5192#[target_feature(enable = "neon,dotprod")]
5193unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5194 unsafe {
5197 use core::arch::aarch64::*;
5198 use core::arch::asm;
5199 let lomask = vdupq_n_u8(0x0F);
5200 let eight = vdupq_n_s8(8);
5201 let mut acc = 0f32;
5202 for gi in 0..gpr {
5203 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5204 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5205 let b = vld1q_u8(t.add(2));
5206 let lo = vandq_u8(b, lomask);
5207 let hi = vshrq_n_u8::<4>(b);
5208 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5209 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5210 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5211 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5212 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5213 asm!(
5214 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5215 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5216 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5217 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5218 options(pure, nomem, nostack),
5219 );
5220 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5221 }
5222 acc
5223 }
5224}
5225
5226#[cfg(target_arch = "x86_64")]
5227#[target_feature(enable = "avx2")]
5228unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5229 unsafe {
5231 use core::arch::x86_64::*;
5232 let lomask = _mm_set1_epi8(0x0F);
5233 let eight = _mm256_set1_epi8(8);
5234 let ones = _mm256_set1_epi16(1);
5235 let mut acc = 0f32;
5236 for gi in 0..gpr {
5237 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5238 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5239 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5240 let lo = _mm_and_si128(b, lomask);
5241 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5242 let w = _mm256_sub_epi8(
5243 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5244 eight,
5245 );
5246 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5247 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5248 let d = _mm256_madd_epi16(p16, ones);
5249 let hi128 = _mm256_extracti128_si256::<1>(d);
5250 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5251 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5252 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5253 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5254 }
5255 acc
5256 }
5257}
5258
5259#[cfg(target_arch = "x86_64")]
5263#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5264unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5265 unsafe {
5267 use core::arch::x86_64::*;
5268 let lomask = _mm_set1_epi8(0x0F);
5269 let eight = _mm256_set1_epi8(8);
5270 let mut acc = 0f32;
5271 for gi in 0..gpr {
5272 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5273 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5274 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5275 let lo = _mm_and_si128(b, lomask);
5276 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5277 let w = _mm256_sub_epi8(
5278 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5279 eight,
5280 );
5281 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5282 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5283 acc += d as f32 * s;
5284 }
5285 acc
5286 }
5287}
5288
5289#[cfg(target_arch = "x86_64")]
5294#[target_feature(enable = "avx2,fma")]
5299unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5300 unsafe {
5302 use core::arch::x86_64::*;
5303 let lomask = _mm_set1_epi8(0x0F);
5304 let eight = _mm256_set1_epi8(8);
5305 let ones = _mm256_set1_epi16(1);
5306 let mut f0 = _mm256_setzero_ps();
5319 let mut f1 = _mm256_setzero_ps();
5320 let mut f2 = _mm256_setzero_ps();
5321 let mut f3 = _mm256_setzero_ps();
5322 for gi in 0..gpr {
5323 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5324 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5325 let sv = _mm256_set1_ps(s);
5326 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5327 let lo = _mm_and_si128(bb, lomask);
5328 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5329 let w = _mm256_sub_epi8(
5330 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5331 eight,
5332 );
5333 let aw = _mm256_abs_epi8(w);
5334 let off = gi * GROUP_SIZE;
5335 let dot = |xq: &[i8]| {
5336 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5337 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5338 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5339 };
5340 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5341 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5342 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5343 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5344 }
5345 [
5346 hsum256_ps(f0),
5347 hsum256_ps(f1),
5348 hsum256_ps(f2),
5349 hsum256_ps(f3),
5350 ]
5351 }
5352}
5353
5354#[cfg(target_arch = "x86_64")]
5357#[target_feature(enable = "avx2")]
5358#[inline]
5359unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5360 unsafe {
5362 use core::arch::x86_64::*;
5363 let hi = _mm256_extractf128_ps::<1>(v);
5364 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5365 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5366 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5367 _mm_cvtss_f32(s)
5368 }
5369}
5370
5371#[cfg(target_arch = "x86_64")]
5373#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5374unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5375 unsafe {
5377 use core::arch::x86_64::*;
5378 let lomask = _mm_set1_epi8(0x0F);
5379 let eight = _mm256_set1_epi8(8);
5380 let mut f0 = _mm256_setzero_ps();
5383 let mut f1 = _mm256_setzero_ps();
5384 let mut f2 = _mm256_setzero_ps();
5385 let mut f3 = _mm256_setzero_ps();
5386 for gi in 0..gpr {
5387 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5388 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5389 let sv = _mm256_set1_ps(s);
5390 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5391 let lo = _mm_and_si128(bb, lomask);
5392 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5393 let w = _mm256_sub_epi8(
5394 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5395 eight,
5396 );
5397 let aw = _mm256_abs_epi8(w);
5398 let off = gi * GROUP_SIZE;
5399 let dot = |xq: &[i8]| {
5400 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5401 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5402 _mm256_setzero_si256(),
5403 aw,
5404 _mm256_sign_epi8(x, w),
5405 ))
5406 };
5407 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5408 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5409 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5410 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5411 }
5412 let acc = [
5413 hsum256_ps(f0),
5414 hsum256_ps(f1),
5415 hsum256_ps(f2),
5416 hsum256_ps(f3),
5417 ];
5418 acc
5419 }
5420}
5421
5422#[cfg(target_arch = "aarch64")]
5427#[target_feature(enable = "neon,dotprod")]
5428unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5429 unsafe {
5431 use core::arch::aarch64::*;
5432 use core::arch::asm;
5433 let lomask = vdupq_n_u8(0x0F);
5434 let eight = vdupq_n_s8(8);
5435 let mut acc = [0f32; 4];
5436 for gi in 0..gpr {
5437 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5438 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5439 let b = vld1q_u8(t.add(2));
5440 let lo = vandq_u8(b, lomask);
5441 let hi = vshrq_n_u8::<4>(b);
5442 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5443 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5444 for (k, xq) in xs.iter().enumerate() {
5445 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5446 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5447 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5448 asm!(
5449 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5450 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5451 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5452 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5453 options(pure, nomem, nostack),
5454 );
5455 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5456 }
5457 }
5458 acc
5459 }
5460}
5461
5462#[inline]
5464fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5465 let gi = j / GROUP_SIZE;
5466 let k = j % GROUP_SIZE;
5467 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5468 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5469 let byte = tile[2 + k / 2];
5470 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5471 ((nib as i32 - 8) as f32, s)
5472}
5473
5474#[inline]
5477fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5478 let mut acc = 0f32;
5479 for gi in 0..gpr {
5480 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5481 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5482 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5483 let mut ga = 0f32;
5484 for (k, &b) in tile[2..].iter().enumerate() {
5485 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5486 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5487 }
5488 acc += ga * s;
5489 }
5490 acc
5491}
5492
5493struct Q4tpView<'a> {
5497 nib: &'a [u8],
5498 params: &'a [u8],
5499 codes: &'a [u8],
5500 stride: usize,
5501 zero_rung: bool,
5503}
5504
5505impl<'a> Q4tpView<'a> {
5506 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5507 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5508 Self {
5509 nib: &bytes[..params_off],
5510 params: &bytes[params_off..codes_off],
5511 codes: &bytes[codes_off..],
5512 stride,
5513 zero_rung: false,
5514 }
5515 }
5516
5517 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5519 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5520 Self {
5521 nib: &bytes[..params_off],
5522 params: &bytes[params_off..codes_off],
5523 codes: &bytes[codes_off..],
5524 stride,
5525 zero_rung: true,
5526 }
5527 }
5528
5529 #[inline]
5545 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5546 let tab = if self.zero_rung {
5547 q2tp_ladder(self.params, r)
5548 } else {
5549 q4tp_ladder(self.params, r)
5550 };
5551 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5552 let out = &mut out[..gpr];
5553 let mut chunks = out.chunks_exact_mut(8);
5554 let mut ci = 0usize;
5555 for c in &mut chunks {
5556 let w = u64::from(codes[ci])
5557 | u64::from(codes[ci + 1]) << 8
5558 | u64::from(codes[ci + 2]) << 16
5559 | u64::from(codes[ci + 3]) << 24
5560 | u64::from(codes[ci + 4]) << 32;
5561 for (k, o) in c.iter_mut().enumerate() {
5562 *o = tab[((w >> (5 * k)) & 31) as usize];
5563 }
5564 ci += 5;
5565 }
5566 let tail = &codes[ci..];
5569 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5570 *o = tab[q4tp_code(tail, k)];
5571 }
5572 }
5573}
5574
5575#[inline]
5576fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5577 #[cfg(target_arch = "aarch64")]
5578 unsafe {
5579 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5580 }
5581 #[cfg(target_arch = "x86_64")]
5582 unsafe {
5583 if vnni_tiles_enabled() {
5584 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5585 }
5586 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5587 }
5588 #[allow(unreachable_code)]
5589 {
5590 let mut acc = 0f32;
5591 for gi in 0..gpr {
5592 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5593 let s = scales[gi];
5594 let mut d = 0i32;
5595 for (k, &b) in tile.iter().enumerate() {
5596 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5597 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5598 }
5599 acc += d as f32 * s;
5600 }
5601 acc
5602 }
5603}
5604
5605#[cfg(target_arch = "aarch64")]
5608#[target_feature(enable = "neon,dotprod")]
5609unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5610 unsafe {
5613 use core::arch::aarch64::*;
5614 use core::arch::asm;
5615 let lomask = vdupq_n_u8(0x0F);
5616 let eight = vdupq_n_s8(8);
5617 let mut acc = 0f32;
5618 for gi in 0..gpr {
5619 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5620 let s = *scales.get_unchecked(gi);
5621 let b = vld1q_u8(t);
5622 let lo = vandq_u8(b, lomask);
5623 let hi = vshrq_n_u8::<4>(b);
5624 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5625 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5626 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5627 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5628 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5629 asm!(
5630 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5631 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5632 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5633 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5634 options(pure, nomem, nostack),
5635 );
5636 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5637 }
5638 acc
5639 }
5640}
5641
5642#[cfg(target_arch = "x86_64")]
5643#[target_feature(enable = "avx2")]
5644unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5645 unsafe {
5647 use core::arch::x86_64::*;
5648 let lomask = _mm_set1_epi8(0x0F);
5649 let eight = _mm256_set1_epi8(8);
5650 let ones = _mm256_set1_epi16(1);
5651 let mut acc = 0f32;
5652 for gi in 0..gpr {
5653 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5654 let s = *scales.get_unchecked(gi);
5655 let b = _mm_loadu_si128(t as *const __m128i);
5656 let lo = _mm_and_si128(b, lomask);
5657 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5658 let w = _mm256_sub_epi8(
5659 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5660 eight,
5661 );
5662 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5663 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5664 let d = _mm256_madd_epi16(p16, ones);
5665 let hi128 = _mm256_extracti128_si256::<1>(d);
5666 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5667 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5668 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5669 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5670 }
5671 acc
5672 }
5673}
5674
5675#[cfg(target_arch = "x86_64")]
5678#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5679unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5680 unsafe {
5682 use core::arch::x86_64::*;
5683 let lomask = _mm_set1_epi8(0x0F);
5684 let eight = _mm256_set1_epi8(8);
5685 let mut acc = 0f32;
5686 for gi in 0..gpr {
5687 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5688 let s = *scales.get_unchecked(gi);
5689 let b = _mm_loadu_si128(t as *const __m128i);
5690 let lo = _mm_and_si128(b, lomask);
5691 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5692 let w = _mm256_sub_epi8(
5693 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5694 eight,
5695 );
5696 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5697 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5698 }
5699 acc
5700 }
5701}
5702
5703#[inline]
5706fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5707 #[cfg(target_arch = "x86_64")]
5708 if avx2_enabled() {
5709 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5711 }
5712 q4tp_row_float_scalar(nib, r, gpr, x, scales)
5713}
5714
5715#[inline]
5716fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5717 let mut acc = 0f32;
5718 for gi in 0..gpr {
5719 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5720 let s = scales[gi];
5721 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5722 let mut ga = 0f32;
5723 for (k, &b) in tile.iter().enumerate() {
5724 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5725 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5726 }
5727 acc += ga * s;
5728 }
5729 acc
5730}
5731
5732#[cfg(target_arch = "x86_64")]
5736#[target_feature(enable = "avx2")]
5737unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5738 unsafe {
5741 use core::arch::x86_64::*;
5742 let mask = _mm_set1_epi8(15);
5743 let eight = _mm_set1_epi8(8);
5744 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5745 let mut acc = 0.0f32;
5746 for gi in 0..gpr {
5747 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5748 let lo = _mm_and_si128(packed, mask);
5749 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5750 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5751 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5752 let xp = x.as_ptr().add(gi * GROUP_SIZE);
5753 let a = _mm256_mul_ps(
5754 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5755 _mm256_loadu_ps(xp),
5756 );
5757 let b = _mm256_mul_ps(
5758 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5759 _mm256_loadu_ps(xp.add(8)),
5760 );
5761 let c = _mm256_mul_ps(
5762 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5763 _mm256_loadu_ps(xp.add(16)),
5764 );
5765 let d = _mm256_mul_ps(
5766 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5767 _mm256_loadu_ps(xp.add(24)),
5768 );
5769 let mut pairs = [0.0f32; 16];
5770 _mm256_storeu_ps(
5771 pairs.as_mut_ptr(),
5772 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5773 );
5774 _mm256_storeu_ps(
5775 pairs.as_mut_ptr().add(8),
5776 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5777 );
5778 let mut ga = 0.0f32;
5779 for v in pairs {
5780 ga += v;
5781 }
5782 acc += ga * scales[gi];
5783 }
5784 acc
5785 }
5786}
5787
5788#[inline]
5791fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5792 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5793 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5794 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5795 ((n as i32 - 8) as f32, scales[gi])
5796}
5797
5798fn q4tp_matvec(
5800 bytes: &[u8],
5801 x: &[f32],
5802 rows: usize,
5803 cols: usize,
5804 out: &mut [f32],
5805 pool: Option<&Pool>,
5806) {
5807 debug_assert_eq!(out.len(), rows);
5808 let gpr = cols / GROUP_SIZE;
5809 let v = Q4tpView::new(bytes, rows, cols);
5810 let out_addr = SendMut(out.as_mut_ptr());
5811 if a8w8_enabled() {
5812 let act = split_act(x);
5813 let run = |start: usize, end: usize| {
5814 with_krow(gpr, |sc| {
5816 for r in start..end {
5817 v.scales_into(r, gpr, sc);
5818 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5819 for &(j, xv) in &act.outliers {
5820 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5821 acc += w * s * xv;
5822 }
5823 unsafe { *out_addr.at(r) = acc };
5825 }
5826 })
5827 };
5828 dispatch_rows(pool, rows, &run);
5829 return;
5830 }
5831 let run = |start: usize, end: usize| {
5832 with_krow(gpr, |sc| {
5833 for r in start..end {
5834 v.scales_into(r, gpr, sc);
5835 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5837 }
5838 })
5839 };
5840 dispatch_rows(pool, rows, &run);
5841}
5842
5843#[allow(clippy::too_many_arguments)]
5846fn q4tp_matvec2(
5847 bytes: &[u8],
5848 x1: &[f32],
5849 x2: &[f32],
5850 rows: usize,
5851 cols: usize,
5852 o1: &mut [f32],
5853 o2: &mut [f32],
5854 pool: Option<&Pool>,
5855) {
5856 let gpr = cols / GROUP_SIZE;
5857 let v = Q4tpView::new(bytes, rows, cols);
5858 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5859 let run = |start: usize, end: usize| {
5860 let mut sc = vec![0f32; gpr];
5861 for r in start..end {
5862 v.scales_into(r, gpr, &mut sc);
5863 unsafe {
5865 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5866 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5867 }
5868 }
5869 };
5870 dispatch_rows(pool, rows, &run);
5871}
5872
5873#[inline]
5876fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5877 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5878 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5879 let c = (byte >> (2 * (k % 4))) & 3;
5880 (c as f32 - 1.5, scales[gi])
5881}
5882
5883#[cfg(target_arch = "x86_64")]
5884const Q2TP_DECODE_U32: [u32; 256] = {
5885 let mut tab = [0u32; 256];
5886 let mut b = 0usize;
5887 while b < 256 {
5888 tab[b] = ((b as u32) & 3)
5889 | ((((b as u32) >> 2) & 3) << 8)
5890 | ((((b as u32) >> 4) & 3) << 16)
5891 | ((((b as u32) >> 6) & 3) << 24);
5892 b += 1;
5893 }
5894 tab
5895};
5896
5897#[cfg(target_arch = "x86_64")]
5901#[target_feature(enable = "avx2")]
5902unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5903 use core::arch::x86_64::*;
5904 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5905 let codes = _mm256_setr_epi32(
5906 Q2TP_DECODE_U32[ch[0] as usize] as i32,
5907 Q2TP_DECODE_U32[ch[1] as usize] as i32,
5908 Q2TP_DECODE_U32[ch[2] as usize] as i32,
5909 Q2TP_DECODE_U32[ch[3] as usize] as i32,
5910 Q2TP_DECODE_U32[ch[4] as usize] as i32,
5911 Q2TP_DECODE_U32[ch[5] as usize] as i32,
5912 Q2TP_DECODE_U32[ch[6] as usize] as i32,
5913 Q2TP_DECODE_U32[ch[7] as usize] as i32,
5914 );
5915 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5916 let pair = _mm256_maddubs_epi16(codes, xv);
5917 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5918 let sum128 = _mm_add_epi32(
5919 _mm256_castsi256_si128(quad),
5920 _mm256_extracti128_si256(quad, 1),
5921 );
5922 let sum64 = _mm_hadd_epi32(sum128, sum128);
5923 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5924}
5925
5926#[cfg(target_arch = "x86_64")]
5927#[inline]
5928fn q2tp_avx2_enabled() -> bool {
5929 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5930 *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
5931}
5932
5933#[inline]
5940fn dot_q2tp_row_i8(
5941 chunks: &[u8],
5942 r: usize,
5943 gpr: usize,
5944 xq: &[i8],
5945 gsum: &[i32],
5946 scales: &[f32],
5947) -> f32 {
5948 let mut acc = 0f32;
5949 let base = r * gpr * Q2TP_CHUNK;
5950 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5951 let mut codes = [0i8; GROUP_SIZE];
5952 #[cfg(target_arch = "x86_64")]
5953 let avx2 = q2tp_avx2_enabled();
5958 for gi in 0..gpr {
5959 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5960 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5961 #[cfg(target_arch = "aarch64")]
5962 let dot = unsafe {
5968 use core::arch::aarch64::*;
5969 let b = vld1_u8(ch.as_ptr());
5970 let three = vdup_n_u8(3);
5971 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5972 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5973 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5974 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5975 let x4 = vld4_s8(xg.as_ptr());
5976 let mut acc4 = vdupq_n_s32(0);
5977 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5978 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5979 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5980 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5981 vaddvq_s32(acc4)
5982 };
5983 #[cfg(target_arch = "x86_64")]
5984 let dot: i32 = if avx2 {
5985 unsafe { q2tp_code_dot_avx2(ch, xg) }
5988 } else {
5989 ch.iter()
5990 .enumerate()
5991 .map(|(k, &b)| {
5992 ((b & 3) as i32) * xg[k * 4] as i32
5993 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5994 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5995 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5996 })
5997 .sum()
5998 };
5999 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6000 let dot: i32 = {
6001 for (k, &b) in ch.iter().enumerate() {
6002 codes[k * 4] = (b & 3) as i8;
6003 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
6004 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
6005 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
6006 }
6007 codes
6008 .iter()
6009 .zip(xg)
6010 .map(|(&c, &x)| c as i32 * x as i32)
6011 .sum()
6012 };
6013 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
6014 }
6015 acc
6016}
6017
6018fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6022 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
6023}
6024
6025#[inline]
6029fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6030 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
6031}
6032
6033#[inline]
6034fn q2tp_row_exact_center(
6035 chunks: &[u8],
6036 r: usize,
6037 gpr: usize,
6038 x: &[f32],
6039 scales: &[f32],
6040 center: f32,
6041) -> f32 {
6042 let mut acc = 0f32;
6043 for gi in 0..gpr {
6044 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
6045 let s = scales[gi];
6046 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6047 let mut g = 0f32;
6048 for (k, &b) in ch.iter().enumerate() {
6049 g += ((b & 3) as f32 - center) * xb[k * 4]
6050 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
6051 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
6052 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
6053 }
6054 acc += s * g;
6055 }
6056 acc
6057}
6058
6059fn q2tp_matvec(
6060 bytes: &[u8],
6061 x: &[f32],
6062 rows: usize,
6063 cols: usize,
6064 out: &mut [f32],
6065 pool: Option<&Pool>,
6066) {
6067 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
6068}
6069
6070fn q2tp_affine_matvec(
6071 bytes: &[u8],
6072 x: &[f32],
6073 rows: usize,
6074 cols: usize,
6075 out: &mut [f32],
6076 pool: Option<&Pool>,
6077) {
6078 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
6079}
6080
6081fn q2tp_matvec_mode(
6082 bytes: &[u8],
6083 x: &[f32],
6084 rows: usize,
6085 cols: usize,
6086 out: &mut [f32],
6087 pool: Option<&Pool>,
6088 affine: bool,
6089) {
6090 debug_assert_eq!(out.len(), rows);
6091 let gpr = cols / GROUP_SIZE;
6092 let v = Q4tpView::new_q2(bytes, rows, cols);
6093 let out_addr = SendMut(out.as_mut_ptr());
6094 if !affine && a8w8_enabled() {
6099 let act = split_act(x);
6100 let gsum = q1_group_sums(&act.xq, gpr);
6101 let (act, gsum) = (&act, &gsum);
6102 let run = move |start: usize, end: usize| {
6103 with_krow(gpr, |sc| {
6104 for r in start..end {
6105 v.scales_into(r, gpr, sc);
6106 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
6107 for &(j, xv) in &act.outliers {
6108 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
6109 acc += w * s * xv;
6110 }
6111 unsafe { *out_addr.at(r) = acc };
6113 }
6114 })
6115 };
6116 dispatch_rows(pool, rows, &run);
6117 return;
6118 }
6119 let run = |start: usize, end: usize| {
6120 with_krow(gpr, |sc| {
6121 for r in start..end {
6122 v.scales_into(r, gpr, sc);
6123 unsafe {
6125 *out_addr.at(r) = if affine {
6126 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
6127 } else {
6128 q2tp_row_exact(v.nib, r, gpr, x, sc)
6129 }
6130 };
6131 }
6132 })
6133 };
6134 dispatch_rows(pool, rows, &run);
6135}
6136
6137#[allow(clippy::too_many_arguments)]
6139fn q2tp_matvec2(
6140 bytes: &[u8],
6141 x1: &[f32],
6142 x2: &[f32],
6143 rows: usize,
6144 cols: usize,
6145 o1: &mut [f32],
6146 o2: &mut [f32],
6147 pool: Option<&Pool>,
6148) {
6149 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
6150}
6151
6152#[allow(clippy::too_many_arguments)]
6153fn q2tp_affine_matvec2(
6154 bytes: &[u8],
6155 x1: &[f32],
6156 x2: &[f32],
6157 rows: usize,
6158 cols: usize,
6159 o1: &mut [f32],
6160 o2: &mut [f32],
6161 pool: Option<&Pool>,
6162) {
6163 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
6164}
6165
6166#[allow(clippy::too_many_arguments)]
6167fn q2tp_matvec2_mode(
6168 bytes: &[u8],
6169 x1: &[f32],
6170 x2: &[f32],
6171 rows: usize,
6172 cols: usize,
6173 o1: &mut [f32],
6174 o2: &mut [f32],
6175 pool: Option<&Pool>,
6176 affine: bool,
6177) {
6178 let gpr = cols / GROUP_SIZE;
6179 let v = Q4tpView::new_q2(bytes, rows, cols);
6180 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6181 let run = |start: usize, end: usize| {
6182 let mut sc = vec![0f32; gpr];
6183 for r in start..end {
6184 v.scales_into(r, gpr, &mut sc);
6185 unsafe {
6187 *p1.at(r) = if affine {
6188 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
6189 } else {
6190 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
6191 };
6192 *p2.at(r) = if affine {
6193 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
6194 } else {
6195 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
6196 };
6197 }
6198 }
6199 };
6200 dispatch_rows(pool, rows, &run);
6201}
6202
6203pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6210 let gpr = cols / GROUP_SIZE;
6215 let v = Q4tpView::new_q2(bytes, rows, cols);
6216 with_krow(gpr, |sc| {
6217 for r in 0..rows {
6218 v.scales_into(r, gpr, sc);
6219 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6220 }
6221 });
6222}
6223
6224pub fn q2tp_affine_matvec_for_test(
6227 bytes: &[u8],
6228 x: &[f32],
6229 rows: usize,
6230 cols: usize,
6231 out: &mut [f32],
6232) {
6233 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6234}
6235
6236pub fn q2tp_matmat_for_test(
6237 bytes: &[u8],
6238 xs_all: &[f32],
6239 b: usize,
6240 rows: usize,
6241 cols: usize,
6242 out: &mut [f32],
6243) {
6244 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6245}
6246
6247fn q2tp_matmat(
6248 bytes: &[u8],
6249 xs_all: &[f32],
6250 b: usize,
6251 rows: usize,
6252 cols: usize,
6253 out: &mut [f32],
6254 pool: Option<&Pool>,
6255) {
6256 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6257}
6258
6259fn q2tp_affine_matmat(
6260 bytes: &[u8],
6261 xs_all: &[f32],
6262 b: usize,
6263 rows: usize,
6264 cols: usize,
6265 out: &mut [f32],
6266 pool: Option<&Pool>,
6267) {
6268 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6269}
6270
6271fn q2tp_matmat_mode(
6272 bytes: &[u8],
6273 xs_all: &[f32],
6274 b: usize,
6275 rows: usize,
6276 cols: usize,
6277 out: &mut [f32],
6278 pool: Option<&Pool>,
6279 affine: bool,
6280) {
6281 debug_assert_eq!(out.len(), b * rows);
6282 let gpr = cols / GROUP_SIZE;
6283 let v = Q4tpView::new_q2(bytes, rows, cols);
6284 let out_addr = SendMut(out.as_mut_ptr());
6285 let run = |start: usize, end: usize| {
6286 let mut sc = vec![0f32; gpr];
6287 for r in start..end {
6288 v.scales_into(r, gpr, &mut sc);
6289 for bi in 0..b {
6290 let x = &xs_all[bi * cols..(bi + 1) * cols];
6291 unsafe {
6293 *out_addr.at(bi * rows + r) = if affine {
6294 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6295 } else {
6296 q2tp_row_exact(v.nib, r, gpr, x, &sc)
6297 }
6298 };
6299 }
6300 }
6301 };
6302 dispatch_rows(pool, rows, &run);
6303}
6304
6305#[cfg(target_arch = "aarch64")]
6313#[target_feature(enable = "neon,dotprod")]
6314unsafe fn dot_q4tp_row_1x4_sdot_v1(
6315 nib: &[u8],
6316 r: usize,
6317 gpr: usize,
6318 xs: [&[i8]; 4],
6319 scales: &[f32],
6320) -> [f32; 4] {
6321 unsafe {
6322 use core::arch::aarch64::*;
6323 use core::arch::asm;
6324 let lomask = vdupq_n_u8(0x0F);
6325 let eight = vdupq_n_s8(8);
6326 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6327 for gi in 0..gpr {
6328 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6329 let s = *scales.get_unchecked(gi);
6330 let bb = vld1q_u8(t);
6331 let lo = vandq_u8(bb, lomask);
6332 let hi = vshrq_n_u8::<4>(bb);
6333 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6334 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6335 let mut d = [0f32; 4];
6336 for (k, dk) in d.iter_mut().enumerate() {
6337 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6338 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6339 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6340 asm!(
6341 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6342 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6343 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6344 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6345 options(pure, nomem, nostack),
6346 );
6347 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6348 }
6349 f0 += d[0];
6350 f1 += d[1];
6351 f2 += d[2];
6352 f3 += d[3];
6353 }
6354 [f0, f1, f2, f3]
6355 }
6356}
6357
6358#[allow(dead_code)]
6365static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6366
6367#[cfg(test)]
6370static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6371
6372#[cfg(target_arch = "x86_64")]
6378fn q4tp_blocked_x86() -> bool {
6379 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6380 1 => false,
6381 2 => avx512vnni_enabled(),
6386 _ => blocked_enabled() && avx512vnni_enabled(),
6390 }
6391}
6392
6393#[cfg(target_arch = "aarch64")]
6395#[allow(dead_code)]
6396fn q4tp_v1() -> bool {
6397 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6398 1 => true,
6399 2 => false,
6400 _ => {
6401 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6402 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6403 }
6404 }
6405}
6406
6407#[cfg(target_arch = "x86_64")]
6419#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6420unsafe fn dot_q4tp_2x8_avx512(
6421 nib: &[u8],
6422 r0: usize,
6423 gpr: usize,
6424 xs: [&[i8]; 8],
6425 sc0: &[f32],
6426 sc1: &[f32],
6427) -> [[f32; 8]; 2] {
6428 unsafe {
6431 use core::arch::x86_64::*;
6432 let lomask = _mm256_set1_epi8(0x0F);
6433 let eight = _mm256_set1_epi8(8);
6434 let zero = _mm512_setzero_si512();
6435 let mut v0 = [_mm512_setzero_ps(); 8];
6436 let mut v1 = [_mm512_setzero_ps(); 8];
6437 let pairs = gpr / 2;
6438 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6439 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6440 let bb = _mm256_loadu_si256(t as *const __m256i);
6441 let lo = _mm256_and_si256(bb, lomask);
6442 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6443 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6444 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6445 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6446 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6447 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6448 };
6449 for gp in 0..pairs {
6450 let gi = gp * 2;
6451 let (wa0, neg0) = unpack(r0, gi);
6452 let (wa1, neg1) = unpack(r0 + 1, gi);
6453 let off = gi * GROUP_SIZE;
6454 let sv = |sc: &[f32]| {
6455 _mm512_insertf32x8::<1>(
6456 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6457 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6458 )
6459 };
6460 let s0 = sv(sc0);
6461 let s1 = sv(sc1);
6462 for k in 0..8 {
6463 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6464 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6465 zero,
6466 wa0,
6467 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6468 ));
6469 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6470 zero,
6471 wa1,
6472 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6473 ));
6474 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6475 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6476 }
6477 }
6478 let mut acc = [[0f32; 8]; 2];
6479 for k in 0..8 {
6480 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6481 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6482 }
6483 if gpr % 2 == 1 {
6484 let off = (gpr - 1) * GROUP_SIZE;
6485 for j in off..off + GROUP_SIZE {
6486 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6487 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6488 for k in 0..8 {
6489 let x = *xs[k].get_unchecked(j) as f32;
6490 acc[0][k] += w0 * sa * x;
6491 acc[1][k] += w1 * sb * x;
6492 }
6493 }
6494 }
6495 acc
6496 }
6497}
6498
6499#[cfg(target_arch = "x86_64")]
6504#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6505unsafe fn dot_q4tp_row_1x8_avx512(
6506 nib: &[u8],
6507 r: usize,
6508 gpr: usize,
6509 xs: [&[i8]; 8],
6510 scales: &[f32],
6511) -> [f32; 8] {
6512 unsafe {
6514 use core::arch::x86_64::*;
6515 let lomask = _mm256_set1_epi8(0x0F);
6516 let eight = _mm256_set1_epi8(8);
6517 let zero = _mm512_setzero_si512();
6518 let (mut v0, mut v1, mut v2, mut v3) = (
6519 _mm512_setzero_ps(),
6520 _mm512_setzero_ps(),
6521 _mm512_setzero_ps(),
6522 _mm512_setzero_ps(),
6523 );
6524 let (mut v4, mut v5, mut v6, mut v7) = (
6525 _mm512_setzero_ps(),
6526 _mm512_setzero_ps(),
6527 _mm512_setzero_ps(),
6528 _mm512_setzero_ps(),
6529 );
6530 let pairs = gpr / 2;
6531 for gp in 0..pairs {
6532 let gi = gp * 2;
6533 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6534 let bb = _mm256_loadu_si256(t as *const __m256i);
6535 let lo = _mm256_and_si256(bb, lomask);
6536 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6537 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6542 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6543 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6544 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6545 let wabs = _mm512_abs_epi8(w);
6546 let neg = _mm512_movepi8_mask(w);
6547 let off = gi * GROUP_SIZE;
6548 let sv = _mm512_insertf32x8::<1>(
6549 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6550 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6551 );
6552 let dot = |x: &[i8]| -> __m512 {
6553 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6554 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6555 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6556 };
6557 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6558 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6559 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6560 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6561 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6562 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6563 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6564 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6565 }
6566 let mut acc = [
6567 _mm512_reduce_add_ps(v0),
6568 _mm512_reduce_add_ps(v1),
6569 _mm512_reduce_add_ps(v2),
6570 _mm512_reduce_add_ps(v3),
6571 _mm512_reduce_add_ps(v4),
6572 _mm512_reduce_add_ps(v5),
6573 _mm512_reduce_add_ps(v6),
6574 _mm512_reduce_add_ps(v7),
6575 ];
6576 if gpr % 2 == 1 {
6579 let off = (gpr - 1) * GROUP_SIZE;
6580 for j in off..off + GROUP_SIZE {
6581 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6582 let ws = w * s;
6583 for k in 0..8 {
6584 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6585 }
6586 }
6587 }
6588 acc
6589 }
6590}
6591
6592#[cfg(target_arch = "x86_64")]
6605#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6606unsafe fn dot_q4tp_row_1x4_avx512(
6607 nib: &[u8],
6608 r: usize,
6609 gpr: usize,
6610 xs: [&[i8]; 4],
6611 scales: &[f32],
6612) -> [f32; 4] {
6613 unsafe {
6615 use core::arch::x86_64::*;
6616 let lomask = _mm256_set1_epi8(0x0F);
6617 let eight = _mm256_set1_epi8(8);
6618 let zero = _mm512_setzero_si512();
6619 let (mut v0, mut v1, mut v2, mut v3) = (
6620 _mm512_setzero_ps(),
6621 _mm512_setzero_ps(),
6622 _mm512_setzero_ps(),
6623 _mm512_setzero_ps(),
6624 );
6625 let pairs = gpr / 2;
6626 for gp in 0..pairs {
6627 let gi = gp * 2;
6628 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6629 let bb = _mm256_loadu_si256(t as *const __m256i);
6630 let lo = _mm256_and_si256(bb, lomask);
6631 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6632 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6637 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6638 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6639 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6640 let wabs = _mm512_abs_epi8(w);
6641 let neg = _mm512_movepi8_mask(w);
6642 let off = gi * GROUP_SIZE;
6643 let sv = _mm512_insertf32x8::<1>(
6644 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6645 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6646 );
6647 let dot = |x: &[i8]| -> __m512 {
6648 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6649 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6650 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6651 };
6652 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6653 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6654 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6655 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6656 }
6657 let mut acc = [
6658 _mm512_reduce_add_ps(v0),
6659 _mm512_reduce_add_ps(v1),
6660 _mm512_reduce_add_ps(v2),
6661 _mm512_reduce_add_ps(v3),
6662 ];
6663 if gpr % 2 == 1 {
6666 let off = (gpr - 1) * GROUP_SIZE;
6667 for j in off..off + GROUP_SIZE {
6668 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6669 let ws = w * s;
6670 for k in 0..4 {
6671 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6672 }
6673 }
6674 }
6675 acc
6676 }
6677}
6678
6679#[cfg(target_arch = "aarch64")]
6683#[target_feature(enable = "neon,dotprod")]
6684unsafe fn dot_q4tp_row_1x4_sdot(
6685 nib: &[u8],
6686 r: usize,
6687 gpr: usize,
6688 xs: [&[i8]; 4],
6689 scales: &[f32],
6690) -> [f32; 4] {
6691 unsafe {
6693 use core::arch::aarch64::*;
6694 use core::arch::asm;
6695 let lomask = vdupq_n_u8(0x0F);
6696 let eight = vdupq_n_s8(8);
6697 let (mut v0, mut v1, mut v2, mut v3) = (
6713 vdupq_n_f32(0.0),
6714 vdupq_n_f32(0.0),
6715 vdupq_n_f32(0.0),
6716 vdupq_n_f32(0.0),
6717 );
6718 for gi in 0..gpr {
6719 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6720 let s = *scales.get_unchecked(gi);
6721 let bb = vld1q_u8(t);
6722 let lo = vandq_u8(bb, lomask);
6723 let hi = vshrq_n_u8::<4>(bb);
6724 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6725 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6726 let off = gi * GROUP_SIZE;
6727 let dot4 = |x: &[i8]| -> int32x4_t {
6728 let x0 = vld1q_s8(x.as_ptr().add(off));
6729 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6730 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6731 asm!(
6732 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6733 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6734 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6735 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6736 options(pure, nomem, nostack),
6737 );
6738 vaddq_s32(a0, a1)
6739 };
6740 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6741 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6742 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6743 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6744 }
6745 [
6746 vaddvq_f32(v0),
6747 vaddvq_f32(v1),
6748 vaddvq_f32(v2),
6749 vaddvq_f32(v3),
6750 ]
6751 }
6752}
6753
6754fn q4tp_matmat(
6762 bytes: &[u8],
6763 xs_all: &[f32],
6764 b: usize,
6765 rows: usize,
6766 cols: usize,
6767 out: &mut [f32],
6768 pool: Option<&Pool>,
6769) {
6770 q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6771}
6772
6773#[allow(clippy::too_many_arguments)]
6777fn q4tp_matmat_with(
6778 bytes: &[u8],
6779 xs_all: &[f32],
6780 b: usize,
6781 rows: usize,
6782 cols: usize,
6783 out: &mut [f32],
6784 pool: Option<&Pool>,
6785 exact: bool,
6786) {
6787 debug_assert_eq!(out.len(), b * rows);
6788 let gpr = cols / GROUP_SIZE;
6789 let v = Q4tpView::new(bytes, rows, cols);
6790
6791 #[cfg(target_os = "macos")]
6795 if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6796 dequant_matmat_accel(
6797 &|r, dst| {
6798 let mut sc = [0f32; 32];
6799 let mut scv;
6800 let s: &[f32] = if gpr <= 32 {
6801 v.scales_into(r, gpr, &mut sc);
6802 &sc[..gpr]
6803 } else {
6804 scv = vec![0f32; gpr];
6805 v.scales_into(r, gpr, &mut scv);
6806 &scv
6807 };
6808 for gi in 0..gpr {
6809 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6810 for (k, &bb) in tile.iter().enumerate() {
6811 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6812 dst[gi * GROUP_SIZE + k * 2 + 1] =
6813 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6814 }
6815 }
6816 },
6817 xs_all,
6818 b,
6819 rows,
6820 cols,
6821 out,
6822 pool,
6823 );
6824 return;
6825 }
6826
6827 let out_addr = SendMut(out.as_mut_ptr());
6828 if a8w8_enabled() {
6829 let acts: Vec<SplitAct> = (0..b)
6830 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6831 .collect();
6832 let acts = &acts;
6833 #[cfg(target_arch = "aarch64")]
6840 let blocked_ok = sdot_enabled() && blocked_enabled();
6841 #[cfg(target_arch = "x86_64")]
6848 let blocked_ok = q4tp_blocked_x86() && !exact;
6849 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6850 let blocked_ok = {
6851 let _ = exact;
6852 false
6853 };
6854 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6864 .ok()
6865 .and_then(|v| v.parse().ok())
6866 .filter(|v| *v > 0)
6867 .unwrap_or(256);
6868 let run = |start: usize, end: usize| {
6869 for abase in (0..acts.len()).step_by(panel_cols) {
6870 let alen = (acts.len() - abase).min(panel_cols);
6871 let mut sc = vec![0f32; gpr];
6872 #[cfg(target_arch = "x86_64")]
6873 let mut r_lo = start;
6874 #[cfg(target_arch = "x86_64")]
6875 if blocked_ok && alen >= 8 {
6876 let mut sc1 = vec![0f32; gpr];
6877 while r_lo + 2 <= end {
6878 v.scales_into(r_lo, gpr, &mut sc);
6879 v.scales_into(r_lo + 1, gpr, &mut sc1);
6880 let mut bi = 0usize;
6881 while bi + 8 <= alen {
6882 let xs = [
6883 acts[abase + bi].xq.as_slice(),
6884 acts[abase + bi + 1].xq.as_slice(),
6885 acts[abase + bi + 2].xq.as_slice(),
6886 acts[abase + bi + 3].xq.as_slice(),
6887 acts[abase + bi + 4].xq.as_slice(),
6888 acts[abase + bi + 5].xq.as_slice(),
6889 acts[abase + bi + 6].xq.as_slice(),
6890 acts[abase + bi + 7].xq.as_slice(),
6891 ];
6892 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6893 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6894 for k in 0..8 {
6895 let act = &acts[abase + bi + k];
6896 let mut acc = dr[k] * act.sx;
6897 for &(j, xv) in &act.outliers {
6898 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6899 acc += w * s * xv;
6900 }
6901 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6903 }
6904 }
6905 bi += 8;
6906 }
6907 for row in [r_lo, r_lo + 1] {
6910 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6911 for b2 in bi..alen {
6912 let act = &acts[abase + b2];
6913 let xs4 = [
6914 act.xq.as_slice(),
6915 act.xq.as_slice(),
6916 act.xq.as_slice(),
6917 act.xq.as_slice(),
6918 ];
6919 let d =
6920 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6921 let mut acc = d[0] * act.sx;
6922 for &(j, xv) in &act.outliers {
6923 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6924 acc += w * s * xv;
6925 }
6926 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6928 }
6929 }
6930 r_lo += 2;
6931 }
6932 }
6933 #[cfg(target_arch = "x86_64")]
6934 let row_start = r_lo;
6935 #[cfg(not(target_arch = "x86_64"))]
6936 let row_start = start;
6937 for r in row_start..end {
6938 v.scales_into(r, gpr, &mut sc);
6939 let mut bi = 0usize;
6940 #[cfg(target_arch = "x86_64")]
6941 if blocked_ok {
6942 while bi + 8 <= alen {
6943 let xs = [
6944 acts[abase + bi].xq.as_slice(),
6945 acts[abase + bi + 1].xq.as_slice(),
6946 acts[abase + bi + 2].xq.as_slice(),
6947 acts[abase + bi + 3].xq.as_slice(),
6948 acts[abase + bi + 4].xq.as_slice(),
6949 acts[abase + bi + 5].xq.as_slice(),
6950 acts[abase + bi + 6].xq.as_slice(),
6951 acts[abase + bi + 7].xq.as_slice(),
6952 ];
6953 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6954 for k in 0..8 {
6955 let act = &acts[abase + bi + k];
6956 let mut acc = d[k] * act.sx;
6957 for &(j, xv) in &act.outliers {
6958 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6959 acc += w * s * xv;
6960 }
6961 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6963 }
6964 bi += 8;
6965 }
6966 while bi + 4 <= alen {
6967 let xs = [
6968 acts[abase + bi].xq.as_slice(),
6969 acts[abase + bi + 1].xq.as_slice(),
6970 acts[abase + bi + 2].xq.as_slice(),
6971 acts[abase + bi + 3].xq.as_slice(),
6972 ];
6973 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6974 for k in 0..4 {
6975 let act = &acts[abase + bi + k];
6976 let mut acc = d[k] * act.sx;
6977 for &(j, xv) in &act.outliers {
6978 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6979 acc += w * s * xv;
6980 }
6981 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6983 }
6984 bi += 4;
6985 }
6986 }
6987 #[cfg(target_arch = "aarch64")]
6988 if blocked_ok {
6989 while bi + 4 <= alen {
6990 let xs = [
6991 acts[abase + bi].xq.as_slice(),
6992 acts[abase + bi + 1].xq.as_slice(),
6993 acts[abase + bi + 2].xq.as_slice(),
6994 acts[abase + bi + 3].xq.as_slice(),
6995 ];
6996 let d = unsafe {
6997 if exact || q4tp_v1() {
6998 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6999 } else {
7000 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
7001 }
7002 };
7003 for k in 0..4 {
7004 let act = &acts[abase + bi + k];
7005 let mut acc = d[k] * act.sx;
7006 for &(j, xv) in &act.outliers {
7007 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7008 acc += w * s * xv;
7009 }
7010 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7012 }
7013 bi += 4;
7014 }
7015 }
7016 let _ = blocked_ok;
7017 while bi < alen {
7018 let act = &acts[abase + bi];
7019 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
7020 for &(j, xv) in &act.outliers {
7021 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7022 acc += w * s * xv;
7023 }
7024 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
7026 bi += 1;
7027 }
7028 }
7029 }
7030 };
7031 dispatch_rows(pool, rows, &run);
7032 return;
7033 }
7034
7035 let run = |start: usize, end: usize| {
7036 let mut sc = vec![0f32; gpr];
7037 for r in start..end {
7038 v.scales_into(r, gpr, &mut sc);
7039 for bi in 0..b {
7040 let x = &xs_all[bi * cols..(bi + 1) * cols];
7041 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
7043 }
7044 }
7045 };
7046 dispatch_rows(pool, rows, &run);
7047}
7048
7049fn q4t_matvec(
7051 bytes: &[u8],
7052 x: &[f32],
7053 rows: usize,
7054 cols: usize,
7055 out: &mut [f32],
7056 pool: Option<&Pool>,
7057) {
7058 debug_assert_eq!(out.len(), rows);
7059 let gpr = cols / GROUP_SIZE;
7060 let out_addr = SendMut(out.as_mut_ptr());
7061 if a8w8_enabled() {
7062 let act = split_act(x);
7063 let run = move |start: usize, end: usize| {
7064 for r in start..end {
7065 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7066 for &(j, xv) in &act.outliers {
7067 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7068 acc += w * s * xv;
7069 }
7070 unsafe { *out_addr.at(r) = acc };
7072 }
7073 };
7074 dispatch_rows(pool, rows, &run);
7075 return;
7076 }
7077 let run = move |start: usize, end: usize| {
7078 for r in start..end {
7079 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
7081 }
7082 };
7083 dispatch_rows(pool, rows, &run);
7084}
7085
7086#[allow(clippy::too_many_arguments)]
7088fn q4t_matvec2(
7089 bytes: &[u8],
7090 x1: &[f32],
7091 x2: &[f32],
7092 rows: usize,
7093 cols: usize,
7094 o1: &mut [f32],
7095 o2: &mut [f32],
7096 pool: Option<&Pool>,
7097) {
7098 let gpr = cols / GROUP_SIZE;
7099 let p1 = SendMut(o1.as_mut_ptr());
7100 let p2 = SendMut(o2.as_mut_ptr());
7101 if a8w8_enabled() {
7102 let a1 = split_act(x1);
7103 let a2 = split_act(x2);
7104 let run = move |start: usize, end: usize| {
7105 for r in start..end {
7106 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
7107 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
7108 for &(j, xv) in &a1.outliers {
7109 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7110 v1 += w * s * xv;
7111 }
7112 for &(j, xv) in &a2.outliers {
7113 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7114 v2 += w * s * xv;
7115 }
7116 unsafe {
7118 *p1.at(r) = v1;
7119 *p2.at(r) = v2;
7120 }
7121 }
7122 };
7123 dispatch_rows(pool, rows, &run);
7124 return;
7125 }
7126 let run = move |start: usize, end: usize| {
7127 for r in start..end {
7128 unsafe {
7130 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
7131 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
7132 }
7133 }
7134 };
7135 dispatch_rows(pool, rows, &run);
7136}
7137
7138#[allow(clippy::too_many_arguments)]
7140#[cfg(target_os = "macos")]
7146fn dequant_matmat_accel(
7147 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
7148 xs_all: &[f32],
7149 b: usize,
7150 rows: usize,
7151 cols: usize,
7152 out: &mut [f32],
7153 pool: Option<&Pool>,
7154) {
7155 const TR: usize = 2048;
7156 thread_local! {
7157 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7158 }
7159 WTILE.with(|wt| {
7160 let mut wtile = wt.borrow_mut();
7161 wtile.resize(TR * cols, 0.0);
7162 let mut r0 = 0usize;
7163 while r0 < rows {
7164 let tr = TR.min(rows - r0);
7165 let wt_addr = SendMut(wtile.as_mut_ptr());
7166 let run = |start: usize, end: usize| {
7167 for r in start..end {
7168 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
7170 dequant_row(r0 + r, dst);
7171 }
7172 };
7173 dispatch_rows(pool, tr, &run);
7174 unsafe {
7175 accel_blas::cblas_sgemm(
7176 101, 111, 112, b as i32,
7180 tr as i32,
7181 cols as i32,
7182 1.0,
7183 xs_all.as_ptr(),
7184 cols as i32,
7185 wtile.as_ptr(),
7186 cols as i32,
7187 0.0,
7188 out.as_mut_ptr().add(r0),
7189 rows as i32,
7190 );
7191 }
7192 r0 += tr;
7193 }
7194 });
7195}
7196
7197fn q4t_matmat(
7198 bytes: &[u8],
7199 xs_all: &[f32],
7200 b: usize,
7201 rows: usize,
7202 cols: usize,
7203 out: &mut [f32],
7204 pool: Option<&Pool>,
7205) {
7206 debug_assert_eq!(out.len(), b * rows);
7207 let gpr = cols / GROUP_SIZE;
7208 #[cfg(target_os = "macos")]
7212 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7213 dequant_matmat_accel(
7214 &|r, dst| {
7215 for gi in 0..gpr {
7216 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7217 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7218 for (k, &bb) in tile[2..].iter().enumerate() {
7219 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7220 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7221 }
7222 }
7223 },
7224 xs_all,
7225 b,
7226 rows,
7227 cols,
7228 out,
7229 pool,
7230 );
7231 return;
7232 }
7233 let out_addr = SendMut(out.as_mut_ptr());
7234 if a8w8_enabled() {
7235 let acts: Vec<SplitAct> = (0..b)
7236 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7237 .collect();
7238 let acts = &acts;
7239 #[cfg(target_arch = "x86_64")]
7240 let blocked_ok = avx2_enabled() && blocked_enabled();
7241 #[cfg(target_arch = "aarch64")]
7242 let blocked_ok = sdot_enabled() && blocked_enabled();
7243 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7244 let blocked_ok = false;
7245 let run = move |start: usize, end: usize| {
7246 for r in start..end {
7247 let mut bi = 0usize;
7248 #[cfg(target_arch = "aarch64")]
7249 if blocked_ok {
7250 while bi + 4 <= acts.len() {
7251 let xs = [
7252 acts[bi].xq.as_slice(),
7253 acts[bi + 1].xq.as_slice(),
7254 acts[bi + 2].xq.as_slice(),
7255 acts[bi + 3].xq.as_slice(),
7256 ];
7257 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7258 for k in 0..4 {
7259 let act = &acts[bi + k];
7260 let mut acc = d[k] * act.sx;
7261 for &(j, xv) in &act.outliers {
7262 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7263 acc += w * sc * xv;
7264 }
7265 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7267 }
7268 bi += 4;
7269 }
7270 }
7271 #[cfg(target_arch = "x86_64")]
7272 if blocked_ok {
7273 while bi + 4 <= acts.len() {
7274 let xs = [
7275 acts[bi].xq.as_slice(),
7276 acts[bi + 1].xq.as_slice(),
7277 acts[bi + 2].xq.as_slice(),
7278 acts[bi + 3].xq.as_slice(),
7279 ];
7280 let d = unsafe {
7281 if vnni_tiles_enabled() {
7282 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7283 } else {
7284 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7285 }
7286 };
7287 for k in 0..4 {
7288 let act = &acts[bi + k];
7289 let mut acc = d[k] * act.sx;
7290 for &(j, xv) in &act.outliers {
7291 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7292 acc += w * sc * xv;
7293 }
7294 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7296 }
7297 bi += 4;
7298 }
7299 }
7300 let _ = blocked_ok;
7301 while bi < acts.len() {
7302 let act = &acts[bi];
7303 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7304 for &(j, xv) in &act.outliers {
7305 let (w, s) = q4t_outlier(bytes, r, gpr, j);
7306 acc += w * s * xv;
7307 }
7308 unsafe { *out_addr.at(bi * rows + r) = acc };
7310 bi += 1;
7311 }
7312 }
7313 };
7314 dispatch_rows(pool, rows, &run);
7315 return;
7316 }
7317 let run = move |start: usize, end: usize| {
7318 for r in start..end {
7319 for bi in 0..b {
7320 let x = &xs_all[bi * cols..(bi + 1) * cols];
7321 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7323 }
7324 }
7325 };
7326 dispatch_rows(pool, rows, &run);
7327}
7328
7329fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7338 (0..gpr)
7339 .map(|gi| {
7340 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7341 .iter()
7342 .map(|&v| v as i32)
7343 .sum()
7344 })
7345 .collect()
7346}
7347
7348#[inline]
7352#[allow(unreachable_code)]
7353#[cfg(target_arch = "x86_64")]
7358#[target_feature(enable = "avx2")]
7359unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7360 unsafe {
7362 use core::arch::x86_64::*;
7363 let expand = _mm256_setr_epi8(
7365 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,
7366 3, 3, 3,
7367 );
7368 let bitsel = _mm256_setr_epi8(
7369 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7370 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7371 );
7372 let ones8 = _mm256_set1_epi8(1);
7373 let ones16 = _mm256_set1_epi16(1);
7374 let mut acc = 0f32;
7375 for gi in 0..gpr {
7376 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7377 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7378 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7379 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7380 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7381 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7382 let sel = _mm256_and_si256(x, mask);
7383 let p16 = _mm256_maddubs_epi16(ones8, sel);
7385 let d32 = _mm256_madd_epi16(p16, ones16);
7386 let hi128 = _mm256_extracti128_si256::<1>(d32);
7387 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7388 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7389 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7390 let msum = _mm_cvtsi128_si32(s32);
7391 let d = 2 * msum - gsum[gi];
7394 acc += d as f32 * s;
7395 }
7396 acc
7397 }
7398}
7399
7400#[cfg(target_arch = "x86_64")]
7403#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7404unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7405 unsafe {
7407 use core::arch::x86_64::*;
7408 let expand = _mm256_setr_epi8(
7409 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,
7410 3, 3, 3,
7411 );
7412 let bitsel = _mm256_setr_epi8(
7413 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7414 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7415 );
7416 let ones8 = _mm256_set1_epi8(1);
7417 let mut acc = 0f32;
7418 for gi in 0..gpr {
7419 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7420 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7421 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7422 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7423 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7424 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7425 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7426 let d = 2 * msum - gsum[gi];
7427 acc += d as f32 * s;
7428 }
7429 acc
7430 }
7431}
7432
7433#[cfg(target_arch = "x86_64")]
7435#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7436unsafe fn dot_q1_row_1x4_vnni(
7437 bytes: &[u8],
7438 r: usize,
7439 gpr: usize,
7440 xs: [&[i8]; 4],
7441 gsums: [&[i32]; 4],
7442) -> [f32; 4] {
7443 unsafe {
7445 use core::arch::x86_64::*;
7446 let expand = _mm256_setr_epi8(
7447 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,
7448 3, 3, 3,
7449 );
7450 let bitsel = _mm256_setr_epi8(
7451 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7452 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7453 );
7454 let ones8 = _mm256_set1_epi8(1);
7455 let mut acc = [0f32; 4];
7456 for gi in 0..gpr {
7457 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7458 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7459 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7460 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7461 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7462 for (k, xq) in xs.iter().enumerate() {
7463 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7464 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7465 let d = 2 * msum - gsums[k][gi];
7466 acc[k] += d as f32 * s;
7467 }
7468 }
7469 acc
7470 }
7471}
7472
7473#[cfg(target_arch = "x86_64")]
7476#[target_feature(enable = "avx2")]
7477unsafe fn dot_q1_row_1x4_avx2(
7478 bytes: &[u8],
7479 r: usize,
7480 gpr: usize,
7481 xs: [&[i8]; 4],
7482 gsums: [&[i32]; 4],
7483) -> [f32; 4] {
7484 unsafe {
7486 use core::arch::x86_64::*;
7487 let expand = _mm256_setr_epi8(
7488 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,
7489 3, 3, 3,
7490 );
7491 let bitsel = _mm256_setr_epi8(
7492 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7493 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7494 );
7495 let ones8 = _mm256_set1_epi8(1);
7496 let ones16 = _mm256_set1_epi16(1);
7497 let mut acc = [0f32; 4];
7498 for gi in 0..gpr {
7499 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7500 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7501 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7502 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7503 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7504 for (k, xq) in xs.iter().enumerate() {
7505 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7506 let sel = _mm256_and_si256(x, mask);
7507 let p16 = _mm256_maddubs_epi16(ones8, sel);
7508 let d32 = _mm256_madd_epi16(p16, ones16);
7509 let hi128 = _mm256_extracti128_si256::<1>(d32);
7510 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7511 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7512 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7513 let msum = _mm_cvtsi128_si32(s32);
7514 let d = 2 * msum - gsums[k][gi];
7515 acc[k] += d as f32 * s;
7516 }
7517 }
7518 acc
7519 }
7520}
7521
7522#[allow(unreachable_code)]
7523fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7524 #[cfg(target_arch = "aarch64")]
7525 unsafe {
7526 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7527 }
7528 #[cfg(target_arch = "x86_64")]
7529 if avx2_enabled() {
7530 unsafe {
7531 if vnni_tiles_enabled() {
7532 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7533 }
7534 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7535 }
7536 }
7537 let _ = gsum;
7538 let mut acc = 0f32;
7539 for gi in 0..gpr {
7540 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7541 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7542 let mut d = 0i32;
7543 for (j, &b) in tile[2..].iter().enumerate() {
7544 for k in 0..8 {
7545 let w = ((b >> k) & 1) as i32 * 2 - 1;
7546 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7547 }
7548 }
7549 acc += d as f32 * s;
7550 }
7551 acc
7552}
7553
7554#[cfg(target_arch = "aarch64")]
7563#[target_feature(enable = "neon,dotprod")]
7564unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7565 unsafe {
7568 use core::arch::aarch64::*;
7569 use core::arch::asm;
7570 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7571 let m = vld1q_u8(MASKS.as_ptr());
7572 macro_rules! tile_dot {
7574 ($t:expr, $x:expr) => {{
7575 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7576 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7577 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7578 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7579 let x0 = vld1q_s8($x);
7580 let x1 = vld1q_s8($x.add(16));
7581 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7582 asm!(
7583 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7584 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7585 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7586 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7587 options(pure, nomem, nostack),
7588 );
7589 vaddq_s32(a0, a1)
7590 }};
7591 }
7592 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7601 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7602 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7603 const IW11: [u8; 16] = [
7604 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7605 ];
7606 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7607 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7608 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7609 let isc = vld1_u8(ISC.as_ptr());
7610 macro_rules! tile_dot_tbl {
7612 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7613 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7614 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7615 let x0 = vld1q_s8($x);
7616 let x1 = vld1q_s8($x.add(16));
7617 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7618 asm!(
7619 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7620 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7621 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7622 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7623 options(pure, nomem, nostack),
7624 );
7625 vaddq_s32(a0, a1)
7626 }};
7627 }
7628 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7629 let row_base = r * gpr * Q1_TILE;
7630 let abs_end = bytes.len();
7631 let xp = xq.as_ptr();
7632 let gp = gsum.as_ptr();
7633 let mut accv = vdupq_n_f32(0.0);
7634 let mut gi = 0;
7635 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7638 let t0 = base.add(gi * Q1_TILE);
7639 let ld_a = vld1q_u8(t0);
7640 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7641 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7642 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7643 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7644 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7645 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7647 let g = vld1q_s32(gp.add(gi));
7648 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7649 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7650 let scf: float32x4_t;
7651 asm!(
7652 "fcvtl {o:v}.4s, {i:v}.4h",
7653 o = out(vreg) scf, i = in(vreg) sc16,
7654 options(pure, nomem, nostack),
7655 );
7656 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7657 gi += 4;
7658 }
7659 let mut acc = vaddvq_f32(accv);
7660 while gi < gpr {
7661 let t = base.add(gi * Q1_TILE);
7662 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7663 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7664 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7665 gi += 1;
7666 }
7667 acc
7668 }
7669}
7670
7671#[cfg(target_arch = "aarch64")]
7676#[target_feature(enable = "neon,dotprod")]
7677unsafe fn dot_q1_row_1x4_sdot(
7678 bytes: &[u8],
7679 r: usize,
7680 gpr: usize,
7681 xs: [&[i8]; 4],
7682 gs: [&[i32]; 4],
7683) -> [f32; 4] {
7684 unsafe {
7686 use core::arch::aarch64::*;
7687 use core::arch::asm;
7688 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7689 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7690 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7691 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7692 const IW11: [u8; 16] = [
7693 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7694 ];
7695 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7696 let m = vld1q_u8(MASKS.as_ptr());
7697 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7698 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7699 let isc = vld1_u8(ISC.as_ptr());
7700 macro_rules! sdot2 {
7701 ($w0:expr, $w1:expr, $x:expr) => {{
7702 let x0 = vld1q_s8($x);
7703 let x1 = vld1q_s8($x.add(16));
7704 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7705 asm!(
7706 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7707 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7708 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7709 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7710 options(pure, nomem, nostack),
7711 );
7712 vaddq_s32(a0, a1)
7713 }};
7714 }
7715 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7716 let row_base = r * gpr * Q1_TILE;
7717 let abs_end = bytes.len();
7718 let mut accv = [vdupq_n_f32(0.0); 4];
7719 let mut gi = 0;
7720 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7721 let t0 = base.add(gi * Q1_TILE);
7722 let ld_a = vld1q_u8(t0);
7723 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7724 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7726 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7727 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7728 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7729 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7730 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7731 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7732 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7733 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7734 let scf: float32x4_t;
7735 asm!(
7736 "fcvtl {o:v}.4s, {i:v}.4h",
7737 o = out(vreg) scf, i = in(vreg) sc16,
7738 options(pure, nomem, nostack),
7739 );
7740 for k in 0..4 {
7741 let xp = xs[k].as_ptr();
7742 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7743 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7744 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7745 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7746 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7747 let g = vld1q_s32(gs[k].as_ptr().add(gi));
7748 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7749 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7750 }
7751 gi += 4;
7752 }
7753 let mut acc = [
7754 vaddvq_f32(accv[0]),
7755 vaddvq_f32(accv[1]),
7756 vaddvq_f32(accv[2]),
7757 vaddvq_f32(accv[3]),
7758 ];
7759 while gi < gpr {
7760 let t = base.add(gi * Q1_TILE);
7761 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7762 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7763 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7764 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7765 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7766 for k in 0..4 {
7767 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7768 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7769 }
7770 gi += 1;
7771 }
7772 acc
7773 }
7774}
7775
7776#[inline]
7778fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7779 let gi = j / GROUP_SIZE;
7780 let k = j % GROUP_SIZE;
7781 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7782 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7783 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7784 ((bit as i32 * 2 - 1) as f32, s)
7785}
7786
7787#[inline]
7789fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7790 let mut acc = 0f32;
7791 for gi in 0..gpr {
7792 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7793 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7794 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7795 let mut ga = 0f32;
7796 for (j, &b) in tile[2..].iter().enumerate() {
7797 for k in 0..8 {
7798 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7799 }
7800 }
7801 acc += ga * s;
7802 }
7803 acc
7804}
7805
7806#[allow(clippy::too_many_arguments)]
7809fn bin_range_a8w8(
7810 w: BinW,
7811 gpr: usize,
7812 act: &SplitAct,
7813 gsum: &[i32],
7814 out: SendMut,
7815 start: usize,
7816 end: usize,
7817) {
7818 for r in start..end {
7819 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
7820 for &(j, xv) in &act.outliers {
7821 let (wv, s) = w.outlier(r, gpr, j);
7822 acc += wv * s * xv;
7823 }
7824 unsafe { *out.at(r) = acc };
7826 }
7827}
7828
7829fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7831 for r in start..end {
7832 unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
7834 }
7835}
7836
7837#[derive(Clone, Copy)]
7845pub(crate) enum BinW<'a> {
7846 Q1(&'a [u8]),
7847 Q1g {
7848 bytes: &'a [u8],
7849 sc_off: usize,
7851 rows: usize,
7854 },
7855}
7856
7857impl<'a> BinW<'a> {
7858 fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
7859 Self::Q1g {
7860 bytes,
7861 sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
7862 rows,
7863 }
7864 }
7865
7866 #[inline]
7867 fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7868 match self {
7869 Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
7870 Self::Q1g {
7871 bytes,
7872 sc_off,
7873 rows,
7874 } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
7875 }
7876 }
7877
7878 #[inline]
7879 fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
7880 match self {
7881 Self::Q1(b) => q1_row_exact(b, r, gpr, x),
7882 Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
7883 }
7884 }
7885
7886 #[inline]
7887 fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
7888 match self {
7889 Self::Q1(b) => q1_outlier(b, r, gpr, j),
7890 Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
7891 }
7892 }
7893
7894 #[cfg(target_arch = "aarch64")]
7899 #[inline]
7900 unsafe fn dot_1x4_sdot(
7901 self,
7902 r: usize,
7903 gpr: usize,
7904 xs: [&[i8]; 4],
7905 gs: [&[i32]; 4],
7906 ) -> [f32; 4] {
7907 unsafe {
7909 match self {
7910 Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
7911 Self::Q1g {
7912 bytes,
7913 sc_off,
7914 rows,
7915 } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
7916 }
7917 }
7918 }
7919
7920 #[cfg(target_arch = "x86_64")]
7925 #[inline]
7926 unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
7927 unsafe {
7929 match (self, vnni_tiles_enabled()) {
7930 (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
7931 (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
7932 (Self::Q1g { bytes, sc_off, .. }, true) => {
7933 dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
7934 }
7935 (Self::Q1g { bytes, sc_off, .. }, false) => {
7936 dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
7937 }
7938 }
7939 }
7940 }
7941}
7942
7943#[inline]
7953fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
7954 u32::from_le_bytes([
7955 bytes[t * 4],
7956 bytes[t * 4 + 1],
7957 bytes[t * 4 + 2],
7958 bytes[t * 4 + 3],
7959 ])
7960}
7961
7962#[inline]
7964fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
7965 let o = sc_off + (t / 4) * 2;
7966 u16::from_le_bytes([bytes[o], bytes[o + 1]])
7967}
7968
7969#[allow(unreachable_code)]
7971#[allow(clippy::too_many_arguments)]
7972fn dot_q1g_row_i8(
7973 bytes: &[u8],
7974 sc_off: usize,
7975 rows: usize,
7976 r: usize,
7977 gpr: usize,
7978 xq: &[i8],
7979 gsum: &[i32],
7980) -> f32 {
7981 #[cfg(target_arch = "aarch64")]
7982 unsafe {
7983 return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
7984 }
7985 #[cfg(target_arch = "x86_64")]
7986 if avx2_enabled() {
7987 unsafe {
7988 if vnni_tiles_enabled() {
7989 return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
7990 }
7991 return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
7992 }
7993 }
7994 let _ = (gsum, rows);
7995 let mut acc = 0f32;
7996 for gi in 0..gpr {
7997 let t = r * gpr + gi;
7998 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
7999 let bits = q1g_bits(bytes, t);
8000 let mut d = 0i32;
8001 for k in 0..GROUP_SIZE {
8002 let w = ((bits >> k) & 1) as i32 * 2 - 1;
8003 d += w * xq[gi * GROUP_SIZE + k] as i32;
8004 }
8005 acc += d as f32 * s;
8006 }
8007 acc
8008}
8009
8010#[cfg(target_arch = "x86_64")]
8015#[target_feature(enable = "avx2")]
8016#[inline]
8017unsafe fn q1g_hsum4(
8018 d0: core::arch::x86_64::__m256i,
8019 d1: core::arch::x86_64::__m256i,
8020 d2: core::arch::x86_64::__m256i,
8021 d3: core::arch::x86_64::__m256i,
8022) -> core::arch::x86_64::__m128i {
8023 unsafe {
8025 use core::arch::x86_64::*;
8026 let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
8027 _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
8028 }
8029}
8030
8031#[cfg(target_arch = "x86_64")]
8035#[target_feature(enable = "avx2")]
8036#[inline]
8037unsafe fn q1g_group_acc(
8038 acc: f32,
8039 msum4: core::arch::x86_64::__m128i,
8040 gsum4: &[i32],
8041 s: f32,
8042) -> f32 {
8043 unsafe {
8045 use core::arch::x86_64::*;
8046 let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
8047 let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
8048 let mut f = [0f32; 4];
8049 _mm_storeu_ps(f.as_mut_ptr(), d);
8050 let mut acc = acc;
8051 for v in f {
8052 acc += v * s;
8053 }
8054 acc
8055 }
8056}
8057
8058#[cfg(target_arch = "x86_64")]
8063#[target_feature(enable = "avx2")]
8064unsafe fn dot_q1g_row_avx2(
8065 bytes: &[u8],
8066 sc_off: usize,
8067 r: usize,
8068 gpr: usize,
8069 xq: &[i8],
8070 gsum: &[i32],
8071) -> f32 {
8072 unsafe {
8075 use core::arch::x86_64::*;
8076 let expand = _mm256_setr_epi8(
8077 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,
8078 3, 3, 3,
8079 );
8080 let bitsel = _mm256_setr_epi8(
8081 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8082 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8083 );
8084 let ones8 = _mm256_set1_epi8(1);
8085 let ones16 = _mm256_set1_epi16(1);
8086 let bp = bytes.as_ptr().add(r * gpr * 4);
8087 let mut acc = 0f32;
8088 let mut gi = 0;
8089 while gi < gpr {
8090 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8091 let mut d = [_mm256_setzero_si256(); 4];
8092 for (k, dk) in d.iter_mut().enumerate() {
8093 let t = gi + k;
8094 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8095 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8096 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8097 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8098 let sel = _mm256_and_si256(x, mask);
8099 *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8100 }
8101 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8102 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8103 gi += 4;
8104 }
8105 acc
8106 }
8107}
8108
8109#[cfg(target_arch = "x86_64")]
8112#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8113unsafe fn dot_q1g_row_vnni(
8114 bytes: &[u8],
8115 sc_off: usize,
8116 r: usize,
8117 gpr: usize,
8118 xq: &[i8],
8119 gsum: &[i32],
8120) -> f32 {
8121 unsafe {
8123 use core::arch::x86_64::*;
8124 let expand = _mm256_setr_epi8(
8125 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,
8126 3, 3, 3,
8127 );
8128 let bitsel = _mm256_setr_epi8(
8129 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8130 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8131 );
8132 let ones8 = _mm256_set1_epi8(1);
8133 let bp = bytes.as_ptr().add(r * gpr * 4);
8134 let mut acc = 0f32;
8135 let mut gi = 0;
8136 while gi < gpr {
8137 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8138 let mut d = [_mm256_setzero_si256(); 4];
8139 for (k, dk) in d.iter_mut().enumerate() {
8140 let t = gi + k;
8141 let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8142 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8143 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8144 let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8145 *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
8146 }
8147 let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8148 acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8149 gi += 4;
8150 }
8151 acc
8152 }
8153}
8154
8155#[cfg(target_arch = "x86_64")]
8159#[target_feature(enable = "avx2")]
8160unsafe fn dot_q1g_row_1x4_avx2(
8161 bytes: &[u8],
8162 sc_off: usize,
8163 r: usize,
8164 gpr: usize,
8165 xs: [&[i8]; 4],
8166 gsums: [&[i32]; 4],
8167) -> [f32; 4] {
8168 unsafe {
8170 use core::arch::x86_64::*;
8171 let expand = _mm256_setr_epi8(
8172 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,
8173 3, 3, 3,
8174 );
8175 let bitsel = _mm256_setr_epi8(
8176 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8177 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8178 );
8179 let ones8 = _mm256_set1_epi8(1);
8180 let ones16 = _mm256_set1_epi16(1);
8181 let bp = bytes.as_ptr().add(r * gpr * 4);
8182 let mut acc = [0f32; 4];
8183 let mut gi = 0;
8184 while gi < gpr {
8185 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8186 let mut d = [[_mm256_setzero_si256(); 4]; 4];
8188 for t in 0..4 {
8189 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8190 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8191 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8192 for (k, xq) in xs.iter().enumerate() {
8193 let x = _mm256_loadu_si256(
8194 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8195 );
8196 let sel = _mm256_and_si256(x, mask);
8197 d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8198 }
8199 }
8200 for k in 0..4 {
8201 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8202 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8203 }
8204 gi += 4;
8205 }
8206 acc
8207 }
8208}
8209
8210#[cfg(target_arch = "x86_64")]
8212#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8213unsafe fn dot_q1g_row_1x4_vnni(
8214 bytes: &[u8],
8215 sc_off: usize,
8216 r: usize,
8217 gpr: usize,
8218 xs: [&[i8]; 4],
8219 gsums: [&[i32]; 4],
8220) -> [f32; 4] {
8221 unsafe {
8223 use core::arch::x86_64::*;
8224 let expand = _mm256_setr_epi8(
8225 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,
8226 3, 3, 3,
8227 );
8228 let bitsel = _mm256_setr_epi8(
8229 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8230 -128, 1, 2, 4, 8, 16, 32, 64, -128,
8231 );
8232 let ones8 = _mm256_set1_epi8(1);
8233 let bp = bytes.as_ptr().add(r * gpr * 4);
8234 let mut acc = [0f32; 4];
8235 let mut gi = 0;
8236 while gi < gpr {
8237 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8238 let mut d = [[_mm256_setzero_si256(); 4]; 4];
8239 for t in 0..4 {
8240 let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8241 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8242 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8243 for (k, xq) in xs.iter().enumerate() {
8244 let x = _mm256_loadu_si256(
8245 xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8246 );
8247 d[k][t] = _mm256_dpbusd_epi32(
8248 _mm256_setzero_si256(),
8249 ones8,
8250 _mm256_and_si256(x, mask),
8251 );
8252 }
8253 }
8254 for k in 0..4 {
8255 let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8256 acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8257 }
8258 gi += 4;
8259 }
8260 acc
8261 }
8262}
8263
8264#[cfg(target_arch = "aarch64")]
8275#[target_feature(enable = "neon,dotprod")]
8276#[allow(clippy::too_many_arguments)]
8277unsafe fn dot_q1g_row_sdot(
8278 bytes: &[u8],
8279 sc_off: usize,
8280 rows: usize,
8281 r: usize,
8282 gpr: usize,
8283 xq: &[i8],
8284 gsum: &[i32],
8285) -> f32 {
8286 unsafe {
8289 use core::arch::aarch64::*;
8290 use core::arch::asm;
8291 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8292 const IW: [[u8; 16]; 8] = [
8294 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8295 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8296 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8297 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8298 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8299 [
8300 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8301 ],
8302 [
8303 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8304 ],
8305 [
8306 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8307 ],
8308 ];
8309 let m = vld1q_u8(MASKS.as_ptr());
8310 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8311 macro_rules! tile_dot {
8312 ($w0:expr, $w1:expr, $x:expr) => {{
8313 let x0 = vld1q_s8($x);
8314 let x1 = vld1q_s8($x.add(16));
8315 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8316 asm!(
8317 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8318 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8319 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8320 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8321 options(pure, nomem, nostack),
8322 );
8323 vaddq_s32(a0, a1)
8324 }};
8325 }
8326 macro_rules! tbl_mask {
8327 ($ld:expr, $i:expr) => {
8328 vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
8329 };
8330 }
8331 let bp = bytes.as_ptr().add(r * gpr * 4);
8332 let xp = xq.as_ptr();
8333 let gp = gsum.as_ptr();
8334 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8335 let mut accv = vdupq_n_f32(0.0);
8336 let mut gi = 0;
8337 while gi + 4 <= vec_end {
8338 let ld = vld1q_u8(bp.add(gi * 4));
8339 let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
8340 let d1 = tile_dot!(
8341 tbl_mask!(ld, 2),
8342 tbl_mask!(ld, 3),
8343 xp.add((gi + 1) * GROUP_SIZE)
8344 );
8345 let d2 = tile_dot!(
8346 tbl_mask!(ld, 4),
8347 tbl_mask!(ld, 5),
8348 xp.add((gi + 2) * GROUP_SIZE)
8349 );
8350 let d3 = tile_dot!(
8351 tbl_mask!(ld, 6),
8352 tbl_mask!(ld, 7),
8353 xp.add((gi + 3) * GROUP_SIZE)
8354 );
8355 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8356 let g = vld1q_s32(gp.add(gi));
8357 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8358 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8359 let scf: float32x4_t;
8360 asm!(
8361 "fcvtl {o:v}.4s, {i:v}.4h",
8362 o = out(vreg) scf, i = in(vreg) h,
8363 options(pure, nomem, nostack),
8364 );
8365 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8366 gi += 4;
8367 }
8368 let mut acc = vaddvq_f32(accv);
8369 while gi < gpr {
8370 let t = bp.add(gi * 4);
8371 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8372 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8373 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8374 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8375 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8376 let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
8377 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8378 gi += 1;
8379 }
8380 acc
8381 }
8382}
8383
8384#[cfg(target_arch = "aarch64")]
8388#[target_feature(enable = "neon,dotprod")]
8389#[allow(clippy::too_many_arguments)]
8390unsafe fn dot_q1g_row_1x4_sdot(
8391 bytes: &[u8],
8392 sc_off: usize,
8393 rows: usize,
8394 r: usize,
8395 gpr: usize,
8396 xs: [&[i8]; 4],
8397 gs: [&[i32]; 4],
8398) -> [f32; 4] {
8399 unsafe {
8401 use core::arch::aarch64::*;
8402 use core::arch::asm;
8403 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8404 const IW: [[u8; 16]; 8] = [
8405 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8406 [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8407 [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8408 [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8409 [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8410 [
8411 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8412 ],
8413 [
8414 12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8415 ],
8416 [
8417 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8418 ],
8419 ];
8420 let m = vld1q_u8(MASKS.as_ptr());
8421 let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8422 macro_rules! sdot2 {
8423 ($w0:expr, $w1:expr, $x:expr) => {{
8424 let x0 = vld1q_s8($x);
8425 let x1 = vld1q_s8($x.add(16));
8426 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8427 asm!(
8428 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8429 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8430 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8431 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8432 options(pure, nomem, nostack),
8433 );
8434 vaddq_s32(a0, a1)
8435 }};
8436 }
8437 let bp = bytes.as_ptr().add(r * gpr * 4);
8438 let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8439 let mut accv = [vdupq_n_f32(0.0); 4];
8440 let mut gi = 0;
8441 while gi + 4 <= vec_end {
8442 let ld = vld1q_u8(bp.add(gi * 4));
8443 let w: [int8x16_t; 8] =
8445 core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
8446 let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8447 let scf: float32x4_t;
8448 asm!(
8449 "fcvtl {o:v}.4s, {i:v}.4h",
8450 o = out(vreg) scf, i = in(vreg) h,
8451 options(pure, nomem, nostack),
8452 );
8453 for k in 0..4 {
8454 let xp = xs[k].as_ptr();
8455 let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
8456 let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
8457 let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
8458 let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
8459 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8460 let g = vld1q_s32(gs[k].as_ptr().add(gi));
8461 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8462 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8463 }
8464 gi += 4;
8465 }
8466 let mut acc = [
8467 vaddvq_f32(accv[0]),
8468 vaddvq_f32(accv[1]),
8469 vaddvq_f32(accv[2]),
8470 vaddvq_f32(accv[3]),
8471 ];
8472 while gi < gpr {
8473 let t = bp.add(gi * 4);
8474 let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8475 let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8476 let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8477 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8478 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8479 for k in 0..4 {
8480 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8481 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8482 }
8483 gi += 1;
8484 }
8485 acc
8486 }
8487}
8488
8489#[inline]
8491fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8492 let t = r * gpr + j / GROUP_SIZE;
8493 let k = j % GROUP_SIZE;
8494 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8495 let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
8496 ((bit as i32 * 2 - 1) as f32, s)
8497}
8498
8499#[inline]
8501fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
8502 let mut acc = 0f32;
8503 for gi in 0..gpr {
8504 let t = r * gpr + gi;
8505 let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8506 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8507 let mut ga = 0f32;
8508 for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
8509 for k in 0..8 {
8510 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8511 }
8512 }
8513 acc += ga * s;
8514 }
8515 acc
8516}
8517
8518fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
8523 let entries = base_len + (rows + 1) * 4;
8524 (base_len, entries, entries <= bytes.len())
8525}
8526
8527#[inline]
8529fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
8530 let o = rp_off + r * 4;
8531 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
8532}
8533
8534const SIGN5: [[f32; 5]; 256] = {
8538 let mut lut = [[0.0f32; 5]; 256];
8539 let pow3 = [1u16, 3, 9, 27, 81];
8540 let mut byte = 0usize;
8541 while byte < 256 {
8542 let mut i = 0usize;
8543 while i < 5 {
8544 let code = (byte as u16 / pow3[i]) % 3;
8545 lut[byte][i] = if code == 1 {
8546 1.0
8547 } else if code == 2 {
8548 -1.0
8549 } else {
8550 0.0
8551 };
8552 i += 1;
8553 }
8554 byte += 1;
8555 }
8556 lut
8557};
8558
8559const SIGN5_I8: [[i8; 5]; 256] = {
8561 let mut lut = [[0i8; 5]; 256];
8562 let pow3 = [1u16, 3, 9, 27, 81];
8563 let mut byte = 0usize;
8564 while byte < 256 {
8565 let mut i = 0usize;
8566 while i < 5 {
8567 let code = (byte as u16 / pow3[i]) % 3;
8568 lut[byte][i] = if code == 1 {
8569 1
8570 } else if code == 2 {
8571 -1
8572 } else {
8573 0
8574 };
8575 i += 1;
8576 }
8577 byte += 1;
8578 }
8579 lut
8580};
8581
8582const SIGN5_U64: [u64; 256] = {
8588 let mut lut = [0u64; 256];
8589 let pow3 = [1u16, 3, 9, 27, 81];
8590 let mut byte = 0usize;
8591 while byte < 256 {
8592 let mut v = 0u64;
8593 let mut i = 0usize;
8594 while i < 5 {
8595 let code = (byte as u16 / pow3[i]) % 3;
8596 let s: u8 = if code == 1 {
8597 1
8598 } else if code == 2 {
8599 0xFF
8600 } else {
8601 0
8602 };
8603 v |= (s as u64) << (i * 8);
8604 i += 1;
8605 }
8606 lut[byte] = v;
8607 byte += 1;
8608 }
8609 lut
8610};
8611
8612#[inline]
8617fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
8618 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8619 let off = (r * gpr + j / GROUP_SIZE) * TILE;
8620 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8621 let within = j % GROUP_SIZE;
8622 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
8623}
8624
8625#[cfg(target_arch = "aarch64")]
8628#[target_feature(enable = "neon,dotprod")]
8629#[inline]
8630unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
8631 unsafe {
8633 use core::arch::aarch64::*;
8634 use core::arch::asm;
8635 let w0 = vld1q_s8(w);
8636 let w1 = vld1q_s8(w.add(16));
8637 let x0 = vld1q_s8(x);
8638 let x1 = vld1q_s8(x.add(16));
8639 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8640 asm!(
8641 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8642 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8643 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8644 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8645 options(pure, nomem, nostack),
8646 );
8647 vaddvq_s32(vaddq_s32(a0, a1))
8648 }
8649}
8650
8651#[cfg(target_arch = "x86_64")]
8654#[target_feature(enable = "avx2")]
8655#[inline]
8656unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
8657 unsafe {
8659 use core::arch::x86_64::*;
8660 let wv = _mm256_loadu_si256(w as *const __m256i);
8661 let xv = _mm256_loadu_si256(x as *const __m256i);
8662 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8663 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
8664 let hi128 = _mm256_extracti128_si256::<1>(d);
8665 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8666 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8667 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8668 _mm_cvtsi128_si32(s32)
8669 }
8670}
8671
8672#[inline]
8677fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
8678 debug_assert!(dst.len() >= 40);
8679 unsafe {
8682 let p = dst.as_mut_ptr();
8683 for bi in 0..7 {
8684 core::ptr::write_unaligned(
8685 p.add(bi * 5) as *mut u64,
8686 SIGN5_U64[*codes.add(bi) as usize],
8687 );
8688 }
8689 }
8690}
8691
8692#[inline]
8697fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
8698 #[cfg(target_arch = "aarch64")]
8699 unsafe {
8700 return sdot32_i8(w, x);
8701 }
8702 #[cfg(target_arch = "x86_64")]
8703 unsafe {
8704 return i8dot32_avx2(w, x);
8705 }
8706 #[allow(unreachable_code)]
8707 unsafe {
8708 let mut s = 0i32;
8709 for k in 0..GROUP_SIZE {
8710 s += *w.add(k) as i32 * *x.add(k) as i32;
8711 }
8712 s
8713 }
8714}
8715
8716#[inline]
8717unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
8718 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
8719 (
8720 SIGN5_U64[*codes as usize],
8721 SIGN5_U64[*codes.add(1) as usize],
8722 SIGN5_U64[*codes.add(2) as usize],
8723 SIGN5_U64[*codes.add(3) as usize],
8724 SIGN5_U64[*codes.add(4) as usize],
8725 SIGN5_U64[*codes.add(5) as usize],
8726 SIGN5_U64[*codes.add(6) as usize],
8727 )
8728 };
8729
8730 let u0 = s0 | (s1 << 40);
8731 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
8732 let u2 = (s3 >> 8) | (s4 << 32);
8733 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
8734
8735 (u0, u1, u2, u3)
8736}
8737
8738#[cfg(target_arch = "aarch64")]
8742#[target_feature(enable = "neon,dotprod")]
8743unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8744 use core::arch::aarch64::*;
8745 use core::arch::asm;
8746 unsafe {
8747 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8748 let mut acc = 0f32;
8749 let bytes_ptr = bytes.as_ptr();
8750 let xq_ptr = xq.as_ptr();
8751 let row_off = r * gpr * TILE;
8752
8753 let gpr2 = gpr & !1;
8754 let mut gi = 0;
8755 while gi < gpr2 {
8756 let off0 = row_off + gi * TILE;
8757 let off1 = off0 + TILE;
8758 let s0 = f16_to_f32(u16::from_le_bytes([
8759 *bytes_ptr.add(off0),
8760 *bytes_ptr.add(off0 + 1),
8761 ]));
8762 let s1 = f16_to_f32(u16::from_le_bytes([
8763 *bytes_ptr.add(off1),
8764 *bytes_ptr.add(off1 + 1),
8765 ]));
8766
8767 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8768 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8769
8770 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8771 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8772 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8773 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8774
8775 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8776 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8777 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
8778 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
8779
8780 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
8781 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8782 asm!(
8783 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
8784 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
8785 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
8786 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
8787 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
8788 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
8789 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
8790 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
8791 options(pure, nomem, nostack),
8792 );
8793 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
8794 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
8795 acc += d0 as f32 * s0 + d1 as f32 * s1;
8796 gi += 2;
8797 }
8798
8799 if gi < gpr {
8800 let off = row_off + gi * TILE;
8801 let s = f16_to_f32(u16::from_le_bytes([
8802 *bytes_ptr.add(off),
8803 *bytes_ptr.add(off + 1),
8804 ]));
8805 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8806 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8807 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8808 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8809 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8810 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8811 asm!(
8812 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8813 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8814 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8815 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8816 options(pure, nomem, nostack),
8817 );
8818 let d = vaddvq_s32(vaddq_s32(a0, a1));
8819 acc += d as f32 * s;
8820 }
8821 acc
8822 }
8823}
8824
8825#[cfg(target_arch = "x86_64")]
8827#[target_feature(enable = "avx2")]
8828unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8829 use core::arch::x86_64::*;
8830 unsafe {
8831 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8832 let mut acc = 0f32;
8833 let bytes_ptr = bytes.as_ptr();
8834 let xq_ptr = xq.as_ptr();
8835 let row_off = r * gpr * TILE;
8836
8837 let ones = _mm256_set1_epi16(1);
8838 for gi in 0..gpr {
8839 let off = row_off + gi * TILE;
8840 let s = f16_to_f32(u16::from_le_bytes([
8841 *bytes_ptr.add(off),
8842 *bytes_ptr.add(off + 1),
8843 ]));
8844 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8845 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8846 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8847 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8848 let d256 = _mm256_madd_epi16(p16, ones);
8849 let d128 = _mm_add_epi32(
8850 _mm256_castsi256_si128(d256),
8851 _mm256_extracti128_si256(d256, 1),
8852 );
8853 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
8854 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
8855 acc += d32 as f32 * s;
8856 }
8857 acc
8858 }
8859}
8860
8861#[cfg(target_arch = "x86_64")]
8863#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8864unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8865 use core::arch::x86_64::*;
8866 unsafe {
8868 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8869 let mut acc = 0f32;
8870 let bytes_ptr = bytes.as_ptr();
8871 let xq_ptr = xq.as_ptr();
8872 let row_off = r * gpr * TILE;
8873 for gi in 0..gpr {
8874 let off = row_off + gi * TILE;
8875 let s = f16_to_f32(u16::from_le_bytes([
8876 *bytes_ptr.add(off),
8877 *bytes_ptr.add(off + 1),
8878 ]));
8879 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8880 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8881 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8882 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8883 acc += d as f32 * s;
8884 }
8885 acc
8886 }
8887}
8888
8889#[inline]
8893fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8894 #[cfg(target_arch = "aarch64")]
8895 unsafe {
8896 return q1t_dot_row_sdot(bytes, r, gpr, xq);
8897 }
8898 #[cfg(target_arch = "x86_64")]
8899 unsafe {
8900 if vnni_tiles_enabled() {
8901 return q1t_dot_row_vnni(bytes, r, gpr, xq);
8902 }
8903 return q1t_dot_row_avx2(bytes, r, gpr, xq);
8904 }
8905 #[allow(unreachable_code)]
8906 {
8907 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8908 let mut acc = 0f32;
8909 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
8911 let off = (r * gpr + gi) * TILE;
8912 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8913 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
8914 let mut d = 0i32;
8915 for k in 0..GROUP_SIZE {
8916 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
8917 }
8918 acc += d as f32 * s;
8919 }
8920 acc
8921 }
8922}
8923
8924fn q1t_row_outlier_correction(
8931 bytes: &[u8],
8932 r: usize,
8933 rp_off: usize,
8934 entries_off: usize,
8935 has_ov: bool,
8936 x: &[f32],
8937) -> f32 {
8938 if !has_ov {
8939 return 0.0;
8940 }
8941 let (c0, c1) = (
8942 q1t_rowptr(bytes, rp_off, r),
8943 q1t_rowptr(bytes, rp_off, r + 1),
8944 );
8945 let mut corr = 0f32;
8946 for p in c0..c1 {
8947 let e = entries_off + p * 4;
8948 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8949 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8950 corr += val * x[col];
8951 }
8952 corr
8953}
8954
8955fn q1t_dequant_row(
8959 bytes: &[u8],
8960 r: usize,
8961 gpr: usize,
8962 rp_off: usize,
8963 entries_off: usize,
8964 has_ov: bool,
8965 buf: &mut [f32],
8966) {
8967 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8968 for g in 0..gpr {
8969 let off = (r * gpr + g) * TILE;
8970 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8971 let codes = &bytes[off + 2..off + TILE];
8972 let bc = g * GROUP_SIZE;
8973 for bi in 0..6 {
8975 let lut = &SIGN5[codes[bi] as usize];
8976 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
8977 for i in 0..5 {
8978 d[i] = lut[i] * s;
8979 }
8980 }
8981 let lut = &SIGN5[codes[6] as usize];
8982 buf[bc + 30] = lut[0] * s;
8983 buf[bc + 31] = lut[1] * s;
8984 }
8985 if !has_ov {
8986 return;
8987 }
8988 let (c0, c1) = (
8989 q1t_rowptr(bytes, rp_off, r),
8990 q1t_rowptr(bytes, rp_off, r + 1),
8991 );
8992 for p in c0..c1 {
8993 let e = entries_off + p * 4;
8994 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8995 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8996 }
8997}
8998
8999fn q1t_add_overlay(
9003 bytes: &[u8],
9004 x: &[f32],
9005 rows: usize,
9006 cols: usize,
9007 out: &mut [f32],
9008 pool: Option<&Pool>,
9009) {
9010 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9011 let gpr = cols / GROUP_SIZE;
9012 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9013 if !has_ov {
9014 return;
9015 }
9016 let out_addr = SendMut(out.as_mut_ptr());
9017 let run = move |start: usize, end: usize| {
9018 for r in start..end {
9019 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9020 unsafe { *out_addr.at(r) += corr };
9022 }
9023 };
9024 dispatch_rows(pool, rows, &run);
9025}
9026
9027#[allow(clippy::too_many_arguments)]
9030fn q1t_range_a8w8(
9031 bytes: &[u8],
9032 gpr: usize,
9033 rp_off: usize,
9034 ent_off: usize,
9035 has_ov: bool,
9036 act: &SplitAct,
9037 x: &[f32],
9038 out: SendMut,
9039 start: usize,
9040 end: usize,
9041) {
9042 for r in start..end {
9043 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9044 for &(j, xv) in &act.outliers {
9045 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9046 }
9047 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9048 unsafe { *out.at(r) = acc };
9050 }
9051}
9052
9053#[allow(clippy::too_many_arguments)]
9056fn q1t_range_f32_batch(
9057 bytes: &[u8],
9058 gpr: usize,
9059 rp_off: usize,
9060 ent_off: usize,
9061 has_ov: bool,
9062 x: &[f32],
9063 out: SendMut,
9064 start: usize,
9065 end: usize,
9066) {
9067 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9068 let mut sg = [0f32; GROUP_SIZE];
9069 for r in start..end {
9070 let mut acc = 0f32;
9071 for g in 0..gpr {
9072 let off = (r * gpr + g) * TILE;
9073 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9074 let codes = &bytes[off + 2..off + TILE];
9075 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9076 for bi in 0..6 {
9077 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9078 }
9079 let lut = &SIGN5[codes[6] as usize];
9080 sg[30] = lut[0];
9081 sg[31] = lut[1];
9082 let mut gsum = 0f32;
9083 for k in 0..GROUP_SIZE {
9084 gsum += sg[k] * xg[k];
9085 }
9086 acc += s * gsum;
9087 }
9088 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9089 unsafe { *out.at(r) = acc };
9091 }
9092}
9093
9094fn q1t_matvec(
9098 bytes: &[u8],
9099 x: &[f32],
9100 rows: usize,
9101 cols: usize,
9102 out: &mut [f32],
9103 pool: Option<&Pool>,
9104) {
9105 debug_assert_eq!(out.len(), rows);
9106 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9107 let gpr = cols / GROUP_SIZE;
9108 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9109 let out_addr = SendMut(out.as_mut_ptr());
9110 if a8w8_enabled() {
9114 let act = split_act(x);
9115 let act = &act;
9116 let run = move |start: usize, end: usize| {
9117 for r in start..end {
9118 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9119 for &(j, xv) in &act.outliers {
9120 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9121 }
9122 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9123 unsafe { *out_addr.at(r) = acc };
9125 }
9126 };
9127 dispatch_rows(pool, rows, &run);
9128 return;
9129 }
9130 let run = move |start: usize, end: usize| {
9131 let mut sg = [0f32; GROUP_SIZE];
9135 for r in start..end {
9136 let mut acc = 0f32;
9137 for g in 0..gpr {
9138 let off = (r * gpr + g) * TILE;
9139 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9140 let codes = &bytes[off + 2..off + TILE];
9141 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9142 for bi in 0..6 {
9143 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9144 }
9145 let lut = &SIGN5[codes[6] as usize];
9146 sg[30] = lut[0];
9147 sg[31] = lut[1];
9148 let mut gsum = 0f32;
9149 for k in 0..GROUP_SIZE {
9150 gsum += sg[k] * xg[k];
9151 }
9152 acc += s * gsum;
9153 }
9154 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9155 unsafe { *out_addr.at(r) = acc };
9156 }
9157 };
9158 dispatch_rows(pool, rows, &run);
9159}
9160
9161#[cfg(target_arch = "aarch64")]
9167#[target_feature(enable = "neon,dotprod")]
9168unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
9169 use core::arch::aarch64::*;
9170 use core::arch::asm;
9171 unsafe {
9173 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9174 let bytes_ptr = bytes.as_ptr();
9175 let row_off = r * gpr * TILE;
9176 let xp = [xa.as_ptr(), xb.as_ptr()];
9177 let mut acc = [0f32; 2];
9178 macro_rules! sdot2 {
9179 ($w0:expr, $w1:expr, $x:expr) => {{
9180 let x0 = vld1q_s8($x);
9181 let x1 = vld1q_s8($x.add(16));
9182 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9183 asm!(
9184 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9185 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9186 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9187 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9188 options(pure, nomem, nostack),
9189 );
9190 vaddvq_s32(vaddq_s32(a0, a1))
9191 }};
9192 }
9193 let gpr2 = gpr & !1;
9194 let mut gi = 0;
9195 while gi < gpr2 {
9196 let off0 = row_off + gi * TILE;
9197 let off1 = off0 + TILE;
9198 let s0 = f16_to_f32(u16::from_le_bytes([
9199 *bytes_ptr.add(off0),
9200 *bytes_ptr.add(off0 + 1),
9201 ]));
9202 let s1 = f16_to_f32(u16::from_le_bytes([
9203 *bytes_ptr.add(off1),
9204 *bytes_ptr.add(off1 + 1),
9205 ]));
9206 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9207 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9208 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9209 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9210 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9211 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9212 for k in 0..2 {
9213 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
9214 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
9215 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
9216 }
9217 gi += 2;
9218 }
9219 if gi < gpr {
9220 let off = row_off + gi * TILE;
9221 let s = f16_to_f32(u16::from_le_bytes([
9222 *bytes_ptr.add(off),
9223 *bytes_ptr.add(off + 1),
9224 ]));
9225 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9226 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9227 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9228 for k in 0..2 {
9229 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
9230 acc[k] += d as f32 * s;
9231 }
9232 }
9233 acc
9234 }
9235}
9236
9237fn q1t_matvec2(
9243 bytes: &[u8],
9244 x1: &[f32],
9245 x2: &[f32],
9246 rows: usize,
9247 cols: usize,
9248 o1: &mut [f32],
9249 o2: &mut [f32],
9250 pool: Option<&Pool>,
9251) {
9252 debug_assert_eq!(o1.len(), rows);
9253 debug_assert_eq!(o2.len(), rows);
9254 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9255 let gpr = cols / GROUP_SIZE;
9256 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9257 let out1 = SendMut(o1.as_mut_ptr());
9258 let out2 = SendMut(o2.as_mut_ptr());
9259 if a8w8_enabled() {
9260 let a1 = split_act(x1);
9261 let a2 = split_act(x2);
9262 let (a1, a2) = (&a1, &a2);
9263 let run = move |start: usize, end: usize| {
9264 for r in start..end {
9265 #[cfg(target_arch = "aarch64")]
9266 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
9269 #[cfg(not(target_arch = "aarch64"))]
9270 let ds = [
9271 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
9272 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
9273 ];
9274 let mut acc1 = ds[0] * a1.sx;
9275 for &(j, xv) in &a1.outliers {
9276 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
9277 }
9278 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9279 let mut acc2 = ds[1] * a2.sx;
9280 for &(j, xv) in &a2.outliers {
9281 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
9282 }
9283 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9284 unsafe {
9286 *out1.at(r) = acc1;
9287 *out2.at(r) = acc2;
9288 }
9289 }
9290 };
9291 dispatch_rows(pool, rows, &run);
9292 return;
9293 }
9294 let run = move |start: usize, end: usize| {
9295 let mut sg = [0f32; GROUP_SIZE];
9298 for r in start..end {
9299 let mut acc1 = 0f32;
9300 let mut acc2 = 0f32;
9301 for g in 0..gpr {
9302 let off = (r * gpr + g) * TILE;
9303 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9304 let codes = &bytes[off + 2..off + TILE];
9305 for bi in 0..6 {
9306 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9307 }
9308 let lut = &SIGN5[codes[6] as usize];
9309 sg[30] = lut[0];
9310 sg[31] = lut[1];
9311 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9312 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9313 let mut gsum1 = 0f32;
9314 for k in 0..GROUP_SIZE {
9315 gsum1 += sg[k] * xg1[k];
9316 }
9317 acc1 += s * gsum1;
9318 let mut gsum2 = 0f32;
9319 for k in 0..GROUP_SIZE {
9320 gsum2 += sg[k] * xg2[k];
9321 }
9322 acc2 += s * gsum2;
9323 }
9324 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9325 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9326 unsafe {
9328 *out1.at(r) = acc1;
9329 *out2.at(r) = acc2;
9330 }
9331 }
9332 };
9333 dispatch_rows(pool, rows, &run);
9334}
9335
9336fn q1t_matmat(
9339 bytes: &[u8],
9340 xs: &[f32],
9341 b: usize,
9342 rows: usize,
9343 cols: usize,
9344 out: &mut [f32],
9345 pool: Option<&Pool>,
9346) {
9347 debug_assert_eq!(out.len(), b * rows);
9348 const TILE: usize = cortiq_core::quant::Q1T_TILE;
9349 let gpr = cols / GROUP_SIZE;
9350 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9351 let out_addr = SendMut(out.as_mut_ptr());
9352 if a8w8_enabled() {
9356 let acts: Vec<SplitAct> = (0..b)
9357 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
9358 .collect();
9359 let acts = &acts;
9360 let run = move |start: usize, end: usize| {
9361 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
9365 for g in 0..gpr {
9366 let off = (r * gpr + g) * TILE;
9367 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9368 q1t_unpack_group_i8(
9369 bytes.as_ptr().wrapping_add(off + 2),
9370 &mut sg[g * GROUP_SIZE..],
9371 );
9372 }
9373 for bi in 0..b {
9374 let act = &acts[bi];
9375 let mut isum = 0f32;
9376 for g in 0..gpr {
9377 let d = q1t_i8dot32(
9378 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
9379 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
9380 );
9381 isum += d as f32 * sc[g];
9382 }
9383 let mut acc = isum * act.sx;
9384 for &(j, xv) in &act.outliers {
9385 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9386 }
9387 accs[bi] = acc;
9388 }
9389 if has_ov {
9393 let (c0, c1) = (
9394 q1t_rowptr(bytes, rp_off, r),
9395 q1t_rowptr(bytes, rp_off, r + 1),
9396 );
9397 for p in c0..c1 {
9398 let e = ent_off + p * 4;
9399 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9400 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9401 for bi in 0..b {
9402 accs[bi] += val * xs[bi * cols + col];
9403 }
9404 }
9405 }
9406 for bi in 0..b {
9407 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
9408 }
9409 }
9410 };
9411 dispatch_rows(pool, rows, &run);
9412 return;
9413 }
9414 let run = move |start: usize, end: usize| {
9415 let mut buf = vec![0f32; cols];
9416 for r in start..end {
9417 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
9418 for bi in 0..b {
9419 let xr = &xs[bi * cols..(bi + 1) * cols];
9420 let mut acc = 0f32;
9421 for j in 0..cols {
9422 acc += buf[j] * xr[j];
9423 }
9424 unsafe { *out_addr.at(bi * rows + r) = acc };
9425 }
9426 }
9427 };
9428 dispatch_rows(pool, rows, &run);
9429}
9430
9431#[cfg(test)]
9434fn q1_matvec(
9435 bytes: &[u8],
9436 x: &[f32],
9437 rows: usize,
9438 cols: usize,
9439 out: &mut [f32],
9440 pool: Option<&Pool>,
9441) {
9442 bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
9443}
9444
9445fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
9448 debug_assert_eq!(out.len(), rows);
9449 let gpr = cols / GROUP_SIZE;
9450 let out_addr = SendMut(out.as_mut_ptr());
9451 if a8w8_enabled() {
9452 let act = split_act(x);
9453 let gsum = q1_group_sums(&act.xq, gpr);
9454 let (act, gsum) = (&act, &gsum);
9455 let run =
9456 move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
9457 dispatch_rows(pool, rows, &run);
9458 return;
9459 }
9460 let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
9461 dispatch_rows(pool, rows, &run);
9462}
9463
9464#[cfg(test)]
9466#[allow(clippy::too_many_arguments)]
9467fn q1_matvec2(
9468 bytes: &[u8],
9469 x1: &[f32],
9470 x2: &[f32],
9471 rows: usize,
9472 cols: usize,
9473 o1: &mut [f32],
9474 o2: &mut [f32],
9475 pool: Option<&Pool>,
9476) {
9477 bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
9478}
9479
9480#[allow(clippy::too_many_arguments)]
9482fn bin_matvec2(
9483 w: BinW,
9484 x1: &[f32],
9485 x2: &[f32],
9486 rows: usize,
9487 cols: usize,
9488 o1: &mut [f32],
9489 o2: &mut [f32],
9490 pool: Option<&Pool>,
9491) {
9492 let gpr = cols / GROUP_SIZE;
9493 let p1 = SendMut(o1.as_mut_ptr());
9494 let p2 = SendMut(o2.as_mut_ptr());
9495 if a8w8_enabled() {
9496 let a1 = split_act(x1);
9497 let a2 = split_act(x2);
9498 let g1 = q1_group_sums(&a1.xq, gpr);
9499 let g2 = q1_group_sums(&a2.xq, gpr);
9500 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
9501 let run = move |start: usize, end: usize| {
9502 for r in start..end {
9503 let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
9504 let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
9505 for &(j, xv) in &a1.outliers {
9506 let (wv, s) = w.outlier(r, gpr, j);
9507 v1 += wv * s * xv;
9508 }
9509 for &(j, xv) in &a2.outliers {
9510 let (wv, s) = w.outlier(r, gpr, j);
9511 v2 += wv * s * xv;
9512 }
9513 unsafe {
9515 *p1.at(r) = v1;
9516 *p2.at(r) = v2;
9517 }
9518 }
9519 };
9520 dispatch_rows(pool, rows, &run);
9521 return;
9522 }
9523 let run = move |start: usize, end: usize| {
9524 for r in start..end {
9525 unsafe {
9527 *p1.at(r) = w.row_exact(r, gpr, x1);
9528 *p2.at(r) = w.row_exact(r, gpr, x2);
9529 }
9530 }
9531 };
9532 dispatch_rows(pool, rows, &run);
9533}
9534
9535#[cfg(test)]
9537#[allow(clippy::too_many_arguments)]
9538fn q1_matmat(
9539 bytes: &[u8],
9540 xs_all: &[f32],
9541 b: usize,
9542 rows: usize,
9543 cols: usize,
9544 out: &mut [f32],
9545 pool: Option<&Pool>,
9546) {
9547 bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
9548}
9549
9550#[allow(clippy::too_many_arguments)]
9552fn bin_matmat(
9553 w: BinW,
9554 xs_all: &[f32],
9555 b: usize,
9556 rows: usize,
9557 cols: usize,
9558 out: &mut [f32],
9559 pool: Option<&Pool>,
9560) {
9561 debug_assert_eq!(out.len(), b * rows);
9562 let gpr = cols / GROUP_SIZE;
9563 let out_addr = SendMut(out.as_mut_ptr());
9564 if a8w8_enabled() {
9565 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
9566 .map(|bi| {
9567 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
9568 let gsum = q1_group_sums(&act.xq, gpr);
9569 (act, gsum)
9570 })
9571 .collect();
9572 let acts = &acts;
9573 #[cfg(target_arch = "x86_64")]
9574 let blocked_ok = avx2_enabled() && blocked_enabled();
9575 #[cfg(target_arch = "aarch64")]
9576 let blocked_ok = sdot_enabled() && blocked_enabled();
9577 let run = move |start: usize, end: usize| {
9578 for r in start..end {
9579 let mut bi = 0usize;
9580 #[cfg(target_arch = "aarch64")]
9583 if blocked_ok {
9584 while bi + 4 <= acts.len() {
9585 let xs = [
9586 acts[bi].0.xq.as_slice(),
9587 acts[bi + 1].0.xq.as_slice(),
9588 acts[bi + 2].0.xq.as_slice(),
9589 acts[bi + 3].0.xq.as_slice(),
9590 ];
9591 let gs = [
9592 acts[bi].1.as_slice(),
9593 acts[bi + 1].1.as_slice(),
9594 acts[bi + 2].1.as_slice(),
9595 acts[bi + 3].1.as_slice(),
9596 ];
9597 let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
9599 for k in 0..4 {
9600 let (act, _) = &acts[bi + k];
9601 let mut acc = d[k] * act.sx;
9602 for &(j, xv) in &act.outliers {
9603 let (wv, sc) = w.outlier(r, gpr, j);
9604 acc += wv * sc * xv;
9605 }
9606 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9608 }
9609 bi += 4;
9610 }
9611 }
9612 #[cfg(target_arch = "x86_64")]
9613 if blocked_ok {
9614 while bi + 4 <= acts.len() {
9615 let xs = [
9616 acts[bi].0.xq.as_slice(),
9617 acts[bi + 1].0.xq.as_slice(),
9618 acts[bi + 2].0.xq.as_slice(),
9619 acts[bi + 3].0.xq.as_slice(),
9620 ];
9621 let gs = [
9622 acts[bi].1.as_slice(),
9623 acts[bi + 1].1.as_slice(),
9624 acts[bi + 2].1.as_slice(),
9625 acts[bi + 3].1.as_slice(),
9626 ];
9627 let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
9629 for k in 0..4 {
9630 let (act, _) = &acts[bi + k];
9631 let mut acc = d[k] * act.sx;
9632 for &(j, xv) in &act.outliers {
9633 let (wv, sc) = w.outlier(r, gpr, j);
9634 acc += wv * sc * xv;
9635 }
9636 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9638 }
9639 bi += 4;
9640 }
9641 }
9642 while bi < acts.len() {
9643 let (act, gsum) = &acts[bi];
9644 let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9645 for &(j, xv) in &act.outliers {
9646 let (wv, s) = w.outlier(r, gpr, j);
9647 acc += wv * s * xv;
9648 }
9649 unsafe { *out_addr.at(bi * rows + r) = acc };
9651 bi += 1;
9652 }
9653 }
9654 };
9655 dispatch_rows(pool, rows, &run);
9656 return;
9657 }
9658 let run = move |start: usize, end: usize| {
9659 for r in start..end {
9660 for bi in 0..b {
9661 let x = &xs_all[bi * cols..(bi + 1) * cols];
9662 unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
9664 }
9665 }
9666 };
9667 dispatch_rows(pool, rows, &run);
9668}
9669
9670fn q4matvec(
9676 bytes: &[u8],
9677 x: &[f32],
9678 rows: usize,
9679 cols: usize,
9680 out: &mut [f32],
9681 pool: Option<&Pool>,
9682) {
9683 debug_assert_eq!(out.len(), rows);
9684 let (packed, scales) = q4_split(bytes, rows, cols);
9685 let gpr = cols / GROUP_SIZE;
9686 let out_addr = SendMut(out.as_mut_ptr());
9687
9688 if a8w8_enabled() {
9689 let act = split_act(x);
9690 let run = move |start: usize, end: usize| {
9691 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
9692 };
9693 dispatch_rows(pool, rows, &run);
9694 return;
9695 }
9696
9697 let run =
9698 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
9699 dispatch_rows(pool, rows, &run);
9700}
9701
9702#[inline]
9705#[allow(unreachable_code)]
9706#[cfg(target_arch = "x86_64")]
9711#[target_feature(enable = "avx2")]
9712unsafe fn dot_q4b_row_1x4_avx2(
9713 buf: &[u8],
9714 scales: &[u8],
9715 g0: usize,
9716 gpr: usize,
9717 xs: [&[i8]; 4],
9718) -> [f32; 4] {
9719 unsafe {
9721 use core::arch::x86_64::*;
9722 let ones = _mm256_set1_epi16(1);
9723 let mut acc = [0f32; 4];
9724 for gi in 0..gpr {
9725 let s = f16_to_f32(u16::from_le_bytes([
9726 scales[(g0 + gi) * 2],
9727 scales[(g0 + gi) * 2 + 1],
9728 ]));
9729 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9730 let aw = _mm256_abs_epi8(w);
9731 for (k, xq) in xs.iter().enumerate() {
9732 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9733 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9734 let d = _mm256_madd_epi16(p16, ones);
9735 let hi128 = _mm256_extracti128_si256::<1>(d);
9736 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9737 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9738 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9739 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
9740 }
9741 }
9742 acc
9743 }
9744}
9745
9746#[cfg(target_arch = "x86_64")]
9748#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9749unsafe fn dot_q4b_row_1x4_vnni(
9750 buf: &[u8],
9751 scales: &[u8],
9752 g0: usize,
9753 gpr: usize,
9754 xs: [&[i8]; 4],
9755) -> [f32; 4] {
9756 unsafe {
9758 use core::arch::x86_64::*;
9759 let mut acc = [0f32; 4];
9760 for gi in 0..gpr {
9761 let s = f16_to_f32(u16::from_le_bytes([
9762 scales[(g0 + gi) * 2],
9763 scales[(g0 + gi) * 2 + 1],
9764 ]));
9765 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9766 let aw = _mm256_abs_epi8(w);
9767 for (k, xq) in xs.iter().enumerate() {
9768 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9769 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9770 acc[k] += d as f32 * s;
9771 }
9772 }
9773 acc
9774 }
9775}
9776
9777#[cfg(target_arch = "x86_64")]
9783#[target_feature(enable = "avx2")]
9784unsafe fn dot_q4b_row_1x4_sx_avx2(
9785 buf: &[u8],
9786 scales: &[u8],
9787 g0: usize,
9788 gpr: usize,
9789 xs: [&[i8]; 4],
9790 sxs: [f32; 4],
9791) -> [f32; 4] {
9792 unsafe {
9794 use core::arch::x86_64::*;
9795 let ones = _mm256_set1_epi16(1);
9796 let mut acc = [0f32; 4];
9797 for gi in 0..gpr {
9798 let s = f16_to_f32(u16::from_le_bytes([
9799 scales[(g0 + gi) * 2],
9800 scales[(g0 + gi) * 2 + 1],
9801 ]));
9802 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9803 let aw = _mm256_abs_epi8(w);
9804 for (k, xq) in xs.iter().enumerate() {
9805 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9806 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9807 let d = _mm256_madd_epi16(p16, ones);
9808 let hi128 = _mm256_extracti128_si256::<1>(d);
9809 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9810 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9811 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9812 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
9813 }
9814 }
9815 acc
9816 }
9817}
9818
9819#[cfg(target_arch = "x86_64")]
9822#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9823unsafe fn dot_q4b_row_1x4_sx_vnni(
9824 buf: &[u8],
9825 scales: &[u8],
9826 g0: usize,
9827 gpr: usize,
9828 xs: [&[i8]; 4],
9829 sxs: [f32; 4],
9830) -> [f32; 4] {
9831 unsafe {
9833 use core::arch::x86_64::*;
9834 let mut acc = [0f32; 4];
9835 for gi in 0..gpr {
9836 let s = f16_to_f32(u16::from_le_bytes([
9837 scales[(g0 + gi) * 2],
9838 scales[(g0 + gi) * 2 + 1],
9839 ]));
9840 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9841 let aw = _mm256_abs_epi8(w);
9842 for (k, xq) in xs.iter().enumerate() {
9843 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9844 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9845 acc[k] += (d as f32 * sxs[k]) * s;
9846 }
9847 }
9848 acc
9849 }
9850}
9851
9852#[allow(unreachable_code)]
9853fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9854 #[cfg(target_arch = "aarch64")]
9855 unsafe {
9856 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
9857 }
9858 #[cfg(target_arch = "x86_64")]
9859 unsafe {
9860 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
9861 }
9862 let mut acc = 0f32;
9863 for gi in 0..gpr {
9864 let g = g0 + gi;
9865 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9866 let mut d = 0i32;
9867 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9868 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
9869 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
9870 }
9871 acc += d as f32 * s;
9872 }
9873 acc
9874}
9875
9876#[inline]
9878#[allow(unreachable_code)]
9879fn dot_q4_row_i8_2(
9880 packed: &[u8],
9881 scales: &[u8],
9882 g0: usize,
9883 gpr: usize,
9884 xq1: &[i8],
9885 xq2: &[i8],
9886) -> (f32, f32) {
9887 #[cfg(target_arch = "aarch64")]
9888 unsafe {
9889 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
9890 }
9891 #[cfg(target_arch = "x86_64")]
9892 unsafe {
9893 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
9894 }
9895 (
9896 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
9897 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
9898 )
9899}
9900
9901#[allow(clippy::too_many_arguments)]
9904fn q4_range_a8w8(
9905 packed: &[u8],
9906 scales: &[u8],
9907 gpr: usize,
9908 cols: usize,
9909 act: &SplitAct,
9910 out: SendMut,
9911 start: usize,
9912 end: usize,
9913) {
9914 for r in start..end {
9915 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
9916 for &(j, xv) in &act.outliers {
9918 let flat = r * cols + j;
9919 let byte = packed[flat / 2];
9920 let nib = if flat & 1 == 0 {
9921 byte & 0x0F
9922 } else {
9923 byte >> 4
9924 };
9925 let s = f16_to_f32(u16::from_le_bytes([
9926 scales[(flat / GROUP_SIZE) * 2],
9927 scales[(flat / GROUP_SIZE) * 2 + 1],
9928 ]));
9929 acc += ((nib as i32 - 8) as f32) * s * xv;
9930 }
9931 unsafe { *out.at(r) = acc };
9933 }
9934}
9935
9936#[allow(clippy::too_many_arguments)]
9939fn q4_range2_a8w8(
9940 packed: &[u8],
9941 scales: &[u8],
9942 gpr: usize,
9943 cols: usize,
9944 a1: &SplitAct,
9945 a2: &SplitAct,
9946 p1: SendMut,
9947 p2: SendMut,
9948 start: usize,
9949 end: usize,
9950) {
9951 for r in start..end {
9952 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
9953 let mut acc1 = s1 * a1.sx;
9954 let mut acc2 = s2 * a2.sx;
9955 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
9957 for &(j, xv) in outliers {
9958 let flat = r * cols + j;
9959 let byte = packed[flat / 2];
9960 let nib = if flat & 1 == 0 {
9961 byte & 0x0F
9962 } else {
9963 byte >> 4
9964 };
9965 let s = f16_to_f32(u16::from_le_bytes([
9966 scales[(flat / GROUP_SIZE) * 2],
9967 scales[(flat / GROUP_SIZE) * 2 + 1],
9968 ]));
9969 *acc += ((nib as i32 - 8) as f32) * s * xv;
9970 }
9971 };
9972 fix(&a1.outliers, &mut acc1);
9973 fix(&a2.outliers, &mut acc2);
9974 unsafe {
9976 *p1.at(r) = acc1;
9977 *p2.at(r) = acc2;
9978 }
9979 }
9980}
9981
9982fn q4_range_f32(
9984 packed: &[u8],
9985 scales: &[u8],
9986 gpr: usize,
9987 x: &[f32],
9988 out: SendMut,
9989 start: usize,
9990 end: usize,
9991) {
9992 for r in start..end {
9993 let mut acc = 0f32;
9994 for gi in 0..gpr {
9995 let g = r * gpr + gi;
9996 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9997 let pk = &packed[g * 16..(g + 1) * 16];
9998 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9999 let mut ga = 0f32;
10000 for (k, &b) in pk.iter().enumerate() {
10001 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
10002 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
10003 }
10004 acc += ga * s;
10005 }
10006 unsafe { *out.at(r) = acc };
10008 }
10009}
10010
10011#[allow(clippy::too_many_arguments)]
10015fn q4matvec2(
10016 bytes: &[u8],
10017 x1: &[f32],
10018 x2: &[f32],
10019 rows: usize,
10020 cols: usize,
10021 o1: &mut [f32],
10022 o2: &mut [f32],
10023 pool: Option<&Pool>,
10024) {
10025 debug_assert_eq!(o1.len(), rows);
10026 debug_assert_eq!(o2.len(), rows);
10027 let (packed, scales) = q4_split(bytes, rows, cols);
10028 let gpr = cols / GROUP_SIZE;
10029
10030 if a8w8_enabled() {
10031 let a1 = split_act(x1);
10032 let a2 = split_act(x2);
10033 let p1 = SendMut(o1.as_mut_ptr());
10034 let p2 = SendMut(o2.as_mut_ptr());
10035 let run = move |start: usize, end: usize| {
10036 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
10037 };
10038 dispatch_rows(pool, rows, &run);
10039 return;
10040 }
10041
10042 let p1 = SendMut(o1.as_mut_ptr());
10043 let p2 = SendMut(o2.as_mut_ptr());
10044 let run = move |start: usize, end: usize| {
10045 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
10046 };
10047 dispatch_rows(pool, rows, &run);
10048}
10049
10050#[allow(clippy::too_many_arguments)]
10052fn q4_range2_f32(
10053 packed: &[u8],
10054 scales: &[u8],
10055 gpr: usize,
10056 x1: &[f32],
10057 x2: &[f32],
10058 p1: SendMut,
10059 p2: SendMut,
10060 start: usize,
10061 end: usize,
10062) {
10063 for r in start..end {
10064 let (mut acc1, mut acc2) = (0f32, 0f32);
10065 for gi in 0..gpr {
10066 let g = r * gpr + gi;
10067 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10068 let pk = &packed[g * 16..(g + 1) * 16];
10069 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10070 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10071 let (mut g1, mut g2) = (0f32, 0f32);
10072 for (k, &b) in pk.iter().enumerate() {
10073 let wl = (b & 0x0F) as f32 - 8.0;
10074 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
10075 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
10076 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
10077 }
10078 acc1 += g1 * s;
10079 acc2 += g2 * s;
10080 }
10081 unsafe {
10083 *p1.at(r) = acc1;
10084 *p2.at(r) = acc2;
10085 }
10086 }
10087}
10088
10089thread_local! {
10090 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
10093 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10094}
10095
10096#[allow(clippy::too_many_arguments)]
10102fn q4matmat(
10103 bytes: &[u8],
10104 xs_all: &[f32],
10105 b: usize,
10106 rows: usize,
10107 cols: usize,
10108 out: &mut [f32],
10109 pool: Option<&Pool>,
10110) {
10111 debug_assert_eq!(xs_all.len(), b * cols);
10112 debug_assert_eq!(out.len(), b * rows);
10113 let (packed, scales) = q4_split(bytes, rows, cols);
10114 let gpr = cols / GROUP_SIZE;
10115 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10116
10117 if a8w8_enabled() {
10118 let acts: Vec<SplitAct> = (0..b)
10119 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10120 .collect();
10121 let acts = &acts;
10122 let out_addr = SendMut(out.as_mut_ptr());
10123 let run = move |start: usize, end: usize| {
10124 ROW_I8.with(|rb| {
10125 let mut buf = rb.borrow_mut();
10126 buf.resize(cols, 0);
10127 for r in start..end {
10128 for gi in 0..gpr {
10132 let g = r * gpr + gi;
10133 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10134 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
10135 buf[gi * GROUP_SIZE + k * 2 + 1] =
10136 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
10137 }
10138 }
10139 let mut bi = 0usize;
10140 #[cfg(target_arch = "x86_64")]
10141 if avx2_enabled() && blocked_enabled() {
10142 while bi + 4 <= acts.len() {
10143 let xs = [
10144 acts[bi].xq.as_slice(),
10145 acts[bi + 1].xq.as_slice(),
10146 acts[bi + 2].xq.as_slice(),
10147 acts[bi + 3].xq.as_slice(),
10148 ];
10149 let d = unsafe {
10150 if vnni_tiles_enabled() {
10151 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
10152 } else {
10153 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
10154 }
10155 };
10156 for k in 0..4 {
10157 let act = &acts[bi + k];
10158 let mut acc = d[k] * act.sx;
10159 for &(j, xv) in &act.outliers {
10160 acc += (buf[j] as i8) as f32
10161 * gscale((r * cols + j) / GROUP_SIZE)
10162 * xv;
10163 }
10164 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10166 }
10167 bi += 4;
10168 }
10169 }
10170 while bi < acts.len() {
10171 let act = &acts[bi];
10172 let mut acc = 0f32;
10173 for gi in 0..gpr {
10174 let d = dot_i8_i8(
10175 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10176 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10177 );
10178 acc += d as f32 * gscale(r * gpr + gi);
10179 }
10180 acc *= act.sx;
10181 for &(j, xv) in &act.outliers {
10183 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
10184 }
10185 unsafe { *out_addr.at(bi * rows + r) = acc };
10187 bi += 1;
10188 }
10189 }
10190 })
10191 };
10192 dispatch_rows(pool, rows, &run);
10193 return;
10194 }
10195
10196 let out_addr = SendMut(out.as_mut_ptr());
10197 let run = move |start: usize, end: usize| {
10198 ROW_F32.with(|rb| {
10199 let mut buf = rb.borrow_mut();
10200 buf.resize(cols, 0.0);
10201 for r in start..end {
10202 for gi in 0..gpr {
10205 let g = r * gpr + gi;
10206 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10207 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
10208 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
10209 }
10210 }
10211 for bi in 0..b {
10212 let x = &xs_all[bi * cols..(bi + 1) * cols];
10213 let mut acc = 0f32;
10214 for gi in 0..gpr {
10215 let mut ga = 0f32;
10216 for k in 0..GROUP_SIZE / 2 {
10221 let e = gi * GROUP_SIZE + k * 2;
10222 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
10223 }
10224 acc += ga * gscale(r * gpr + gi);
10225 }
10226 unsafe { *out_addr.at(bi * rows + r) = acc };
10228 }
10229 }
10230 })
10231 };
10232 dispatch_rows(pool, rows, &run);
10233}
10234
10235#[allow(clippy::too_many_arguments)]
10240fn vbitmatmat(
10241 bytes: &[u8],
10242 offsets: &[usize],
10243 xs_all: &[f32],
10244 b: usize,
10245 rows: usize,
10246 cols: usize,
10247 out: &mut [f32],
10248 pool: Option<&Pool>,
10249) {
10250 debug_assert_eq!(xs_all.len(), b * cols);
10251 debug_assert_eq!(out.len(), b * rows);
10252 debug_assert_eq!(offsets.len(), rows + 1);
10253 let ng = cols / GROUP_SIZE;
10254 let bits = &bytes[..rows];
10255 let sc_off = rows;
10256 let gscale = |r: usize, g: usize| {
10257 let so = (r * ng + g) * 2;
10258 f16_to_f32(u16::from_le_bytes([
10259 bytes[sc_off + so],
10260 bytes[sc_off + so + 1],
10261 ]))
10262 };
10263
10264 let decode_f32 = |r: usize, dst: &mut [f32]| {
10266 let bw = bits[r] as usize;
10267 let l = ((1i32 << (bw - 1)) - 1) as f32;
10268 let data = &bytes[offsets[r]..offsets[r + 1]];
10269 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
10270 for d in dst.iter_mut() {
10271 while nbits < bw {
10272 acc = (acc << 8) | data[idx] as u64;
10273 idx += 1;
10274 nbits += 8;
10275 }
10276 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
10277 nbits -= bw;
10278 *d = u - l;
10279 }
10280 };
10281
10282 if a8w8_enabled() {
10283 let acts: Vec<SplitAct> = (0..b)
10284 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10285 .collect();
10286 let acts = &acts;
10287 let out_addr = SendMut(out.as_mut_ptr());
10288 let run = move |start: usize, end: usize| {
10289 for r in start..end {
10290 let bw = bits[r] as usize;
10291 if bw == 8 {
10292 ROW_F32.with(|rb| {
10295 let mut buf = rb.borrow_mut();
10296 buf.resize(cols, 0.0);
10297 decode_f32(r, &mut buf);
10298 for bi in 0..b {
10299 let x = &xs_all[bi * cols..(bi + 1) * cols];
10300 let mut dot = 0f32;
10301 for g in 0..ng {
10302 let mut gd = 0f32;
10303 for k in 0..GROUP_SIZE {
10304 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10305 }
10306 dot += gd * gscale(r, g);
10307 }
10308 unsafe { *out_addr.at(bi * rows + r) = dot };
10310 }
10311 });
10312 continue;
10313 }
10314 let l = (1i32 << (bw - 1)) - 1;
10315 let data = &bytes[offsets[r]..offsets[r + 1]];
10316 ROW_I8.with(|rb| {
10317 let mut buf = rb.borrow_mut();
10318 buf.resize(cols, 0);
10319 #[inline(always)]
10320 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
10321 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
10322 let u = unpack8::<B>(&data[blk * B..]);
10323 for k in 0..8 {
10324 chunk[k] = (u[k] - l) as i8 as u8;
10325 }
10326 }
10327 }
10328 match bw {
10329 3 => fill::<3>(data, l, &mut buf),
10330 4 => vbit_fill4(data, &mut buf),
10331 5 => fill::<5>(data, l, &mut buf),
10332 6 => fill::<6>(data, l, &mut buf),
10333 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
10334 }
10335 let mut bi = 0usize;
10336 #[cfg(target_arch = "x86_64")]
10340 if avx2_enabled() && blocked_enabled() {
10341 while bi + 4 <= acts.len() {
10342 let xs = [
10343 acts[bi].xq.as_slice(),
10344 acts[bi + 1].xq.as_slice(),
10345 acts[bi + 2].xq.as_slice(),
10346 acts[bi + 3].xq.as_slice(),
10347 ];
10348 let sxs = [
10349 acts[bi].sx,
10350 acts[bi + 1].sx,
10351 acts[bi + 2].sx,
10352 acts[bi + 3].sx,
10353 ];
10354 let d = unsafe {
10355 if vnni_tiles_enabled() {
10356 dot_q4b_row_1x4_sx_vnni(
10357 &buf,
10358 &bytes[sc_off..],
10359 r * ng,
10360 ng,
10361 xs,
10362 sxs,
10363 )
10364 } else {
10365 dot_q4b_row_1x4_sx_avx2(
10366 &buf,
10367 &bytes[sc_off..],
10368 r * ng,
10369 ng,
10370 xs,
10371 sxs,
10372 )
10373 }
10374 };
10375 for k in 0..4 {
10376 let act = &acts[bi + k];
10377 let mut dot = d[k];
10378 for &(j, xv) in &act.outliers {
10379 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10380 }
10381 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
10383 }
10384 bi += 4;
10385 }
10386 }
10387 while bi < acts.len() {
10388 let act = &acts[bi];
10389 let mut dot = 0f32;
10390 for g in 0..ng {
10391 let d = dot_i8_i8(
10392 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10393 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10394 ) as f32
10395 * act.sx;
10396 dot += d * gscale(r, g);
10397 }
10398 for &(j, xv) in &act.outliers {
10399 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10400 }
10401 unsafe { *out_addr.at(bi * rows + r) = dot };
10403 bi += 1;
10404 }
10405 });
10406 }
10407 };
10408 dispatch_rows(pool, rows, &run);
10409 return;
10410 }
10411
10412 let out_addr = SendMut(out.as_mut_ptr());
10413 let run = move |start: usize, end: usize| {
10414 ROW_F32.with(|rb| {
10415 let mut buf = rb.borrow_mut();
10416 buf.resize(cols, 0.0);
10417 for r in start..end {
10418 decode_f32(r, &mut buf);
10419 for bi in 0..b {
10420 let x = &xs_all[bi * cols..(bi + 1) * cols];
10421 let mut dot = 0f32;
10422 for g in 0..ng {
10423 let mut gd = 0f32;
10424 for k in 0..GROUP_SIZE {
10425 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10426 }
10427 dot += gd * gscale(r, g);
10428 }
10429 unsafe { *out_addr.at(bi * rows + r) = dot };
10431 }
10432 }
10433 })
10434 };
10435 dispatch_rows(pool, rows, &run);
10436}
10437
10438pub(crate) fn gpu_batch_job<'a>(
10442 t: &'a QTensor,
10443 x: &[f32],
10444) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
10445 match t {
10446 QTensor::Mapped {
10447 model,
10448 idx,
10449 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
10450 rows,
10451 cols,
10452 row_scale,
10453 col_field,
10454 ..
10455 } => Some((
10456 model.clone(),
10457 crate::gpu::BatchJob {
10458 idx: *idx,
10459 rows: *rows,
10460 cols: *cols,
10461 row_scale,
10462 xs: prescale(x, col_field, *dt).into_owned(),
10463 layout: crate::gpu::BatchLayout::Q8,
10464 },
10465 )),
10466 QTensor::Mapped {
10468 model,
10469 idx,
10470 dtype: TensorDtype::Q1,
10471 rows,
10472 cols,
10473 ..
10474 } => Some((
10475 model.clone(),
10476 crate::gpu::BatchJob {
10477 idx: *idx,
10478 rows: *rows,
10479 cols: *cols,
10480 row_scale: &[],
10481 xs: x.to_vec(),
10482 layout: crate::gpu::BatchLayout::Q1,
10483 },
10484 )),
10485 QTensor::Mapped {
10487 model,
10488 idx,
10489 dtype: TensorDtype::Q1G,
10490 rows,
10491 cols,
10492 ..
10493 } => Some((
10494 model.clone(),
10495 crate::gpu::BatchJob {
10496 idx: *idx,
10497 rows: *rows,
10498 cols: *cols,
10499 row_scale: &[],
10500 xs: x.to_vec(),
10501 layout: crate::gpu::BatchLayout::Q1g,
10502 },
10503 )),
10504 QTensor::Mapped {
10509 model,
10510 idx,
10511 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
10512 rows,
10513 cols,
10514 ..
10515 } => Some((
10516 model.clone(),
10517 crate::gpu::BatchJob {
10518 idx: *idx,
10519 rows: *rows,
10520 cols: *cols,
10521 row_scale: &[],
10522 xs: x.to_vec(),
10523 layout: if *dt == TensorDtype::Q4Tiled {
10524 crate::gpu::BatchLayout::Q4t
10525 } else {
10526 crate::gpu::BatchLayout::Q4tp
10527 },
10528 },
10529 )),
10530 _ => None,
10531 }
10532}
10533
10534thread_local! {
10535 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10536 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10537}
10538
10539pub(crate) fn prescale<'a>(
10540 x: &'a [f32],
10541 col_field: &[f32],
10542 dtype: TensorDtype,
10543) -> std::borrow::Cow<'a, [f32]> {
10544 if dtype == TensorDtype::Q8_2f {
10545 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
10546 } else {
10547 std::borrow::Cow::Borrowed(x)
10548 }
10549}
10550
10551pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
10554 x: &[f32],
10555 col_field: &[f32],
10556 dtype: TensorDtype,
10557 buf_id: u8,
10558 f: F,
10559) -> R {
10560 if dtype == TensorDtype::Q8_2f {
10561 if buf_id == 1 {
10562 PRESCALE_BUF1.with(|b| {
10563 let mut buf = b.borrow_mut();
10564 buf.clear();
10565 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10566 f(&buf)
10567 })
10568 } else {
10569 PRESCALE_BUF2.with(|b| {
10570 let mut buf = b.borrow_mut();
10571 buf.clear();
10572 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10573 f(&buf)
10574 })
10575 }
10576 } else {
10577 f(x)
10578 }
10579}
10580
10581#[cfg(target_arch = "x86_64")]
10586pub(crate) fn avx2_enabled() -> bool {
10587 use std::sync::OnceLock;
10588 static ON: OnceLock<bool> = OnceLock::new();
10589 *ON.get_or_init(|| {
10590 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
10591 && std::arch::is_x86_feature_detected!("avx2")
10592 && std::arch::is_x86_feature_detected!("fma")
10593 })
10594}
10595
10596#[cfg(target_arch = "x86_64")]
10601fn avx2_a8w8_enabled() -> bool {
10602 if FLOAT_ACTIVATIONS.get() {
10603 return false;
10604 }
10605 use std::sync::OnceLock;
10606 static ON: OnceLock<bool> = OnceLock::new();
10607 *ON.get_or_init(|| {
10608 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
10609 })
10610}
10611
10612thread_local! {
10613 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10614}
10615
10616pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
10619 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
10620 impl Drop for Restore {
10621 fn drop(&mut self) {
10622 FULL_GPU_Q8.set(self.0);
10623 }
10624 }
10625 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
10626}
10627
10628thread_local! {
10632 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10633}
10634
10635pub(crate) fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
10636 struct Restore(bool);
10637 impl Drop for Restore {
10638 fn drop(&mut self) {
10639 FLOAT_ACTIVATIONS.set(self.0);
10640 }
10641 }
10642 let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
10643 f()
10644}
10645
10646static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
10658
10659pub(crate) fn row_exact() -> bool {
10660 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
10661}
10662
10663fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
10664 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
10665 impl Drop for Restore<'_> {
10666 fn drop(&mut self) {
10667 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
10668 }
10669 }
10670 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
10671 let _restore = Restore(active);
10672 f()
10673}
10674
10675pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
10677 counted_row_exact_scope(&ROW_EXACT, f)
10678}
10679
10680#[inline]
10684pub(crate) fn a8w8_enabled() -> bool {
10685 #[cfg(target_arch = "aarch64")]
10686 {
10687 sdot_enabled()
10688 }
10689 #[cfg(target_arch = "x86_64")]
10690 {
10691 avx2_a8w8_enabled()
10692 }
10693 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
10694 {
10695 false
10696 }
10697}
10698
10699#[inline]
10702#[allow(unreachable_code)]
10703fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
10704 #[cfg(target_arch = "aarch64")]
10705 unsafe {
10706 return dot_i8_sdot(w, xq);
10707 }
10708 #[cfg(target_arch = "x86_64")]
10709 unsafe {
10710 if avx512vnni_enabled() {
10711 return dot_i8_i8_vnni(w, xq);
10712 }
10713 return dot_i8_i8_avx2(w, xq);
10714 }
10715 w.iter()
10716 .zip(xq)
10717 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
10718 .sum()
10719}
10720
10721#[cfg(target_arch = "x86_64")]
10725fn avx512vnni_enabled() -> bool {
10726 use std::sync::OnceLock;
10727 static ON: OnceLock<bool> = OnceLock::new();
10728 *ON.get_or_init(|| {
10729 std::env::var("CMF_AVX512")
10730 .map(|v| v != "0")
10731 .unwrap_or(true)
10732 && std::arch::is_x86_feature_detected!("avx512f")
10733 && std::arch::is_x86_feature_detected!("avx512bw")
10734 && std::arch::is_x86_feature_detected!("avx512vl")
10735 && std::arch::is_x86_feature_detected!("avx512vnni")
10736 })
10737}
10738
10739#[cfg(target_arch = "x86_64")]
10747fn vnni_tiles_enabled() -> bool {
10748 use std::sync::OnceLock;
10749 static ON: OnceLock<bool> = OnceLock::new();
10750 *ON.get_or_init(|| {
10751 std::env::var("CMF_VNNI_TILES")
10752 .map(|v| v != "0")
10753 .unwrap_or(true)
10754 && avx512vnni_enabled()
10755 })
10756}
10757
10758#[cfg(target_arch = "x86_64")]
10763#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10764#[inline]
10765unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
10766 unsafe {
10768 use core::arch::x86_64::*;
10769 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
10770 let hi128 = _mm256_extracti128_si256::<1>(d);
10771 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10772 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10773 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10774 _mm_cvtsi128_si32(s32)
10775 }
10776}
10777
10778#[cfg(target_arch = "x86_64")]
10783#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10784unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
10785 unsafe {
10787 use core::arch::x86_64::*;
10788 let n = w.len();
10789 let mut j = 0usize;
10790 let mut total: i32;
10791 {
10796 #[inline(always)]
10797 unsafe fn step(
10798 w: *const u8,
10799 x: *const i8,
10800 acc: core::arch::x86_64::__m512i,
10801 ) -> core::arch::x86_64::__m512i {
10802 unsafe {
10803 use core::arch::x86_64::*;
10804 let wv = _mm512_loadu_si512(w as *const _);
10805 let xv = _mm512_loadu_si512(x as *const _);
10806 let aw = _mm512_abs_epi8(wv);
10807 let neg = _mm512_movepi8_mask(wv);
10808 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
10809 _mm512_dpbusd_epi32(acc, aw, sx)
10810 }
10811 }
10812 let (mut a0, mut a1, mut a2, mut a3) = (
10813 _mm512_setzero_si512(),
10814 _mm512_setzero_si512(),
10815 _mm512_setzero_si512(),
10816 _mm512_setzero_si512(),
10817 );
10818 while j + 256 <= n {
10819 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10820 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
10821 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
10822 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
10823 j += 256;
10824 }
10825 while j + 64 <= n {
10826 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10827 j += 64;
10828 }
10829 let s01 = _mm512_add_epi32(a0, a1);
10830 let s23 = _mm512_add_epi32(a2, a3);
10831 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
10832 }
10833 if j + 32 <= n {
10835 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10836 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10837 let d = _mm256_dpbusd_epi32(
10838 _mm256_setzero_si256(),
10839 _mm256_abs_epi8(wv),
10840 _mm256_sign_epi8(xv, wv),
10841 );
10842 let hi128 = _mm256_extracti128_si256::<1>(d);
10843 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10844 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10845 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10846 total += _mm_cvtsi128_si32(s32);
10847 j += 32;
10848 }
10849 while j < n {
10850 total += (w[j] as i8) as i32 * xq[j] as i32;
10851 j += 1;
10852 }
10853 total
10854 }
10855}
10856
10857#[cfg(target_arch = "x86_64")]
10859#[target_feature(enable = "avx2,fma")]
10860unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
10861 unsafe {
10863 use core::arch::x86_64::*;
10864 let n = x.len();
10865 let wp = w.as_ptr();
10866 let xp = x.as_ptr();
10867 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
10868 let mut j = 0usize;
10869 while j + 16 <= n {
10870 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
10871 let lo = _mm256_cvtepi8_epi32(wb);
10872 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
10873 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
10874 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
10875 j += 16;
10876 }
10877 let acc = _mm256_add_ps(a0, a1);
10878 let hi128 = _mm256_extractf128_ps::<1>(acc);
10879 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
10880 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
10881 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
10882 let mut sum = _mm_cvtss_f32(s32);
10883 while j < n {
10884 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
10885 j += 1;
10886 }
10887 sum
10888 }
10889}
10890
10891#[cfg(target_arch = "x86_64")]
10896#[target_feature(enable = "avx2")]
10897unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
10898 unsafe {
10900 use core::arch::x86_64::*;
10901 let n = w.len();
10902 let ones = _mm256_set1_epi16(1);
10903 let mut acc = _mm256_setzero_si256();
10904 let mut j = 0usize;
10905 while j + 32 <= n {
10906 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10907 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10908 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10909 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
10910 j += 32;
10911 }
10912 let hi128 = _mm256_extracti128_si256::<1>(acc);
10913 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
10914 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10915 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10916 let mut s = _mm_cvtsi128_si32(s32);
10917 while j < n {
10918 s += (w[j] as i8) as i32 * xq[j] as i32;
10919 j += 1;
10920 }
10921 s
10922 }
10923}
10924
10925#[cfg(target_arch = "aarch64")]
10929#[target_feature(enable = "neon,i8mm")]
10930unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10931 unsafe {
10933 use core::arch::aarch64::*;
10934 use core::arch::asm;
10935 let n = w0.len();
10936 let w0p = w0.as_ptr() as *const i8;
10937 let w1p = w1.as_ptr() as *const i8;
10938 let mut acc01 = vdupq_n_s32(0);
10941 let mut acc23 = vdupq_n_s32(0);
10942 let mut i = 0usize;
10943 while i + 8 <= n {
10944 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
10945 let xb01 = vcombine_s8(
10946 vld1_s8(xs[0].as_ptr().add(i)),
10947 vld1_s8(xs[1].as_ptr().add(i)),
10948 );
10949 let xb23 = vcombine_s8(
10950 vld1_s8(xs[2].as_ptr().add(i)),
10951 vld1_s8(xs[3].as_ptr().add(i)),
10952 );
10953 asm!(
10954 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
10955 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
10956 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
10957 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
10958 options(pure, nomem, nostack),
10959 );
10960 i += 8;
10961 }
10962 let mut out = [[0i32; 4]; 2];
10963 let a01: [i32; 4] = core::mem::transmute(acc01);
10964 let a23: [i32; 4] = core::mem::transmute(acc23);
10965 out[0][0] = a01[0];
10966 out[0][1] = a01[1];
10967 out[1][0] = a01[2];
10968 out[1][1] = a01[3];
10969 out[0][2] = a23[0];
10970 out[0][3] = a23[1];
10971 out[1][2] = a23[2];
10972 out[1][3] = a23[3];
10973 if i < n {
10974 for (k, x) in xs.iter().enumerate() {
10975 for j in i..n {
10976 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
10977 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
10978 }
10979 }
10980 }
10981 out
10982 }
10983}
10984
10985#[cfg(target_arch = "aarch64")]
10989#[target_feature(enable = "neon,dotprod")]
10990unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10991 unsafe {
10993 use core::arch::aarch64::*;
10994 use core::arch::asm;
10995 let n = w0.len();
10996 let w0p = w0.as_ptr() as *const i8;
10997 let w1p = w1.as_ptr() as *const i8;
10998 let mut acc = [[vdupq_n_s32(0); 4]; 2];
10999 let mut i = 0usize;
11000 while i + 16 <= n {
11001 let wv0 = vld1q_s8(w0p.add(i));
11002 let wv1 = vld1q_s8(w1p.add(i));
11003 for (k, x) in xs.iter().enumerate() {
11004 let xv = vld1q_s8(x.as_ptr().add(i));
11005 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
11006 asm!(
11007 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
11008 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
11009 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11010 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
11011 options(pure, nomem, nostack),
11012 );
11013 acc[0][k] = a0;
11014 acc[1][k] = a1;
11015 }
11016 i += 16;
11017 }
11018 let mut out = [[0i32; 4]; 2];
11019 for r in 0..2 {
11020 for k in 0..4 {
11021 out[r][k] = vaddvq_s32(acc[r][k]);
11022 }
11023 }
11024 if i < n {
11025 for (k, x) in xs.iter().enumerate() {
11026 for j in i..n {
11027 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11028 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11029 }
11030 }
11031 }
11032 out
11033 }
11034}
11035
11036#[cfg(target_arch = "x86_64")]
11042#[target_feature(enable = "avx2")]
11043unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11044 unsafe {
11046 use core::arch::x86_64::*;
11047 let n = w0.len();
11048 let ones = _mm256_set1_epi16(1);
11049 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
11050 let mut j = 0usize;
11051 while j + 32 <= n {
11052 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
11053 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
11054 let aw0 = _mm256_abs_epi8(wv0);
11055 let aw1 = _mm256_abs_epi8(wv1);
11056 for (k, x) in xs.iter().enumerate() {
11057 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
11058 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
11059 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
11060 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
11061 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
11062 }
11063 j += 32;
11064 }
11065 let mut out = [[0i32; 4]; 2];
11066 for r in 0..2 {
11067 for k in 0..4 {
11068 let a = acc[r][k];
11069 let hi128 = _mm256_extracti128_si256::<1>(a);
11070 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
11071 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11072 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11073 out[r][k] = _mm_cvtsi128_si32(s32);
11074 }
11075 }
11076 if j < n {
11077 for (k, x) in xs.iter().enumerate() {
11078 for i in j..n {
11079 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
11080 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
11081 }
11082 }
11083 }
11084 out
11085 }
11086}
11087
11088#[cfg(target_arch = "x86_64")]
11093#[inline]
11094fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
11095 let dot = if avx512vnni_enabled() && row.len() >= 64 {
11096 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
11097 } else {
11098 unsafe { dot_i8_i8_avx2(row, &act.xq) }
11099 };
11100 let mut acc = dot as f32 * act.sx;
11101 for &(j, xv) in &act.outliers {
11102 acc += (row[j] as i8) as f32 * xv;
11103 }
11104 acc
11105}
11106
11107#[cfg(target_arch = "x86_64")]
11111#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11112unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11113 unsafe {
11115 use core::arch::x86_64::*;
11116 let n = w.len();
11117 let flip = _mm512_set1_epi8(-128); #[inline(always)]
11119 unsafe fn step(
11120 w: *const u8,
11121 x: *const i8,
11122 flip: core::arch::x86_64::__m512i,
11123 acc: core::arch::x86_64::__m512i,
11124 ) -> core::arch::x86_64::__m512i {
11125 unsafe {
11126 use core::arch::x86_64::*;
11127 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
11128 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
11129 }
11130 }
11131 let (mut a0, mut a1, mut a2, mut a3) = (
11132 _mm512_setzero_si512(),
11133 _mm512_setzero_si512(),
11134 _mm512_setzero_si512(),
11135 _mm512_setzero_si512(),
11136 );
11137 let mut j = 0usize;
11138 while j + 256 <= n {
11139 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11140 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
11141 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
11142 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
11143 j += 256;
11144 }
11145 while j + 64 <= n {
11146 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11147 j += 64;
11148 }
11149 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
11150 _mm512_add_epi32(a0, a1),
11151 _mm512_add_epi32(a2, a3),
11152 ));
11153 while j < n {
11155 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
11156 j += 1;
11157 }
11158 total
11159 }
11160}
11161
11162#[cfg(target_arch = "x86_64")]
11168#[target_feature(enable = "avx2")]
11169unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11170 unsafe {
11173 use core::arch::x86_64::*;
11174 let lomask = _mm_set1_epi8(0x0F);
11175 let eight = _mm256_set1_epi8(8);
11176 let ones = _mm256_set1_epi16(1);
11177 let mut acc = 0f32;
11178 for gi in 0..gpr {
11179 let g = g0 + gi;
11180 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11181 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11182 let lo = _mm_and_si128(b, lomask);
11183 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11184 let w = _mm256_sub_epi8(
11185 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11186 eight,
11187 );
11188 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11189 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
11190 let d = _mm256_madd_epi16(p16, ones);
11191 let hi128 = _mm256_extracti128_si256::<1>(d);
11192 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11193 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11194 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11195 acc += _mm_cvtsi128_si32(s32) as f32 * s;
11196 }
11197 acc
11198 }
11199}
11200
11201#[cfg(target_arch = "x86_64")]
11204#[target_feature(enable = "avx2")]
11205unsafe fn dot_q4_row_avx2_2(
11206 packed: &[u8],
11207 scales: &[u8],
11208 g0: usize,
11209 gpr: usize,
11210 xq1: &[i8],
11211 xq2: &[i8],
11212) -> (f32, f32) {
11213 unsafe {
11215 use core::arch::x86_64::*;
11216 let lomask = _mm_set1_epi8(0x0F);
11217 let eight = _mm256_set1_epi8(8);
11218 let ones = _mm256_set1_epi16(1);
11219 let (mut acc1, mut acc2) = (0f32, 0f32);
11220 #[inline(always)]
11221 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
11222 unsafe {
11223 use core::arch::x86_64::*;
11224 let hi128 = _mm256_extracti128_si256::<1>(d);
11225 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11226 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11227 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11228 _mm_cvtsi128_si32(s32)
11229 }
11230 }
11231 for gi in 0..gpr {
11232 let g = g0 + gi;
11233 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11234 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11235 let lo = _mm_and_si128(b, lomask);
11236 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11237 let w = _mm256_sub_epi8(
11238 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11239 eight,
11240 );
11241 let aw = _mm256_abs_epi8(w);
11242 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11243 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11244 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
11245 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
11246 acc1 += hsum(d1) as f32 * s;
11247 acc2 += hsum(d2) as f32 * s;
11248 }
11249 (acc1, acc2)
11250 }
11251}
11252
11253#[cfg(target_arch = "x86_64")]
11255fn q8_range_avx2(
11256 q: &[u8],
11257 row_scale: &[f32],
11258 act: &SplitAct,
11259 cols: usize,
11260 out_addr: SendMut,
11261 start: usize,
11262 end: usize,
11263) {
11264 for o in start..end {
11265 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11266 unsafe { *out_addr.at(o) = v };
11268 }
11269}
11270
11271#[cfg(target_arch = "x86_64")]
11273#[allow(clippy::too_many_arguments)]
11274fn q8_range2_avx2(
11275 q: &[u8],
11276 row_scale: &[f32],
11277 a1: &SplitAct,
11278 a2: &SplitAct,
11279 cols: usize,
11280 p1: SendMut,
11281 p2: SendMut,
11282 start: usize,
11283 end: usize,
11284) {
11285 for o in start..end {
11286 let row = &q[o * cols..(o + 1) * cols];
11287 unsafe {
11289 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
11290 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
11291 }
11292 }
11293}
11294
11295#[cfg(target_arch = "aarch64")]
11306fn i8mm_enabled() -> bool {
11307 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11308 *ON.get_or_init(|| {
11309 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
11310 && std::arch::is_aarch64_feature_detected!("i8mm")
11311 })
11312}
11313
11314#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
11318fn sdot_enabled() -> bool {
11319 if FLOAT_ACTIVATIONS.get() {
11320 return false;
11321 }
11322 use std::sync::OnceLock;
11323 static ON: OnceLock<bool> = OnceLock::new();
11324 *ON.get_or_init(|| {
11325 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
11326 if !want {
11327 return false;
11328 }
11329
11330 #[cfg(target_arch = "aarch64")]
11331 {
11332 if std::arch::is_aarch64_feature_detected!("dotprod") {
11333 return true;
11334 }
11335 #[cfg(target_os = "android")]
11336 {
11337 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
11338 if cpuinfo.lines().any(|l| {
11339 (l.starts_with("Features") || l.starts_with("features"))
11340 && l.contains("asimddp")
11341 }) {
11342 return true;
11343 }
11344 }
11345 }
11346 false
11347 }
11348 #[cfg(not(target_arch = "aarch64"))]
11349 {
11350 false
11351 }
11352 })
11353}
11354
11355struct SplitAct {
11360 xq: Vec<i8>,
11361 sx: f32,
11362 outliers: Vec<(usize, f32)>,
11363 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
11366 xsum: i32,
11367}
11368
11369thread_local! {
11370 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
11373 const { std::cell::RefCell::new(Vec::new()) };
11374}
11375
11376impl Drop for SplitAct {
11377 fn drop(&mut self) {
11378 let buf = std::mem::take(&mut self.xq);
11379 if buf.capacity() > 0 {
11380 XQ_FREE.with(|f| {
11381 let mut f = f.borrow_mut();
11382 if f.len() < 16 {
11383 f.push(buf);
11384 }
11385 });
11386 }
11387 }
11388}
11389
11390thread_local! {
11391 static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
11398 static KROWS: UnsafeCell<[Vec<f32>; 2]> =
11403 const { UnsafeCell::new([Vec::new(), Vec::new()]) };
11404}
11405
11406#[inline]
11409fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
11410 KROW.with(|s| {
11411 let b = unsafe { &mut *s.get() };
11414 if b.len() < n {
11415 b.resize(n, 0.0);
11416 }
11417 f(&mut b[..n])
11418 })
11419}
11420
11421#[inline(always)]
11432fn q8_round(t: f32) -> i8 {
11433 let t = t.clamp(-127.0, 127.0);
11434 let i = t as i32;
11435 let f = t - i as f32;
11436 let r = if f >= 0.5 {
11437 i + 1
11438 } else if f <= -0.5 {
11439 i - 1
11440 } else {
11441 i
11442 };
11443 r as i8
11444}
11445
11446#[inline]
11447fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
11448 KROWS.with(|s| {
11449 let b = unsafe { &mut *s.get() };
11452 for row in b.iter_mut() {
11453 if row.len() < n {
11454 row.resize(n, 0.0);
11455 }
11456 }
11457 let (a, b) = b.split_at_mut(1);
11458 f(&mut a[0][..n], &mut b[0][..n])
11459 })
11460}
11461
11462#[inline]
11463fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
11464 if limit > 0.0 {
11465 up = up.clamp(-limit, limit);
11466 gate = gate.min(limit);
11467 }
11468 gate / (1.0 + (-gate).exp()) * up
11469}
11470
11471fn split_act(x: &[f32]) -> SplitAct {
11472 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11473 let n = x.len();
11474 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
11475 let thr = 8.0 * rms;
11476 let mut outliers: Vec<(usize, f32)> = Vec::new();
11480 let mut amax = 0f32;
11481 for (j, &v) in x.iter().enumerate() {
11482 let a = v.abs();
11483 if a > thr {
11484 outliers.push((j, v));
11485 } else if a > amax {
11486 amax = a;
11487 }
11488 }
11489 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11490 let inv = 1.0 / sx;
11491 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11492 xq.clear();
11493 xq.reserve(n);
11494 if outliers.is_empty() {
11495 xq.extend(
11496 x.iter()
11497 .map(|&v| q8_round(v * inv)),
11498 );
11499 } else {
11500 xq.extend(x.iter().map(|&v| {
11502 if v.abs() > thr {
11503 0
11504 } else {
11505 q8_round(v * inv)
11506 }
11507 }));
11508 }
11509 let xsum = xq.iter().map(|&v| v as i32).sum();
11510 SplitAct {
11511 xq,
11512 sx,
11513 outliers,
11514 xsum,
11515 }
11516}
11517
11518fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
11519 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11520 let n = x.len();
11521 let rms = (x
11522 .iter()
11523 .zip(col)
11524 .map(|(&a, &c)| {
11525 let v = a * c;
11526 (v * v) as f64
11527 })
11528 .sum::<f64>()
11529 / n.max(1) as f64)
11530 .sqrt() as f32;
11531 let thr = 8.0 * rms;
11532
11533 let mut outliers = Vec::new();
11534 let mut amax = 0f32;
11535 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
11536 let v = a * c;
11537 let s = v.abs();
11538 if s > thr {
11539 outliers.push((j, v));
11540 } else if s > amax {
11541 amax = s;
11542 }
11543 }
11544
11545 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11546 let inv = 1.0 / sx;
11547 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11548 xq.clear();
11549 xq.reserve(n);
11550 if outliers.is_empty() {
11551 xq.extend(
11552 x.iter()
11553 .zip(col)
11554 .map(|(&a, &c)| q8_round((a * c) * inv)),
11555 );
11556 } else {
11557 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
11558 let v = a * c;
11559 if v.abs() > thr {
11560 0
11561 } else {
11562 q8_round(v * inv)
11563 }
11564 }));
11565 }
11566 let xsum = xq.iter().map(|&v| v as i32).sum();
11567 SplitAct {
11568 xq,
11569 sx,
11570 outliers,
11571 xsum,
11572 }
11573}
11574
11575#[cfg(target_arch = "aarch64")]
11578#[target_feature(enable = "neon,dotprod")]
11579unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
11580 unsafe {
11582 use core::arch::aarch64::*;
11583 use core::arch::asm;
11584 let wp = w.as_ptr() as *const i8;
11585 let n = w.len();
11586 let (mut a0, mut a1, mut a2, mut a3) = (
11587 vdupq_n_s32(0),
11588 vdupq_n_s32(0),
11589 vdupq_n_s32(0),
11590 vdupq_n_s32(0),
11591 );
11592 let mut i = 0;
11593 while i + 64 <= n {
11594 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11595 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
11596 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
11597 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
11598 asm!(
11599 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
11600 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
11601 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
11602 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
11603 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11604 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
11605 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
11606 options(pure, nomem, nostack),
11607 );
11608 i += 64;
11609 }
11610 while i + 16 <= n {
11611 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11612 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
11613 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
11614 i += 16;
11615 }
11616 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
11617 while i < n {
11618 s += (*wp.add(i)) as i32 * xq[i] as i32;
11619 i += 1;
11620 }
11621 s
11622 }
11623}
11624
11625#[cfg(target_arch = "aarch64")]
11629#[target_feature(enable = "neon,dotprod")]
11630unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
11631 unsafe {
11633 use core::arch::aarch64::*;
11634 use core::arch::asm;
11635 let n = xq.len();
11636 let px = xq.as_ptr();
11637 let (p0, p1, p2, p3) = (
11638 w0.as_ptr() as *const i8,
11639 w1.as_ptr() as *const i8,
11640 w2.as_ptr() as *const i8,
11641 w3.as_ptr() as *const i8,
11642 );
11643 let (mut a0, mut a1, mut a2, mut a3) = (
11644 vdupq_n_s32(0),
11645 vdupq_n_s32(0),
11646 vdupq_n_s32(0),
11647 vdupq_n_s32(0),
11648 );
11649 let mut i = 0;
11650 while i + 16 <= n {
11651 let x = vld1q_s8(px.add(i));
11652 let v0 = vld1q_s8(p0.add(i));
11653 let v1 = vld1q_s8(p1.add(i));
11654 let v2 = vld1q_s8(p2.add(i));
11655 let v3 = vld1q_s8(p3.add(i));
11656 asm!(
11657 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11658 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11659 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11660 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11661 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11662 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11663 options(pure, nomem, nostack),
11664 );
11665 i += 16;
11666 }
11667 let mut r = [
11668 vaddvq_s32(a0),
11669 vaddvq_s32(a1),
11670 vaddvq_s32(a2),
11671 vaddvq_s32(a3),
11672 ];
11673 while i < n {
11674 let xi = *px.add(i) as i32;
11675 r[0] += (*p0.add(i)) as i32 * xi;
11676 r[1] += (*p1.add(i)) as i32 * xi;
11677 r[2] += (*p2.add(i)) as i32 * xi;
11678 r[3] += (*p3.add(i)) as i32 * xi;
11679 i += 1;
11680 }
11681 r
11682 }
11683}
11684
11685#[cfg(target_arch = "aarch64")]
11692#[target_feature(enable = "neon,dotprod")]
11693unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
11694 unsafe {
11697 use core::arch::aarch64::*;
11698 use core::arch::asm;
11699 let n = xq.len();
11700 let px = xq.as_ptr();
11701 let pg = g.as_ptr() as *const i8;
11702 let (mut a0, mut a1, mut a2, mut a3) = (
11703 vdupq_n_s32(0),
11704 vdupq_n_s32(0),
11705 vdupq_n_s32(0),
11706 vdupq_n_s32(0),
11707 );
11708 let mut i = 0;
11709 while i + 16 <= n {
11710 let x = vld1q_s8(px.add(i));
11711 let base = pg.add(4 * i);
11712 let v0 = vld1q_s8(base);
11713 let v1 = vld1q_s8(base.add(16));
11714 let v2 = vld1q_s8(base.add(32));
11715 let v3 = vld1q_s8(base.add(48));
11716 asm!(
11717 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11718 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11719 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11720 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11721 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11722 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11723 options(pure, nomem, nostack),
11724 );
11725 i += 16;
11726 }
11727 [
11728 vaddvq_s32(a0),
11729 vaddvq_s32(a1),
11730 vaddvq_s32(a2),
11731 vaddvq_s32(a3),
11732 ]
11733 }
11734}
11735
11736#[cfg(target_arch = "aarch64")]
11742fn q8_range_sdot(
11743 q: &[u8],
11744 rep: &[u8],
11745 row_scale: &[f32],
11746 act: &SplitAct,
11747 cols: usize,
11748 out_addr: SendMut,
11749 start: usize,
11750 end: usize,
11751) {
11752 let mut o = start;
11753 if !rep.is_empty() {
11756 while o < end && o % 4 != 0 {
11757 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11758 unsafe { *out_addr.at(o) = v };
11759 o += 1;
11760 }
11761 }
11762 while o + 4 <= end {
11763 let r = if rep.is_empty() {
11764 unsafe {
11765 dot_i8_sdot_4rows(
11766 &q[o * cols..(o + 1) * cols],
11767 &q[(o + 1) * cols..(o + 2) * cols],
11768 &q[(o + 2) * cols..(o + 3) * cols],
11769 &q[(o + 3) * cols..(o + 4) * cols],
11770 &act.xq,
11771 )
11772 }
11773 } else {
11774 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
11775 };
11776 for k in 0..4 {
11777 let mut acc = r[k] as f32 * act.sx;
11778 for &(j, xv) in &act.outliers {
11779 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
11780 }
11781 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
11783 }
11784 o += 4;
11785 }
11786 while o < end {
11787 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11788 unsafe { *out_addr.at(o) = v };
11789 o += 1;
11790 }
11791}
11792
11793#[cfg(target_arch = "aarch64")]
11796#[allow(clippy::too_many_arguments)]
11797fn q8_range2_sdot(
11798 q: &[u8],
11799 row_scale: &[f32],
11800 a1: &SplitAct,
11801 a2: &SplitAct,
11802 cols: usize,
11803 p1: SendMut,
11804 p2: SendMut,
11805 start: usize,
11806 end: usize,
11807) {
11808 for o in start..end {
11809 let row = &q[o * cols..(o + 1) * cols];
11810 unsafe {
11812 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
11813 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
11814 }
11815 }
11816}
11817
11818#[allow(clippy::too_many_arguments)]
11820fn q8_range2_f32(
11821 q: &[u8],
11822 row_scale: &[f32],
11823 x1: &[f32],
11824 x2: &[f32],
11825 cols: usize,
11826 p1: SendMut,
11827 p2: SendMut,
11828 start: usize,
11829 end: usize,
11830) {
11831 for o in start..end {
11832 let row = &q[o * cols..(o + 1) * cols];
11833 unsafe {
11835 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
11836 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
11837 }
11838 }
11839}
11840
11841fn q8_range_f32(
11843 q: &[u8],
11844 row_scale: &[f32],
11845 xs: &[f32],
11846 cols: usize,
11847 out_addr: SendMut,
11848 start: usize,
11849 end: usize,
11850) {
11851 for o in start..end {
11852 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
11853 unsafe { *out_addr.at(o) = v };
11855 }
11856}
11857
11858#[inline]
11862fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
11863 #[cfg(target_arch = "aarch64")]
11864 return row_dot_sdot(row, act);
11865 #[cfg(target_arch = "x86_64")]
11866 return row_dot_avx2(row, act);
11867 #[allow(unreachable_code)]
11868 q8_row_dot_scalar(row, act)
11869}
11870
11871#[allow(dead_code)]
11872fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
11873 let mut acc = 0i32;
11874 for (k, &b) in row.iter().enumerate() {
11875 acc += (b as i8) as i32 * act.xq[k] as i32;
11876 }
11877 let mut acc = acc as f32 * act.sx;
11878 for &(j, xv) in &act.outliers {
11879 acc += (row[j] as i8) as f32 * xv;
11880 }
11881 acc
11882}
11883
11884#[cfg(target_arch = "aarch64")]
11887#[inline]
11888fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
11889 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
11890 for &(j, xv) in &act.outliers {
11891 acc += (row[j] as i8) as f32 * xv;
11892 }
11893 acc
11894}
11895
11896#[cfg(target_arch = "aarch64")]
11904#[target_feature(enable = "neon,dotprod")]
11905unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11906 unsafe {
11909 use core::arch::aarch64::*;
11910 use core::arch::asm;
11911 let lomask = vdupq_n_u8(0x0F);
11912 let eight = vdupq_n_s8(8);
11913 let mut acc = 0f32;
11914 for gi in 0..gpr {
11915 let g = g0 + gi;
11916 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11917 let b = vld1q_u8(packed.as_ptr().add(g * 16));
11918 let lo = vandq_u8(b, lomask);
11919 let hi = vshrq_n_u8::<4>(b);
11920 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11921 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11922 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
11923 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
11924 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
11925 asm!(
11926 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
11927 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
11928 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11929 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
11930 options(pure, nomem, nostack),
11931 );
11932 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11933 }
11934 acc
11935 }
11936}
11937
11938#[cfg(target_arch = "aarch64")]
11943#[target_feature(enable = "neon,dotprod")]
11944unsafe fn dot_q4_row_sdot2(
11945 packed: &[u8],
11946 scales: &[u8],
11947 g0: usize,
11948 gpr: usize,
11949 xq1: &[i8],
11950 xq2: &[i8],
11951) -> (f32, f32) {
11952 unsafe {
11955 use core::arch::aarch64::*;
11956 use core::arch::asm;
11957 let lomask = vdupq_n_u8(0x0F);
11958 let eight = vdupq_n_s8(8);
11959 let (mut acc1, mut acc2) = (0f32, 0f32);
11960 for gi in 0..gpr {
11961 let g = g0 + gi;
11962 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11963 let b = vld1q_u8(packed.as_ptr().add(g * 16));
11964 let lo = vandq_u8(b, lomask);
11965 let hi = vshrq_n_u8::<4>(b);
11966 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11967 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11968 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
11969 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
11970 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
11971 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
11972 let (mut a0, mut a1, mut b0, mut b1) = (
11973 vdupq_n_s32(0),
11974 vdupq_n_s32(0),
11975 vdupq_n_s32(0),
11976 vdupq_n_s32(0),
11977 );
11978 asm!(
11979 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
11980 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
11981 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
11982 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
11983 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11984 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
11985 e0 = in(vreg) e0, e1 = in(vreg) e1,
11986 x10 = in(vreg) x10, x11 = in(vreg) x11,
11987 x20 = in(vreg) x20, x21 = in(vreg) x21,
11988 options(pure, nomem, nostack),
11989 );
11990 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11991 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
11992 }
11993 (acc1, acc2)
11994 }
11995}
11996
11997#[inline]
12002pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
12003 #[cfg(target_arch = "aarch64")]
12004 unsafe {
12005 return axpy_i8_f32_neon(acc, row, w);
12006 }
12007 #[cfg(target_arch = "x86_64")]
12008 if avx2_enabled() {
12009 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
12010 }
12011 #[allow(unreachable_code)]
12012 {
12013 for (a, &b) in acc.iter_mut().zip(row) {
12014 *a += w * b as f32;
12015 }
12016 }
12017}
12018
12019#[cfg(target_arch = "x86_64")]
12021#[target_feature(enable = "avx2,fma")]
12022unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
12023 unsafe {
12025 use core::arch::x86_64::*;
12026 let n = acc.len().min(row.len());
12027 let ap = acc.as_mut_ptr();
12028 let rp = row.as_ptr();
12029 let wv = _mm256_set1_ps(w);
12030 let mut j = 0usize;
12031 while j + 16 <= n {
12032 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
12033 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
12034 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
12035 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
12036 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
12037 _mm256_storeu_ps(ap.add(j), v0);
12038 _mm256_storeu_ps(ap.add(j + 8), v1);
12039 j += 16;
12040 }
12041 while j < n {
12042 *ap.add(j) += w * (*rp.add(j)) as f32;
12043 j += 1;
12044 }
12045 }
12046}
12047
12048#[cfg(target_arch = "aarch64")]
12049#[target_feature(enable = "neon")]
12050unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
12051 unsafe {
12053 use core::arch::aarch64::*;
12054 let n = acc.len().min(row.len());
12055 let ap = acc.as_mut_ptr();
12056 let rp = row.as_ptr();
12057 let wv = vdupq_n_f32(w);
12058 let mut j = 0usize;
12059 while j + 16 <= n {
12060 let rb = vld1q_s8(rp.add(j));
12061 let lo = vmovl_s8(vget_low_s8(rb));
12062 let hi = vmovl_s8(vget_high_s8(rb));
12063 for (off, half) in [(0, lo), (8, hi)] {
12064 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
12065 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
12066 let o = j + off;
12067 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
12068 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
12069 }
12070 j += 16;
12071 }
12072 while j < n {
12073 *ap.add(j) += w * (*rp.add(j)) as f32;
12074 j += 1;
12075 }
12076 }
12077}
12078
12079#[inline]
12082pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
12083 #[cfg(target_arch = "aarch64")]
12084 unsafe {
12085 return dot_i8_f32_neon(w, x);
12086 }
12087 #[cfg(target_arch = "x86_64")]
12088 if avx2_enabled() {
12089 return unsafe { dot_i8_f32_avx2(w, x) };
12090 }
12091 #[allow(unreachable_code)]
12092 {
12093 let mut sum = 0.0f32;
12094 for (j, &b) in w.iter().enumerate() {
12095 sum += (b as i8) as f32 * x[j];
12096 }
12097 sum
12098 }
12099}
12100
12101#[inline]
12105fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12106 #[cfg(target_arch = "aarch64")]
12107 unsafe {
12108 return dot_i8_col_f32_neon(w, x, col);
12109 }
12110 #[allow(unreachable_code)]
12111 {
12112 let mut sum = 0.0f32;
12113 for (j, &b) in w.iter().enumerate() {
12114 sum += (b as i8) as f32 * x[j] * col[j];
12115 }
12116 sum
12117 }
12118}
12119
12120#[cfg(target_arch = "aarch64")]
12121#[target_feature(enable = "neon")]
12122unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12123 unsafe {
12125 use core::arch::aarch64::*;
12126 let n = x.len();
12127 let wp = w.as_ptr() as *const i8;
12128 let xp = x.as_ptr();
12129 let cp = col.as_ptr();
12130 let (mut a0, mut a1, mut a2, mut a3) = (
12131 vdupq_n_f32(0.0),
12132 vdupq_n_f32(0.0),
12133 vdupq_n_f32(0.0),
12134 vdupq_n_f32(0.0),
12135 );
12136 let mut j = 0usize;
12137 while j + 16 <= n {
12138 let wb = vld1q_s8(wp.add(j));
12139 let lo = vmovl_s8(vget_low_s8(wb));
12140 let hi = vmovl_s8(vget_high_s8(wb));
12141 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12142 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12143 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12144 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12145 a0 = vfmaq_f32(
12146 a0,
12147 w0,
12148 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
12149 );
12150 a1 = vfmaq_f32(
12151 a1,
12152 w1,
12153 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
12154 );
12155 a2 = vfmaq_f32(
12156 a2,
12157 w2,
12158 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
12159 );
12160 a3 = vfmaq_f32(
12161 a3,
12162 w3,
12163 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
12164 );
12165 j += 16;
12166 }
12167 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12168 while j < n {
12169 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
12170 j += 1;
12171 }
12172 sum
12173 }
12174}
12175
12176#[cfg(target_arch = "aarch64")]
12177#[target_feature(enable = "neon")]
12178unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
12179 unsafe {
12181 use core::arch::aarch64::*;
12182 let n = x.len();
12183 let wp = w.as_ptr() as *const i8;
12184 let xp = x.as_ptr();
12185 let (mut a0, mut a1, mut a2, mut a3) = (
12186 vdupq_n_f32(0.0),
12187 vdupq_n_f32(0.0),
12188 vdupq_n_f32(0.0),
12189 vdupq_n_f32(0.0),
12190 );
12191 let mut j = 0usize;
12192 while j + 16 <= n {
12193 let wb = vld1q_s8(wp.add(j));
12194 let lo = vmovl_s8(vget_low_s8(wb));
12195 let hi = vmovl_s8(vget_high_s8(wb));
12196 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12197 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12198 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12199 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12200 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
12201 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
12202 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
12203 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
12204 j += 16;
12205 }
12206 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12207 while j < n {
12208 sum += (*wp.add(j)) as f32 * *xp.add(j);
12209 j += 1;
12210 }
12211 sum
12212 }
12213}
12214
12215#[allow(clippy::too_many_arguments)]
12216fn qmatvec(
12217 q: &[u8],
12218 rep: &[u8],
12219 row_scale: &[f32],
12220 x: &[f32],
12221 col_field: &[f32],
12222 dtype: TensorDtype,
12223 rows: usize,
12224 cols: usize,
12225 out: &mut [f32],
12226 pool: Option<&Pool>,
12227) {
12228 debug_assert_eq!(out.len(), rows);
12229 #[cfg(not(target_arch = "aarch64"))]
12230 let _ = rep;
12231
12232 #[cfg(target_arch = "aarch64")]
12233 if sdot_enabled() {
12234 let act = if dtype == TensorDtype::Q8_2f {
12235 split_act_q8_2f(x, col_field)
12236 } else {
12237 split_act(x)
12238 };
12239 let out_addr = SendMut(out.as_mut_ptr());
12240 let run_range = |start: usize, end: usize| {
12241 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
12242 };
12243 match pool {
12244 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12245 _ => run_range(0, rows),
12246 }
12247 return;
12248 }
12249 #[cfg(target_arch = "x86_64")]
12252 if avx2_a8w8_enabled() {
12253 let act = if dtype == TensorDtype::Q8_2f {
12254 split_act_q8_2f(x, col_field)
12255 } else {
12256 split_act(x)
12257 };
12258 let out_addr = SendMut(out.as_mut_ptr());
12259 let run_range = |start: usize, end: usize| {
12260 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
12261 };
12262 match pool {
12263 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12264 _ => run_range(0, rows),
12265 }
12266 return;
12267 }
12268
12269 prescale_with(x, col_field, dtype, 1, |xs| {
12270 let out_addr = SendMut(out.as_mut_ptr());
12271 let run_range = move |start: usize, end: usize| {
12272 for o in start..end {
12273 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12274 unsafe { *out_addr.at(o) = v };
12276 }
12277 };
12278 match pool {
12279 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12280 _ => run_range(0, rows),
12281 }
12282 });
12283}
12284
12285#[allow(clippy::too_many_arguments)]
12286fn qmatvec2(
12287 q: &[u8],
12288 row_scale: &[f32],
12289 x1: &[f32],
12290 x2: &[f32],
12291 col_field: &[f32],
12292 dtype: TensorDtype,
12293 rows: usize,
12294 cols: usize,
12295 o1: &mut [f32],
12296 o2: &mut [f32],
12297 pool: Option<&Pool>,
12298) {
12299 #[cfg(target_arch = "aarch64")]
12300 if sdot_enabled() {
12301 let a1s = if dtype == TensorDtype::Q8_2f {
12302 split_act_q8_2f(x1, col_field)
12303 } else {
12304 split_act(x1)
12305 };
12306 let a2s = if dtype == TensorDtype::Q8_2f {
12307 split_act_q8_2f(x2, col_field)
12308 } else {
12309 split_act(x2)
12310 };
12311 let p1 = SendMut(o1.as_mut_ptr());
12312 let p2 = SendMut(o2.as_mut_ptr());
12313 let run_range = |start: usize, end: usize| {
12314 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12315 };
12316 match pool {
12317 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12318 _ => run_range(0, rows),
12319 }
12320 return;
12321 }
12322 #[cfg(target_arch = "x86_64")]
12323 if avx2_a8w8_enabled() {
12324 let a1s = if dtype == TensorDtype::Q8_2f {
12325 split_act_q8_2f(x1, col_field)
12326 } else {
12327 split_act(x1)
12328 };
12329 let a2s = if dtype == TensorDtype::Q8_2f {
12330 split_act_q8_2f(x2, col_field)
12331 } else {
12332 split_act(x2)
12333 };
12334 let p1 = SendMut(o1.as_mut_ptr());
12335 let p2 = SendMut(o2.as_mut_ptr());
12336 let run_range = |start: usize, end: usize| {
12337 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12338 };
12339 match pool {
12340 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12341 _ => run_range(0, rows),
12342 }
12343 return;
12344 }
12345
12346 prescale_with(x1, col_field, dtype, 1, |x1s| {
12347 prescale_with(x2, col_field, dtype, 2, |x2s| {
12348 let p1 = SendMut(o1.as_mut_ptr());
12349 let p2 = SendMut(o2.as_mut_ptr());
12350 let run_range = move |start: usize, end: usize| {
12351 for o in start..end {
12352 let row = &q[o * cols..(o + 1) * cols];
12353 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
12354 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
12355 unsafe {
12357 *p1.at(o) = s1;
12358 *p2.at(o) = s2;
12359 }
12360 }
12361 };
12362 match pool {
12363 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12364 _ => run_range(0, rows),
12365 }
12366 });
12367 });
12368}
12369
12370#[derive(Clone, Copy)]
12371struct SendMut(*mut f32);
12372unsafe impl Send for SendMut {}
12373unsafe impl Sync for SendMut {}
12374
12375impl SendMut {
12376 #[inline]
12377 fn at(self, i: usize) -> *mut f32 {
12378 unsafe { self.0.add(i) }
12379 }
12380}
12381
12382#[cfg(test)]
12383mod tests {
12384 #[test]
12388 fn q8_round_is_round_clamp() {
12389 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
12390 let mut probe = vec![
12391 0.0f32,
12392 -0.0,
12393 f32::NAN,
12394 f32::INFINITY,
12395 f32::NEG_INFINITY,
12396 f32::MAX,
12397 f32::MIN,
12398 1e30,
12399 -1e30,
12400 f32::MIN_POSITIVE,
12401 -f32::MIN_POSITIVE,
12402 ];
12403 for k in -300i32..=300 {
12404 let h = k as f32 * 0.5;
12405 let up = f32::from_bits(h.to_bits() + 1);
12406 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
12407 for t in [h, up, down] {
12408 probe.push(t);
12409 probe.push(-t);
12410 }
12411 }
12412 let mut t = -140.0f32;
12413 while t < 140.0 {
12414 probe.push(t);
12415 t += 0.000_731;
12416 }
12417 for t in probe {
12418 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
12419 }
12420 }
12421
12422 use super::*;
12423
12424 #[test]
12428 fn f32_matmat_equals_matvec_bitwise() {
12429 let (rows, cols, b) = (5usize, 37usize, 7usize);
12430 let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
12431 let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
12432 let t = QTensor::from_f32(w, rows, cols);
12433 let mut out = vec![0f32; b * rows];
12434 t.matmat(&xs, b, &mut out, None);
12435 for bi in 0..b {
12436 let mut one = vec![0f32; rows];
12437 t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
12438 for o in 0..rows {
12439 assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
12440 }
12441 }
12442 }
12443
12444 #[test]
12445 fn q2tp_i8_dot_matches_exact_on_grid() {
12446 let (rows, cols) = (5, 64);
12450 let gpr = cols / GROUP_SIZE;
12451 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
12455 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
12456 .collect();
12457 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
12458 let x: Vec<f32> = (0..cols)
12459 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
12460 .collect();
12461 let act = split_act(&x);
12462 assert!(
12463 act.outliers.is_empty(),
12464 "on-grid input must have no outliers"
12465 );
12466 let gsum = q1_group_sums(&act.xq, gpr);
12467 for r in 0..rows {
12468 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
12469 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
12470 assert!(
12471 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
12472 "row {r}: exact {exact} vs i8 {fast}"
12473 );
12474 }
12475 }
12476
12477 #[cfg(target_arch = "x86_64")]
12478 #[test]
12479 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
12480 if !std::arch::is_x86_feature_detected!("avx2") {
12485 return;
12486 }
12487 let mut seed = 0x9e3779b9u32;
12488 let mut next = || {
12489 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
12490 seed
12491 };
12492 for _ in 0..20_000 {
12493 let mut ch = [0u8; Q2TP_CHUNK];
12494 let mut x = [0i8; GROUP_SIZE];
12495 for b in &mut ch {
12496 *b = next() as u8;
12497 }
12498 for v in &mut x {
12499 *v = (next() >> 24) as i8;
12500 }
12501 let mut reference = 0i32;
12502 for (k, &b) in ch.iter().enumerate() {
12503 reference += (b & 3) as i32 * x[k * 4] as i32;
12504 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
12505 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
12506 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
12507 }
12508 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
12511 assert_eq!(got, reference, "packed q2 lane mismatch");
12512 }
12513 }
12514
12515 #[test]
12516 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
12517 let (rows, cols) = (1usize, GROUP_SIZE);
12518 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
12519 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
12523 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
12524 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
12526 let mut raw = vec![0.0f32; rows];
12527 let mut affine = vec![0.0f32; rows];
12528 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
12529 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
12530 assert_eq!(raw, vec![-24.0]);
12531 assert_eq!(affine, vec![-8.0]);
12532 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
12533 }
12534
12535 #[test]
12536 fn q8_row_dot_fast_matches_scalar() {
12537 let cols = 96;
12540 let row: Vec<u8> = (0..cols)
12541 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
12542 .collect();
12543 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
12544 let act = split_act(&x);
12545 let fast = q8_row_dot(&row, &act);
12546 let scalar = q8_row_dot_scalar(&row, &act);
12547 assert!(
12548 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
12549 "fast {fast} vs scalar {scalar}"
12550 );
12551 }
12552
12553 #[test]
12554 fn f32_matvec_matches_matvec_rows_bitexact() {
12555 let (rows, cols) = (300, 40);
12556 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
12557 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
12558 let qt = QTensor::from_f32(w.clone(), rows, cols);
12559
12560 let mut a = vec![0.0f32; rows];
12561 matvec_rows(None, &w, &x, &mut a);
12562 let mut b = vec![0.0f32; rows];
12563 qt.matvec(&x, &mut b, None);
12564 assert_eq!(a, b);
12565 }
12566
12567 #[test]
12568 fn sdot_kernel_exact_on_grid() {
12569 eprintln!("sdot_enabled = {}", sdot_enabled());
12574 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
12576 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
12577 .collect();
12578 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
12579 let x: Vec<f32> = (0..cols)
12580 .map(|i| match i % 3 {
12581 0 => 1.0,
12582 1 => -1.0,
12583 _ => 0.0,
12584 })
12585 .collect();
12586 let mut a = vec![0.0f32; rows];
12587 qmatvec(
12588 &w,
12589 &[],
12590 &scales,
12591 &x,
12592 &[],
12593 TensorDtype::Q8Row,
12594 rows,
12595 cols,
12596 &mut a,
12597 None,
12598 );
12599 for o in 0..rows {
12600 let mut acc = 0.0f32;
12601 for j in 0..cols {
12602 acc += (w[o * cols + j] as i8) as f32 * x[j];
12603 }
12604 let expect = acc * scales[o];
12605 assert!(
12606 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12607 "row {o}: {} vs {expect}",
12608 a[o]
12609 );
12610 }
12611 }
12612
12613 #[test]
12614 fn q1_tbl_fast_path_matches_reference() {
12615 let (rows, cols) = (5, 256);
12620 let gpr = cols / GROUP_SIZE;
12621 let mut bytes = Vec::new();
12622 for t in 0..rows * gpr {
12623 let s = 0.007 + (t % 11) as f32 * 0.004;
12624 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12625 for j in 0..4 {
12626 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
12627 }
12628 }
12629 let x: Vec<f32> = (0..cols)
12630 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
12631 .collect();
12632 let mut w = vec![0.0f32; rows * cols];
12633 cortiq_core::quant::dequant_q1(&bytes, &mut w);
12634 let mut got = vec![0.0f32; rows];
12635 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12636 for o in 0..rows {
12637 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12638 assert!(
12639 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12640 "row {o}: {} vs {expect}",
12641 got[o]
12642 );
12643 }
12644 let b = 5usize;
12647 let mut xs_all = Vec::new();
12648 for bi in 0..b {
12649 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
12650 }
12651 let mut mm = vec![0.0f32; b * rows];
12652 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
12653 for bi in 0..b {
12654 let mut single = vec![0.0f32; rows];
12655 q1_matvec(
12656 &bytes,
12657 &xs_all[bi * cols..(bi + 1) * cols],
12658 rows,
12659 cols,
12660 &mut single,
12661 None,
12662 );
12663 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
12664 }
12665 }
12666
12667 #[test]
12668 fn q1_kernels_match_exact_reference() {
12669 let (rows, cols) = (7, 96);
12671 let gpr = cols / GROUP_SIZE;
12672 let mut bytes = Vec::new();
12673 for t in 0..rows * gpr {
12674 let s = 0.01 + (t % 13) as f32 * 0.003;
12675 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12676 for j in 0..4 {
12677 bytes.push(((t * 31 + j * 97) % 251) as u8);
12678 }
12679 }
12680 let x: Vec<f32> = (0..cols)
12682 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
12683 .collect();
12684 let mut w = vec![0.0f32; rows * cols];
12686 cortiq_core::quant::dequant_q1(&bytes, &mut w);
12687 let mut expect = vec![0.0f32; rows];
12688 for o in 0..rows {
12689 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12690 }
12691 let mut got = vec![0.0f32; rows];
12692 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12693 for o in 0..rows {
12694 assert!(
12695 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
12696 "row {o}: {} vs {}",
12697 got[o],
12698 expect[o]
12699 );
12700 }
12701 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
12703 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
12704 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12705 assert_eq!(a1, got);
12706 let mut xs = x.clone();
12707 xs.extend_from_slice(&x2);
12708 let mut mm = vec![0.0f32; 2 * rows];
12709 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
12710 assert_eq!(&mm[..rows], got.as_slice());
12711 assert_eq!(&mm[rows..], a2.as_slice());
12712 }
12713
12714 fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
12717 let tiles = rows * cols / GROUP_SIZE;
12718 let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
12719 for t in 0..tiles {
12720 let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
12721 q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12722 let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
12723 q1.extend_from_slice(&bits.to_le_bytes());
12724 }
12725 let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
12726 assert_eq!(rp.differing, 0, "uniform scales repack exactly");
12727 (q1, rp.data)
12728 }
12729
12730 fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
12733 let mut x: Vec<f32> = (0..cols)
12734 .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
12735 .collect();
12736 x[3 % cols] = 41.0;
12737 x[(cols / 2 + 5) % cols] = -37.5;
12738 x
12739 }
12740
12741 #[test]
12748 fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
12749 for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
12750 let gpr = cols / GROUP_SIZE;
12751 let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
12752 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12753 let x = q1g_test_x(cols, rows);
12754 let act = split_act(&x);
12755 assert!(!act.outliers.is_empty(), "test wants the outlier path");
12756 let gsum = q1_group_sums(&act.xq, gpr);
12757 for r in 0..rows {
12758 let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
12759 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
12760 let (a, b) = (
12761 wa.dot_i8(r, gpr, &act.xq, &gsum),
12762 wb.dot_i8(r, gpr, &act.xq, &gsum),
12763 );
12764 assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
12765 for j in [0, 31, 32, 127, cols - 1] {
12766 let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
12767 assert_eq!(
12768 (a.0.to_bits(), a.1.to_bits()),
12769 (b.0.to_bits(), b.1.to_bits())
12770 );
12771 }
12772 #[cfg(target_arch = "x86_64")]
12773 if avx2_enabled() {
12774 let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
12775 unsafe {
12777 let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
12778 let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
12779 assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
12780 }
12781 if std::is_x86_feature_detected!("avx512vnni")
12782 && std::is_x86_feature_detected!("avx512bw")
12783 && std::is_x86_feature_detected!("avx512vl")
12784 {
12785 unsafe {
12787 let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
12788 let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
12789 assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
12790 }
12791 }
12792 }
12793 }
12794 let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
12796 bin_matvec(wa, &x, rows, cols, &mut ga, None);
12797 bin_matvec(wb, &x, rows, cols, &mut gb, None);
12798 assert_eq!(ga, gb, "{rows}x{cols} matvec");
12799 let x2 = q1g_test_x(cols, rows + 1);
12800 let (mut a1, mut a2, mut b1, mut b2) = (
12801 vec![0f32; rows],
12802 vec![0f32; rows],
12803 vec![0f32; rows],
12804 vec![0f32; rows],
12805 );
12806 bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12807 bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
12808 assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
12809 let b = 5usize;
12810 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
12811 let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12812 bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
12813 bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
12814 assert_eq!(ma, mb, "{rows}x{cols} matmat");
12815 let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
12818 .map(|k| {
12819 let a = split_act(&xs[k * cols..(k + 1) * cols]);
12820 let g = q1_group_sums(&a.xq, gpr);
12821 (a, g)
12822 })
12823 .collect();
12824 let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
12825 let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
12826 for r in 0..rows {
12827 #[cfg(target_arch = "x86_64")]
12828 if avx2_enabled() {
12829 let (a, b) = unsafe {
12831 (
12832 wa.dot_1x4_x86(r, gpr, xq4, gs4),
12833 wb.dot_1x4_x86(r, gpr, xq4, gs4),
12834 )
12835 };
12836 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12837 }
12838 #[cfg(target_arch = "aarch64")]
12839 if sdot_enabled() {
12840 let (a, b) = unsafe {
12842 (
12843 wa.dot_1x4_sdot(r, gpr, xq4, gs4),
12844 wb.dot_1x4_sdot(r, gpr, xq4, gs4),
12845 )
12846 };
12847 assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12848 }
12849 let _ = (r, &xq4, &gs4);
12850 }
12851 }
12852 }
12853
12854 #[test]
12860 #[ignore]
12861 fn q1g_cpu_matvec_speed() {
12862 use crate::pool::Pool;
12863 let threads = std::env::var("CMF_THREADS")
12864 .ok()
12865 .and_then(|v| v.parse().ok())
12866 .unwrap_or(8usize);
12867 let pool = Pool::new(threads);
12868 for &(name, rows, cols, reps) in &[
12869 ("gate/up_proj", 6144usize, 2048usize, 41usize),
12870 ("down_proj", 2048, 6144, 41),
12871 ("lm_head", 151_669, 2048, 9),
12872 ] {
12873 let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
12874 let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12875 let x = q1g_test_x(cols, 1);
12876 let mut out = vec![0f32; rows];
12877 let mut time = |w: BinW| {
12878 let t0 = std::time::Instant::now();
12879 bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
12880 t0.elapsed().as_secs_f64()
12881 };
12882 let _ = (time(wa), time(wb));
12883 let (mut ta, mut tb) = (Vec::new(), Vec::new());
12884 for r in 0..reps {
12885 if r % 2 == 0 {
12886 ta.push(time(wa));
12887 tb.push(time(wb));
12888 } else {
12889 tb.push(time(wb));
12890 ta.push(time(wa));
12891 }
12892 }
12893 ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
12894 tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
12895 let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
12896 eprintln!(
12897 "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
12898 a8w8_enabled(),
12899 ma * 1e6,
12900 mb * 1e6,
12901 mb / ma
12902 );
12903 }
12904 }
12905
12906 #[test]
12911 fn q1g_mapped_tensor_paths_equal_q1() {
12912 use crate::pool::Pool;
12913 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
12914 let (rows, cols) = (300usize, 256usize);
12915 let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
12916 let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
12917 let header: CmfHeader = serde_json::from_value(serde_json::json!({
12918 "format": "cmf", "version": cortiq_core::CMF_VERSION,
12919 "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
12920 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
12921 "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
12922 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
12923 "max_position_embeddings": 8},
12924 "quant_type": "VBIT"
12925 }))
12926 .unwrap();
12927 let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
12928 name: name.into(),
12929 dtype,
12930 shape: vec![rows, cols],
12931 data: data.clone(),
12932 };
12933 let specs = [
12934 spec("g1", TensorDtype::Q1, &gq1),
12935 spec("u1", TensorDtype::Q1, &uq1),
12936 spec("gg", TensorDtype::Q1G, &gq1g),
12937 spec("ug", TensorDtype::Q1G, &uq1g),
12938 ];
12939 let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
12940 std::fs::create_dir_all(&dir).unwrap();
12941 let path = dir.join("m.cmf");
12942 CmfModel::write(&path, &header, &specs, None, None).unwrap();
12943 let model = Arc::new(CmfModel::open(&path).unwrap());
12944 assert!(model.verify().is_empty(), "{:?}", model.verify());
12945 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
12946 let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
12947 assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
12948 assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
12949 let pool = Pool::new(3);
12950 let x = q1g_test_x(cols, 4);
12951 let run = |a: &QTensor| {
12952 let mut o = vec![0f32; rows];
12953 a.matvec(&x, &mut o, Some(&pool));
12954 o
12955 };
12956 assert_eq!(run(&g1), run(&gg), "matvec");
12957 let many = |a: &QTensor, b: &QTensor| {
12958 let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
12959 QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
12960 (oa, ob)
12961 };
12962 let want = many(&g1, &u1);
12963 assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
12964 assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
12965 assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
12966 let fused = |g: &QTensor, u: &QTensor| {
12967 let mut o = vec![0f32; rows];
12968 let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
12969 (ok, o)
12970 };
12971 let (ok_a, fa) = fused(&g1, &u1);
12972 let (ok_b, fb) = fused(&gg, &ug);
12973 assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
12974 assert_eq!(fa, fb, "fused gate/up");
12975 let b = 6usize;
12976 let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
12977 let mm = |a: &QTensor| {
12978 let mut o = vec![0f32; b * rows];
12979 a.matmat(&xs, b, &mut o, Some(&pool));
12980 o
12981 };
12982 assert_eq!(mm(&g1), mm(&gg), "matmat");
12983 let x2 = q1g_test_x(cols, 77);
12984 let pair = |a: &QTensor| {
12985 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
12986 a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
12987 (o1, o2)
12988 };
12989 assert_eq!(pair(&g1), pair(&gg), "matvec2");
12990 for r in [0, 1, rows / 2, rows - 1] {
12991 let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
12992 g1.row_f32(r, &mut ra);
12993 gg.row_f32(r, &mut rb);
12994 assert_eq!(ra, rb, "embedding row {r}");
12995 }
12996 drop((g1, u1, gg, ug, model));
12997 let _ = std::fs::remove_dir_all(&dir);
12998 }
12999
13000 #[test]
13007 fn blr_mapped_tensor_paths_equal_dense() {
13008 use crate::pool::Pool;
13009 use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13010 let (rows, cols, r) = (96usize, 128usize, 16usize);
13011 let mut rng = crate::blr::tests::Rng(41);
13012 let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
13013 let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
13014 let mut q8: Vec<u8> = (0..rows * cols)
13016 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13017 .collect();
13018 for o in 0..rows {
13019 q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
13020 }
13021 let header: CmfHeader = serde_json::from_value(serde_json::json!({
13022 "format": "cmf", "version": cortiq_core::CMF_VERSION,
13023 "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
13024 "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13025 "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
13026 "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13027 "max_position_embeddings": 8},
13028 "quant_type": "Q8_2F"
13029 }))
13030 .unwrap();
13031 let specs = [
13032 TensorSpec {
13033 name: "g".into(),
13034 dtype: TensorDtype::BinaryLowRank,
13035 shape: vec![rows, cols],
13036 data: gate.clone(),
13037 },
13038 TensorSpec {
13039 name: "u".into(),
13040 dtype: TensorDtype::BinaryLowRank,
13041 shape: vec![rows, cols],
13042 data: up.clone(),
13043 },
13044 TensorSpec {
13045 name: "q".into(),
13046 dtype: TensorDtype::Q8Row,
13047 shape: vec![rows, cols],
13048 data: q8,
13049 },
13050 ];
13051 let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
13052 std::fs::create_dir_all(&dir).unwrap();
13053 let path = dir.join("m.cmf");
13054 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13055 let model = Arc::new(CmfModel::open(&path).unwrap());
13056 assert!(model.verify().is_empty(), "{:?}", model.verify());
13057 let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13058 let (g, u, q) = (t("g"), t("u"), t("q"));
13059 assert!(g.is_blr() && u.is_blr() && !q.is_blr());
13060 assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
13061 assert!(g.as_f32().is_none() && g.model_idx().is_none());
13062 assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
13063 assert!(!g.sparse_col_ok());
13064 assert_eq!((g.rows(), g.cols()), (rows, cols));
13065 let deq = |bytes: &[u8]| {
13066 let mut w = vec![0f32; rows * cols];
13067 cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
13068 w
13069 };
13070 let (wg, wu) = (deq(&gate), deq(&up));
13071 let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
13072 (0..rows)
13073 .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
13074 .collect()
13075 };
13076 let close = |a: &[f32], b: &[f32], what: &str| {
13077 let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
13078 let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
13079 assert!(e / s < 1e-5, "{what}: rel {}", e / s);
13080 };
13081 let pool = Pool::new(3);
13082 let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13083 let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13084 let mut o = vec![0f32; rows];
13085 g.matvec(&x, &mut o, Some(&pool));
13086 close(&o, &mv(&wg, &x), "matvec");
13087 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13088 g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13089 close(&o1, &mv(&wg, &x), "matvec2 a");
13090 close(&o2, &mv(&wg, &x2), "matvec2 b");
13091 for b in [2usize, 5] {
13092 let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
13093 let mut om = vec![0f32; b * rows];
13094 g.matmat(&xs, b, &mut om, Some(&pool));
13095 for t in 0..b {
13096 close(
13097 &om[t * rows..(t + 1) * rows],
13098 &mv(&wg, &xs[t * cols..(t + 1) * cols]),
13099 "matmat",
13100 );
13101 }
13102 }
13103 let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
13106 QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
13107 close(&a, &mv(&wg, &x), "matvec_many blr gate");
13108 close(&c, &mv(&wu, &x), "matvec_many blr up");
13109 let mut qo = vec![0f32; rows];
13110 q.matvec(&x, &mut qo, Some(&pool));
13111 assert_eq!(bq, qo, "matvec_many q8_row sibling");
13112 let mut f = vec![0f32; rows];
13113 assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
13114 let mut row = vec![0f32; cols];
13115 g.row_f32(rows / 3, &mut row);
13116 close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
13117 let mut scratch = vec![0f32; cols];
13118 let d = g.row_dot(7, &x, &mut scratch);
13119 let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
13120 assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
13121 let mut col = vec![0f32; rows];
13122 g.add_col_scaled(cols / 2, 1.5, &mut col);
13123 let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
13124 close(&col, &want, "add_col_scaled");
13125 drop((g, u, q, model));
13126 let _ = std::fs::remove_dir_all(&dir);
13127 }
13128
13129 #[test]
13130 fn repack_is_bit_identical() {
13131 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
13137 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13138 .collect();
13139 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
13140 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
13141 let rep = q8_repack_layout(&w, rows, cols);
13142 for g in 0..rows / 4 {
13144 for c in 0..cols / 16 {
13145 for lane in 0..4 {
13146 assert_eq!(
13147 &rep[g * 4 * cols + c * 64 + lane * 16
13148 ..g * 4 * cols + c * 64 + lane * 16 + 16],
13149 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
13150 );
13151 }
13152 }
13153 }
13154 let mut a = vec![0.0f32; rows];
13155 qmatvec(
13156 &w,
13157 &[],
13158 &scales,
13159 &x,
13160 &[],
13161 TensorDtype::Q8Row,
13162 rows,
13163 cols,
13164 &mut a,
13165 None,
13166 );
13167 let mut b = vec![0.0f32; rows];
13168 qmatvec(
13169 &w,
13170 &rep,
13171 &scales,
13172 &x,
13173 &[],
13174 TensorDtype::Q8Row,
13175 rows,
13176 cols,
13177 &mut b,
13178 None,
13179 );
13180 assert_eq!(a, b, "full-range repack output diverged");
13181
13182 #[cfg(target_arch = "aarch64")]
13183 if sdot_enabled() {
13184 let act = split_act(&x);
13186 let mut c1 = vec![0.0f32; rows];
13187 let mut c2 = vec![0.0f32; rows];
13188 q8_range_sdot(
13189 &w,
13190 &[],
13191 &scales,
13192 &act,
13193 cols,
13194 SendMut(c1.as_mut_ptr()),
13195 3,
13196 rows - 2,
13197 );
13198 q8_range_sdot(
13199 &w,
13200 &rep,
13201 &scales,
13202 &act,
13203 cols,
13204 SendMut(c2.as_mut_ptr()),
13205 3,
13206 rows - 2,
13207 );
13208 assert_eq!(c1, c2, "unaligned-range repack output diverged");
13209 }
13210 }
13211
13212 #[test]
13213 fn sdot_a8w8_noise_is_bounded() {
13214 let (rows, cols) = (16, 512);
13218 let w: Vec<u8> = (0..rows * cols)
13219 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13220 .collect();
13221 let scales = vec![0.01f32; rows];
13222 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
13223 let mut a = vec![0.0f32; rows];
13224 qmatvec(
13225 &w,
13226 &[],
13227 &scales,
13228 &x,
13229 &[],
13230 TensorDtype::Q8Row,
13231 rows,
13232 cols,
13233 &mut a,
13234 None,
13235 );
13236 let (mut num, mut den) = (0f64, 0f64);
13237 for o in 0..rows {
13238 let mut acc = 0.0f32;
13239 for j in 0..cols {
13240 acc += (w[o * cols + j] as i8) as f32 * x[j];
13241 }
13242 let expect = acc * scales[o];
13243 num += ((a[o] - expect) as f64).powi(2);
13244 den += (expect as f64).powi(2);
13245 }
13246 let rel = (num / den.max(1e-12)).sqrt();
13247 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
13248 }
13249
13250 #[test]
13251 fn i8_dot_neon_matches_scalar() {
13252 let n = 100;
13253 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
13254 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
13255 let mut scalar = 0.0f32;
13256 for j in 0..n {
13257 scalar += (w[j] as i8) as f32 * x[j];
13258 }
13259 let fast = dot_i8_f32(&w, &x);
13260 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
13261 }
13262
13263 #[test]
13265 fn vbitmatvec_matches_full_dequant() {
13266 let (rows, cols) = (6, 64);
13267 let ng = cols / GROUP_SIZE;
13268 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13270 let mut bytes = bits.clone();
13271 for g in 0..rows * ng {
13272 let s = 0.02 + 0.001 * g as f32;
13273 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13274 }
13275 for r in 0..rows {
13276 let b = bits[r] as usize;
13277 let (mut acc, mut nb) = (0u64, 0usize);
13278 let mut rowbytes = Vec::new();
13279 for i in 0..cols {
13280 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13281 acc = (acc << b) | v;
13282 nb += b;
13283 while nb >= 8 {
13284 nb -= 8;
13285 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13286 }
13287 }
13288 if nb > 0 {
13289 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13290 }
13291 bytes.extend_from_slice(&rowbytes);
13292 }
13293 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13294
13295 let mut reference = vec![0f32; rows * cols];
13296 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
13297 let mut expect = vec![0f32; rows];
13298 for r in 0..rows {
13299 expect[r] = reference[r * cols..(r + 1) * cols]
13300 .iter()
13301 .zip(&x)
13302 .map(|(w, xv)| w * xv)
13303 .sum();
13304 }
13305 let mut got = vec![0f32; rows];
13306 let offsets = vbit_row_offsets(&bytes, rows, cols);
13307 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
13308 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13312 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13313 for r in 0..rows {
13314 assert!(
13315 (got[r] - expect[r]).abs() < tol * scale,
13316 "row {r}: {} vs {}",
13317 got[r],
13318 expect[r]
13319 );
13320 }
13321 }
13322
13323 #[test]
13328 #[cfg(target_arch = "x86_64")]
13329 fn vbit_matmat_blocked_matches_per_row() {
13330 let (rows, cols, b) = (64usize, 128usize, 9usize);
13331 let ng = cols / GROUP_SIZE;
13332 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
13333 let mut bytes = bits.clone();
13334 for g in 0..rows * ng {
13335 let sc = 0.02 + 0.0005 * g as f32;
13336 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13337 }
13338 for r in 0..rows {
13339 let bw = bits[r] as usize;
13340 let (mut acc, mut nb) = (0u64, 0usize);
13341 let mut rowbytes = Vec::new();
13342 for i in 0..cols {
13343 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13344 acc = (acc << bw) | v;
13345 nb += bw;
13346 while nb >= 8 {
13347 nb -= 8;
13348 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13349 }
13350 }
13351 if nb > 0 {
13352 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13353 }
13354 bytes.extend_from_slice(&rowbytes);
13355 }
13356 let x: Vec<f32> = (0..b * cols)
13357 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13358 .collect();
13359 let offsets = vbit_row_offsets(&bytes, rows, cols);
13360 let mut y_a = vec![0f32; b * rows];
13361 let mut y_b = vec![0f32; b * rows];
13362 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13363 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
13364 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13365 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
13366 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13367 let max_d = y_a
13368 .iter()
13369 .zip(&y_b)
13370 .map(|(p, q)| (p - q).abs())
13371 .fold(0.0f32, f32::max);
13372 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
13373 }
13374
13375 #[test]
13383 fn q4t_matmat_blocked_matches_per_row() {
13384 let (rows, cols, b) = (16usize, 64usize, 9usize);
13385 let gpr = cols / GROUP_SIZE;
13386 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13387 for r in 0..rows {
13388 for g in 0..gpr {
13389 let t = (r * gpr + g) * Q4_TILE;
13390 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
13391 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13392 for k in 0..16 {
13393 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13394 }
13395 }
13396 }
13397 let x: Vec<f32> = (0..b * cols)
13398 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13399 .collect();
13400 let mut y_blk = vec![0f32; b * rows];
13401 let mut y_row = vec![0f32; b * rows];
13402 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13403 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
13404 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13405 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
13406 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13407 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
13408 }
13409
13410 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
13417 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
13418 let gpr = cols / GROUP_SIZE;
13419 let stride = q4tp_code_stride(gpr);
13420 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
13421 let mut b = vec![0u8; codes_off + rows * stride];
13422 for r in 0..rows {
13423 for g in 0..gpr {
13424 let t = (r * gpr + g) * Q4TP_NIB;
13425 for k in 0..16 {
13426 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13427 }
13428 }
13429 let lo = -6.0 - 0.03 * (r % 17) as f32;
13430 let step = 0.01 + 0.004 * (r % 11) as f32;
13431 let p = params_off + r * 4;
13432 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
13433 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
13434 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
13435 for g in 0..gpr {
13436 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
13437 }
13438 }
13439 b
13440 }
13441
13442 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
13446 let gpr = cols / GROUP_SIZE;
13447 let v = Q4tpView::new(bytes, rows, cols);
13448 let mut out = vec![0u8; rows * gpr * Q4_TILE];
13449 let mut sc = vec![0f32; gpr];
13450 for r in 0..rows {
13451 v.scales_into(r, gpr, &mut sc);
13452 for g in 0..gpr {
13453 let t = (r * gpr + g) * Q4_TILE;
13454 let s = sc[g];
13455 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13456 let src = (r * gpr + g) * Q4TP_NIB;
13457 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
13458 }
13459 }
13460 out
13461 }
13462
13463 #[test]
13469 fn q4tp_exact_path_matches_dequant_reference() {
13470 let (rows, cols) = (256usize, 512usize);
13471 let gpr = cols / GROUP_SIZE;
13472 let bytes = synth_q4tp(rows, cols);
13473 let mut w = vec![0f32; rows * cols];
13474 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13475
13476 let x: Vec<f32> = (0..cols)
13477 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13478 .collect();
13479 let v = Q4tpView::new(&bytes, rows, cols);
13480 let mut sc = vec![0f32; gpr];
13481 for r in 0..rows {
13482 v.scales_into(r, gpr, &mut sc);
13483 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
13484 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
13485 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13489 assert!(
13490 (got - want).abs() <= 1e-5 * mag,
13491 "row {r}: kernel {got} vs dequant {want}"
13492 );
13493 }
13494 }
13495
13496 #[test]
13502 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
13503 let (rows, cols) = (256usize, 512usize);
13504 let bytes = synth_q4tp(rows, cols);
13505 let twin = q4tp_as_q4t(&bytes, rows, cols);
13506 let x: Vec<f32> = (0..cols)
13507 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13508 .collect();
13509
13510 let mut got = vec![0f32; rows];
13511 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
13512 let mut want = vec![0f32; rows];
13513 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
13514
13515 let mut w = vec![0f32; rows * cols];
13518 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13519 for r in 0..rows {
13520 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13521 assert!(
13522 (got[r] - want[r]).abs() <= 1e-3 * mag,
13523 "row {r}: q4tp {} vs q4t {}",
13524 got[r],
13525 want[r]
13526 );
13527 }
13528 }
13529
13530 #[test]
13535 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
13536 let (rows, cols, b) = (256usize, 512usize, 5usize);
13537 let bytes = synth_q4tp(rows, cols);
13538 let twin = q4tp_as_q4t(&bytes, rows, cols);
13539 let xs: Vec<f32> = (0..b * cols)
13540 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13541 .collect();
13542
13543 let mut got = vec![0f32; b * rows];
13544 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
13545 let mut want = vec![0f32; b * rows];
13546 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
13547
13548 let mut w = vec![0f32; rows * cols];
13549 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13550 for t in 0..b {
13551 for r in 0..rows {
13552 let mag: f32 = (0..cols)
13553 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
13554 .sum();
13555 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
13556 assert!(
13557 (g - wa).abs() <= 1e-3 * mag,
13558 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
13559 );
13560 }
13561 }
13562 }
13563
13564 #[test]
13565 fn q4tp_matvec2_matches_the_single_stream_kernel() {
13566 let (rows, cols) = (128usize, 256usize);
13567 let gpr = cols / GROUP_SIZE;
13568 let bytes = synth_q4tp(rows, cols);
13569 let xs: Vec<f32> = (0..2 * cols)
13570 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13571 .collect();
13572
13573 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13574 q4tp_matvec2(
13575 &bytes,
13576 &xs[..cols],
13577 &xs[cols..],
13578 rows,
13579 cols,
13580 &mut o1,
13581 &mut o2,
13582 None,
13583 );
13584
13585 let v = Q4tpView::new(&bytes, rows, cols);
13588 let mut sc = vec![0f32; gpr];
13589 for r in 0..rows {
13590 v.scales_into(r, gpr, &mut sc);
13591 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
13592 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
13593 }
13594 }
13595
13596 #[test]
13603 fn q4tp_matvec_keeps_pace_with_q4t() {
13604 let (rows, cols) = (4096usize, 3072usize);
13605 let bytes = synth_q4tp(rows, cols);
13606 let twin = q4tp_as_q4t(&bytes, rows, cols);
13607 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
13608 let mut o = vec![0f32; rows];
13609 let n = 12;
13610 let mut best = (f64::MAX, f64::MAX);
13611 for _ in 0..3 {
13614 let t0 = std::time::Instant::now();
13615 for _ in 0..n {
13616 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
13617 }
13618 best.0 = best.0.min(t0.elapsed().as_secs_f64());
13619 let t0 = std::time::Instant::now();
13620 for _ in 0..n {
13621 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
13622 }
13623 best.1 = best.1.min(t0.elapsed().as_secs_f64());
13624 }
13625 let ratio = best.1 / best.0;
13626 println!(
13627 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
13628 best.0 * 1e3 / n as f64,
13629 best.1 * 1e3 / n as f64
13630 );
13631 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
13632 }
13633
13634 #[cfg(target_os = "macos")]
13635 #[test]
13636 fn q4t_matmat_accel_matches_dequant_reference() {
13637 if !accel_gemm_enabled() {
13638 return; }
13640 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
13642 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13643 for r in 0..rows {
13644 for g in 0..gpr {
13645 let t = (r * gpr + g) * Q4_TILE;
13646 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
13647 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13648 for k in 0..16 {
13649 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13650 }
13651 }
13652 }
13653 let x: Vec<f32> = (0..b * cols)
13654 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13655 .collect();
13656 let mut got = vec![0f32; b * rows];
13657 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
13658 let mut w = vec![0f32; rows * cols];
13660 for r in 0..rows {
13661 for g in 0..gpr {
13662 let t = (r * gpr + g) * Q4_TILE;
13663 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
13664 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
13665 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
13666 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
13667 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
13668 }
13669 }
13670 }
13671 for bi in 0..b {
13672 for r in 0..rows {
13673 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
13674 let d = (got[bi * rows + r] - want).abs();
13675 assert!(
13676 d <= want.abs().max(1.0) * 1e-4,
13677 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
13678 got[bi * rows + r]
13679 );
13680 }
13681 }
13682 }
13683
13684 #[test]
13685 fn q4matvec_matches_full_dequant() {
13686 let (rows, cols) = (8, 64);
13687 let groups = rows * cols / GROUP_SIZE;
13688 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13690 for i in 0..groups * 16 {
13691 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13692 }
13693 for g in 0..groups {
13694 let s = 0.01 + 0.003 * g as f32;
13695 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13696 }
13697 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13698
13699 let mut reference = vec![0.0f32; rows * cols];
13700 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13701 let mut expect = vec![0.0f32; rows];
13702 for r in 0..rows {
13703 expect[r] = reference[r * cols..(r + 1) * cols]
13704 .iter()
13705 .zip(&x)
13706 .map(|(w, xv)| w * xv)
13707 .sum();
13708 }
13709
13710 let mut got = vec![0.0f32; rows];
13711 q4matvec(&bytes, &x, rows, cols, &mut got, None);
13712 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13716 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13717 for r in 0..rows {
13718 assert!(
13719 (got[r] - expect[r]).abs() < tol * scale,
13720 "row {r}: {} vs {}",
13721 got[r],
13722 expect[r]
13723 );
13724 }
13725 }
13726
13727 #[test]
13730 fn vbitmatvec2_equals_two_singles() {
13731 let (rows, cols) = (6, 64);
13732 let ng = cols / GROUP_SIZE;
13733 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13734 let mut bytes = bits.clone();
13735 for g in 0..rows * ng {
13736 let s = 0.02 + 0.001 * g as f32;
13737 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13738 }
13739 for r in 0..rows {
13740 let b = bits[r] as usize;
13741 let (mut acc, mut nb) = (0u64, 0usize);
13742 let mut rowbytes = Vec::new();
13743 for i in 0..cols {
13744 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13745 acc = (acc << b) | v;
13746 nb += b;
13747 while nb >= 8 {
13748 nb -= 8;
13749 rowbytes.push(((acc >> nb) & 0xFF) as u8);
13750 }
13751 }
13752 if nb > 0 {
13753 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13754 }
13755 bytes.extend_from_slice(&rowbytes);
13756 }
13757 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13758 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
13759 let offsets = vbit_row_offsets(&bytes, rows, cols);
13760
13761 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13762 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
13763 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
13764 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13765 vbitmatvec2(
13766 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
13767 );
13768 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
13769 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
13770 }
13771
13772 #[test]
13774 fn q4matvec2_equals_two_singles() {
13775 let (rows, cols) = (8, 128);
13776 let groups = rows * cols / GROUP_SIZE;
13777 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13778 for i in 0..groups * 16 {
13779 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13780 }
13781 for g in 0..groups {
13782 let s = 0.01 + 0.003 * g as f32;
13783 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13784 }
13785 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13788 x1[9] = 250.0;
13789 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
13790
13791 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13792 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
13793 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
13794 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13795 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
13796 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
13797 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
13798 }
13799
13800 #[test]
13803 fn matvec_many_equals_separate_matvecs() {
13804 use crate::pool::Pool;
13805 let (r1, r2, cols) = (300, 200, 64);
13806 let mk = |salt: usize, rows: usize| {
13807 QTensor::from_f32(
13808 (0..rows * cols)
13809 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
13810 .collect(),
13811 rows,
13812 cols,
13813 )
13814 };
13815 let (a, b) = (mk(1, r1), mk(5, r2));
13816 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
13817 let pool = Pool::new(3);
13818
13819 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
13820 a.matvec(&x, &mut ea, Some(&pool));
13821 b.matvec(&x, &mut eb, Some(&pool));
13822 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
13823 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13824 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
13825 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
13826 }
13827
13828 #[test]
13833 fn q4tp_matvec_many_equals_separate_matvecs() {
13834 use crate::pool::Pool;
13835 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13836
13837 let (r1, r2, cols) = (300usize, 200usize, 64usize);
13838 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13839 "arch_name": "tiny-q4tp",
13840 "hidden_size": cols,
13841 "intermediate_size": cols * 2,
13842 "num_layers": 1,
13843 "num_attention_heads": 2,
13844 "num_kv_heads": 1,
13845 "head_dim": 32,
13846 "vocab_size": r1,
13847 "layer_types": ["FullAttention"],
13848 "rms_norm_eps": 1e-6,
13849 "max_position_embeddings": 8,
13850 "linear_conv_kernel_dim": 0,
13851 "linear_num_key_heads": 0,
13852 "linear_num_value_heads": 0
13853 }))
13854 .unwrap();
13855 let header = CmfHeader {
13856 format: "cmf".into(),
13857 version: CMF_VERSION,
13858 arch,
13859 quant_type: QuantType::Q4Block,
13860 provenance: None,
13861 tokenizer_config: None,
13862 section_hashes: None,
13863 skills: Vec::new(),
13864 shard: None,
13865 calibration: None,
13866 routing: None,
13867 genome: None,
13868 lineage: Vec::new(),
13869 router: None,
13870 segments: Vec::new(),
13871 };
13872 let specs = [
13873 TensorSpec {
13874 name: "q".into(),
13875 dtype: TensorDtype::Q4TiledP,
13876 shape: vec![r1, cols],
13877 data: synth_q4tp(r1, cols),
13878 },
13879 TensorSpec {
13880 name: "kv".into(),
13881 dtype: TensorDtype::Q4TiledP,
13882 shape: vec![r2, cols],
13883 data: synth_q4tp(r2, cols),
13884 },
13885 ];
13886 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
13887 std::fs::create_dir_all(&dir).unwrap();
13888 let path = dir.join("m.cmf");
13889 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13890 let model = Arc::new(CmfModel::open(&path).unwrap());
13891 let (a, b) = (
13892 QTensor::from_model(&model, "q").unwrap(),
13893 QTensor::from_model(&model, "kv").unwrap(),
13894 );
13895 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
13896 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
13897 let x: Vec<f32> = (0..cols)
13898 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
13899 .collect();
13900 let pool = Pool::new(3);
13901 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
13902 a.matvec(&x, &mut ea, Some(&pool));
13903 b.matvec(&x, &mut eb, Some(&pool));
13904 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
13905 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13906 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
13907 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
13908 let _ = std::fs::remove_dir_all(&dir);
13909 }
13910
13911 #[test]
13917 fn multi_token_moe_rows_equal_single_token_decode() {
13918 use crate::pool::Pool;
13919 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13920
13921 let (h, inter, ne) = (64usize, 128usize, 3usize);
13922 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13923 "arch_name": "tiny-q4tp-moe",
13924 "hidden_size": h,
13925 "intermediate_size": inter,
13926 "num_layers": 1,
13927 "num_attention_heads": 2,
13928 "num_kv_heads": 1,
13929 "head_dim": 32,
13930 "vocab_size": 8,
13931 "layer_types": ["FullAttention"],
13932 "rms_norm_eps": 1e-6,
13933 "max_position_embeddings": 8,
13934 "linear_conv_kernel_dim": 0,
13935 "linear_num_key_heads": 0,
13936 "linear_num_value_heads": 0
13937 }))
13938 .unwrap();
13939 let header = CmfHeader {
13940 format: "cmf".into(),
13941 version: CMF_VERSION,
13942 arch,
13943 quant_type: QuantType::Q4Block,
13944 provenance: None,
13945 tokenizer_config: None,
13946 section_hashes: None,
13947 skills: Vec::new(),
13948 shard: None,
13949 calibration: None,
13950 routing: None,
13951 genome: None,
13952 lineage: Vec::new(),
13953 router: None,
13954 segments: Vec::new(),
13955 };
13956 let mut specs = Vec::new();
13957 for e in 0..ne {
13958 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
13959 .into_iter()
13960 .enumerate()
13961 {
13962 let mut data = synth_q4tp(r, c);
13965 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
13966 .iter_mut()
13967 .enumerate()
13968 {
13969 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
13970 }
13971 specs.push(TensorSpec {
13972 name: format!("{n}{e}"),
13973 dtype: TensorDtype::Q4TiledP,
13974 shape: vec![r, c],
13975 data,
13976 });
13977 }
13978 }
13979 let dir = std::env::temp_dir().join(format!(
13980 "cmf-moe-rows-{}-{}",
13981 std::process::id(),
13982 FLOAT_ACTIVATIONS.get()
13983 ));
13984 std::fs::create_dir_all(&dir).unwrap();
13985 let path = dir.join("m.cmf");
13986 CmfModel::write(&path, &header, &specs, None, None).unwrap();
13987 let model = Arc::new(CmfModel::open(&path).unwrap());
13988 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
13989 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
13990 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
13991 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
13992 let b = 4usize;
13993 let mut xs: Vec<f32> = (0..b * h)
13994 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
13995 .collect();
13996 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
13999 (vec![2, 0], vec![0.6, 0.4]),
14000 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
14001 (vec![1], vec![1.0]),
14002 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
14003 ];
14004 let pool = Pool::new(3);
14005 let mut want = vec![0f32; b * h];
14007 for (tk, (idx, w)) in routes.iter().enumerate() {
14008 let x = &xs[tk * h..(tk + 1) * h];
14009 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
14010 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
14011 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
14012 if FLOAT_ACTIVATIONS.get() {
14013 for (slot, &e) in idx.iter().enumerate() {
14014 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
14015 g[e].matvec(x, &mut gate, Some(&pool));
14016 u[e].matvec(x, &mut up, Some(&pool));
14017 for (v, u) in gate.iter_mut().zip(up) {
14018 *v = (*v / (1.0 + (-*v).exp())) * u;
14019 }
14020 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
14021 }
14022 }
14023 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
14024 assert!(QTensor::moe_down_many(
14025 &downs,
14026 &gs,
14027 w,
14028 &mut want[tk * h..(tk + 1) * h],
14029 Some(&pool)
14030 ));
14031 }
14032 if FLOAT_ACTIVATIONS.get() {
14033 for (tk, (idx, w)) in routes.iter().enumerate() {
14034 let mut scalar = vec![0.0; h];
14035 for (&e, &weight) in idx.iter().zip(w) {
14036 let (mut gate, mut up, mut down) =
14037 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
14038 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
14039 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
14040 for (v, u) in gate.iter_mut().zip(up) {
14041 *v = (*v / (1.0 + (-*v).exp())) * u;
14042 }
14043 d[e].matvec(&gate, &mut down, Some(&pool));
14044 for (v, d) in scalar.iter_mut().zip(down) {
14045 *v += weight * d;
14046 }
14047 }
14048 assert_eq!(
14049 &want[tk * h..(tk + 1) * h],
14050 scalar,
14051 "float many equals scalar experts"
14052 );
14053 }
14054 }
14055 let mut experts: Vec<usize> = Vec::new();
14057 let mut groups: Vec<Vec<usize>> = Vec::new();
14058 for (tk, (idx, _)) in routes.iter().enumerate() {
14059 for &e in idx {
14060 match experts.iter().position(|&x| x == e) {
14061 Some(k) => groups[k].push(tk),
14062 None => {
14063 experts.push(e);
14064 groups.push(vec![tk]);
14065 }
14066 }
14067 }
14068 }
14069 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
14070 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
14071 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
14072 assert!(QTensor::moe_gate_up_rows(
14073 &pairs,
14074 &groups,
14075 &xs,
14076 &mut gs,
14077 Some(&pool)
14078 ));
14079 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
14080 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
14081 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
14082 assert!(QTensor::moe_down_rows(
14083 &downs,
14084 &lens,
14085 &gs,
14086 &mut ds,
14087 Some(&pool)
14088 ));
14089 let slot = |tk: usize, e: usize| {
14090 let k = experts.iter().position(|&x| x == e).unwrap();
14091 groups[..k].iter().map(|g| g.len()).sum::<usize>()
14092 + groups[k].iter().position(|&x| x == tk).unwrap()
14093 };
14094 let mut got = vec![0f32; b * h];
14095 for (tk, (idx, w)) in routes.iter().enumerate() {
14096 for i in 0..h {
14097 let mut acc = 0f32;
14098 for (&e, &we) in idx.iter().zip(w) {
14099 acc += we * ds[slot(tk, e)][i];
14100 }
14101 got[tk * h + i] = acc;
14102 }
14103 }
14104 assert!(want.iter().any(|v| *v != 0.0));
14105 assert_eq!(
14106 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14107 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14108 "multi-token MoE must equal decode bit for bit"
14109 );
14110
14111 let b5 = 5usize;
14114 let x5: Vec<f32> = (0..b5 * h)
14115 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
14116 .collect();
14117 let mut mm = vec![0f32; b5 * inter];
14118 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
14119 for tk in 0..b5 {
14120 let mut mv = vec![0f32; inter];
14121 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
14122 assert_eq!(
14123 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14124 mm[tk * inter..(tk + 1) * inter]
14125 .iter()
14126 .map(|v| v.to_bits())
14127 .collect::<Vec<_>>(),
14128 "row-exact matmat token {tk}"
14129 );
14130 }
14131 let _ = std::fs::remove_dir_all(&dir);
14134 }
14135
14136 #[test]
14144 fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
14145 use crate::pool::Pool;
14146 use std::sync::atomic::Ordering::Relaxed;
14147 let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
14148 Q4TP_ALT.store(2, Relaxed);
14151 let pool = Pool::new(3);
14152 for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
14156 let bytes = synth_q4tp(rows, cols);
14157 let mut xs: Vec<f32> = (0..b * cols)
14158 .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
14159 .collect();
14160 xs[3] = 7.5; let run = |exact: bool| {
14162 let mut out = vec![0f32; b * rows];
14163 q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
14164 out
14165 };
14166 let (fast, exact) = (run(false), run(true));
14167 #[cfg(not(target_arch = "aarch64"))]
14168 let _ = fast;
14169 let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14170 let mut matvecs = vec![0f32; b * rows];
14171 for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
14172 q4tp_matvec(
14173 &bytes,
14174 &xs[bi * cols..(bi + 1) * cols],
14175 rows,
14176 cols,
14177 o,
14178 Some(&pool),
14179 );
14180 }
14181 assert!(matvecs.iter().any(|v| *v != 0.0));
14182 assert_eq!(
14183 bits(&exact),
14184 bits(&matvecs),
14185 "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
14186 );
14187 #[cfg(target_arch = "aarch64")]
14188 {
14189 let gpr = cols / GROUP_SIZE;
14191 let v = Q4tpView::new(&bytes, rows, cols);
14192 let mut old = vec![0f32; b * rows];
14193 let mut sc = vec![0f32; gpr];
14194 let a8w8 = a8w8_enabled();
14195 let blocked = sdot_enabled() && blocked_enabled();
14196 let acts: Vec<SplitAct> = (0..b)
14197 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
14198 .collect();
14199 for r in 0..rows {
14200 v.scales_into(r, gpr, &mut sc);
14201 if !a8w8 {
14202 for bi in 0..b {
14203 let x = &xs[bi * cols..(bi + 1) * cols];
14204 old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
14205 }
14206 continue;
14207 }
14208 let finish = |d: f32, act: &SplitAct| {
14209 let mut acc = d * act.sx;
14210 for &(j, xv) in &act.outliers {
14211 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
14212 acc += w * s * xv;
14213 }
14214 acc
14215 };
14216 let mut bi = 0usize;
14217 while blocked && bi + 4 <= b {
14218 let xs4 = [
14219 acts[bi].xq.as_slice(),
14220 acts[bi + 1].xq.as_slice(),
14221 acts[bi + 2].xq.as_slice(),
14222 acts[bi + 3].xq.as_slice(),
14223 ];
14224 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
14225 for k in 0..4 {
14226 old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
14227 }
14228 bi += 4;
14229 }
14230 for (bi, act) in acts.iter().enumerate().skip(bi) {
14231 let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
14232 old[bi * rows + r] = finish(d, act);
14233 }
14234 }
14235 assert_eq!(
14236 bits(&fast),
14237 bits(&old),
14238 "{rows}x{cols} b={b}: the fast path changed outside row_exact"
14239 );
14240 if blocked {
14243 assert_ne!(
14244 bits(&fast),
14245 bits(&matvecs),
14246 "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
14247 );
14248 }
14249 }
14250 }
14251 Q4TP_ALT.store(0, Relaxed);
14252 }
14253
14254 #[test]
14255 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
14256 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
14257 let active = AtomicUsize::new(0);
14260 counted_row_exact_scope(&active, || {
14261 assert_eq!(active.load(Ordering::Acquire), 1);
14262 counted_row_exact_scope(&active, || {
14263 assert_eq!(active.load(Ordering::Acquire), 2);
14264 });
14265 assert_eq!(active.load(Ordering::Acquire), 1);
14266 });
14267 assert_eq!(active.load(Ordering::Acquire), 0);
14268
14269 let both_entered = Barrier::new(2);
14270 let release_last = Barrier::new(2);
14271 std::thread::scope(|s| {
14272 let first = s.spawn(|| counted_row_exact_scope(&active, || {
14273 both_entered.wait();
14274 }));
14275 let last = s.spawn(|| counted_row_exact_scope(&active, || {
14276 both_entered.wait();
14277 release_last.wait();
14278 }));
14279 first.join().unwrap();
14280 let after_first = active.load(Ordering::Acquire);
14281 release_last.wait();
14282 last.join().unwrap();
14283 assert_eq!(after_first, 1, "second request must remain exact");
14284 });
14285 assert_eq!(active.load(Ordering::Acquire), 0);
14286 let panic = std::panic::catch_unwind(|| {
14287 counted_row_exact_scope(&active, || panic!("scope unwind"));
14288 });
14289 assert!(panic.is_err());
14290 assert_eq!(active.load(Ordering::Acquire), 0);
14291 }
14292
14293 #[test]
14294 #[cfg(target_arch = "x86_64")]
14295 fn q4tp_float_avx2_is_bitwise_scalar() {
14296 if !avx2_enabled() {
14297 return;
14298 }
14299 for cols in [32, 64, 96, 2048, 4096] {
14300 let rows = 9;
14301 let bytes = synth_q4tp(rows, cols);
14302 let v = Q4tpView::new(&bytes, rows, cols);
14303 let gpr = cols / GROUP_SIZE;
14304 let mut sc = vec![0.0; gpr];
14305 for seed in 1..=5 {
14306 let xs: Vec<f32> = (0..cols)
14307 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
14308 .collect();
14309 for r in 0..rows {
14310 v.scales_into(r, gpr, &mut sc);
14311 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
14312 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
14313 assert_eq!(
14314 scalar.to_bits(),
14315 vector.to_bits(),
14316 "cols={cols} row={r} seed={seed}"
14317 );
14318 }
14319 }
14320 }
14321 }
14322
14323 #[test]
14324 fn multi_token_moe_rows_float_equal_single_token_decode() {
14325 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
14326 }
14327
14328 #[test]
14329 fn full_gpu_q8_scope_is_nested_and_thread_local() {
14330 assert!(!FULL_GPU_Q8.get());
14331 let before = gpu_split_frac();
14332 {
14333 let _guard = enter_full_gpu_q8_scope();
14334 assert_eq!(gpu_split_frac(), 1.0);
14335 {
14336 let _nested = enter_full_gpu_q8_scope();
14337 }
14338 assert_eq!(gpu_split_frac(), 1.0);
14339 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
14340 }
14341 assert!(!FULL_GPU_Q8.get());
14342 assert_eq!(gpu_split_frac(), before);
14343 }
14344
14345 #[test]
14346 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
14347 assert!(!FLOAT_ACTIVATIONS.get());
14348 let before = a8w8_enabled();
14349 float_activations_scope(|| {
14350 assert!(!a8w8_enabled());
14351 float_activations_scope(|| assert!(!a8w8_enabled()));
14352 assert!(FLOAT_ACTIVATIONS.get());
14353 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
14354 .join()
14355 .unwrap();
14356 });
14357 assert!(!FLOAT_ACTIVATIONS.get());
14358 assert_eq!(a8w8_enabled(), before);
14359 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
14360 assert!(!FLOAT_ACTIVATIONS.get());
14361 }
14362
14363 #[test]
14366 fn batched_matmat_equals_per_position_matvec() {
14367 let (rows, cols, b) = (8, 64, 5);
14368 let groups = rows * cols / GROUP_SIZE;
14370 let mut q4 = Vec::new();
14371 for i in 0..groups * 16 {
14372 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14373 }
14374 for g in 0..groups {
14375 q4.extend_from_slice(
14376 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14377 );
14378 }
14379 let ng = cols / GROUP_SIZE;
14381 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
14382 let mut vb = bits.clone();
14383 for g in 0..rows * ng {
14384 vb.extend_from_slice(
14385 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
14386 );
14387 }
14388 for r in 0..rows {
14389 let bw = bits[r] as usize;
14390 let (mut acc, mut nb) = (0u64, 0usize);
14391 let mut rowbytes = Vec::new();
14392 for i in 0..cols {
14393 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14394 acc = (acc << bw) | v;
14395 nb += bw;
14396 while nb >= 8 {
14397 nb -= 8;
14398 rowbytes.push(((acc >> nb) & 0xFF) as u8);
14399 }
14400 }
14401 if nb > 0 {
14402 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14403 }
14404 vb.extend_from_slice(&rowbytes);
14405 }
14406 let offsets = vbit_row_offsets(&vb, rows, cols);
14407
14408 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14409
14410 let mut got = vec![0f32; b * rows];
14412 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
14413 for bi in 0..b {
14414 let mut expect = vec![0f32; rows];
14415 q4matvec(
14416 &q4,
14417 &xs[bi * cols..(bi + 1) * cols],
14418 rows,
14419 cols,
14420 &mut expect,
14421 None,
14422 );
14423 assert_eq!(
14424 &got[bi * rows..(bi + 1) * rows],
14425 &expect[..],
14426 "q4 batch pos {bi}"
14427 );
14428 }
14429
14430 let mut got = vec![0f32; b * rows];
14432 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
14433 for bi in 0..b {
14434 let mut expect = vec![0f32; rows];
14435 vbitmatvec(
14436 &vb,
14437 &offsets,
14438 &xs[bi * cols..(bi + 1) * cols],
14439 rows,
14440 cols,
14441 &mut expect,
14442 None,
14443 );
14444 assert_eq!(
14445 &got[bi * rows..(bi + 1) * rows],
14446 &expect[..],
14447 "vbit batch pos {bi}"
14448 );
14449 }
14450 }
14451
14452 #[test]
14456 fn q4_tiled_matches_q4_block_bitexact() {
14457 let (rows, cols, b) = (8usize, 128usize, 3usize);
14458 let groups = rows * cols / GROUP_SIZE;
14459 let mut split = Vec::with_capacity(groups * 18);
14460 for i in 0..groups * 16 {
14461 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14462 }
14463 for g in 0..groups {
14464 split.extend_from_slice(
14465 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14466 );
14467 }
14468 let (packed, scales) = split.split_at(groups * 16);
14470 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
14471 for g in 0..groups {
14472 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
14473 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
14474 }
14475
14476 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14477 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14479
14480 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
14481 q4matvec(&split, &x1, rows, cols, &mut a, None);
14482 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
14483 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
14484
14485 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14486 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
14487 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
14488 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
14489 assert_eq!(a1, t1);
14490 assert_eq!(a2, t2);
14491
14492 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14493 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14494 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
14495 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
14496 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
14497 }
14498
14499 #[test]
14506 fn q4matvec_sdot_outlier_exact() {
14507 let (rows, cols) = (4, 128);
14508 let groups = rows * cols / GROUP_SIZE;
14509 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14510 for i in 0..groups * 16 {
14511 bytes.push(((i * 11 + 5) % 256) as u8);
14512 }
14513 for g in 0..groups {
14514 let s = 0.02 + 0.002 * g as f32;
14515 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14516 }
14517 let mut x: Vec<f32> = (0..cols)
14518 .map(|i| match i % 3 {
14519 0 => 1.0,
14520 1 => -1.0,
14521 _ => 0.0,
14522 })
14523 .collect();
14524 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
14527 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14528 let mut expect = vec![0.0f32; rows];
14529 for r in 0..rows {
14530 expect[r] = reference[r * cols..(r + 1) * cols]
14531 .iter()
14532 .zip(&x)
14533 .map(|(w, xv)| w * xv)
14534 .sum();
14535 }
14536 let mut got = vec![0.0f32; rows];
14537 q4matvec(&bytes, &x, rows, cols, &mut got, None);
14538 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14539 for r in 0..rows {
14540 assert!(
14541 (got[r] - expect[r]).abs() < 2e-3 * scale,
14542 "row {r}: {} vs {} (outlier term must be exact)",
14543 got[r],
14544 expect[r]
14545 );
14546 }
14547 }
14548
14549 #[test]
14553 fn q1t_matvec_matches_reference() {
14554 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
14555 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
14557 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
14558 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
14560 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
14561 let mut bytes = Vec::new();
14562 for r in 0..rows {
14563 for g in 0..gpr {
14564 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
14565 let mut c = [0u8; 7];
14566 for k in 0..GROUP_SIZE {
14567 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
14569 0
14570 } else {
14571 ((k + r * 3 + g) % 3) as u8 };
14573 cortiq_core::quant::q1t_pack(&mut c, k, code);
14574 }
14575 bytes.extend_from_slice(&c);
14576 }
14577 }
14578 let mut row_ptr = vec![0u32; rows + 1];
14581 for &(idx, _) in &outliers {
14582 row_ptr[idx as usize / cols + 1] += 1;
14583 }
14584 for r in 0..rows {
14585 row_ptr[r + 1] += row_ptr[r];
14586 }
14587 for &p in &row_ptr {
14588 bytes.extend_from_slice(&p.to_le_bytes());
14589 }
14590 for &(idx, v) in &outliers {
14591 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
14592 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
14593 }
14594
14595 let mut refw = vec![0f32; rows * cols];
14596 dequant_q1t(&bytes, rows, cols, &mut refw);
14597 let x: Vec<f32> = (0..cols)
14600 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14601 .collect();
14602 let mut expect = vec![0f32; rows];
14603 for r in 0..rows {
14604 let mut a = 0.0f32;
14605 for j in 0..cols {
14606 a += refw[r * cols + j] * x[j];
14607 }
14608 expect[r] = a;
14609 }
14610 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
14611 let mut got = vec![0f32; rows];
14612 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
14613 for r in 0..rows {
14614 assert!(
14615 (got[r] - expect[r]).abs() < tol(expect[r]),
14616 "row {r}: {} vs {}",
14617 got[r],
14618 expect[r]
14619 );
14620 }
14621 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
14623 let mut gm = vec![0f32; 2 * rows];
14624 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
14625 for r in 0..rows {
14626 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
14627 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
14628 }
14629 let xb: Vec<f32> = (0..cols)
14633 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
14634 .collect();
14635 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14636 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
14637 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
14638 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14639 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
14640 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
14641 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
14642 }
14643
14644 #[test]
14647 fn q1t_matvec2_odd_gpr_matches_singles() {
14648 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14649 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
14651 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14652 for r in 0..rows {
14653 for g in 0..gpr {
14654 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
14655 let mut c = [0u8; 7];
14656 for k in 0..GROUP_SIZE {
14657 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
14658 }
14659 bytes.extend_from_slice(&c);
14660 }
14661 }
14662 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14663 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14664 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14665 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14666 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14667 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14668 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14669 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
14670 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
14671 }
14672
14673 #[test]
14677 #[ignore]
14678 fn q1t_matvec2_speed() {
14679 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14680 use std::time::Instant;
14681 let (rows, cols) = (8192usize, 4096usize);
14682 let gpr = cols / GROUP_SIZE;
14683 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14684 for r in 0..rows {
14685 for g in 0..gpr {
14686 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14687 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14688 let mut c = [0u8; 7];
14689 for k in 0..GROUP_SIZE {
14690 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14691 }
14692 bytes.extend_from_slice(&c);
14693 }
14694 }
14695 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14696 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14697 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14698 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14699 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14701 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14702 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
14703 for _ in 0..8 {
14704 let t0 = Instant::now();
14705 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14706 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
14707 let t1 = Instant::now();
14708 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14709 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14710 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
14711 }
14712 assert_eq!(p1, s1);
14713 assert_eq!(p2, s2);
14714 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
14715 }
14716
14717 #[test]
14721 #[ignore]
14722 fn q1t_matvec_speed() {
14723 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
14724 use std::time::Instant;
14725 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
14727 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
14728 for r in 0..rows {
14729 for g in 0..gpr {
14730 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14731 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14732 let mut c = [0u8; 7];
14733 for k in 0..GROUP_SIZE {
14734 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14735 }
14736 bytes.extend_from_slice(&c);
14737 }
14738 }
14739 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
14741 let mut idx = 0usize;
14742 while idx < n {
14743 row_ptr[idx / cols + 1] += 1;
14744 idx += stride;
14745 }
14746 for r in 0..rows {
14747 row_ptr[r + 1] += row_ptr[r];
14748 }
14749 for &p in &row_ptr {
14750 bytes.extend_from_slice(&p.to_le_bytes());
14751 }
14752 let mut idx = 0usize;
14753 while idx < n {
14754 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
14755 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
14756 idx += stride;
14757 }
14758 let x: Vec<f32> = (0..cols)
14761 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14762 .collect();
14763 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
14764
14765 let slow = |out: &mut [f32]| {
14767 let mut buf = vec![0f32; cols];
14768 for r in 0..rows {
14769 for g in 0..gpr {
14770 let off = (r * gpr + g) * Q1T_TILE;
14771 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
14772 let codes = &bytes[off + 2..off + Q1T_TILE];
14773 for k in 0..GROUP_SIZE {
14774 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
14775 1 => s,
14776 2 => -s,
14777 _ => 0.0,
14778 };
14779 }
14780 }
14781 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
14782 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
14783 }
14784 };
14785 let iters = 5;
14786 let mut a = vec![0f32; rows];
14787 slow(&mut a); let t = Instant::now();
14789 for _ in 0..iters {
14790 slow(&mut a);
14791 }
14792 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14793
14794 let mut b = vec![0f32; rows];
14795 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
14797 for _ in 0..iters {
14798 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
14799 }
14800 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14801
14802 for r in 0..rows {
14803 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
14804 }
14805 println!(
14806 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
14807 slow_ms / fast_ms
14808 );
14809 }
14810}
14811
14812#[cfg(test)]
14813mod gemm_bench {
14814 #[test]
14825 #[ignore]
14826 fn q4tp_matmat_throughput() {
14827 let _alt = super::Q4TP_ALT_TEST_LOCK
14828 .lock()
14829 .unwrap_or_else(|e| e.into_inner());
14830 let b: usize = std::env::var("CMF_BENCH_B")
14834 .ok()
14835 .and_then(|v| v.parse().ok())
14836 .unwrap_or(296);
14837 let (rows, cols) = (9216usize, 2304usize);
14838 let (_, _, _) = (rows, cols, b);
14839 let total =
14840 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
14841 .unwrap();
14842 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
14847 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
14848 let lo = cortiq_core::quant::f32_to_f16(-4.0);
14849 let step = cortiq_core::quant::f32_to_f16(0.1);
14850 for r in 0..rows {
14851 let o = params_off + r * 4;
14852 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
14853 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
14854 }
14855 let _ = codes_off;
14856 let xs: Vec<f32> = (0..b * cols)
14857 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
14858 .collect();
14859 let mut out = vec![0f32; b * rows];
14860 let pool = crate::pool::Pool::from_env();
14861 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
14867 let reps: usize = std::env::var("CMF_BENCH_REPS")
14868 .ok()
14869 .and_then(|v| v.parse().ok())
14870 .unwrap_or(10);
14871 let mut best = [f64::MAX; 2];
14872 let mut sums = [0f32; 2];
14873 for _ in 0..reps {
14874 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
14875 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
14876 let t = std::time::Instant::now();
14877 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
14878 best[k] = best[k].min(t.elapsed().as_secs_f64());
14879 sums[k] = out.iter().take(64).sum::<f32>();
14880 }
14881 }
14882 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
14883 for (k, name) in ["previous", "tuned "].iter().enumerate() {
14884 println!(
14885 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
14886 best[k] * 1e3,
14887 flops / best[k] / 1e9,
14888 sums[k]
14889 );
14890 }
14891 assert!(
14892 (sums[0] - sums[1]).abs() < 1e-2,
14893 "the tuned kernel changed the result: {} vs {}",
14894 sums[0],
14895 sums[1]
14896 );
14897 }
14898
14899 #[test]
14905 fn q4tp_matmat_blocked_matches_scalar() {
14906 use std::sync::atomic::Ordering::Relaxed;
14907 let _alt = super::Q4TP_ALT_TEST_LOCK
14908 .lock()
14909 .unwrap_or_else(|e| e.into_inner());
14910 for &(rows, cols, b) in &[
14918 (64usize, 128usize, 7usize),
14919 (33, 96, 4),
14920 (16, 256, 9),
14921 (192, 2304, 37),
14922 ] {
14923 let total = cortiq_core::quant::expected_nbytes(
14924 cortiq_core::TensorDtype::Q4TiledP,
14925 &[rows, cols],
14926 )
14927 .unwrap();
14928 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
14929 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
14930 let lo = cortiq_core::quant::f32_to_f16(-4.0);
14931 let step = cortiq_core::quant::f32_to_f16(0.1);
14932 for r in 0..rows {
14933 let o = params_off + r * 4;
14934 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
14935 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
14936 }
14937 let xs: Vec<f32> = (0..b * cols)
14938 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
14939 .collect();
14940 let mut got = vec![0f32; b * rows];
14941 let mut want = vec![0f32; b * rows];
14942 let gpr = cols / 32;
14943 let view = super::Q4tpView::new(&bytes, rows, cols);
14944 let pool = crate::pool::Pool::from_env();
14945 super::Q4TP_ALT.store(2, Relaxed);
14946 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
14947 super::Q4TP_ALT.store(1, Relaxed);
14948 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
14949 super::Q4TP_ALT.store(0, Relaxed);
14950 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
14957 let (mut worst, mut at) = (0f32, 0usize);
14958 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
14959 if (g - w).abs() > worst {
14960 worst = (g - w).abs();
14961 at = i;
14962 }
14963 }
14964 assert!(
14965 worst <= 1e-4 * scale,
14966 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
14967 (scale {scale:.3e}) at cell {at}: {} vs {}",
14968 got[at],
14969 want[at]
14970 );
14971
14972 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
14979 for bi in 0..b {
14980 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
14981 for r in 0..rows {
14982 let mut sc = vec![0f32; gpr];
14983 view.scales_into(r, gpr, &mut sc);
14984 let mut exact = 0f64;
14985 for j in 0..cols {
14986 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
14987 exact += w as f64 * sq as f64 * act.xq[j] as f64;
14988 }
14989 exact *= act.sx as f64;
14990 for &(j, xv) in &act.outliers {
14991 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
14992 exact += w as f64 * sq as f64 * xv as f64;
14993 }
14994 let i = bi * rows + r;
14995 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
14996 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
14997 }
14998 }
14999 println!(
15000 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
15001 per-column {e_scalar:.3e}"
15002 );
15003 assert!(
15008 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
15009 "{rows}x{cols} b={b}: error against f64 too large — blocked \
15010 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
15011 );
15012 }
15013 }
15014
15015 #[test]
15021 #[ignore]
15022 fn q4tp_matmat_row_exact_cost() {
15023 let _alt = super::Q4TP_ALT_TEST_LOCK
15024 .lock()
15025 .unwrap_or_else(|e| e.into_inner());
15026 let pool = crate::pool::Pool::from_env();
15027 let reps: usize = std::env::var("CMF_BENCH_REPS")
15028 .ok()
15029 .and_then(|v| v.parse().ok())
15030 .unwrap_or(30);
15031 for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
15032 let total = cortiq_core::quant::expected_nbytes(
15033 cortiq_core::TensorDtype::Q4TiledP,
15034 &[rows, cols],
15035 )
15036 .unwrap();
15037 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15038 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15039 let lo = cortiq_core::quant::f32_to_f16(-4.0);
15040 let step = cortiq_core::quant::f32_to_f16(0.1);
15041 for r in 0..rows {
15042 let o = params_off + r * 4;
15043 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15044 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15045 }
15046 for &b in &[2usize, 4, 5, 8] {
15047 let xs: Vec<f32> = (0..b * cols)
15048 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15049 .collect();
15050 let mut out = vec![0f32; b * rows];
15051 let mut best = [f64::MAX; 3];
15052 for _ in 0..reps {
15053 for (k, best_k) in best.iter_mut().enumerate() {
15054 let t = std::time::Instant::now();
15055 match k {
15056 0 | 1 => super::q4tp_matmat_with(
15057 &bytes,
15058 &xs,
15059 b,
15060 rows,
15061 cols,
15062 &mut out,
15063 pool.as_deref(),
15064 k == 1,
15065 ),
15066 _ => {
15067 for (bi, o) in out.chunks_mut(rows).enumerate() {
15068 super::q4tp_matvec(
15069 &bytes,
15070 &xs[bi * cols..(bi + 1) * cols],
15071 rows,
15072 cols,
15073 o,
15074 pool.as_deref(),
15075 );
15076 }
15077 }
15078 }
15079 *best_k = best_k.min(t.elapsed().as_secs_f64());
15080 }
15081 }
15082 println!(
15083 "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
15084 {b} matvecs {:.3} ms",
15085 best[0] * 1e3,
15086 best[1] * 1e3,
15087 best[2] * 1e3
15088 );
15089 }
15090 }
15091 }
15092
15093 #[test]
15097 #[ignore]
15098 fn q4t_matmat_throughput() {
15099 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
15100 let total =
15101 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
15102 .unwrap();
15103 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15106 let sc = cortiq_core::quant::f32_to_f16(0.02);
15107 for t in bytes.chunks_mut(super::Q4_TILE) {
15108 t[..2].copy_from_slice(&sc.to_le_bytes());
15109 }
15110 let xs: Vec<f32> = (0..b * cols)
15111 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15112 .collect();
15113 let mut out = vec![0f32; b * rows];
15114 let pool = crate::pool::Pool::from_env();
15115 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15116 let reps: usize = std::env::var("CMF_BENCH_REPS")
15117 .ok()
15118 .and_then(|v| v.parse().ok())
15119 .unwrap_or(10);
15120 let mut best = f64::MAX;
15121 for _ in 0..reps {
15122 let t = std::time::Instant::now();
15123 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15124 best = best.min(t.elapsed().as_secs_f64());
15125 }
15126 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15127 println!(
15128 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
15129 best * 1e3,
15130 flops / best / 1e9,
15131 out.iter().take(64).sum::<f32>()
15132 );
15133 }
15134}