1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16 GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17 q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23 F32 {
24 data: Vec<f32>,
25 rows: usize,
26 cols: usize,
27 },
28 Mapped {
29 model: Arc<CmfModel>,
30 idx: usize,
32 dtype: TensorDtype,
33 rows: usize,
34 cols: usize,
35 row_scale: Vec<f32>,
37 col_field: Vec<f32>,
39 vbit_offsets: Vec<usize>,
43 repack: Vec<u8>,
51 },
52}
53
54fn repack_enabled() -> bool {
61 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62 *ON.get_or_init(|| {
63 std::env::var("CMF_REPACK")
64 .map(|v| v == "1")
65 .unwrap_or(cfg!(target_os = "android"))
66 })
67}
68
69fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73 #[cfg(target_arch = "aarch64")]
74 let arch_ok = sdot_enabled();
75 #[cfg(not(target_arch = "aarch64"))]
76 let arch_ok = false;
77 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78 return Vec::new();
79 }
80 q8_repack_layout(bytes, rows, cols)
81}
82
83fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86 let groups = rows / 4;
87 let mut rep = vec![0u8; groups * 4 * cols];
88 for g in 0..groups {
89 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90 for c in 0..cols / 16 {
91 for lane in 0..4 {
92 let src = (g * 4 + lane) * cols + c * 16;
93 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94 .copy_from_slice(&bytes[src..src + 16]);
95 }
96 }
97 }
98 rep
99}
100
101fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104 let ng = cols / GROUP_SIZE;
105 let bits = &bytes[..rows];
106 let mut offsets = Vec::with_capacity(rows + 1);
107 let mut off = rows + rows * ng * 2;
108 for r in 0..rows {
109 offsets.push(off);
110 off += (cols * bits[r] as usize).div_ceil(8);
111 }
112 offsets.push(off);
113 offsets
114}
115
116fn blocked_enabled() -> bool {
122 use std::sync::atomic::Ordering::Relaxed;
123 match BLOCKED_OVERRIDE.load(Relaxed) {
124 1 => false,
125 2 => true,
126 _ => {
127 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128 *ON.get_or_init(|| {
129 std::env::var("CMF_X86_BLOCKED")
130 .map(|v| v != "0")
131 .unwrap_or(true)
132 })
133 }
134 }
135}
136
137static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
138
139pub fn set_blocked_override(on: Option<bool>) {
146 let v = match on {
147 None => 0,
148 Some(false) => 1,
149 Some(true) => 2,
150 };
151 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
152}
153
154fn gpu_lmhead_enabled() -> bool {
155 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
156 *ON.get_or_init(|| {
157 std::env::var("CMF_GPU_LMHEAD")
158 .map(|v| v != "0")
159 .unwrap_or(true)
160 })
161}
162
163fn gpu_split_frac() -> f32 {
164 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
165 *F.get_or_init(|| {
166 std::env::var("CMF_GPU_SPLIT")
167 .ok()
168 .and_then(|v| v.parse::<f32>().ok())
169 .unwrap_or(0.5)
170 .clamp(0.0, 1.0)
171 })
172}
173
174impl QTensor {
175 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
176 debug_assert_eq!(data.len(), rows * cols);
177 Self::F32 { data, rows, cols }
178 }
179
180 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
183 let idx = model
186 .tensor_index(name)
187 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
188 let entry = &model.tensors[idx];
189 if entry.shape.len() != 2 {
190 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
191 }
192 let (rows, cols) = (entry.shape[0], entry.shape[1]);
193 let bytes = model.entry_bytes(entry);
194
195 match entry.dtype {
196 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
197 let n = rows * cols;
198 let scales_off = n;
199 let row_scale: Vec<f32> = (0..rows)
200 .map(|o| {
201 f16_to_f32(u16::from_le_bytes([
202 bytes[scales_off + o * 2],
203 bytes[scales_off + o * 2 + 1],
204 ]))
205 })
206 .collect();
207 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
208 let col_off = n + rows * 2;
209 (0..cols)
210 .map(|i| {
211 f16_to_f32(u16::from_le_bytes([
212 bytes[col_off + i * 2],
213 bytes[col_off + i * 2 + 1],
214 ]))
215 })
216 .collect()
217 } else {
218 Vec::new()
219 };
220 Ok(Self::Mapped {
221 model: model.clone(),
222 idx,
223 dtype: entry.dtype,
224 rows,
225 cols,
226 row_scale,
227 col_field,
228 vbit_offsets: Vec::new(),
229 repack: q8_repack(bytes, rows, cols),
230 })
231 }
232 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
234 model: model.clone(),
235 idx,
236 dtype: entry.dtype,
237 rows,
238 cols,
239 row_scale: Vec::new(),
240 col_field: Vec::new(),
241 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
242 repack: Vec::new(),
243 }),
244 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
248 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
249 let offsets: Vec<usize> = (0..=rows)
250 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
251 .collect();
252 Ok(Self::Mapped {
253 model: model.clone(),
254 idx,
255 dtype: entry.dtype,
256 rows,
257 cols,
258 row_scale: Vec::new(),
259 col_field: Vec::new(),
260 vbit_offsets: offsets,
261 repack: Vec::new(),
262 })
263 }
264 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
270 model: model.clone(),
271 idx,
272 dtype: entry.dtype,
273 rows,
274 cols,
275 row_scale: Vec::new(),
276 col_field: Vec::new(),
277 vbit_offsets: Vec::new(),
278 repack: Vec::new(),
279 }),
280 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
283 model: model.clone(),
284 idx,
285 dtype: entry.dtype,
286 rows,
287 cols,
288 row_scale: Vec::new(),
289 col_field: Vec::new(),
290 vbit_offsets: Vec::new(),
291 repack: Vec::new(),
292 }),
293 TensorDtype::Q2TiledP 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::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
306 model: model.clone(),
307 idx,
308 dtype: entry.dtype,
309 rows,
310 cols,
311 row_scale: Vec::new(),
312 col_field: Vec::new(),
313 vbit_offsets: Vec::new(),
314 repack: Vec::new(),
315 }),
316 TensorDtype::Q1 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::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
332 model: model.clone(),
333 idx,
334 dtype: entry.dtype,
335 rows,
336 cols,
337 row_scale: Vec::new(),
338 col_field: Vec::new(),
339 vbit_offsets: Vec::new(),
340 repack: Vec::new(),
341 }),
342 _ => {
344 let mut data = vec![0.0f32; rows * cols];
345 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
346 Ok(Self::from_f32(data, rows, cols))
347 }
348 }
349 }
350
351 pub(crate) fn is_q1(&self) -> bool {
354 matches!(
355 self,
356 Self::Mapped {
357 dtype: TensorDtype::Q1,
358 ..
359 }
360 )
361 }
362
363 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
366 match self {
367 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
368 _ => None,
369 }
370 }
371
372 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
379 if self.has_prism_contract() {
380 return None;
381 }
382 match self {
383 #[cfg(target_os = "macos")]
384 Self::Mapped {
385 dtype: TensorDtype::Q1T,
386 ..
387 } if !crate::gpu::metal_q1t_enabled() => None,
388 Self::Mapped {
389 idx,
390 dtype:
391 TensorDtype::Q1
392 | TensorDtype::Q1T
393 | TensorDtype::Q4Block
394 | TensorDtype::Q4Tiled
395 | TensorDtype::Q4TiledP
399 | TensorDtype::Q8Row
400 | TensorDtype::Q8_2f,
401 rows,
402 cols,
403 ..
404 } => Some((*idx, *rows, *cols)),
405 _ => None,
406 }
407 }
408
409 #[cfg(target_os = "macos")]
416 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
417 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
418 let name = &model.tensors[idx].name;
419 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
420 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
421 if forward && affine {
422 let e = model.tensors.get(idx)?;
423 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
424 }
425 }
426 self.q1_parts()
427 }
428
429 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
435 if self.has_prism_contract() {
436 return None;
437 }
438 match self {
439 Self::Mapped {
440 idx,
441 dtype: TensorDtype::Q4Tiled,
442 rows,
443 cols,
444 ..
445 } => Some((*idx, *rows, *cols)),
446 _ => None,
447 }
448 }
449
450 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
454 if self.has_prism_contract() {
455 return None;
456 }
457 match self {
458 Self::Mapped {
459 idx,
460 dtype: TensorDtype::Q4TiledP,
461 rows,
462 cols,
463 ..
464 } => Some((*idx, *rows, *cols)),
465 _ => None,
466 }
467 }
468
469 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
474 if self.has_prism_contract() {
475 return None;
476 }
477 match self {
478 Self::Mapped {
479 idx,
480 dtype: TensorDtype::Q8Row,
481 rows,
482 cols,
483 row_scale,
484 col_field,
485 ..
486 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
487 _ => None,
488 }
489 }
490
491 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
495 match self {
496 Self::Mapped { dtype, .. } => Some(*dtype),
497 _ => None,
498 }
499 }
500
501 pub fn model_idx(&self) -> Option<usize> {
506 match self {
507 Self::Mapped { idx, .. } => Some(*idx),
508 _ => None,
509 }
510 }
511
512 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
517 match self {
518 Self::Mapped { model, .. } => Some(model.clone()),
519 _ => None,
520 }
521 }
522
523 pub(crate) fn has_prism_contract(&self) -> bool {
528 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
529 }
530
531 pub fn rows(&self) -> usize {
532 match self {
533 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
534 }
535 }
536
537 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
540 if self.has_prism_contract() {
541 return None;
542 }
543 match self {
544 Self::Mapped {
545 model,
546 idx,
547 dtype: TensorDtype::Q4Tiled,
548 ..
549 } => Some((model, *idx)),
550 _ => None,
551 }
552 }
553
554 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
557 if self.has_prism_contract() {
558 return None;
559 }
560 match self {
561 Self::Mapped {
562 model,
563 idx,
564 dtype: TensorDtype::Q4TiledP,
565 ..
566 } => Some((model, *idx)),
567 _ => None,
568 }
569 }
570
571 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
579 if self.has_prism_contract() {
580 return None;
581 }
582 match self {
583 Self::Mapped {
584 model,
585 idx,
586 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
587 ..
588 } => Some((model, *idx)),
589 _ => None,
590 }
591 }
592
593 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
596 if self.has_prism_contract() {
597 return None;
598 }
599 match self {
600 Self::Mapped {
601 model,
602 idx,
603 dtype: TensorDtype::Q2TiledP,
604 ..
605 } => Some((model, *idx)),
606 _ => None,
607 }
608 }
609
610 pub fn cols(&self) -> usize {
611 match self {
612 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
613 }
614 }
615
616 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
619 if self.has_prism_contract() {
620 return None;
621 }
622 match self {
623 Self::Mapped {
624 model,
625 idx,
626 dtype: TensorDtype::Q1,
627 ..
628 } => Some((model, *idx)),
629 _ => None,
630 }
631 }
632
633 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
643 if self.has_prism_contract() {
644 return None;
645 }
646 self.graph_weight_descriptor()
647 }
648
649 pub(crate) fn graph_weight_descriptor(
654 &self,
655 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
656 match self {
657 Self::Mapped {
658 model,
659 idx,
660 dtype: TensorDtype::Q8Row,
661 row_scale,
662 ..
663 } => Some((model, *idx, 0, row_scale.as_slice())),
664 Self::Mapped {
665 model,
666 idx,
667 dtype: TensorDtype::Q1,
668 ..
669 } => Some((model, *idx, 1, &[])),
670 Self::Mapped {
675 model,
676 idx,
677 dtype: TensorDtype::Q4Tiled,
678 ..
679 } => Some((model, *idx, 5, &[])),
680 Self::Mapped {
684 model,
685 idx,
686 dtype: TensorDtype::Q4TiledP,
687 ..
688 } => Some((model, *idx, 6, &[])),
689 Self::Mapped {
690 model,
691 idx,
692 dtype: TensorDtype::Q4Block,
693 ..
694 } => Some((model, *idx, 2, &[])),
695 Self::Mapped {
700 model,
701 idx,
702 dtype: TensorDtype::Q8_2f,
703 ..
704 } => Some((model, *idx, 7, &[])),
705 Self::Mapped {
706 model,
707 idx,
708 dtype: TensorDtype::Q1T,
709 ..
710 } => Some((model, *idx, 3, &[])),
711 Self::Mapped {
715 model,
716 idx,
717 dtype: TensorDtype::Q2TiledP,
718 ..
719 } => Some((model, *idx, 9, &[])),
720 _ => None,
721 }
722 }
723
724 pub fn as_f32(&self) -> Option<&[f32]> {
727 match self {
728 Self::F32 { data, .. } => Some(data),
729 Self::Mapped { .. } => None,
730 }
731 }
732
733 fn quant_bytes(&self) -> &[u8] {
734 match self {
735 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
736 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
737 }
738 }
739
740 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
742 let cols = self.cols();
743 debug_assert_eq!(dst.len(), cols);
744 match self {
745 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
746 Self::Mapped {
747 model,
748 idx,
749 dtype,
750 row_scale,
751 col_field,
752 vbit_offsets,
753 ..
754 } => {
755 if *dtype == TensorDtype::Q4Tiled {
756 let bytes = self.quant_bytes();
757 let gpr = cols / GROUP_SIZE;
758 for gi in 0..gpr {
759 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
760 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
761 for (k, &b) in tile[2..].iter().enumerate() {
762 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
763 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
764 }
765 }
766 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
767 crate::prism::inverse_embedding(model, dst);
768 }
769 return;
770 }
771 if *dtype == TensorDtype::Q4TiledP {
772 let bytes = self.quant_bytes();
773 let gpr = cols / GROUP_SIZE;
774 let v = Q4tpView::new(bytes, self.rows(), cols);
775 let mut sc = vec![0f32; gpr];
776 v.scales_into(r, gpr, &mut sc);
777 for gi in 0..gpr {
778 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
779 let s = sc[gi];
780 for (k, &b) in tile.iter().enumerate() {
781 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
782 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
783 }
784 }
785 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
786 crate::prism::inverse_embedding(model, dst);
787 }
788 return;
789 }
790 if *dtype == TensorDtype::Q2TiledP {
791 let bytes = self.quant_bytes();
792 let gpr = cols / GROUP_SIZE;
793 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
794 let mut sc = vec![0f32; gpr];
795 v.scales_into(r, gpr, &mut sc);
796 for gi in 0..gpr {
797 let ch =
798 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
799 let s = sc[gi];
800 for (k, &b) in ch.iter().enumerate() {
801 for j in 0..4 {
802 let center = if crate::prism::is_affine_target(
803 model,
804 &model.tensors[*idx].name,
805 ) {
806 1.0
807 } else {
808 1.5
809 };
810 dst[gi * GROUP_SIZE + k * 4 + j] =
811 (((b >> (2 * j)) & 3) as f32 - center) * s;
812 }
813 }
814 }
815 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
816 crate::prism::inverse_embedding(model, dst);
817 }
818 return;
819 }
820 if *dtype == TensorDtype::Q4Block {
821 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
822 let gpr = cols / GROUP_SIZE;
823 for gi in 0..gpr {
824 let g = r * gpr + gi;
825 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
826 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
827 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
828 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
829 }
830 }
831 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
832 crate::prism::inverse_embedding(model, dst);
833 }
834 return;
835 }
836 if *dtype == TensorDtype::Q1 {
837 let bytes = self.quant_bytes();
838 let gpr = cols / GROUP_SIZE;
839 for gi in 0..gpr {
840 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
841 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
842 for (j, &b) in tile[2..].iter().enumerate() {
843 for k in 0..8 {
844 dst[gi * GROUP_SIZE + j * 8 + k] =
845 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
846 }
847 }
848 }
849 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
850 crate::prism::inverse_embedding(model, dst);
851 }
852 return;
853 }
854 if *dtype == TensorDtype::Q1T {
855 let bytes = self.quant_bytes();
856 let gpr = cols / GROUP_SIZE;
857 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
858 for gi in 0..gpr {
859 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
860 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
861 bytes[off],
862 bytes[off + 1],
863 ]));
864 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
865 for k in 0..GROUP_SIZE {
866 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
867 {
868 1 => s,
869 2 => -s,
870 _ => 0.0,
871 };
872 }
873 }
874 let rows = self.rows();
876 let entries = base_len + (rows + 1) * 4;
877 if entries <= bytes.len() {
878 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
879 let r0 = u32::from_le_bytes([
880 ptrs[r * 4],
881 ptrs[r * 4 + 1],
882 ptrs[r * 4 + 2],
883 ptrs[r * 4 + 3],
884 ]) as usize;
885 let r1 = u32::from_le_bytes([
886 ptrs[(r + 1) * 4],
887 ptrs[(r + 1) * 4 + 1],
888 ptrs[(r + 1) * 4 + 2],
889 ptrs[(r + 1) * 4 + 3],
890 ]) as usize;
891 let off = entries + r0 * 4;
892 for i in 0..r1 - r0 {
893 let item = &bytes[off + i * 4..off + i * 4 + 4];
894 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
895 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
896 item[2], item[3],
897 ]));
898 if c < cols {
899 dst[c] = v;
900 }
901 }
902 }
903 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
904 crate::prism::inverse_embedding(model, dst);
905 }
906 return;
907 }
908 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
909 let bytes = self.quant_bytes();
910 let rows = self.rows();
911 let ng = cols / GROUP_SIZE;
912 let bits = &bytes[..rows];
913 let sc_off = rows;
914 let off = vbit_offsets[r];
917 let b = bits[r] as usize;
918 let l = ((1usize << (b - 1)) - 1) as f32;
919 let data = &bytes[off..];
920 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
921 for (i, d) in dst.iter_mut().enumerate() {
922 while nbits < b {
923 acc = (acc << 8) | data[byte_idx] as u64;
924 byte_idx += 1;
925 nbits += 8;
926 }
927 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
928 nbits -= b;
929 let so = (r * ng + i / GROUP_SIZE) * 2;
930 let sv = f16_to_f32(u16::from_le_bytes([
931 bytes[sc_off + so],
932 bytes[sc_off + so + 1],
933 ]));
934 *d = (u - l) * sv;
935 }
936 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
937 crate::prism::inverse_embedding(model, dst);
938 }
939 return;
940 }
941 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
942 let s = row_scale[r];
943 match dtype {
944 TensorDtype::Q8Row => {
945 for (d, &b) in dst.iter_mut().zip(q) {
946 *d = (b as i8) as f32 * s;
947 }
948 }
949 TensorDtype::Q8_2f => {
950 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
951 *d = (b as i8) as f32 * s * col_field[i];
952 }
953 }
954 _ => unreachable!(),
955 }
956 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
957 crate::prism::inverse_embedding(model, dst);
958 }
959 }
960 }
961 }
962
963 pub fn sparse_col_ok(&self) -> bool {
968 match self {
969 Self::F32 { .. } => true,
970 Self::Mapped { dtype, .. } => {
971 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
972 }
973 }
974 }
975
976 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
980 let inter = self.cols();
981 let hidden = self.rows();
982 debug_assert_eq!(out.len(), hidden);
983 match self {
984 Self::F32 { data, .. } => {
985 for (k, o) in out.iter_mut().enumerate() {
986 *o += w * data[k * inter + c];
987 }
988 }
989 Self::Mapped {
990 dtype,
991 row_scale,
992 col_field,
993 ..
994 } => {
995 let q = self.quant_bytes();
996 let colf = if *dtype == TensorDtype::Q8_2f {
997 col_field[c]
998 } else {
999 1.0
1000 };
1001 let wc = w * colf;
1002 for (k, o) in out.iter_mut().enumerate() {
1003 let b = q[k * inter + c] as i8 as f32;
1004 *o += wc * b * row_scale[k];
1005 }
1006 }
1007 }
1008 }
1009
1010 #[inline]
1019 pub fn prefetch_row(&self, r: usize) {
1020 let Self::Mapped { dtype, .. } = self else {
1021 return;
1022 };
1023 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1024 return;
1025 }
1026 let cols = self.cols();
1027 let q = self.quant_bytes();
1028 let (a, b) = (r * cols, (r + 1) * cols);
1029 if b > q.len() {
1030 return;
1031 }
1032 let mut j = a;
1033 while j < b {
1034 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1035 j += 512;
1036 }
1037 }
1038
1039 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1048 let cols = self.cols();
1049 debug_assert_eq!(out.len(), cols);
1050 match self {
1051 Self::F32 { data, .. } => {
1052 let row = &data[r * cols..(r + 1) * cols];
1053 for (o, v) in out.iter_mut().zip(row) {
1054 *o += w * v;
1055 }
1056 }
1057 Self::Mapped {
1058 dtype,
1059 row_scale,
1060 col_field,
1061 ..
1062 } => match dtype {
1063 TensorDtype::Q8Row => {
1064 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1065 let ws = w * row_scale[r];
1066 let row: &[i8] =
1067 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1068 axpy_i8_f32(out, row, ws);
1069 }
1070 TensorDtype::Q8_2f => {
1071 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1072 let ws = w * row_scale[r];
1073 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1074 *o += ws * c * (*b as i8 as f32);
1075 }
1076 }
1077 _ => {
1078 self.row_f32(r, scratch);
1079 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1080 *o += w * v;
1081 }
1082 }
1083 },
1084 }
1085 }
1086
1087 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1091 let cols = self.cols();
1092 match self {
1093 Self::F32 { data, .. } => {
1094 let row = &data[r * cols..(r + 1) * cols];
1095 row.iter().zip(x).map(|(w, v)| w * v).sum()
1096 }
1097 Self::Mapped {
1098 model,
1099 idx,
1100 dtype,
1101 row_scale,
1102 col_field,
1103 ..
1104 } => {
1105 let prism_forward =
1106 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1107 if prism_forward {
1108 let transformed = crate::prism::forward(model, &x[..cols]);
1109 let gpr = cols / GROUP_SIZE;
1110 match dtype {
1111 TensorDtype::Q2TiledP => {
1112 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1113 let mut sc = vec![0f32; gpr];
1114 v.scales_into(r, gpr, &mut sc);
1115 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1116 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1117 }
1118 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1119 }
1120 _ => {
1121 self.row_f32(r, scratch);
1122 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1123 }
1124 }
1125 }
1126 match dtype {
1127 TensorDtype::Q8Row => {
1128 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1129 dot_i8_f32(q, x) * row_scale[r]
1130 }
1131 TensorDtype::Q8_2f => {
1132 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1133 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1134 }
1135 _ => {
1136 self.row_f32(r, scratch);
1137 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1138 }
1139 }
1140 }
1141 }
1142 }
1143
1144 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1147 match self {
1148 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
1152 Self::Mapped {
1153 model,
1154 idx,
1155 dtype,
1156 rows,
1157 cols,
1158 row_scale,
1159 col_field,
1160 vbit_offsets,
1161 repack,
1162 } => {
1163 let _ = (model, idx);
1164 assert!(
1174 out.len() >= *rows && x.len() >= *cols,
1175 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1176 out.len(),
1177 x.len(),
1178 );
1179 let prism_forward =
1180 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1181 if *dtype == TensorDtype::Q2TiledP
1182 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1183 {
1184 use std::sync::atomic::{AtomicUsize, Ordering};
1185 static N: AtomicUsize = AtomicUsize::new(0);
1186 let n = N.fetch_add(1, Ordering::Relaxed);
1187 if n < 128 {
1188 eprintln!(
1189 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1190 model.tensors[*idx].name,
1191 prism_forward,
1192 rows,
1193 cols,
1194 crate::gpu::enabled_here(),
1195 crate::gpu::q2tp_gpu_opt_in(),
1196 crate::gpu::cur_layer(),
1197 );
1198 }
1199 }
1200 if prism_forward {
1205 let transformed = crate::prism::forward(model, &x[..*cols]);
1206 match dtype {
1207 TensorDtype::Q4Block => {
1208 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1209 }
1210 TensorDtype::Q4Tiled => {
1211 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1212 }
1213 TensorDtype::Q4TiledP => {
1214 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1215 }
1216 TensorDtype::Q2TiledP => {
1217 let affine =
1218 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1219 if *rows * *cols >= 8_388_608
1220 && crate::gpu::enabled_here()
1221 && crate::gpu::q2tp_gpu_opt_in()
1222 {
1223 let gpu_ok = if affine {
1224 crate::gpu::q2tp_affine_matvec(
1225 model,
1226 *idx,
1227 &transformed,
1228 *rows,
1229 *cols,
1230 out,
1231 )
1232 } else {
1233 crate::gpu::q2tp_matvec(
1234 model,
1235 *idx,
1236 &transformed,
1237 *rows,
1238 *cols,
1239 out,
1240 )
1241 };
1242 if gpu_ok {
1243 return;
1244 }
1245 }
1246 if affine {
1247 q2tp_affine_matvec(
1248 self.quant_bytes(),
1249 &transformed,
1250 *rows,
1251 *cols,
1252 out,
1253 pool,
1254 )
1255 } else {
1256 q2tp_matvec(
1257 self.quant_bytes(),
1258 &transformed,
1259 *rows,
1260 *cols,
1261 out,
1262 pool,
1263 )
1264 }
1265 }
1266 TensorDtype::Q1 => {
1267 q1_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1268 }
1269 TensorDtype::Q1T => {
1270 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1271 }
1272 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1273 self.quant_bytes(),
1274 vbit_offsets,
1275 &transformed,
1276 *rows,
1277 *cols,
1278 out,
1279 pool,
1280 ),
1281 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1282 self.quant_bytes(),
1283 repack,
1284 row_scale,
1285 &transformed,
1286 col_field,
1287 *dtype,
1288 *rows,
1289 *cols,
1290 out,
1291 pool,
1292 ),
1293 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1294 }
1295 return;
1296 }
1297 if *dtype == TensorDtype::Q4Block {
1298 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1302 let t0 = std::time::Instant::now();
1303 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1304 crate::gpu::ProbeArm::Gpu => {
1305 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1306 crate::gpu::probe_record(
1307 crate::gpu::OpClass::Matvec,
1308 true,
1309 t0.elapsed(),
1310 );
1311 return;
1312 }
1313 }
1314 crate::gpu::ProbeArm::CpuTimed => {
1315 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1316 crate::gpu::probe_record(
1317 crate::gpu::OpClass::Matvec,
1318 false,
1319 t0.elapsed(),
1320 );
1321 return;
1322 }
1323 crate::gpu::ProbeArm::Cpu => {}
1324 }
1325 }
1326 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1327 return;
1328 }
1329 if *dtype == TensorDtype::Q4Tiled {
1330 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1335 let t0 = std::time::Instant::now();
1336 let cls = crate::gpu::matvec_class(*rows, *cols);
1337 match crate::gpu::probe_arm(cls) {
1338 crate::gpu::ProbeArm::Gpu => {
1339 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1340 crate::gpu::probe_record(cls, true, t0.elapsed());
1341 return;
1342 }
1343 }
1344 crate::gpu::ProbeArm::CpuTimed => {
1345 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1346 crate::gpu::probe_record(cls, false, t0.elapsed());
1347 return;
1348 }
1349 crate::gpu::ProbeArm::Cpu => {}
1350 }
1351 }
1352 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1353 return;
1354 }
1355 if *dtype == TensorDtype::Q4TiledP {
1356 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1362 let t0 = std::time::Instant::now();
1363 let cls = crate::gpu::matvec_class(*rows, *cols);
1364 match crate::gpu::probe_arm(cls) {
1365 crate::gpu::ProbeArm::Gpu => {
1366 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1367 crate::gpu::probe_record(cls, true, t0.elapsed());
1368 return;
1369 }
1370 }
1371 crate::gpu::ProbeArm::CpuTimed => {
1372 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1373 crate::gpu::probe_record(cls, false, t0.elapsed());
1374 return;
1375 }
1376 crate::gpu::ProbeArm::Cpu => {}
1377 }
1378 }
1379 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1380 return;
1381 }
1382 if *dtype == TensorDtype::Q2TiledP {
1383 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1384 return;
1385 }
1386 if *dtype == TensorDtype::Q1 {
1387 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1392 let t0 = std::time::Instant::now();
1393 let arm = if crate::gpu::q1_force() {
1394 crate::gpu::ProbeArm::Gpu
1395 } else {
1396 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1397 };
1398 match arm {
1399 crate::gpu::ProbeArm::Gpu => {
1400 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1401 crate::gpu::probe_record(
1402 crate::gpu::OpClass::Matvec,
1403 true,
1404 t0.elapsed(),
1405 );
1406 return;
1407 }
1408 }
1409 crate::gpu::ProbeArm::CpuTimed => {
1410 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1411 crate::gpu::probe_record(
1412 crate::gpu::OpClass::Matvec,
1413 false,
1414 t0.elapsed(),
1415 );
1416 return;
1417 }
1418 crate::gpu::ProbeArm::Cpu => {}
1419 }
1420 }
1421 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1422 return;
1423 }
1424 if *dtype == TensorDtype::Q1T {
1425 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1429 let t0 = std::time::Instant::now();
1430 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1431 crate::gpu::ProbeArm::Gpu => {
1432 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1433 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1434 crate::gpu::probe_record(
1435 crate::gpu::OpClass::Matvec,
1436 true,
1437 t0.elapsed(),
1438 );
1439 return;
1440 }
1441 }
1442 crate::gpu::ProbeArm::CpuTimed => {
1443 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1444 crate::gpu::probe_record(
1445 crate::gpu::OpClass::Matvec,
1446 false,
1447 t0.elapsed(),
1448 );
1449 return;
1450 }
1451 crate::gpu::ProbeArm::Cpu => {}
1452 }
1453 }
1454 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1455 return;
1456 }
1457 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1458 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1459 return;
1460 }
1461 let xs = prescale(x, col_field, *dtype);
1462 if *rows >= crate::gpu::min_rows()
1467 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1468 && gpu_lmhead_enabled()
1469 && crate::gpu::enabled_here()
1470 {
1471 let t0 = std::time::Instant::now();
1474 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1475 crate::gpu::ProbeArm::Gpu => {}
1476 crate::gpu::ProbeArm::CpuTimed => {
1477 qmatvec(
1478 self.quant_bytes(),
1479 repack,
1480 row_scale,
1481 x,
1482 col_field,
1483 *dtype,
1484 *rows,
1485 *cols,
1486 out,
1487 pool,
1488 );
1489 crate::gpu::probe_record(
1490 crate::gpu::OpClass::Matvec,
1491 false,
1492 t0.elapsed(),
1493 );
1494 return;
1495 }
1496 crate::gpu::ProbeArm::Cpu => {
1497 qmatvec(
1498 self.quant_bytes(),
1499 repack,
1500 row_scale,
1501 x,
1502 col_field,
1503 *dtype,
1504 *rows,
1505 *cols,
1506 out,
1507 pool,
1508 );
1509 return;
1510 }
1511 }
1512 let frac = gpu_split_frac();
1513 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1514 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1515 let bytes = self.quant_bytes();
1516 let ok = std::thread::scope(|sc| {
1517 let g = sc.spawn(|| {
1518 crate::gpu::q8_matvec_range(
1519 model,
1520 *idx,
1521 cpu_rows,
1522 &row_scale[cpu_rows..],
1523 &xs,
1524 *rows - cpu_rows,
1525 *cols,
1526 out_gpu,
1527 )
1528 });
1529 if cpu_rows > 0 {
1530 let rep_cpu = if repack.is_empty() {
1533 &[][..]
1534 } else {
1535 &repack[..(cpu_rows / 4) * 4 * *cols]
1536 };
1537 qmatvec(
1538 &bytes[..cpu_rows * *cols],
1539 rep_cpu,
1540 &row_scale[..cpu_rows],
1541 x,
1542 col_field,
1543 *dtype,
1544 cpu_rows,
1545 *cols,
1546 out_cpu,
1547 pool,
1548 );
1549 }
1550 g.join().unwrap_or(false)
1551 });
1552 if ok {
1553 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1554 return;
1555 }
1556 qmatvec(
1559 &bytes[cpu_rows * *cols..(*rows) * *cols],
1560 &[],
1561 &row_scale[cpu_rows..],
1562 x,
1563 col_field,
1564 *dtype,
1565 *rows - cpu_rows,
1566 *cols,
1567 out_gpu,
1568 pool,
1569 );
1570 return;
1571 }
1572 qmatvec(
1573 self.quant_bytes(),
1574 repack,
1575 row_scale,
1576 x,
1577 col_field,
1578 *dtype,
1579 *rows,
1580 *cols,
1581 out,
1582 pool,
1583 );
1584 }
1585 }
1586 }
1587
1588 pub fn matvec2(
1590 &self,
1591 x1: &[f32],
1592 x2: &[f32],
1593 o1: &mut [f32],
1594 o2: &mut [f32],
1595 pool: Option<&Pool>,
1596 ) {
1597 match self {
1598 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1599 Self::Mapped {
1600 model,
1601 idx,
1602 dtype,
1603 rows,
1604 cols,
1605 row_scale,
1606 col_field,
1607 vbit_offsets,
1608 ..
1609 } => {
1610 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1611 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1612 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1613 match dtype {
1614 TensorDtype::Q4Block => {
1615 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1616 }
1617 TensorDtype::Q4Tiled => {
1618 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1619 }
1620 TensorDtype::Q4TiledP => {
1621 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1622 }
1623 TensorDtype::Q2TiledP => {
1624 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1625 q2tp_affine_matvec2(
1626 self.quant_bytes(),
1627 &tx1,
1628 &tx2,
1629 *rows,
1630 *cols,
1631 o1,
1632 o2,
1633 pool,
1634 )
1635 } else {
1636 q2tp_matvec2(
1637 self.quant_bytes(),
1638 &tx1,
1639 &tx2,
1640 *rows,
1641 *cols,
1642 o1,
1643 o2,
1644 pool,
1645 )
1646 }
1647 }
1648 TensorDtype::Q1 => {
1649 q1_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1650 }
1651 TensorDtype::Q1T => {
1652 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1653 }
1654 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1655 self.quant_bytes(),
1656 vbit_offsets,
1657 &tx1,
1658 &tx2,
1659 *rows,
1660 *cols,
1661 o1,
1662 o2,
1663 pool,
1664 ),
1665 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1666 self.quant_bytes(),
1667 row_scale,
1668 &tx1,
1669 &tx2,
1670 col_field,
1671 *dtype,
1672 *rows,
1673 *cols,
1674 o1,
1675 o2,
1676 pool,
1677 ),
1678 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1679 }
1680 return;
1681 }
1682 if *dtype == TensorDtype::Q4Block {
1683 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1684 return;
1685 }
1686 if *dtype == TensorDtype::Q4Tiled {
1687 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1688 return;
1689 }
1690 if *dtype == TensorDtype::Q4TiledP {
1691 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1692 return;
1693 }
1694 if *dtype == TensorDtype::Q2TiledP {
1695 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1696 return;
1697 }
1698 if *dtype == TensorDtype::Q1 {
1699 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1700 return;
1701 }
1702 if *dtype == TensorDtype::Q1T {
1703 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1709 return;
1710 }
1711 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1712 vbitmatvec2(
1713 self.quant_bytes(),
1714 vbit_offsets,
1715 x1,
1716 x2,
1717 *rows,
1718 *cols,
1719 o1,
1720 o2,
1721 pool,
1722 );
1723 return;
1724 }
1725 qmatvec2(
1726 self.quant_bytes(),
1727 row_scale,
1728 x1,
1729 x2,
1730 col_field,
1731 *dtype,
1732 *rows,
1733 *cols,
1734 o1,
1735 o2,
1736 pool,
1737 );
1738 }
1739 }
1740 }
1741}
1742
1743impl QTensor {
1744 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1752 if self.has_prism_contract() {
1753 return None;
1754 }
1755 match self {
1756 Self::Mapped {
1757 model, idx, dtype, ..
1758 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1759 _ => None,
1760 }
1761 }
1762
1763 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1764 let cols = self.cols();
1765 let rows = self.rows();
1766 debug_assert_eq!(xs_all.len(), b * cols);
1767 debug_assert_eq!(out.len(), b * rows);
1768 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1769 if crate::gptq_capture::capturing() {
1773 if let Self::Mapped { model, idx, .. } = self {
1774 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1775 }
1776 }
1777 match self {
1778 Self::F32 { data, .. } => {
1779 let out_addr = SendMut(out.as_mut_ptr());
1780 let run = |start: usize, end: usize| {
1781 for o in start..end {
1782 let row = &data[o * cols..(o + 1) * cols];
1783 for bi in 0..b {
1784 let x = &xs_all[bi * cols..(bi + 1) * cols];
1785 let mut acc = 0f32;
1786 for j in 0..cols {
1787 acc += row[j] * x[j];
1788 }
1789 unsafe { *out_addr.at(bi * rows + o) = acc };
1790 }
1791 }
1792 };
1793 dispatch_rows(pool, rows, &run);
1794 }
1795 Self::Mapped {
1796 model,
1797 idx,
1798 dtype,
1799 row_scale,
1800 col_field,
1801 vbit_offsets,
1802 ..
1803 } => {
1804 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1805 let mut transformed = Vec::with_capacity(xs_all.len());
1806 for bi in 0..b {
1807 transformed.extend_from_slice(&crate::prism::forward(
1808 model,
1809 &xs_all[bi * cols..(bi + 1) * cols],
1810 ));
1811 }
1812 match dtype {
1813 TensorDtype::Q4Block => {
1814 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1815 }
1816 TensorDtype::Q4Tiled => {
1817 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1818 }
1819 TensorDtype::Q4TiledP => {
1820 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1821 }
1822 TensorDtype::Q2TiledP => {
1823 let affine =
1824 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1825 let gpu_batch_ok = if affine {
1831 b >= 2
1832 } else {
1833 b >= 32 && b * rows * cols >= 128_000_000
1834 };
1835 if gpu_batch_ok
1836 && cols % 32 == 0
1837 && crate::gpu::enabled_here()
1838 && crate::gpu::q2tp_gpu_opt_in()
1839 {
1840 let gpu_ok = if affine {
1841 crate::gpu::q2tp_affine_matmat(
1842 model,
1843 *idx,
1844 &transformed,
1845 b,
1846 rows,
1847 cols,
1848 out,
1849 )
1850 } else {
1851 crate::gpu::q2tp_matmat(
1852 model,
1853 *idx,
1854 &transformed,
1855 b,
1856 rows,
1857 cols,
1858 out,
1859 )
1860 };
1861 if gpu_ok {
1862 return;
1863 }
1864 }
1865 if affine
1869 && b == 1
1870 && cols % 32 == 0
1871 && crate::gpu::enabled_here()
1872 && crate::gpu::q2tp_gpu_opt_in()
1873 && crate::gpu::q2tp_affine_matvec(
1874 model,
1875 *idx,
1876 &transformed[..cols],
1877 rows,
1878 cols,
1879 &mut out[..rows],
1880 )
1881 {
1882 return;
1883 }
1884 if affine {
1885 q2tp_affine_matmat(
1886 self.quant_bytes(),
1887 &transformed,
1888 b,
1889 rows,
1890 cols,
1891 out,
1892 pool,
1893 )
1894 } else {
1895 q2tp_matmat(
1896 self.quant_bytes(),
1897 &transformed,
1898 b,
1899 rows,
1900 cols,
1901 out,
1902 pool,
1903 )
1904 }
1905 }
1906 TensorDtype::Q1 => {
1907 q1_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1908 }
1909 TensorDtype::Q1T => {
1910 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1911 }
1912 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
1913 self.quant_bytes(),
1914 vbit_offsets,
1915 &transformed,
1916 b,
1917 rows,
1918 cols,
1919 out,
1920 pool,
1921 ),
1922 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1923 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1924 .map(|bi| {
1925 prescale(
1926 &transformed[bi * cols..(bi + 1) * cols],
1927 col_field,
1928 *dtype,
1929 )
1930 })
1931 .collect();
1932 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
1933 }
1934 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1935 }
1936 return;
1937 }
1938 if *dtype == TensorDtype::Q4Block {
1939 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1940 return;
1941 }
1942 if *dtype == TensorDtype::Q4TiledP {
1943 if b >= 32
1955 && b * rows * cols >= 128_000_000
1956 && cols % 32 == 0
1957 && !crate::gpu::mm_killed()
1958 && crate::gpu::enabled_here()
1959 {
1960 let class = if b >= 128 {
1961 crate::gpu::OpClass::MatmatWide
1962 } else {
1963 crate::gpu::OpClass::Matmat
1964 };
1965 if let Self::Mapped { model, idx, .. } = self {
1966 if crate::mm_ab::on() {
1978 let mut g = vec![0f32; b * rows];
1979 let t = std::time::Instant::now();
1980 let took = crate::gpu::q4tp_matmat(
1981 model, *idx, xs_all, b, rows, cols, &mut g,
1982 );
1983 let dg = t.elapsed();
1984 let t = std::time::Instant::now();
1985 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1986 let dc = t.elapsed();
1987 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1988 return;
1989 }
1990 let t0 = std::time::Instant::now();
1991 let resident = crate::gpu::weight_is_resident(model, *idx);
1995 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1996 crate::gpu::ProbeArm::Gpu => {
1997 if crate::gpu::q4tp_matmat(
1998 model, *idx, xs_all, b, rows, cols, out,
1999 ) {
2000 let el = t0.elapsed();
2001 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2011 let budget = std::time::Duration::from_secs_f64(
2012 flops / 1.5e12 * 8.0 + 0.020,
2013 );
2014 crate::gpu::mm_budget_check(
2015 "q4tp matmat",
2016 el,
2017 budget,
2018 crate::gpu::probe_was_cold() || !resident,
2019 );
2020 crate::gpu::probe_record(class, true, el);
2021 return;
2022 }
2023 }
2024 crate::gpu::ProbeArm::CpuTimed => {
2025 q4tp_matmat(
2026 self.quant_bytes(),
2027 xs_all,
2028 b,
2029 rows,
2030 cols,
2031 out,
2032 pool,
2033 );
2034 crate::gpu::probe_record(class, false, t0.elapsed());
2035 return;
2036 }
2037 crate::gpu::ProbeArm::Cpu => {}
2038 }
2039 }
2040 }
2041 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2042 return;
2043 }
2044 if *dtype == TensorDtype::Q2TiledP {
2045 if b >= 32
2051 && b * rows * cols >= 128_000_000
2052 && cols % 32 == 0
2053 && !crate::gpu::mm_killed()
2054 && crate::gpu::enabled_here()
2055 {
2056 let class = if b >= 128 {
2057 crate::gpu::OpClass::MatmatWide
2058 } else {
2059 crate::gpu::OpClass::Matmat
2060 };
2061 if let Self::Mapped { model, idx, .. } = self {
2062 let t0 = std::time::Instant::now();
2063 match crate::gpu::probe_arm(class) {
2064 crate::gpu::ProbeArm::Gpu => {
2065 if crate::gpu::q2tp_matmat(
2066 model, *idx, xs_all, b, rows, cols, out,
2067 ) {
2068 crate::gpu::probe_record(class, true, t0.elapsed());
2069 return;
2070 }
2071 }
2072 crate::gpu::ProbeArm::CpuTimed => {
2073 q2tp_matmat(
2074 self.quant_bytes(),
2075 xs_all,
2076 b,
2077 rows,
2078 cols,
2079 out,
2080 pool,
2081 );
2082 crate::gpu::probe_record(class, false, t0.elapsed());
2083 return;
2084 }
2085 crate::gpu::ProbeArm::Cpu => {}
2086 }
2087 }
2088 }
2089 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2094 return;
2095 }
2096 if *dtype == TensorDtype::Q4Tiled {
2097 if b >= 32
2109 && b * rows * cols >= 128_000_000
2110 && cols % 32 == 0
2111 && !crate::gpu::mm_killed()
2112 && crate::gpu::enabled_here()
2113 {
2114 let class = if b >= 128 {
2115 crate::gpu::OpClass::MatmatWide
2116 } else {
2117 crate::gpu::OpClass::Matmat
2118 };
2119 if let Self::Mapped { model, idx, .. } = self {
2120 let t0 = std::time::Instant::now();
2121 match crate::gpu::probe_arm(class) {
2122 crate::gpu::ProbeArm::Gpu => {
2123 if crate::gpu::q4t_matmat(
2124 model, *idx, xs_all, b, rows, cols, out,
2125 ) {
2126 let el = t0.elapsed();
2127 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2137 let budget = std::time::Duration::from_secs_f64(
2138 flops / 1.5e12 * 8.0 + 0.020,
2139 );
2140 crate::gpu::mm_budget_check(
2141 "q4t matmat",
2142 el,
2143 budget,
2144 crate::gpu::probe_was_cold(),
2145 );
2146 crate::gpu::probe_record(class, true, el);
2147 return;
2148 }
2149 }
2150 crate::gpu::ProbeArm::CpuTimed => {
2151 q4t_matmat(
2152 self.quant_bytes(),
2153 xs_all,
2154 b,
2155 rows,
2156 cols,
2157 out,
2158 pool,
2159 );
2160 crate::gpu::probe_record(class, false, t0.elapsed());
2161 return;
2162 }
2163 crate::gpu::ProbeArm::Cpu => {}
2164 }
2165 }
2166 }
2167 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2168 return;
2169 }
2170 if *dtype == TensorDtype::Q1 {
2171 if b >= 32
2174 && b * rows * cols >= 128_000_000
2175 && cols % 64 == 0
2176 && crate::gpu::enabled_here()
2177 {
2178 if let Self::Mapped { model, idx, .. } = self {
2179 let t0 = std::time::Instant::now();
2180 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2181 crate::gpu::ProbeArm::Gpu => {
2182 if crate::gpu::q1_matmat(
2183 model, *idx, xs_all, b, rows, cols, out,
2184 ) {
2185 crate::gpu::probe_record(
2186 crate::gpu::OpClass::Matmat,
2187 true,
2188 t0.elapsed(),
2189 );
2190 return;
2191 }
2192 }
2193 crate::gpu::ProbeArm::CpuTimed => {
2194 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2195 crate::gpu::probe_record(
2196 crate::gpu::OpClass::Matmat,
2197 false,
2198 t0.elapsed(),
2199 );
2200 return;
2201 }
2202 crate::gpu::ProbeArm::Cpu => {}
2203 }
2204 }
2205 }
2206 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2207 return;
2208 }
2209 if *dtype == TensorDtype::Q1T {
2210 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2213 if let Self::Mapped { model, idx, .. } = self {
2214 let t0 = std::time::Instant::now();
2215 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2216 crate::gpu::ProbeArm::Gpu => {
2217 if crate::gpu::q1t_matmat(
2218 model, *idx, xs_all, b, rows, cols, out,
2219 ) {
2220 crate::gpu::probe_record(
2221 crate::gpu::OpClass::Matmat,
2222 true,
2223 t0.elapsed(),
2224 );
2225 return;
2226 }
2227 }
2228 crate::gpu::ProbeArm::CpuTimed => {
2229 q1t_matmat(
2230 self.quant_bytes(),
2231 xs_all,
2232 b,
2233 rows,
2234 cols,
2235 out,
2236 pool,
2237 );
2238 crate::gpu::probe_record(
2239 crate::gpu::OpClass::Matmat,
2240 false,
2241 t0.elapsed(),
2242 );
2243 return;
2244 }
2245 crate::gpu::ProbeArm::Cpu => {}
2246 }
2247 }
2248 }
2249 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2250 return;
2251 }
2252 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2253 vbitmatmat(
2254 self.quant_bytes(),
2255 vbit_offsets,
2256 xs_all,
2257 b,
2258 rows,
2259 cols,
2260 out,
2261 pool,
2262 );
2263 return;
2264 }
2265 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2266 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2267 .collect();
2268 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2274 if let Self::Mapped { model, idx, .. } = self {
2275 let t0 = std::time::Instant::now();
2276 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2277 crate::gpu::ProbeArm::Gpu
2278 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2279 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2280 {
2281 let q = self.quant_bytes();
2285 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2286 return;
2287 }
2288 crate::gpu::ProbeArm::Gpu => {
2289 let flat: Vec<f32> =
2290 pre.iter().flat_map(|v| v.iter().copied()).collect();
2291 if crate::gpu::q8_matmat(
2292 model, *idx, row_scale, &flat, b, rows, cols, out,
2293 ) {
2294 crate::gpu::probe_record(
2295 crate::gpu::OpClass::Matmat,
2296 true,
2297 t0.elapsed(),
2298 );
2299 return;
2300 }
2301 }
2302 crate::gpu::ProbeArm::CpuTimed => {
2303 let q = self.quant_bytes();
2304 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2305 crate::gpu::probe_record(
2306 crate::gpu::OpClass::Matmat,
2307 false,
2308 t0.elapsed(),
2309 );
2310 return;
2311 }
2312 crate::gpu::ProbeArm::Cpu => {}
2313 }
2314 }
2315 }
2316 let q = self.quant_bytes();
2317 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2318 }
2319 }
2320 }
2321}
2322
2323impl QTensor {
2324 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2334 let (rows, cols) = (self.rows(), self.cols());
2335 let Self::Mapped {
2336 model,
2337 idx,
2338 dtype,
2339 row_scale,
2340 col_field,
2341 ..
2342 } = self
2343 else {
2344 return false;
2345 };
2346 if crate::prism::has_contract(model) {
2347 return false;
2348 }
2349 match *dtype {
2350 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2351 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2355 if *dtype == TensorDtype::Q8_2f
2358 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2359 && crate::gpu::q8_matmat_2f(
2360 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2361 )
2362 {
2363 return true;
2364 }
2365 let flat: Vec<f32> = (0..b)
2366 .flat_map(|bi| {
2367 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2368 })
2369 .collect();
2370 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2371 }
2372 _ => false,
2373 }
2374 }
2375
2376 pub fn matvec_many<const N: usize>(
2383 ts: [&QTensor; N],
2384 x: &[f32],
2385 mut outs: [&mut [f32]; N],
2386 pool: Option<&Pool>,
2387 ) {
2388 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2389 if ts.iter().any(|t| t.has_prism_contract()) {
2390 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2394 t.matvec(x, o, pool);
2395 }
2396 return;
2397 }
2398 let uniform_q8 = ts.iter().all(|t| {
2399 matches!(
2400 t,
2401 Self::Mapped {
2402 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2403 ..
2404 }
2405 )
2406 });
2407 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2408 let uniform_q4 = ts.iter().all(|t| {
2409 matches!(
2410 t,
2411 Self::Mapped {
2412 dtype: TensorDtype::Q4Block,
2413 ..
2414 }
2415 )
2416 });
2417 let uniform_vbit = ts.iter().all(|t| {
2418 matches!(
2419 t,
2420 Self::Mapped {
2421 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2422 ..
2423 }
2424 )
2425 });
2426 let uniform_q1 = ts.iter().all(|t| {
2427 matches!(
2428 t,
2429 Self::Mapped {
2430 dtype: TensorDtype::Q1,
2431 ..
2432 }
2433 )
2434 });
2435 let uniform_q1t = ts.iter().all(|t| {
2436 matches!(
2437 t,
2438 Self::Mapped {
2439 dtype: TensorDtype::Q1T,
2440 ..
2441 }
2442 )
2443 });
2444 let uniform_q4tp = ts.iter().all(|t| {
2449 matches!(
2450 t,
2451 Self::Mapped {
2452 dtype: TensorDtype::Q4TiledP,
2453 ..
2454 }
2455 )
2456 }) && ts
2457 .iter()
2458 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2459 let Some(pool) = pool else {
2460 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2461 t.matvec(x, o, None);
2462 }
2463 return;
2464 };
2465 if total_rows < 256
2466 || !(uniform_q8
2467 || uniform_f32
2468 || uniform_q4
2469 || uniform_vbit
2470 || uniform_q1
2471 || uniform_q1t
2472 || uniform_q4tp)
2473 {
2474 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2475 t.matvec(x, o, Some(pool));
2476 }
2477 return;
2478 }
2479
2480 if uniform_q4tp {
2481 let cols = ts[0].cols();
2487 let gpr = cols / GROUP_SIZE;
2488 let views: Vec<Q4tpView> = ts
2489 .iter()
2490 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2491 .collect();
2492 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2493 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2494 let locate = |flat: usize| -> (usize, usize) {
2496 let mut acc = 0;
2497 for (i, &r) in rows_of.iter().enumerate() {
2498 if flat < acc + r {
2499 return (i, flat - acc);
2500 }
2501 acc += r;
2502 }
2503 (rows_of.len() - 1, 0)
2504 };
2505 let (views, outs_addr) = (&views, &outs_addr);
2506 if a8w8_enabled() {
2507 let act = split_act(x);
2508 let act = &act;
2509 let run = |start: usize, end: usize| {
2510 let mut sc = vec![0f32; gpr];
2511 for flat in start..end {
2512 let (t, r) = locate(flat);
2513 let v = &views[t];
2514 v.scales_into(r, gpr, &mut sc);
2515 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2516 for &(j, xv) in &act.outliers {
2517 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2518 acc += w * s * xv;
2519 }
2520 unsafe { *outs_addr[t].at(r) = acc };
2522 }
2523 };
2524 pool.run_rows(total_rows, &run);
2525 } else {
2526 let run = |start: usize, end: usize| {
2527 let mut sc = vec![0f32; gpr];
2528 for flat in start..end {
2529 let (t, r) = locate(flat);
2530 let v = &views[t];
2531 v.scales_into(r, gpr, &mut sc);
2532 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2534 }
2535 };
2536 pool.run_rows(total_rows, &run);
2537 }
2538 return;
2539 }
2540
2541 if uniform_q1 {
2542 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2545 if a8w8_enabled() {
2546 let act = split_act(x);
2547 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2548 let (act, gsum) = (&act, &gsum);
2549 let closures: [_; N] = std::array::from_fn(|i| {
2550 let (bytes, gpr, out) =
2551 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2552 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2553 });
2554 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2555 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2556 pool.run_many(&parts);
2557 } else {
2558 let closures: [_; N] = std::array::from_fn(|i| {
2559 let (bytes, gpr, out) =
2560 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2561 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2562 });
2563 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2564 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2565 pool.run_many(&parts);
2566 }
2567 return;
2568 }
2569
2570 if uniform_q1t {
2571 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2575 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2576 if a8w8_enabled() {
2577 let act = split_act(x);
2578 let act = &act;
2579 let x_ref = x;
2580 let closures: [_; N] = std::array::from_fn(|i| {
2581 let bytes = ts[i].quant_bytes();
2582 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2583 let gpr = cols / GROUP_SIZE;
2584 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2585 let out = outs_addr[i];
2586 move |s: usize, e: usize| {
2587 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2588 }
2589 });
2590 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2591 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2592 pool.run_many(&parts);
2593 } else {
2594 let x_ref = x;
2595 let closures: [_; N] = std::array::from_fn(|i| {
2596 let bytes = ts[i].quant_bytes();
2597 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2598 let gpr = cols / GROUP_SIZE;
2599 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2600 let out = outs_addr[i];
2601 move |s: usize, e: usize| {
2602 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2603 }
2604 });
2605 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2606 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2607 pool.run_many(&parts);
2608 }
2609 return;
2610 }
2611
2612 if uniform_q4 || uniform_vbit {
2613 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2614 if a8w8_enabled() {
2616 let act = split_act(x);
2617 let act = &act;
2618 if uniform_q4 {
2619 let closures: [_; N] = std::array::from_fn(|i| {
2620 let (packed, scales) =
2621 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2622 let (gpr, cols, out) =
2623 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2624 move |s: usize, e: usize| {
2625 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2626 }
2627 });
2628 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2629 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2630 pool.run_many(&parts);
2631 } else {
2632 let closures: [_; N] = std::array::from_fn(|i| {
2633 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2634 unreachable!()
2635 };
2636 let (bytes, rows, cols, out) = (
2637 ts[i].quant_bytes(),
2638 ts[i].rows(),
2639 ts[i].cols(),
2640 outs_addr[i],
2641 );
2642 move |s: usize, e: usize| {
2643 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2644 }
2645 });
2646 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2647 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2648 pool.run_many(&parts);
2649 }
2650 return;
2651 }
2652 if uniform_q4 {
2653 let closures: [_; N] = std::array::from_fn(|i| {
2654 let (packed, scales) =
2655 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2656 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2657 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2658 });
2659 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2660 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2661 pool.run_many(&parts);
2662 } else {
2663 let closures: [_; N] = std::array::from_fn(|i| {
2664 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2665 unreachable!()
2666 };
2667 let (bytes, rows, cols, out) = (
2668 ts[i].quant_bytes(),
2669 ts[i].rows(),
2670 ts[i].cols(),
2671 outs_addr[i],
2672 );
2673 move |s: usize, e: usize| {
2674 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2675 }
2676 });
2677 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2678 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2679 pool.run_many(&parts);
2680 }
2681 return;
2682 }
2683
2684 if uniform_f32 {
2685 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2686 let closures: [_; N] = std::array::from_fn(|i| {
2687 let Self::F32 { data, cols, .. } = ts[i] else {
2688 unreachable!()
2689 };
2690 let out = outs_addr[i];
2691 move |start: usize, end: usize| {
2692 for o in start..end {
2693 let row = &data[o * cols..(o + 1) * cols];
2694 let mut sum = 0.0f32;
2695 for j in 0..*cols {
2696 sum += row[j] * x[j];
2697 }
2698 unsafe { *out.at(o) = sum };
2700 }
2701 }
2702 });
2703 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2704 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2705 pool.run_many(&parts);
2706 return;
2707 }
2708
2709 struct Ctx<'a> {
2712 bytes: &'a [u8],
2713 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2714 rep: &'a [u8],
2715 row_scale: &'a [f32],
2716 cols: usize,
2717 xs: std::borrow::Cow<'a, [f32]>,
2718 }
2719 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2720 let Self::Mapped {
2721 dtype,
2722 cols,
2723 row_scale,
2724 col_field,
2725 repack,
2726 ..
2727 } = ts[i]
2728 else {
2729 unreachable!()
2730 };
2731 Ctx {
2732 bytes: ts[i].quant_bytes(),
2733 rep: repack,
2734 row_scale,
2735 cols: *cols,
2736 xs: prescale(x, col_field, *dtype),
2737 }
2738 });
2739 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2740 #[cfg(target_arch = "aarch64")]
2741 if sdot_enabled() {
2742 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2743 let closures: [_; N] = std::array::from_fn(|i| {
2744 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2745 move |start: usize, end: usize| {
2746 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2747 }
2748 });
2749 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2750 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2751 pool.run_many(&parts);
2752 return;
2753 }
2754 #[cfg(target_arch = "x86_64")]
2755 if avx2_a8w8_enabled() {
2756 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2757 let closures: [_; N] = std::array::from_fn(|i| {
2758 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2759 move |start: usize, end: usize| {
2760 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2761 }
2762 });
2763 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2764 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2765 pool.run_many(&parts);
2766 return;
2767 }
2768 let closures: [_; N] = std::array::from_fn(|i| {
2769 let (c, out) = (&ctxs[i], outs_addr[i]);
2770 move |start: usize, end: usize| {
2771 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2772 }
2773 });
2774 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2775 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2776 pool.run_many(&parts);
2777 }
2778}
2779
2780impl QTensor {
2781 #[allow(clippy::needless_range_loop)]
2786 pub fn matvec2_many<const N: usize>(
2787 ts: [&QTensor; N],
2788 x1: &[f32],
2789 x2: &[f32],
2790 mut o1s: [&mut [f32]; N],
2791 mut o2s: [&mut [f32]; N],
2792 pool: Option<&Pool>,
2793 ) {
2794 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2795 if ts.iter().any(|t| t.has_prism_contract()) {
2796 for i in 0..N {
2797 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2798 }
2799 return;
2800 }
2801 let uniform_q8 = ts.iter().all(|t| {
2802 matches!(
2803 t,
2804 Self::Mapped {
2805 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2806 ..
2807 }
2808 )
2809 });
2810 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2811 let uniform_q4 = ts.iter().all(|t| {
2812 matches!(
2813 t,
2814 Self::Mapped {
2815 dtype: TensorDtype::Q4Block,
2816 ..
2817 }
2818 )
2819 });
2820 let uniform_vbit = ts.iter().all(|t| {
2821 matches!(
2822 t,
2823 Self::Mapped {
2824 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2825 ..
2826 }
2827 )
2828 });
2829 let fusable = pool.is_some()
2830 && total_rows >= 256
2831 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2832 if !fusable {
2833 for i in 0..N {
2834 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2835 }
2836 return;
2837 }
2838 let pool = pool.unwrap();
2839
2840 if uniform_q4 || uniform_vbit {
2841 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2842 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2843 if a8w8_enabled() {
2845 let a1 = split_act(x1);
2846 let a2 = split_act(x2);
2847 let (a1, a2) = (&a1, &a2);
2848 if uniform_q4 {
2849 let closures: [_; N] = std::array::from_fn(|i| {
2850 let (packed, scales) =
2851 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2852 let (gpr, cols, o1, o2) =
2853 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2854 move |s: usize, e: usize| {
2855 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2856 }
2857 });
2858 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2859 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2860 pool.run_many(&parts);
2861 } else {
2862 let closures: [_; N] = std::array::from_fn(|i| {
2863 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2864 unreachable!()
2865 };
2866 let (bytes, rows, cols, o1, o2) = (
2867 ts[i].quant_bytes(),
2868 ts[i].rows(),
2869 ts[i].cols(),
2870 p1[i],
2871 p2[i],
2872 );
2873 move |s: usize, e: usize| {
2874 vbit_range2_a8w8(
2875 bytes,
2876 vbit_offsets,
2877 x1,
2878 x2,
2879 a1,
2880 a2,
2881 rows,
2882 cols,
2883 o1,
2884 o2,
2885 s,
2886 e,
2887 )
2888 }
2889 });
2890 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2891 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2892 pool.run_many(&parts);
2893 }
2894 return;
2895 }
2896 if uniform_q4 {
2897 let closures: [_; N] = std::array::from_fn(|i| {
2898 let (packed, scales) =
2899 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2900 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2901 move |s: usize, e: usize| {
2902 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2903 }
2904 });
2905 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2906 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2907 pool.run_many(&parts);
2908 } else {
2909 let closures: [_; N] = std::array::from_fn(|i| {
2910 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2911 unreachable!()
2912 };
2913 let (bytes, rows, cols, o1, o2) = (
2914 ts[i].quant_bytes(),
2915 ts[i].rows(),
2916 ts[i].cols(),
2917 p1[i],
2918 p2[i],
2919 );
2920 move |s: usize, e: usize| {
2921 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2922 }
2923 });
2924 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2925 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2926 pool.run_many(&parts);
2927 }
2928 return;
2929 }
2930
2931 if uniform_f32 {
2932 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2933 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2934 let closures: [_; N] = std::array::from_fn(|i| {
2935 let Self::F32 { data, cols, .. } = ts[i] else {
2936 unreachable!()
2937 };
2938 let (o1, o2) = (p1[i], p2[i]);
2939 move |start: usize, end: usize| {
2940 for o in start..end {
2941 let row = &data[o * cols..(o + 1) * cols];
2942 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2943 for j in 0..*cols {
2944 s1 += row[j] * x1[j];
2945 s2 += row[j] * x2[j];
2946 }
2947 unsafe {
2949 *o1.at(o) = s1;
2950 *o2.at(o) = s2;
2951 }
2952 }
2953 }
2954 });
2955 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2956 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2957 pool.run_many(&parts);
2958 return;
2959 }
2960
2961 struct Ctx<'a> {
2962 bytes: &'a [u8],
2963 row_scale: &'a [f32],
2964 cols: usize,
2965 xs1: std::borrow::Cow<'a, [f32]>,
2966 xs2: std::borrow::Cow<'a, [f32]>,
2967 }
2968 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2969 let Self::Mapped {
2970 dtype,
2971 cols,
2972 row_scale,
2973 col_field,
2974 ..
2975 } = ts[i]
2976 else {
2977 unreachable!()
2978 };
2979 Ctx {
2980 bytes: ts[i].quant_bytes(),
2981 row_scale,
2982 cols: *cols,
2983 xs1: prescale(x1, col_field, *dtype),
2984 xs2: prescale(x2, col_field, *dtype),
2985 }
2986 });
2987 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2988 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2989 #[cfg(target_arch = "aarch64")]
2990 if sdot_enabled() {
2991 let acts: [(SplitAct, SplitAct); N] =
2992 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2993 let closures: [_; N] = std::array::from_fn(|i| {
2994 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2995 move |start: usize, end: usize| {
2996 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2997 }
2998 });
2999 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3000 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3001 pool.run_many(&parts);
3002 return;
3003 }
3004 #[cfg(target_arch = "x86_64")]
3005 if avx2_a8w8_enabled() {
3006 let acts: [(SplitAct, SplitAct); N] =
3007 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3008 let closures: [_; N] = std::array::from_fn(|i| {
3009 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3010 move |start: usize, end: usize| {
3011 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3012 }
3013 });
3014 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3015 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3016 pool.run_many(&parts);
3017 return;
3018 }
3019 let closures: [_; N] = std::array::from_fn(|i| {
3020 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3021 move |start: usize, end: usize| {
3022 q8_range2_f32(
3023 c.bytes,
3024 c.row_scale,
3025 &c.xs1,
3026 &c.xs2,
3027 c.cols,
3028 o1,
3029 o2,
3030 start,
3031 end,
3032 )
3033 }
3034 });
3035 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3036 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3037 pool.run_many(&parts);
3038 }
3039
3040 pub fn matvec_silu_mul(
3045 gate: &QTensor,
3046 up: &QTensor,
3047 x: &[f32],
3048 out: &mut [f32],
3049 pool: Option<&Pool>,
3050 ) -> bool {
3051 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3052 }
3053
3054 pub fn matvec_silu_mul_limited(
3060 gate: &QTensor,
3061 up: &QTensor,
3062 x: &[f32],
3063 out: &mut [f32],
3064 limit: f32,
3065 pool: Option<&Pool>,
3066 ) -> bool {
3067 if gate.has_prism_contract() || up.has_prism_contract() {
3068 return false;
3072 }
3073 let inter = gate.rows();
3074 debug_assert_eq!(up.rows(), inter);
3075 debug_assert_eq!(out.len(), inter);
3076 debug_assert_eq!(gate.cols(), up.cols());
3077 if !a8w8_enabled() {
3078 return false;
3079 }
3080 let act = split_act(x);
3081 let act = &act;
3082 let x_ref = x;
3083 let out_addr = SendMut(out.as_mut_ptr());
3084
3085 match (gate, up) {
3086 (
3088 Self::Mapped {
3089 dtype: TensorDtype::Q4Block,
3090 ..
3091 },
3092 Self::Mapped {
3093 dtype: TensorDtype::Q4Block,
3094 ..
3095 },
3096 ) => {
3097 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3098 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3099 let gpr = gate.cols() / GROUP_SIZE;
3100 let cols = gate.cols();
3101 let run = move |start: usize, end: usize| {
3102 for r in start..end {
3103 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3104 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3105 for &(j, xv) in &act.outliers {
3106 let flat = r * cols + j;
3107 let gb = gp[flat / 2];
3108 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3109 let gsc = f16_to_f32(u16::from_le_bytes([
3110 gs[(flat / GROUP_SIZE) * 2],
3111 gs[(flat / GROUP_SIZE) * 2 + 1],
3112 ]));
3113 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3114 let ub = up_p[flat / 2];
3115 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3116 let usc = f16_to_f32(u16::from_le_bytes([
3117 up_s[(flat / GROUP_SIZE) * 2],
3118 up_s[(flat / GROUP_SIZE) * 2 + 1],
3119 ]));
3120 uv += ((un as i32 - 8) as f32) * usc * xv;
3121 }
3122 let silu_g = gv / (1.0 + (-gv).exp());
3123 unsafe { *out_addr.at(r) = silu_g * uv };
3125 }
3126 };
3127 dispatch_rows(pool, inter, &run);
3128 true
3129 }
3130 (
3134 Self::Mapped {
3135 dtype: TensorDtype::Q4Tiled,
3136 ..
3137 },
3138 Self::Mapped {
3139 dtype: TensorDtype::Q4Tiled,
3140 ..
3141 },
3142 ) => {
3143 let g_bytes = gate.quant_bytes();
3144 let u_bytes = up.quant_bytes();
3145 let gpr = gate.cols() / GROUP_SIZE;
3146 let run = move |start: usize, end: usize| {
3147 for r in start..end {
3148 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3149 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3150 for &(j, xv) in &act.outliers {
3151 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3152 gv += w * s * xv;
3153 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3154 uv += w * s * xv;
3155 }
3156 let silu_g = gv / (1.0 + (-gv).exp());
3157 unsafe { *out_addr.at(r) = silu_g * uv };
3159 }
3160 };
3161 dispatch_rows(pool, inter, &run);
3162 true
3163 }
3164 (
3167 Self::Mapped {
3168 dtype: TensorDtype::Q4TiledP,
3169 ..
3170 },
3171 Self::Mapped {
3172 dtype: TensorDtype::Q4TiledP,
3173 ..
3174 },
3175 ) => {
3176 let cols = gate.cols();
3177 let gpr = cols / GROUP_SIZE;
3178 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3179 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3180 let run = |start: usize, end: usize| {
3181 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3182 for r in start..end {
3183 gv_view.scales_into(r, gpr, &mut gsc);
3184 uv_view.scales_into(r, gpr, &mut usc);
3185 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3186 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3187 for &(j, xv) in &act.outliers {
3188 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3189 gv += w * s * xv;
3190 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3191 uv += w * s * xv;
3192 }
3193 let silu_g = gv / (1.0 + (-gv).exp());
3194 unsafe { *out_addr.at(r) = silu_g * uv };
3196 }
3197 };
3198 dispatch_rows(pool, inter, &run);
3199 true
3200 }
3201 (
3207 Self::Mapped {
3208 dtype: TensorDtype::Q1,
3209 ..
3210 },
3211 Self::Mapped {
3212 dtype: TensorDtype::Q1,
3213 ..
3214 },
3215 ) => {
3216 let g_bytes = gate.quant_bytes();
3217 let u_bytes = up.quant_bytes();
3218 let gpr = gate.cols() / GROUP_SIZE;
3219 let gsum = q1_group_sums(&act.xq, gpr);
3220 let gsum = &gsum;
3221 let run = move |start: usize, end: usize| {
3222 for r in start..end {
3223 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
3224 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
3225 for &(j, xv) in &act.outliers {
3226 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
3227 gv += w * s * xv;
3228 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
3229 uv += w * s * xv;
3230 }
3231 let silu_g = gv / (1.0 + (-gv).exp());
3232 unsafe { *out_addr.at(r) = silu_g * uv };
3234 }
3235 };
3236 dispatch_rows(pool, inter, &run);
3237 true
3238 }
3239 (
3243 Self::Mapped {
3244 dtype: TensorDtype::Q2TiledP,
3245 ..
3246 },
3247 Self::Mapped {
3248 dtype: TensorDtype::Q2TiledP,
3249 ..
3250 },
3251 ) => {
3252 let cols = gate.cols();
3253 let gpr = cols / GROUP_SIZE;
3254 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3255 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3256 let gsum = q1_group_sums(&act.xq, gpr);
3257 let gsum = &gsum;
3258 let run = move |start: usize, end: usize| {
3259 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3260 for r in start..end {
3261 gv_view.scales_into(r, gpr, &mut gsc);
3262 uv_view.scales_into(r, gpr, &mut usc);
3263 let mut gv =
3264 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
3265 let mut uv =
3266 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
3267 for &(j, xv) in &act.outliers {
3268 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3269 gv += w * s * xv;
3270 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
3271 uv += w * s * xv;
3272 }
3273 let silu_g = gv / (1.0 + (-gv).exp());
3274 unsafe { *out_addr.at(r) = silu_g * uv };
3276 }
3277 };
3278 dispatch_rows(pool, inter, &run);
3279 true
3280 }
3281 (
3286 Self::Mapped {
3287 dtype: TensorDtype::Q8Row,
3288 row_scale: g_rs,
3289 ..
3290 },
3291 Self::Mapped {
3292 dtype: TensorDtype::Q8Row,
3293 row_scale: u_rs,
3294 ..
3295 },
3296 ) => {
3297 let g_bytes = gate.quant_bytes();
3298 let u_bytes = up.quant_bytes();
3299 let cols = gate.cols();
3300 let run = move |start: usize, end: usize| {
3301 for r in start..end {
3302 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3303 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3304 let silu_g = gv / (1.0 + (-gv).exp());
3305 unsafe { *out_addr.at(r) = silu_g * uv };
3307 }
3308 };
3309 dispatch_rows(pool, inter, &run);
3310 true
3311 }
3312 (
3314 Self::Mapped {
3315 dtype: TensorDtype::Q1T,
3316 ..
3317 },
3318 Self::Mapped {
3319 dtype: TensorDtype::Q1T,
3320 ..
3321 },
3322 ) => {
3323 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3324 let g_bytes = gate.quant_bytes();
3325 let u_bytes = up.quant_bytes();
3326 let gpr = gate.cols() / GROUP_SIZE;
3327 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3328 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3329 let run = move |start: usize, end: usize| {
3330 for r in start..end {
3331 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3332 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3333 for &(j, xv) in &act.outliers {
3334 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3335 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3336 }
3337 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3338 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3339 let silu_g = gv / (1.0 + (-gv).exp());
3340 unsafe { *out_addr.at(r) = silu_g * uv };
3342 }
3343 };
3344 dispatch_rows(pool, inter, &run);
3345 true
3346 }
3347 _ => false,
3348 }
3349 }
3350
3351 pub fn moe_gate_up_many(
3365 pairs: &[(&QTensor, &QTensor)],
3366 x: &[f32],
3367 outs: &mut [Vec<f32>],
3368 pool: Option<&Pool>,
3369 ) -> bool {
3370 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
3371 return false;
3372 }
3373 let inter = pairs[0].0.rows();
3374 let cols = pairs[0].0.cols();
3375 if cols % GROUP_SIZE != 0 {
3376 return false;
3377 }
3378 let gpr = cols / GROUP_SIZE;
3379 let q2 = matches!(
3382 pairs[0].0,
3383 Self::Mapped {
3384 dtype: TensorDtype::Q2TiledP,
3385 ..
3386 }
3387 );
3388 let want = if q2 {
3389 TensorDtype::Q2TiledP
3390 } else {
3391 TensorDtype::Q4TiledP
3392 };
3393 let mut views = Vec::with_capacity(pairs.len() * 2);
3394 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3395 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3396 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3397 if !both
3398 || g.rows() != inter
3399 || u.rows() != inter
3400 || g.cols() != cols
3401 || u.cols() != cols
3402 || o.len() != inter
3403 {
3404 return false;
3405 }
3406 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3407 views.push(mk(g.quant_bytes(), inter, cols));
3408 views.push(mk(u.quant_bytes(), inter, cols));
3409 }
3410 let act = split_act(x);
3411 let gsum = if q2 {
3412 q1_group_sums(&act.xq, gpr)
3413 } else {
3414 Vec::new()
3415 };
3416 let (act, gsum) = (&act, &gsum);
3417 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3418 let (views, ptrs) = (&views, &ptrs);
3419 let run = |start: usize, end: usize| {
3420 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3421 for flat in start..end {
3422 let (e, r) = (flat / inter, flat % inter);
3423 let gv_view = &views[e * 2];
3424 let uv_view = &views[e * 2 + 1];
3425 gv_view.scales_into(r, gpr, &mut gsc);
3426 uv_view.scales_into(r, gpr, &mut usc);
3427 let (mut gv, mut uv) = if q2 {
3428 (
3429 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
3430 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
3431 )
3432 } else {
3433 (
3434 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
3435 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
3436 )
3437 };
3438 for &(j, xv) in &act.outliers {
3439 let (og, ou) = if q2 {
3440 (
3441 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3442 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
3443 )
3444 } else {
3445 (
3446 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3447 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
3448 )
3449 };
3450 gv += og.0 * og.1 * xv;
3451 uv += ou.0 * ou.1 * xv;
3452 }
3453 let silu_g = gv / (1.0 + (-gv).exp());
3454 unsafe { *ptrs[e].at(r) = silu_g * uv };
3456 }
3457 };
3458 dispatch_rows(pool, pairs.len() * inter, &run);
3459 true
3460 }
3461
3462 pub fn moe_down_many(
3471 downs: &[&QTensor],
3472 gs: &[Vec<f32>],
3473 weights: &[f32],
3474 out: &mut [f32],
3475 pool: Option<&Pool>,
3476 ) -> bool {
3477 if downs.is_empty()
3478 || downs.len() != gs.len()
3479 || downs.len() != weights.len()
3480 || !a8w8_enabled()
3481 {
3482 return false;
3483 }
3484 let rows = out.len();
3485 let cols = downs[0].cols();
3486 if cols % GROUP_SIZE != 0 {
3487 return false;
3488 }
3489 let gpr = cols / GROUP_SIZE;
3490 let mut views = Vec::with_capacity(downs.len());
3491 for (d, g) in downs.iter().zip(gs.iter()) {
3492 if !matches!(
3493 d,
3494 Self::Mapped {
3495 dtype: TensorDtype::Q4TiledP,
3496 ..
3497 }
3498 ) || d.rows() != rows
3499 || d.cols() != cols
3500 || g.len() != cols
3501 {
3502 return false;
3503 }
3504 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3505 }
3506 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3508 let out_addr = SendMut(out.as_mut_ptr());
3516 let (views, acts, weights) = (&views, &acts, &weights);
3517 let run = |start: usize, end: usize| {
3518 let mut sc = vec![0f32; gpr];
3519 for r in start..end {
3520 let mut acc = 0f32;
3521 for (e, v) in views.iter().enumerate() {
3522 v.scales_into(r, gpr, &mut sc);
3523 let a = &acts[e];
3524 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3525 for &(j, xv) in &a.outliers {
3526 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3527 d += w * s * xv;
3528 }
3529 acc += weights[e] * d;
3530 }
3531 unsafe { *out_addr.at(r) = acc };
3533 }
3534 };
3535 dispatch_rows(pool, rows, &run);
3536 true
3537 }
3538}
3539
3540#[cfg(target_os = "macos")]
3545mod accel_blas {
3546 #[link(name = "Accelerate", kind = "framework")]
3547 unsafe extern "C" {
3548 pub fn cblas_sgemm(
3549 order: i32,
3550 trans_a: i32,
3551 trans_b: i32,
3552 m: i32,
3553 n: i32,
3554 k: i32,
3555 alpha: f32,
3556 a: *const f32,
3557 lda: i32,
3558 b: *const f32,
3559 ldb: i32,
3560 beta: f32,
3561 c: *mut f32,
3562 ldc: i32,
3563 );
3564 }
3565}
3566
3567#[cfg(target_os = "macos")]
3568pub(crate) fn accel_gemm_enabled() -> bool {
3569 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3570 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3571}
3572
3573#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3576pub(crate) fn accel_gemm_enabled() -> bool {
3577 true
3578}
3579
3580#[cfg(target_arch = "aarch64")]
3587#[allow(clippy::too_many_arguments)]
3588pub(crate) fn neon_gemm_rm(
3589 m: usize,
3590 n: usize,
3591 k: usize,
3592 alpha: f32,
3593 a: &[f32],
3594 lda: usize,
3595 b_mat: &[f32],
3596 ldb: usize,
3597 b_rows_are_n: bool,
3598 c: &mut [f32],
3599 ldc: usize,
3600) {
3601 debug_assert!(a.len() >= (m - 1) * lda + k);
3602 debug_assert!(c.len() >= (m - 1) * ldc + n);
3603 unsafe {
3605 use core::arch::aarch64::*;
3606 let mut i = 0usize;
3607 while i < m {
3608 let mi = (m - i).min(4);
3609 let mut j = 0usize;
3610 while j < n {
3611 let nj = (n - j).min(8);
3612 if mi == 4 && nj == 8 {
3613 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3614 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3615 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3616 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3617 for p in 0..k {
3618 let (b0, b1) = if b_rows_are_n {
3619 let base = b_mat.as_ptr().add(j * ldb + p);
3622 let g = |o: usize| *base.add(o * ldb);
3623 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3624 } else {
3625 let base = b_mat.as_ptr().add(p * ldb + j);
3626 (
3627 [*base, *base.add(1), *base.add(2), *base.add(3)],
3628 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3629 )
3630 };
3631 let bv0 = vld1q_f32(b0.as_ptr());
3632 let bv1 = vld1q_f32(b1.as_ptr());
3633 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3634 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3635 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3636 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3637 c0a = vfmaq_f32(c0a, a0, bv0);
3638 c0b = vfmaq_f32(c0b, a0, bv1);
3639 c1a = vfmaq_f32(c1a, a1, bv0);
3640 c1b = vfmaq_f32(c1b, a1, bv1);
3641 c2a = vfmaq_f32(c2a, a2, bv0);
3642 c2b = vfmaq_f32(c2b, a2, bv1);
3643 c3a = vfmaq_f32(c3a, a3, bv0);
3644 c3b = vfmaq_f32(c3b, a3, bv1);
3645 }
3646 let al = vdupq_n_f32(alpha);
3647 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3648 .iter()
3649 .enumerate()
3650 {
3651 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3652 vst1q_f32(dst, vmulq_f32(*ca, al));
3653 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3654 }
3655 } else {
3656 for r in 0..mi {
3657 for q in 0..nj {
3658 let mut acc = 0f32;
3659 for p in 0..k {
3660 let bv = if b_rows_are_n {
3661 b_mat[(j + q) * ldb + p]
3662 } else {
3663 b_mat[p * ldb + j + q]
3664 };
3665 acc += a[(i + r) * lda + p] * bv;
3666 }
3667 c[(i + r) * ldc + j + q] = acc * alpha;
3668 }
3669 }
3670 }
3671 j += nj;
3672 }
3673 i += mi;
3674 }
3675 }
3676}
3677
3678#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3680#[allow(clippy::too_many_arguments)]
3681pub(crate) fn sgemm_rm(
3682 m: usize,
3683 n: usize,
3684 k: usize,
3685 alpha: f32,
3686 a: &[f32],
3687 lda: usize,
3688 b_mat: &[f32],
3689 ldb: usize,
3690 b_rows_are_n: bool,
3691 c: &mut [f32],
3692 ldc: usize,
3693) {
3694 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3695}
3696
3697#[allow(clippy::too_many_arguments)]
3701pub fn sgemm_public(
3702 m: usize,
3703 n: usize,
3704 k: usize,
3705 alpha: f32,
3706 a: &[f32],
3707 lda: usize,
3708 b_mat: &[f32],
3709 ldb: usize,
3710 b_rows_are_n: bool,
3711 c: &mut [f32],
3712 ldc: usize,
3713) {
3714 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3715 {
3716 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3717 }
3718 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3723 {
3724 for i in 0..m {
3725 for j in 0..n {
3726 let mut acc = 0f32;
3727 for p in 0..k {
3728 let bv = if b_rows_are_n {
3729 b_mat[j * ldb + p]
3730 } else {
3731 b_mat[p * ldb + j]
3732 };
3733 acc += a[i * lda + p] * bv;
3734 }
3735 c[i * ldc + j] = alpha * acc;
3736 }
3737 }
3738 }
3739}
3740
3741#[cfg(target_os = "macos")]
3744#[allow(clippy::too_many_arguments)]
3745pub(crate) fn sgemm_rm(
3746 m: usize,
3747 n: usize,
3748 k: usize,
3749 alpha: f32,
3750 a: &[f32],
3751 lda: usize,
3752 b_mat: &[f32],
3753 ldb: usize,
3754 b_rows_are_n: bool,
3755 c: &mut [f32],
3756 ldc: usize,
3757) {
3758 debug_assert!(a.len() >= (m - 1) * lda + k);
3759 debug_assert!(c.len() >= (m - 1) * ldc + n);
3760 #[cfg(target_arch = "aarch64")]
3765 if std::env::var("CMF_FORCE_NEON_GEMM")
3766 .map(|v| v == "1")
3767 .unwrap_or(false)
3768 {
3769 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3770 }
3771 unsafe {
3772 accel_blas::cblas_sgemm(
3773 101, 111, if b_rows_are_n { 112 } else { 111 },
3776 m as i32,
3777 n as i32,
3778 k as i32,
3779 alpha,
3780 a.as_ptr(),
3781 lda as i32,
3782 b_mat.as_ptr(),
3783 ldb as i32,
3784 0.0,
3785 c.as_mut_ptr(),
3786 ldc as i32,
3787 );
3788 }
3789}
3790
3791#[cfg(target_os = "macos")]
3798fn qmatmat_accel(
3799 q: &[u8],
3800 row_scale: &[f32],
3801 pre: &[std::borrow::Cow<'_, [f32]>],
3802 rows: usize,
3803 cols: usize,
3804 out: &mut [f32],
3805 pool: Option<&Pool>,
3806) {
3807 const TR: usize = 2048;
3812 let b = pre.len();
3813 thread_local! {
3814 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3815 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3816 }
3817 XPANEL.with(|xp| {
3818 WTILE.with(|wt| {
3819 let mut xpanel = xp.borrow_mut();
3820 xpanel.clear();
3821 for x in pre {
3822 xpanel.extend_from_slice(x);
3823 }
3824 let mut wtile = wt.borrow_mut();
3825 wtile.resize(TR * cols, 0.0);
3826 let mut r0 = 0usize;
3827 while r0 < rows {
3828 let tr = TR.min(rows - r0);
3829 let wt_addr = SendMut(wtile.as_mut_ptr());
3831 let run = |start: usize, end: usize| {
3832 for r in start..end {
3833 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3834 let s = row_scale[r0 + r];
3835 let dst =
3837 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3838 for (d, &v) in dst.iter_mut().zip(row) {
3839 *d = (v as i8) as f32 * s;
3840 }
3841 }
3842 };
3843 dispatch_rows(pool, tr, &run);
3844 unsafe {
3846 accel_blas::cblas_sgemm(
3847 101, 111, 112, b as i32,
3851 tr as i32,
3852 cols as i32,
3853 1.0,
3854 xpanel.as_ptr(),
3855 cols as i32,
3856 wtile.as_ptr(),
3857 cols as i32,
3858 0.0,
3859 out.as_mut_ptr().add(r0),
3860 rows as i32,
3861 );
3862 }
3863 r0 += tr;
3864 }
3865 })
3866 });
3867}
3868
3869fn qmatmat(
3870 q: &[u8],
3871 row_scale: &[f32],
3872 pre: &[std::borrow::Cow<'_, [f32]>],
3873 rows: usize,
3874 cols: usize,
3875 out: &mut [f32],
3876 pool: Option<&Pool>,
3877) {
3878 let b = pre.len();
3879 debug_assert_eq!(out.len(), b * rows);
3880 #[cfg(target_os = "macos")]
3885 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3886 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3887 return;
3888 }
3889 #[cfg(target_arch = "aarch64")]
3890 if sdot_enabled() {
3891 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3892 let out_addr = SendMut(out.as_mut_ptr());
3893 let blocked_ok = blocked_enabled();
3896 let use_i8mm = i8mm_enabled();
3897 if blocked_ok {
3898 let run = |start: usize, end: usize| {
3899 let mut o = start;
3900 while o < end {
3901 if o + 2 <= end {
3902 let r0 = &q[o * cols..(o + 1) * cols];
3903 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3904 let mut bi = 0usize;
3905 while bi + 4 <= acts.len() {
3906 let xs = [
3907 acts[bi].xq.as_slice(),
3908 acts[bi + 1].xq.as_slice(),
3909 acts[bi + 2].xq.as_slice(),
3910 acts[bi + 3].xq.as_slice(),
3911 ];
3912 let d = if use_i8mm {
3913 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3914 } else {
3915 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3916 };
3917 for (r, row) in [r0, r1].into_iter().enumerate() {
3918 for k in 0..4 {
3919 let act = &acts[bi + k];
3920 let mut v = d[r][k] as f32 * act.sx;
3921 for &(j, xv) in &act.outliers {
3922 v += (row[j] as i8) as f32 * xv;
3923 }
3924 unsafe {
3925 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3926 };
3927 }
3928 }
3929 bi += 4;
3930 }
3931 while bi < acts.len() {
3932 for (r, row) in [r0, r1].into_iter().enumerate() {
3933 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3934 unsafe { *out_addr.at(bi * rows + o + r) = v };
3935 }
3936 bi += 1;
3937 }
3938 o += 2;
3939 } else {
3940 let row = &q[o * cols..(o + 1) * cols];
3941 for (bi, act) in acts.iter().enumerate() {
3942 let v = row_dot_sdot(row, act) * row_scale[o];
3943 unsafe { *out_addr.at(bi * rows + o) = v };
3944 }
3945 o += 1;
3946 }
3947 }
3948 };
3949 dispatch_rows(pool, rows, &run);
3950 return;
3951 }
3952 let run = |start: usize, end: usize| {
3953 for o in start..end {
3954 let row = &q[o * cols..(o + 1) * cols];
3955 for (bi, act) in acts.iter().enumerate() {
3956 let v = row_dot_sdot(row, act) * row_scale[o];
3957 unsafe { *out_addr.at(bi * rows + o) = v };
3958 }
3959 }
3960 };
3961 dispatch_rows(pool, rows, &run);
3962 return;
3963 }
3964 #[cfg(target_arch = "x86_64")]
3969 if avx2_a8w8_enabled() {
3970 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3971 let out_addr = SendMut(out.as_mut_ptr());
3972 let blocked_ok = blocked_enabled();
3975 if !avx512vnni_enabled() && blocked_ok {
3976 let run = |start: usize, end: usize| {
3977 let mut o = start;
3978 while o < end {
3979 if o + 2 <= end {
3980 let r0 = &q[o * cols..(o + 1) * cols];
3981 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3982 let mut bi = 0usize;
3983 while bi + 4 <= acts.len() {
3984 let xs = [
3985 acts[bi].xq.as_slice(),
3986 acts[bi + 1].xq.as_slice(),
3987 acts[bi + 2].xq.as_slice(),
3988 acts[bi + 3].xq.as_slice(),
3989 ];
3990 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3991 for (r, row) in [r0, r1].into_iter().enumerate() {
3992 for k in 0..4 {
3993 let act = &acts[bi + k];
3994 let mut v = d[r][k] as f32 * act.sx;
3995 for &(j, xv) in &act.outliers {
3996 v += (row[j] as i8) as f32 * xv;
3997 }
3998 unsafe {
3999 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4000 };
4001 }
4002 }
4003 bi += 4;
4004 }
4005 while bi < acts.len() {
4006 for (r, row) in [r0, r1].into_iter().enumerate() {
4007 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4008 unsafe { *out_addr.at(bi * rows + o + r) = v };
4009 }
4010 bi += 1;
4011 }
4012 o += 2;
4013 } else {
4014 let row = &q[o * cols..(o + 1) * cols];
4015 for (bi, act) in acts.iter().enumerate() {
4016 let v = row_dot_avx2(row, act) * row_scale[o];
4017 unsafe { *out_addr.at(bi * rows + o) = v };
4018 }
4019 o += 1;
4020 }
4021 }
4022 };
4023 dispatch_rows(pool, rows, &run);
4024 return;
4025 }
4026 let run = |start: usize, end: usize| {
4027 for o in start..end {
4028 let row = &q[o * cols..(o + 1) * cols];
4029 for (bi, act) in acts.iter().enumerate() {
4030 let v = row_dot_avx2(row, act) * row_scale[o];
4031 unsafe { *out_addr.at(bi * rows + o) = v };
4032 }
4033 }
4034 };
4035 dispatch_rows(pool, rows, &run);
4036 return;
4037 }
4038 let out_addr = SendMut(out.as_mut_ptr());
4039 let run = |start: usize, end: usize| {
4040 for o in start..end {
4041 let row = &q[o * cols..(o + 1) * cols];
4042 for (bi, x) in pre.iter().enumerate() {
4043 let mut acc = 0f32;
4044 for j in 0..cols {
4045 acc += (row[j] as i8) as f32 * x[j];
4046 }
4047 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4048 }
4049 }
4050 };
4051 dispatch_rows(pool, rows, &run);
4052}
4053
4054fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4057 match pool {
4058 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4059 _ => run(0, rows),
4060 }
4061}
4062
4063fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4065 let groups = rows * cols / GROUP_SIZE;
4066 bytes.split_at(groups * 16)
4067}
4068
4069#[inline]
4074fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4075 #[cfg(target_arch = "aarch64")]
4076 unsafe {
4077 return vbit_fill4_neon(data, buf);
4078 }
4079 #[cfg(target_arch = "x86_64")]
4080 if avx2_enabled() {
4081 return unsafe { vbit_fill4_avx2(data, buf) };
4082 }
4083 #[allow(unreachable_code)]
4084 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4085 let u = unpack8::<4>(&data[blk * 4..]);
4086 for k in 0..8 {
4087 chunk[k] = (u[k] - 7) as i8 as u8;
4088 }
4089 }
4090}
4091
4092#[cfg(target_arch = "aarch64")]
4093#[target_feature(enable = "neon")]
4094unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4095 unsafe {
4098 use core::arch::aarch64::*;
4099 let n = buf.len();
4100 let mask = vdupq_n_u8(0x0F);
4101 let seven = vdupq_n_s8(7);
4102 let mut g = 0usize;
4103 while g * 32 + 32 <= n {
4104 let b = vld1q_u8(data.as_ptr().add(g * 16));
4105 let hi = vshrq_n_u8::<4>(b);
4106 let lo = vandq_u8(b, mask);
4107 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4108 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4109 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4110 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4111 g += 1;
4112 }
4113 }
4114}
4115
4116#[cfg(target_arch = "x86_64")]
4117#[target_feature(enable = "avx2")]
4118unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4119 unsafe {
4121 use core::arch::x86_64::*;
4122 let n = buf.len();
4123 let mask = _mm_set1_epi8(0x0F);
4124 let seven = _mm256_set1_epi8(7);
4125 let mut g = 0usize;
4126 while g * 32 + 32 <= n {
4127 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4128 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4129 let lo = _mm_and_si128(b, mask);
4130 let z = _mm256_sub_epi8(
4131 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4132 seven,
4133 );
4134 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4135 g += 1;
4136 }
4137 }
4138}
4139
4140#[inline(always)]
4145fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4146 let mut acc = 0u64;
4147 for i in 0..B {
4148 acc = (acc << 8) | data[i] as u64;
4149 }
4150 let mask = (1u64 << B) - 1;
4151 let mut out = [0i32; 8];
4152 for (k, o) in out.iter_mut().enumerate() {
4153 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4154 }
4155 out
4156}
4157
4158#[allow(clippy::too_many_arguments)]
4164fn vbitmatvec(
4165 bytes: &[u8],
4166 offsets: &[usize],
4167 x: &[f32],
4168 rows: usize,
4169 cols: usize,
4170 out: &mut [f32],
4171 pool: Option<&Pool>,
4172) {
4173 debug_assert_eq!(out.len(), rows);
4174 debug_assert_eq!(offsets.len(), rows + 1);
4175
4176 if a8w8_enabled() {
4180 let act = split_act(x);
4181 let out_addr = SendMut(out.as_mut_ptr());
4182 let run = move |start: usize, end: usize| {
4183 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4184 };
4185 dispatch_rows(pool, rows, &run);
4186 return;
4187 }
4188
4189 let out_addr = SendMut(out.as_mut_ptr());
4190 let run = move |start: usize, end: usize| {
4191 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4192 };
4193 dispatch_rows(pool, rows, &run);
4194}
4195
4196#[allow(clippy::too_many_arguments)]
4200fn vbit_range_a8w8(
4201 bytes: &[u8],
4202 offsets: &[usize],
4203 x: &[f32],
4204 act: &SplitAct,
4205 rows: usize,
4206 cols: usize,
4207 out: SendMut,
4208 start: usize,
4209 end: usize,
4210) {
4211 let ng = cols / GROUP_SIZE;
4212 let bits = &bytes[..rows];
4213 let sc_off = rows;
4214 let row_dot = |r: usize| -> f32 {
4215 let b = bits[r] as usize;
4216 let l = (1i32 << (b - 1)) - 1;
4217 let mask = (1u64 << b) - 1;
4218 let data = &bytes[offsets[r]..offsets[r + 1]];
4219 if b == 8 {
4220 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4222 let mut dot = 0f32;
4223 for g in 0..ng {
4224 let so = (r * ng + g) * 2;
4225 let sgf = f16_to_f32(u16::from_le_bytes([
4226 bytes[sc_off + so],
4227 bytes[sc_off + so + 1],
4228 ]));
4229 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4230 let mut gd = 0f32;
4231 for &xv in xg.iter() {
4232 if nbits < 8 {
4233 acc = (acc << 8) | data[idx] as u64;
4234 idx += 1;
4235 nbits += 8;
4236 }
4237 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4238 nbits -= 8;
4239 gd += (u - l) as f32 * xv;
4240 }
4241 dot += gd * sgf;
4242 }
4243 return dot;
4244 }
4245 thread_local! {
4249 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4250 const { std::cell::RefCell::new(Vec::new()) };
4251 }
4252 #[inline(always)]
4253 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4254 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4255 let u = unpack8::<B>(&data[blk * B..]);
4256 for k in 0..8 {
4257 chunk[k] = (u[k] - l) as i8 as u8;
4258 }
4259 }
4260 }
4261 let _ = mask;
4262 VBIT_SCRATCH.with(|scratch| {
4263 let mut buf = scratch.borrow_mut();
4264 buf.resize(cols, 0);
4265 match b {
4266 3 => fill::<3>(data, l, &mut buf),
4267 4 => vbit_fill4(data, &mut buf),
4268 5 => fill::<5>(data, l, &mut buf),
4269 6 => fill::<6>(data, l, &mut buf),
4270 _ => unreachable!(),
4271 }
4272 let mut dot = 0f32;
4273 for g in 0..ng {
4274 let so = (r * ng + g) * 2;
4275 let s = f16_to_f32(u16::from_le_bytes([
4276 bytes[sc_off + so],
4277 bytes[sc_off + so + 1],
4278 ]));
4279 let d = dot_i8_i8(
4280 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4281 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4282 ) as f32
4283 * act.sx;
4284 dot += d * s;
4285 }
4286 for &(j, xv) in &act.outliers {
4287 let so = (r * ng + j / GROUP_SIZE) * 2;
4288 let s = f16_to_f32(u16::from_le_bytes([
4289 bytes[sc_off + so],
4290 bytes[sc_off + so + 1],
4291 ]));
4292 dot += (buf[j] as i8) as f32 * s * xv;
4294 }
4295 dot
4296 })
4297 };
4298 for r in start..end {
4299 unsafe { *out.at(r) = row_dot(r) };
4301 }
4302}
4303
4304#[allow(clippy::too_many_arguments)]
4306fn vbit_range_f32(
4307 bytes: &[u8],
4308 offsets: &[usize],
4309 x: &[f32],
4310 rows: usize,
4311 cols: usize,
4312 out: SendMut,
4313 start: usize,
4314 end: usize,
4315) {
4316 let ng = cols / GROUP_SIZE;
4317 let bits = &bytes[..rows];
4318 let sc_off = rows;
4319 #[inline(always)]
4323 fn dot_row<const B: usize>(
4324 data: &[u8],
4325 bytes: &[u8],
4326 sc_off: usize,
4327 r: usize,
4328 ng: usize,
4329 x: &[f32],
4330 ) -> f32 {
4331 let l = ((1i32 << (B - 1)) - 1) as f32;
4332 let gbytes = GROUP_SIZE * B / 8;
4333 let mut dot = 0f32;
4334 for g in 0..ng {
4335 let so = (r * ng + g) * 2;
4336 let s = f16_to_f32(u16::from_le_bytes([
4337 bytes[sc_off + so],
4338 bytes[sc_off + so + 1],
4339 ]));
4340 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4341 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4342 let mut gd = 0f32;
4343 for blk in 0..GROUP_SIZE / 8 {
4344 let u = unpack8::<B>(&gd0[blk * B..]);
4345 let xb = &xg[blk * 8..blk * 8 + 8];
4346 for k in 0..8 {
4347 gd += (u[k] as f32 - l) * xb[k];
4348 }
4349 }
4350 dot += gd * s;
4351 }
4352 dot
4353 }
4354 for r in start..end {
4355 let data = &bytes[offsets[r]..offsets[r + 1]];
4356 let v = match bits[r] {
4357 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4358 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4359 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4360 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4361 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4362 b => unreachable!("vbit bit-width {b} (validated at load)"),
4363 };
4364 unsafe { *out.at(r) = v };
4366 }
4367}
4368
4369#[allow(clippy::too_many_arguments)]
4374fn vbitmatvec2(
4375 bytes: &[u8],
4376 offsets: &[usize],
4377 x1: &[f32],
4378 x2: &[f32],
4379 rows: usize,
4380 cols: usize,
4381 o1: &mut [f32],
4382 o2: &mut [f32],
4383 pool: Option<&Pool>,
4384) {
4385 debug_assert_eq!(o1.len(), rows);
4386 debug_assert_eq!(o2.len(), rows);
4387
4388 if a8w8_enabled() {
4389 let a1 = split_act(x1);
4390 let a2 = split_act(x2);
4391 let p1 = SendMut(o1.as_mut_ptr());
4392 let p2 = SendMut(o2.as_mut_ptr());
4393 let run = move |start: usize, end: usize| {
4394 vbit_range2_a8w8(
4395 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4396 )
4397 };
4398 dispatch_rows(pool, rows, &run);
4399 return;
4400 }
4401
4402 let p1 = SendMut(o1.as_mut_ptr());
4403 let p2 = SendMut(o2.as_mut_ptr());
4404 let run = move |start: usize, end: usize| {
4405 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4406 };
4407 dispatch_rows(pool, rows, &run);
4408}
4409
4410#[allow(clippy::too_many_arguments)]
4414fn vbit_range2_a8w8(
4415 bytes: &[u8],
4416 offsets: &[usize],
4417 x1: &[f32],
4418 x2: &[f32],
4419 a1: &SplitAct,
4420 a2: &SplitAct,
4421 rows: usize,
4422 cols: usize,
4423 p1: SendMut,
4424 p2: SendMut,
4425 start: usize,
4426 end: usize,
4427) {
4428 let ng = cols / GROUP_SIZE;
4429 let bits = &bytes[..rows];
4430 let sc_off = rows;
4431 let row_dots = |r: usize| -> (f32, f32) {
4432 let b = bits[r] as usize;
4433 let l = (1i32 << (b - 1)) - 1;
4434 let data = &bytes[offsets[r]..offsets[r + 1]];
4435 if b == 8 {
4436 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4439 let (mut d1, mut d2) = (0f32, 0f32);
4440 for g in 0..ng {
4441 let so = (r * ng + g) * 2;
4442 let sgf = f16_to_f32(u16::from_le_bytes([
4443 bytes[sc_off + so],
4444 bytes[sc_off + so + 1],
4445 ]));
4446 let (mut g1, mut g2) = (0f32, 0f32);
4447 for k in 0..GROUP_SIZE {
4448 if nbits < 8 {
4449 acc = (acc << 8) | data[idx] as u64;
4450 idx += 1;
4451 nbits += 8;
4452 }
4453 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4454 nbits -= 8;
4455 let w = (u - l) as f32;
4456 g1 += w * x1[g * GROUP_SIZE + k];
4457 g2 += w * x2[g * GROUP_SIZE + k];
4458 }
4459 d1 += g1 * sgf;
4460 d2 += g2 * sgf;
4461 }
4462 return (d1, d2);
4463 }
4464 thread_local! {
4465 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4466 const { std::cell::RefCell::new(Vec::new()) };
4467 }
4468 #[inline(always)]
4469 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4470 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4471 let u = unpack8::<B>(&data[blk * B..]);
4472 for k in 0..8 {
4473 chunk[k] = (u[k] - l) as i8 as u8;
4474 }
4475 }
4476 }
4477 VBIT_SCRATCH2.with(|scratch| {
4478 let mut buf = scratch.borrow_mut();
4479 buf.resize(cols, 0);
4480 match b {
4481 3 => fill::<3>(data, l, &mut buf),
4482 4 => vbit_fill4(data, &mut buf),
4483 5 => fill::<5>(data, l, &mut buf),
4484 6 => fill::<6>(data, l, &mut buf),
4485 _ => unreachable!(),
4486 }
4487 let (mut d1, mut d2) = (0f32, 0f32);
4488 for g in 0..ng {
4489 let so = (r * ng + g) * 2;
4490 let s = f16_to_f32(u16::from_le_bytes([
4491 bytes[sc_off + so],
4492 bytes[sc_off + so + 1],
4493 ]));
4494 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4495 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
4496 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
4497 d1 += v1 * s;
4498 d2 += v2 * s;
4499 }
4500 for &(j, xv) in &a1.outliers {
4501 let so = (r * ng + j / GROUP_SIZE) * 2;
4502 let s = f16_to_f32(u16::from_le_bytes([
4503 bytes[sc_off + so],
4504 bytes[sc_off + so + 1],
4505 ]));
4506 d1 += (buf[j] as i8) as f32 * s * xv;
4507 }
4508 for &(j, xv) in &a2.outliers {
4509 let so = (r * ng + j / GROUP_SIZE) * 2;
4510 let s = f16_to_f32(u16::from_le_bytes([
4511 bytes[sc_off + so],
4512 bytes[sc_off + so + 1],
4513 ]));
4514 d2 += (buf[j] as i8) as f32 * s * xv;
4515 }
4516 (d1, d2)
4517 })
4518 };
4519 for r in start..end {
4520 let (v1, v2) = row_dots(r);
4521 unsafe {
4523 *p1.at(r) = v1;
4524 *p2.at(r) = v2;
4525 }
4526 }
4527}
4528
4529#[allow(clippy::too_many_arguments)]
4533fn vbit_range2_f32(
4534 bytes: &[u8],
4535 offsets: &[usize],
4536 x1: &[f32],
4537 x2: &[f32],
4538 rows: usize,
4539 cols: usize,
4540 p1: SendMut,
4541 p2: SendMut,
4542 start: usize,
4543 end: usize,
4544) {
4545 let ng = cols / GROUP_SIZE;
4546 let bits = &bytes[..rows];
4547 let sc_off = rows;
4548 #[inline(always)]
4549 #[allow(clippy::too_many_arguments)]
4550 fn dot_row2<const B: usize>(
4551 data: &[u8],
4552 bytes: &[u8],
4553 sc_off: usize,
4554 r: usize,
4555 ng: usize,
4556 x1: &[f32],
4557 x2: &[f32],
4558 ) -> (f32, f32) {
4559 let l = ((1i32 << (B - 1)) - 1) as f32;
4560 let gbytes = GROUP_SIZE * B / 8;
4561 let (mut d1, mut d2) = (0f32, 0f32);
4562 for g in 0..ng {
4563 let so = (r * ng + g) * 2;
4564 let s = f16_to_f32(u16::from_le_bytes([
4565 bytes[sc_off + so],
4566 bytes[sc_off + so + 1],
4567 ]));
4568 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4569 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4570 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4571 let (mut g1, mut g2) = (0f32, 0f32);
4572 for blk in 0..GROUP_SIZE / 8 {
4573 let u = unpack8::<B>(&gd0[blk * B..]);
4574 for k in 0..8 {
4575 let w = u[k] as f32 - l;
4576 g1 += w * x1g[blk * 8 + k];
4577 g2 += w * x2g[blk * 8 + k];
4578 }
4579 }
4580 d1 += g1 * s;
4581 d2 += g2 * s;
4582 }
4583 (d1, d2)
4584 }
4585 for r in start..end {
4586 let data = &bytes[offsets[r]..offsets[r + 1]];
4587 let (v1, v2) = match bits[r] {
4588 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4589 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4590 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4591 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4592 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4593 b => unreachable!("vbit bit-width {b} (validated at load)"),
4594 };
4595 unsafe {
4597 *p1.at(r) = v1;
4598 *p2.at(r) = v2;
4599 }
4600 }
4601}
4602
4603#[inline]
4610#[allow(unreachable_code)]
4611fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4612 #[cfg(target_arch = "aarch64")]
4613 unsafe {
4614 return dot_q4t_row_sdot(bytes, r, gpr, xq);
4615 }
4616 #[cfg(target_arch = "x86_64")]
4617 unsafe {
4618 if vnni_tiles_enabled() {
4619 return dot_q4t_row_vnni(bytes, r, gpr, xq);
4620 }
4621 return dot_q4t_row_avx2(bytes, r, gpr, xq);
4622 }
4623 let mut acc = 0f32;
4624 for gi in 0..gpr {
4625 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4626 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4627 let mut d = 0i32;
4628 for (k, &b) in tile[2..].iter().enumerate() {
4629 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4630 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4631 }
4632 acc += d as f32 * s;
4633 }
4634 acc
4635}
4636
4637#[cfg(target_arch = "aarch64")]
4638#[target_feature(enable = "neon,dotprod")]
4639unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4640 unsafe {
4643 use core::arch::aarch64::*;
4644 use core::arch::asm;
4645 let lomask = vdupq_n_u8(0x0F);
4646 let eight = vdupq_n_s8(8);
4647 let mut acc = 0f32;
4648 for gi in 0..gpr {
4649 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4650 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4651 let b = vld1q_u8(t.add(2));
4652 let lo = vandq_u8(b, lomask);
4653 let hi = vshrq_n_u8::<4>(b);
4654 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4655 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4656 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4657 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4658 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4659 asm!(
4660 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4661 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4662 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4663 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4664 options(pure, nomem, nostack),
4665 );
4666 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4667 }
4668 acc
4669 }
4670}
4671
4672#[cfg(target_arch = "x86_64")]
4673#[target_feature(enable = "avx2")]
4674unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4675 unsafe {
4677 use core::arch::x86_64::*;
4678 let lomask = _mm_set1_epi8(0x0F);
4679 let eight = _mm256_set1_epi8(8);
4680 let ones = _mm256_set1_epi16(1);
4681 let mut acc = 0f32;
4682 for gi in 0..gpr {
4683 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4684 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4685 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4686 let lo = _mm_and_si128(b, lomask);
4687 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4688 let w = _mm256_sub_epi8(
4689 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4690 eight,
4691 );
4692 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4693 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4694 let d = _mm256_madd_epi16(p16, ones);
4695 let hi128 = _mm256_extracti128_si256::<1>(d);
4696 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4697 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4698 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4699 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4700 }
4701 acc
4702 }
4703}
4704
4705#[cfg(target_arch = "x86_64")]
4709#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4710unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4711 unsafe {
4713 use core::arch::x86_64::*;
4714 let lomask = _mm_set1_epi8(0x0F);
4715 let eight = _mm256_set1_epi8(8);
4716 let mut acc = 0f32;
4717 for gi in 0..gpr {
4718 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4719 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4720 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4721 let lo = _mm_and_si128(b, lomask);
4722 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4723 let w = _mm256_sub_epi8(
4724 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4725 eight,
4726 );
4727 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4728 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4729 acc += d as f32 * s;
4730 }
4731 acc
4732 }
4733}
4734
4735#[cfg(target_arch = "x86_64")]
4740#[target_feature(enable = "avx2,fma")]
4745unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4746 unsafe {
4748 use core::arch::x86_64::*;
4749 let lomask = _mm_set1_epi8(0x0F);
4750 let eight = _mm256_set1_epi8(8);
4751 let ones = _mm256_set1_epi16(1);
4752 let mut f0 = _mm256_setzero_ps();
4765 let mut f1 = _mm256_setzero_ps();
4766 let mut f2 = _mm256_setzero_ps();
4767 let mut f3 = _mm256_setzero_ps();
4768 for gi in 0..gpr {
4769 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4770 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4771 let sv = _mm256_set1_ps(s);
4772 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4773 let lo = _mm_and_si128(bb, lomask);
4774 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4775 let w = _mm256_sub_epi8(
4776 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4777 eight,
4778 );
4779 let aw = _mm256_abs_epi8(w);
4780 let off = gi * GROUP_SIZE;
4781 let dot = |xq: &[i8]| {
4782 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4783 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4784 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4785 };
4786 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4787 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4788 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4789 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4790 }
4791 [
4792 hsum256_ps(f0),
4793 hsum256_ps(f1),
4794 hsum256_ps(f2),
4795 hsum256_ps(f3),
4796 ]
4797 }
4798}
4799
4800#[cfg(target_arch = "x86_64")]
4803#[target_feature(enable = "avx2")]
4804#[inline]
4805unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4806 unsafe {
4808 use core::arch::x86_64::*;
4809 let hi = _mm256_extractf128_ps::<1>(v);
4810 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4811 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4812 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4813 _mm_cvtss_f32(s)
4814 }
4815}
4816
4817#[cfg(target_arch = "x86_64")]
4819#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4820unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4821 unsafe {
4823 use core::arch::x86_64::*;
4824 let lomask = _mm_set1_epi8(0x0F);
4825 let eight = _mm256_set1_epi8(8);
4826 let mut f0 = _mm256_setzero_ps();
4829 let mut f1 = _mm256_setzero_ps();
4830 let mut f2 = _mm256_setzero_ps();
4831 let mut f3 = _mm256_setzero_ps();
4832 for gi in 0..gpr {
4833 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4834 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4835 let sv = _mm256_set1_ps(s);
4836 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4837 let lo = _mm_and_si128(bb, lomask);
4838 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4839 let w = _mm256_sub_epi8(
4840 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4841 eight,
4842 );
4843 let aw = _mm256_abs_epi8(w);
4844 let off = gi * GROUP_SIZE;
4845 let dot = |xq: &[i8]| {
4846 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4847 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4848 _mm256_setzero_si256(),
4849 aw,
4850 _mm256_sign_epi8(x, w),
4851 ))
4852 };
4853 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4854 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4855 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4856 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4857 }
4858 let acc = [
4859 hsum256_ps(f0),
4860 hsum256_ps(f1),
4861 hsum256_ps(f2),
4862 hsum256_ps(f3),
4863 ];
4864 acc
4865 }
4866}
4867
4868#[cfg(target_arch = "aarch64")]
4873#[target_feature(enable = "neon,dotprod")]
4874unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4875 unsafe {
4877 use core::arch::aarch64::*;
4878 use core::arch::asm;
4879 let lomask = vdupq_n_u8(0x0F);
4880 let eight = vdupq_n_s8(8);
4881 let mut acc = [0f32; 4];
4882 for gi in 0..gpr {
4883 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4884 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4885 let b = vld1q_u8(t.add(2));
4886 let lo = vandq_u8(b, lomask);
4887 let hi = vshrq_n_u8::<4>(b);
4888 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4889 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4890 for (k, xq) in xs.iter().enumerate() {
4891 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4892 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4893 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4894 asm!(
4895 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4896 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4897 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4898 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4899 options(pure, nomem, nostack),
4900 );
4901 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4902 }
4903 }
4904 acc
4905 }
4906}
4907
4908#[inline]
4910fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4911 let gi = j / GROUP_SIZE;
4912 let k = j % GROUP_SIZE;
4913 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4914 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4915 let byte = tile[2 + k / 2];
4916 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4917 ((nib as i32 - 8) as f32, s)
4918}
4919
4920#[inline]
4923fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4924 let mut acc = 0f32;
4925 for gi in 0..gpr {
4926 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4927 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4928 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4929 let mut ga = 0f32;
4930 for (k, &b) in tile[2..].iter().enumerate() {
4931 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4932 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4933 }
4934 acc += ga * s;
4935 }
4936 acc
4937}
4938
4939struct Q4tpView<'a> {
4943 nib: &'a [u8],
4944 params: &'a [u8],
4945 codes: &'a [u8],
4946 stride: usize,
4947 zero_rung: bool,
4949}
4950
4951impl<'a> Q4tpView<'a> {
4952 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4953 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4954 Self {
4955 nib: &bytes[..params_off],
4956 params: &bytes[params_off..codes_off],
4957 codes: &bytes[codes_off..],
4958 stride,
4959 zero_rung: false,
4960 }
4961 }
4962
4963 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4965 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4966 Self {
4967 nib: &bytes[..params_off],
4968 params: &bytes[params_off..codes_off],
4969 codes: &bytes[codes_off..],
4970 stride,
4971 zero_rung: true,
4972 }
4973 }
4974
4975 #[inline]
4991 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4992 let tab = if self.zero_rung {
4993 q2tp_ladder(self.params, r)
4994 } else {
4995 q4tp_ladder(self.params, r)
4996 };
4997 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4998 let out = &mut out[..gpr];
4999 let mut chunks = out.chunks_exact_mut(8);
5000 let mut ci = 0usize;
5001 for c in &mut chunks {
5002 let w = u64::from(codes[ci])
5003 | u64::from(codes[ci + 1]) << 8
5004 | u64::from(codes[ci + 2]) << 16
5005 | u64::from(codes[ci + 3]) << 24
5006 | u64::from(codes[ci + 4]) << 32;
5007 for (k, o) in c.iter_mut().enumerate() {
5008 *o = tab[((w >> (5 * k)) & 31) as usize];
5009 }
5010 ci += 5;
5011 }
5012 let tail = &codes[ci..];
5015 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5016 *o = tab[q4tp_code(tail, k)];
5017 }
5018 }
5019}
5020
5021#[inline]
5022fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5023 #[cfg(target_arch = "aarch64")]
5024 unsafe {
5025 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5026 }
5027 #[cfg(target_arch = "x86_64")]
5028 unsafe {
5029 if vnni_tiles_enabled() {
5030 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5031 }
5032 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5033 }
5034 #[allow(unreachable_code)]
5035 {
5036 let mut acc = 0f32;
5037 for gi in 0..gpr {
5038 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5039 let s = scales[gi];
5040 let mut d = 0i32;
5041 for (k, &b) in tile.iter().enumerate() {
5042 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5043 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5044 }
5045 acc += d as f32 * s;
5046 }
5047 acc
5048 }
5049}
5050
5051#[cfg(target_arch = "aarch64")]
5054#[target_feature(enable = "neon,dotprod")]
5055unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5056 unsafe {
5059 use core::arch::aarch64::*;
5060 use core::arch::asm;
5061 let lomask = vdupq_n_u8(0x0F);
5062 let eight = vdupq_n_s8(8);
5063 let mut acc = 0f32;
5064 for gi in 0..gpr {
5065 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5066 let s = *scales.get_unchecked(gi);
5067 let b = vld1q_u8(t);
5068 let lo = vandq_u8(b, lomask);
5069 let hi = vshrq_n_u8::<4>(b);
5070 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5071 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5072 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5073 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5074 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5075 asm!(
5076 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5077 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5078 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5079 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5080 options(pure, nomem, nostack),
5081 );
5082 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5083 }
5084 acc
5085 }
5086}
5087
5088#[cfg(target_arch = "x86_64")]
5089#[target_feature(enable = "avx2")]
5090unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5091 unsafe {
5093 use core::arch::x86_64::*;
5094 let lomask = _mm_set1_epi8(0x0F);
5095 let eight = _mm256_set1_epi8(8);
5096 let ones = _mm256_set1_epi16(1);
5097 let mut acc = 0f32;
5098 for gi in 0..gpr {
5099 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5100 let s = *scales.get_unchecked(gi);
5101 let b = _mm_loadu_si128(t as *const __m128i);
5102 let lo = _mm_and_si128(b, lomask);
5103 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5104 let w = _mm256_sub_epi8(
5105 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5106 eight,
5107 );
5108 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5109 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5110 let d = _mm256_madd_epi16(p16, ones);
5111 let hi128 = _mm256_extracti128_si256::<1>(d);
5112 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5113 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5114 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5115 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5116 }
5117 acc
5118 }
5119}
5120
5121#[cfg(target_arch = "x86_64")]
5124#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5125unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5126 unsafe {
5128 use core::arch::x86_64::*;
5129 let lomask = _mm_set1_epi8(0x0F);
5130 let eight = _mm256_set1_epi8(8);
5131 let mut acc = 0f32;
5132 for gi in 0..gpr {
5133 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5134 let s = *scales.get_unchecked(gi);
5135 let b = _mm_loadu_si128(t as *const __m128i);
5136 let lo = _mm_and_si128(b, lomask);
5137 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5138 let w = _mm256_sub_epi8(
5139 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5140 eight,
5141 );
5142 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5143 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5144 }
5145 acc
5146 }
5147}
5148
5149#[inline]
5152fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5153 let mut acc = 0f32;
5154 for gi in 0..gpr {
5155 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5156 let s = scales[gi];
5157 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5158 let mut ga = 0f32;
5159 for (k, &b) in tile.iter().enumerate() {
5160 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5161 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5162 }
5163 acc += ga * s;
5164 }
5165 acc
5166}
5167
5168#[inline]
5171fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5172 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5173 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5174 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5175 ((n as i32 - 8) as f32, scales[gi])
5176}
5177
5178fn q4tp_matvec(
5180 bytes: &[u8],
5181 x: &[f32],
5182 rows: usize,
5183 cols: usize,
5184 out: &mut [f32],
5185 pool: Option<&Pool>,
5186) {
5187 debug_assert_eq!(out.len(), rows);
5188 let gpr = cols / GROUP_SIZE;
5189 let v = Q4tpView::new(bytes, rows, cols);
5190 let out_addr = SendMut(out.as_mut_ptr());
5191 if a8w8_enabled() {
5192 let act = split_act(x);
5193 let run = |start: usize, end: usize| {
5194 with_krow(gpr, |sc| {
5196 for r in start..end {
5197 v.scales_into(r, gpr, sc);
5198 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5199 for &(j, xv) in &act.outliers {
5200 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5201 acc += w * s * xv;
5202 }
5203 unsafe { *out_addr.at(r) = acc };
5205 }
5206 })
5207 };
5208 dispatch_rows(pool, rows, &run);
5209 return;
5210 }
5211 let run = |start: usize, end: usize| {
5212 with_krow(gpr, |sc| {
5213 for r in start..end {
5214 v.scales_into(r, gpr, sc);
5215 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5217 }
5218 })
5219 };
5220 dispatch_rows(pool, rows, &run);
5221}
5222
5223#[allow(clippy::too_many_arguments)]
5226fn q4tp_matvec2(
5227 bytes: &[u8],
5228 x1: &[f32],
5229 x2: &[f32],
5230 rows: usize,
5231 cols: usize,
5232 o1: &mut [f32],
5233 o2: &mut [f32],
5234 pool: Option<&Pool>,
5235) {
5236 let gpr = cols / GROUP_SIZE;
5237 let v = Q4tpView::new(bytes, rows, cols);
5238 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5239 let run = |start: usize, end: usize| {
5240 let mut sc = vec![0f32; gpr];
5241 for r in start..end {
5242 v.scales_into(r, gpr, &mut sc);
5243 unsafe {
5245 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5246 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5247 }
5248 }
5249 };
5250 dispatch_rows(pool, rows, &run);
5251}
5252
5253#[inline]
5256fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5257 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5258 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5259 let c = (byte >> (2 * (k % 4))) & 3;
5260 (c as f32 - 1.5, scales[gi])
5261}
5262
5263#[cfg(target_arch = "x86_64")]
5264const Q2TP_DECODE_U32: [u32; 256] = {
5265 let mut tab = [0u32; 256];
5266 let mut b = 0usize;
5267 while b < 256 {
5268 tab[b] = ((b as u32) & 3)
5269 | ((((b as u32) >> 2) & 3) << 8)
5270 | ((((b as u32) >> 4) & 3) << 16)
5271 | ((((b as u32) >> 6) & 3) << 24);
5272 b += 1;
5273 }
5274 tab
5275};
5276
5277#[cfg(target_arch = "x86_64")]
5281#[target_feature(enable = "avx2")]
5282unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5283 use core::arch::x86_64::*;
5284 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5285 let codes = _mm256_setr_epi32(
5286 Q2TP_DECODE_U32[ch[0] as usize] as i32,
5287 Q2TP_DECODE_U32[ch[1] as usize] as i32,
5288 Q2TP_DECODE_U32[ch[2] as usize] as i32,
5289 Q2TP_DECODE_U32[ch[3] as usize] as i32,
5290 Q2TP_DECODE_U32[ch[4] as usize] as i32,
5291 Q2TP_DECODE_U32[ch[5] as usize] as i32,
5292 Q2TP_DECODE_U32[ch[6] as usize] as i32,
5293 Q2TP_DECODE_U32[ch[7] as usize] as i32,
5294 );
5295 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5296 let pair = _mm256_maddubs_epi16(codes, xv);
5297 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5298 let sum128 = _mm_add_epi32(
5299 _mm256_castsi256_si128(quad),
5300 _mm256_extracti128_si256(quad, 1),
5301 );
5302 let sum64 = _mm_hadd_epi32(sum128, sum128);
5303 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5304}
5305
5306#[inline]
5313fn dot_q2tp_row_i8(
5314 chunks: &[u8],
5315 r: usize,
5316 gpr: usize,
5317 xq: &[i8],
5318 gsum: &[i32],
5319 scales: &[f32],
5320) -> f32 {
5321 let mut acc = 0f32;
5322 let base = r * gpr * Q2TP_CHUNK;
5323 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5324 let mut codes = [0i8; GROUP_SIZE];
5325 #[cfg(target_arch = "x86_64")]
5326 let avx2 = std::arch::is_x86_feature_detected!("avx2");
5327 for gi in 0..gpr {
5328 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5329 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5330 #[cfg(target_arch = "aarch64")]
5331 let dot = unsafe {
5337 use core::arch::aarch64::*;
5338 let b = vld1_u8(ch.as_ptr());
5339 let three = vdup_n_u8(3);
5340 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5341 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5342 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5343 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5344 let x4 = vld4_s8(xg.as_ptr());
5345 let mut acc4 = vdupq_n_s32(0);
5346 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5347 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5348 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5349 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5350 vaddvq_s32(acc4)
5351 };
5352 #[cfg(target_arch = "x86_64")]
5353 let dot: i32 = if avx2 {
5354 unsafe { q2tp_code_dot_avx2(ch, xg) }
5357 } else {
5358 ch.iter()
5359 .enumerate()
5360 .map(|(k, &b)| {
5361 ((b & 3) as i32) * xg[k * 4] as i32
5362 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5363 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5364 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5365 })
5366 .sum()
5367 };
5368 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5369 let dot: i32 = {
5370 for (k, &b) in ch.iter().enumerate() {
5371 codes[k * 4] = (b & 3) as i8;
5372 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5373 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5374 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5375 }
5376 codes
5377 .iter()
5378 .zip(xg)
5379 .map(|(&c, &x)| c as i32 * x as i32)
5380 .sum()
5381 };
5382 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5383 }
5384 acc
5385}
5386
5387fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5391 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5392}
5393
5394#[inline]
5398fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5399 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
5400}
5401
5402#[inline]
5403fn q2tp_row_exact_center(
5404 chunks: &[u8],
5405 r: usize,
5406 gpr: usize,
5407 x: &[f32],
5408 scales: &[f32],
5409 center: f32,
5410) -> f32 {
5411 let mut acc = 0f32;
5412 for gi in 0..gpr {
5413 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
5414 let s = scales[gi];
5415 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5416 let mut g = 0f32;
5417 for (k, &b) in ch.iter().enumerate() {
5418 g += ((b & 3) as f32 - center) * xb[k * 4]
5419 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
5420 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
5421 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
5422 }
5423 acc += s * g;
5424 }
5425 acc
5426}
5427
5428fn q2tp_matvec(
5429 bytes: &[u8],
5430 x: &[f32],
5431 rows: usize,
5432 cols: usize,
5433 out: &mut [f32],
5434 pool: Option<&Pool>,
5435) {
5436 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
5437}
5438
5439fn q2tp_affine_matvec(
5440 bytes: &[u8],
5441 x: &[f32],
5442 rows: usize,
5443 cols: usize,
5444 out: &mut [f32],
5445 pool: Option<&Pool>,
5446) {
5447 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
5448}
5449
5450fn q2tp_matvec_mode(
5451 bytes: &[u8],
5452 x: &[f32],
5453 rows: usize,
5454 cols: usize,
5455 out: &mut [f32],
5456 pool: Option<&Pool>,
5457 affine: bool,
5458) {
5459 debug_assert_eq!(out.len(), rows);
5460 let gpr = cols / GROUP_SIZE;
5461 let v = Q4tpView::new_q2(bytes, rows, cols);
5462 let out_addr = SendMut(out.as_mut_ptr());
5463 if !affine && a8w8_enabled() {
5468 let act = split_act(x);
5469 let gsum = q1_group_sums(&act.xq, gpr);
5470 let (act, gsum) = (&act, &gsum);
5471 let run = move |start: usize, end: usize| {
5472 with_krow(gpr, |sc| {
5473 for r in start..end {
5474 v.scales_into(r, gpr, sc);
5475 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
5476 for &(j, xv) in &act.outliers {
5477 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
5478 acc += w * s * xv;
5479 }
5480 unsafe { *out_addr.at(r) = acc };
5482 }
5483 })
5484 };
5485 dispatch_rows(pool, rows, &run);
5486 return;
5487 }
5488 let run = |start: usize, end: usize| {
5489 with_krow(gpr, |sc| {
5490 for r in start..end {
5491 v.scales_into(r, gpr, sc);
5492 unsafe {
5494 *out_addr.at(r) = if affine {
5495 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
5496 } else {
5497 q2tp_row_exact(v.nib, r, gpr, x, sc)
5498 }
5499 };
5500 }
5501 })
5502 };
5503 dispatch_rows(pool, rows, &run);
5504}
5505
5506#[allow(clippy::too_many_arguments)]
5508fn q2tp_matvec2(
5509 bytes: &[u8],
5510 x1: &[f32],
5511 x2: &[f32],
5512 rows: usize,
5513 cols: usize,
5514 o1: &mut [f32],
5515 o2: &mut [f32],
5516 pool: Option<&Pool>,
5517) {
5518 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
5519}
5520
5521#[allow(clippy::too_many_arguments)]
5522fn q2tp_affine_matvec2(
5523 bytes: &[u8],
5524 x1: &[f32],
5525 x2: &[f32],
5526 rows: usize,
5527 cols: usize,
5528 o1: &mut [f32],
5529 o2: &mut [f32],
5530 pool: Option<&Pool>,
5531) {
5532 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
5533}
5534
5535#[allow(clippy::too_many_arguments)]
5536fn q2tp_matvec2_mode(
5537 bytes: &[u8],
5538 x1: &[f32],
5539 x2: &[f32],
5540 rows: usize,
5541 cols: usize,
5542 o1: &mut [f32],
5543 o2: &mut [f32],
5544 pool: Option<&Pool>,
5545 affine: bool,
5546) {
5547 let gpr = cols / GROUP_SIZE;
5548 let v = Q4tpView::new_q2(bytes, rows, cols);
5549 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5550 let run = |start: usize, end: usize| {
5551 let mut sc = vec![0f32; gpr];
5552 for r in start..end {
5553 v.scales_into(r, gpr, &mut sc);
5554 unsafe {
5556 *p1.at(r) = if affine {
5557 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
5558 } else {
5559 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
5560 };
5561 *p2.at(r) = if affine {
5562 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
5563 } else {
5564 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
5565 };
5566 }
5567 }
5568 };
5569 dispatch_rows(pool, rows, &run);
5570}
5571
5572pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
5579 let gpr = cols / GROUP_SIZE;
5584 let v = Q4tpView::new_q2(bytes, rows, cols);
5585 with_krow(gpr, |sc| {
5586 for r in 0..rows {
5587 v.scales_into(r, gpr, sc);
5588 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
5589 }
5590 });
5591}
5592
5593pub fn q2tp_affine_matvec_for_test(
5596 bytes: &[u8],
5597 x: &[f32],
5598 rows: usize,
5599 cols: usize,
5600 out: &mut [f32],
5601) {
5602 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
5603}
5604
5605pub fn q2tp_matmat_for_test(
5606 bytes: &[u8],
5607 xs_all: &[f32],
5608 b: usize,
5609 rows: usize,
5610 cols: usize,
5611 out: &mut [f32],
5612) {
5613 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
5614}
5615
5616fn q2tp_matmat(
5617 bytes: &[u8],
5618 xs_all: &[f32],
5619 b: usize,
5620 rows: usize,
5621 cols: usize,
5622 out: &mut [f32],
5623 pool: Option<&Pool>,
5624) {
5625 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
5626}
5627
5628fn q2tp_affine_matmat(
5629 bytes: &[u8],
5630 xs_all: &[f32],
5631 b: usize,
5632 rows: usize,
5633 cols: usize,
5634 out: &mut [f32],
5635 pool: Option<&Pool>,
5636) {
5637 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
5638}
5639
5640fn q2tp_matmat_mode(
5641 bytes: &[u8],
5642 xs_all: &[f32],
5643 b: usize,
5644 rows: usize,
5645 cols: usize,
5646 out: &mut [f32],
5647 pool: Option<&Pool>,
5648 affine: bool,
5649) {
5650 debug_assert_eq!(out.len(), b * rows);
5651 let gpr = cols / GROUP_SIZE;
5652 let v = Q4tpView::new_q2(bytes, rows, cols);
5653 let out_addr = SendMut(out.as_mut_ptr());
5654 let run = |start: usize, end: usize| {
5655 let mut sc = vec![0f32; gpr];
5656 for r in start..end {
5657 v.scales_into(r, gpr, &mut sc);
5658 for bi in 0..b {
5659 let x = &xs_all[bi * cols..(bi + 1) * cols];
5660 unsafe {
5662 *out_addr.at(bi * rows + r) = if affine {
5663 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
5664 } else {
5665 q2tp_row_exact(v.nib, r, gpr, x, &sc)
5666 }
5667 };
5668 }
5669 }
5670 };
5671 dispatch_rows(pool, rows, &run);
5672}
5673
5674#[cfg(target_arch = "aarch64")]
5678#[target_feature(enable = "neon,dotprod")]
5679unsafe fn dot_q4tp_row_1x4_sdot_v1(
5680 nib: &[u8],
5681 r: usize,
5682 gpr: usize,
5683 xs: [&[i8]; 4],
5684 scales: &[f32],
5685) -> [f32; 4] {
5686 unsafe {
5687 use core::arch::aarch64::*;
5688 use core::arch::asm;
5689 let lomask = vdupq_n_u8(0x0F);
5690 let eight = vdupq_n_s8(8);
5691 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
5692 for gi in 0..gpr {
5693 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5694 let s = *scales.get_unchecked(gi);
5695 let bb = vld1q_u8(t);
5696 let lo = vandq_u8(bb, lomask);
5697 let hi = vshrq_n_u8::<4>(bb);
5698 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5699 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5700 let mut d = [0f32; 4];
5701 for (k, dk) in d.iter_mut().enumerate() {
5702 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5703 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5704 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5705 asm!(
5706 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5707 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5708 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5709 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5710 options(pure, nomem, nostack),
5711 );
5712 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5713 }
5714 f0 += d[0];
5715 f1 += d[1];
5716 f2 += d[2];
5717 f3 += d[3];
5718 }
5719 [f0, f1, f2, f3]
5720 }
5721}
5722
5723#[allow(dead_code)]
5730static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
5731
5732#[cfg(target_arch = "x86_64")]
5738fn q4tp_blocked_x86() -> bool {
5739 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5740 1 => false,
5741 2 => avx512vnni_enabled(),
5746 _ => blocked_enabled() && avx512vnni_enabled(),
5750 }
5751}
5752
5753#[cfg(target_arch = "aarch64")]
5755#[allow(dead_code)]
5756fn q4tp_v1() -> bool {
5757 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5758 1 => true,
5759 2 => false,
5760 _ => {
5761 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5762 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
5763 }
5764 }
5765}
5766
5767#[cfg(target_arch = "x86_64")]
5779#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5780unsafe fn dot_q4tp_2x8_avx512(
5781 nib: &[u8],
5782 r0: usize,
5783 gpr: usize,
5784 xs: [&[i8]; 8],
5785 sc0: &[f32],
5786 sc1: &[f32],
5787) -> [[f32; 8]; 2] {
5788 unsafe {
5791 use core::arch::x86_64::*;
5792 let lomask = _mm256_set1_epi8(0x0F);
5793 let eight = _mm256_set1_epi8(8);
5794 let zero = _mm512_setzero_si512();
5795 let mut v0 = [_mm512_setzero_ps(); 8];
5796 let mut v1 = [_mm512_setzero_ps(); 8];
5797 let pairs = gpr / 2;
5798 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
5799 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5800 let bb = _mm256_loadu_si256(t as *const __m256i);
5801 let lo = _mm256_and_si256(bb, lomask);
5802 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5803 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5804 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5805 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5806 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5807 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5808 };
5809 for gp in 0..pairs {
5810 let gi = gp * 2;
5811 let (wa0, neg0) = unpack(r0, gi);
5812 let (wa1, neg1) = unpack(r0 + 1, gi);
5813 let off = gi * GROUP_SIZE;
5814 let sv = |sc: &[f32]| {
5815 _mm512_insertf32x8::<1>(
5816 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5817 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5818 )
5819 };
5820 let s0 = sv(sc0);
5821 let s1 = sv(sc1);
5822 for k in 0..8 {
5823 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5824 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5825 zero,
5826 wa0,
5827 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5828 ));
5829 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5830 zero,
5831 wa1,
5832 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5833 ));
5834 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5835 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5836 }
5837 }
5838 let mut acc = [[0f32; 8]; 2];
5839 for k in 0..8 {
5840 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5841 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5842 }
5843 if gpr % 2 == 1 {
5844 let off = (gpr - 1) * GROUP_SIZE;
5845 for j in off..off + GROUP_SIZE {
5846 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5847 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5848 for k in 0..8 {
5849 let x = *xs[k].get_unchecked(j) as f32;
5850 acc[0][k] += w0 * sa * x;
5851 acc[1][k] += w1 * sb * x;
5852 }
5853 }
5854 }
5855 acc
5856 }
5857}
5858
5859#[cfg(target_arch = "x86_64")]
5864#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5865unsafe fn dot_q4tp_row_1x8_avx512(
5866 nib: &[u8],
5867 r: usize,
5868 gpr: usize,
5869 xs: [&[i8]; 8],
5870 scales: &[f32],
5871) -> [f32; 8] {
5872 unsafe {
5874 use core::arch::x86_64::*;
5875 let lomask = _mm256_set1_epi8(0x0F);
5876 let eight = _mm256_set1_epi8(8);
5877 let zero = _mm512_setzero_si512();
5878 let (mut v0, mut v1, mut v2, mut v3) = (
5879 _mm512_setzero_ps(),
5880 _mm512_setzero_ps(),
5881 _mm512_setzero_ps(),
5882 _mm512_setzero_ps(),
5883 );
5884 let (mut v4, mut v5, mut v6, mut v7) = (
5885 _mm512_setzero_ps(),
5886 _mm512_setzero_ps(),
5887 _mm512_setzero_ps(),
5888 _mm512_setzero_ps(),
5889 );
5890 let pairs = gpr / 2;
5891 for gp in 0..pairs {
5892 let gi = gp * 2;
5893 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5894 let bb = _mm256_loadu_si256(t as *const __m256i);
5895 let lo = _mm256_and_si256(bb, lomask);
5896 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5897 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5902 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5903 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5904 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5905 let wabs = _mm512_abs_epi8(w);
5906 let neg = _mm512_movepi8_mask(w);
5907 let off = gi * GROUP_SIZE;
5908 let sv = _mm512_insertf32x8::<1>(
5909 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5910 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5911 );
5912 let dot = |x: &[i8]| -> __m512 {
5913 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5914 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5915 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5916 };
5917 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5918 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5919 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5920 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5921 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5922 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5923 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5924 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5925 }
5926 let mut acc = [
5927 _mm512_reduce_add_ps(v0),
5928 _mm512_reduce_add_ps(v1),
5929 _mm512_reduce_add_ps(v2),
5930 _mm512_reduce_add_ps(v3),
5931 _mm512_reduce_add_ps(v4),
5932 _mm512_reduce_add_ps(v5),
5933 _mm512_reduce_add_ps(v6),
5934 _mm512_reduce_add_ps(v7),
5935 ];
5936 if gpr % 2 == 1 {
5939 let off = (gpr - 1) * GROUP_SIZE;
5940 for j in off..off + GROUP_SIZE {
5941 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5942 let ws = w * s;
5943 for k in 0..8 {
5944 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5945 }
5946 }
5947 }
5948 acc
5949 }
5950}
5951
5952#[cfg(target_arch = "x86_64")]
5965#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5966unsafe fn dot_q4tp_row_1x4_avx512(
5967 nib: &[u8],
5968 r: usize,
5969 gpr: usize,
5970 xs: [&[i8]; 4],
5971 scales: &[f32],
5972) -> [f32; 4] {
5973 unsafe {
5975 use core::arch::x86_64::*;
5976 let lomask = _mm256_set1_epi8(0x0F);
5977 let eight = _mm256_set1_epi8(8);
5978 let zero = _mm512_setzero_si512();
5979 let (mut v0, mut v1, mut v2, mut v3) = (
5980 _mm512_setzero_ps(),
5981 _mm512_setzero_ps(),
5982 _mm512_setzero_ps(),
5983 _mm512_setzero_ps(),
5984 );
5985 let pairs = gpr / 2;
5986 for gp in 0..pairs {
5987 let gi = gp * 2;
5988 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5989 let bb = _mm256_loadu_si256(t as *const __m256i);
5990 let lo = _mm256_and_si256(bb, lomask);
5991 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5992 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5997 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5998 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5999 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6000 let wabs = _mm512_abs_epi8(w);
6001 let neg = _mm512_movepi8_mask(w);
6002 let off = gi * GROUP_SIZE;
6003 let sv = _mm512_insertf32x8::<1>(
6004 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6005 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6006 );
6007 let dot = |x: &[i8]| -> __m512 {
6008 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6009 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6010 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6011 };
6012 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6013 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6014 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6015 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6016 }
6017 let mut acc = [
6018 _mm512_reduce_add_ps(v0),
6019 _mm512_reduce_add_ps(v1),
6020 _mm512_reduce_add_ps(v2),
6021 _mm512_reduce_add_ps(v3),
6022 ];
6023 if gpr % 2 == 1 {
6026 let off = (gpr - 1) * GROUP_SIZE;
6027 for j in off..off + GROUP_SIZE {
6028 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6029 let ws = w * s;
6030 for k in 0..4 {
6031 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6032 }
6033 }
6034 }
6035 acc
6036 }
6037}
6038
6039#[cfg(target_arch = "aarch64")]
6043#[target_feature(enable = "neon,dotprod")]
6044unsafe fn dot_q4tp_row_1x4_sdot(
6045 nib: &[u8],
6046 r: usize,
6047 gpr: usize,
6048 xs: [&[i8]; 4],
6049 scales: &[f32],
6050) -> [f32; 4] {
6051 unsafe {
6053 use core::arch::aarch64::*;
6054 use core::arch::asm;
6055 let lomask = vdupq_n_u8(0x0F);
6056 let eight = vdupq_n_s8(8);
6057 let (mut v0, mut v1, mut v2, mut v3) = (
6073 vdupq_n_f32(0.0),
6074 vdupq_n_f32(0.0),
6075 vdupq_n_f32(0.0),
6076 vdupq_n_f32(0.0),
6077 );
6078 for gi in 0..gpr {
6079 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6080 let s = *scales.get_unchecked(gi);
6081 let bb = vld1q_u8(t);
6082 let lo = vandq_u8(bb, lomask);
6083 let hi = vshrq_n_u8::<4>(bb);
6084 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6085 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6086 let off = gi * GROUP_SIZE;
6087 let dot4 = |x: &[i8]| -> int32x4_t {
6088 let x0 = vld1q_s8(x.as_ptr().add(off));
6089 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6090 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6091 asm!(
6092 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6093 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6094 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6095 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6096 options(pure, nomem, nostack),
6097 );
6098 vaddq_s32(a0, a1)
6099 };
6100 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6101 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6102 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6103 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6104 }
6105 [
6106 vaddvq_f32(v0),
6107 vaddvq_f32(v1),
6108 vaddvq_f32(v2),
6109 vaddvq_f32(v3),
6110 ]
6111 }
6112}
6113
6114fn q4tp_matmat(
6118 bytes: &[u8],
6119 xs_all: &[f32],
6120 b: usize,
6121 rows: usize,
6122 cols: usize,
6123 out: &mut [f32],
6124 pool: Option<&Pool>,
6125) {
6126 debug_assert_eq!(out.len(), b * rows);
6127 let gpr = cols / GROUP_SIZE;
6128 let v = Q4tpView::new(bytes, rows, cols);
6129
6130 #[cfg(target_os = "macos")]
6132 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6133 dequant_matmat_accel(
6134 &|r, dst| {
6135 let mut sc = [0f32; 32];
6136 let mut scv;
6137 let s: &[f32] = if gpr <= 32 {
6138 v.scales_into(r, gpr, &mut sc);
6139 &sc[..gpr]
6140 } else {
6141 scv = vec![0f32; gpr];
6142 v.scales_into(r, gpr, &mut scv);
6143 &scv
6144 };
6145 for gi in 0..gpr {
6146 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6147 for (k, &bb) in tile.iter().enumerate() {
6148 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6149 dst[gi * GROUP_SIZE + k * 2 + 1] =
6150 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6151 }
6152 }
6153 },
6154 xs_all,
6155 b,
6156 rows,
6157 cols,
6158 out,
6159 pool,
6160 );
6161 return;
6162 }
6163
6164 let out_addr = SendMut(out.as_mut_ptr());
6165 if a8w8_enabled() {
6166 let acts: Vec<SplitAct> = (0..b)
6167 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6168 .collect();
6169 let acts = &acts;
6170 #[cfg(target_arch = "aarch64")]
6171 let blocked_ok = sdot_enabled() && blocked_enabled();
6172 #[cfg(target_arch = "x86_64")]
6179 let blocked_ok = q4tp_blocked_x86();
6180 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6181 let blocked_ok = false;
6182 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6192 .ok()
6193 .and_then(|v| v.parse().ok())
6194 .filter(|v| *v > 0)
6195 .unwrap_or(256);
6196 let run = |start: usize, end: usize| {
6197 for abase in (0..acts.len()).step_by(panel_cols) {
6198 let alen = (acts.len() - abase).min(panel_cols);
6199 let mut sc = vec![0f32; gpr];
6200 #[cfg(target_arch = "x86_64")]
6201 let mut r_lo = start;
6202 #[cfg(target_arch = "x86_64")]
6203 if blocked_ok && alen >= 8 {
6204 let mut sc1 = vec![0f32; gpr];
6205 while r_lo + 2 <= end {
6206 v.scales_into(r_lo, gpr, &mut sc);
6207 v.scales_into(r_lo + 1, gpr, &mut sc1);
6208 let mut bi = 0usize;
6209 while bi + 8 <= alen {
6210 let xs = [
6211 acts[abase + bi].xq.as_slice(),
6212 acts[abase + bi + 1].xq.as_slice(),
6213 acts[abase + bi + 2].xq.as_slice(),
6214 acts[abase + bi + 3].xq.as_slice(),
6215 acts[abase + bi + 4].xq.as_slice(),
6216 acts[abase + bi + 5].xq.as_slice(),
6217 acts[abase + bi + 6].xq.as_slice(),
6218 acts[abase + bi + 7].xq.as_slice(),
6219 ];
6220 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6221 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6222 for k in 0..8 {
6223 let act = &acts[abase + bi + k];
6224 let mut acc = dr[k] * act.sx;
6225 for &(j, xv) in &act.outliers {
6226 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6227 acc += w * s * xv;
6228 }
6229 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6231 }
6232 }
6233 bi += 8;
6234 }
6235 for row in [r_lo, r_lo + 1] {
6238 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6239 for b2 in bi..alen {
6240 let act = &acts[abase + b2];
6241 let xs4 = [
6242 act.xq.as_slice(),
6243 act.xq.as_slice(),
6244 act.xq.as_slice(),
6245 act.xq.as_slice(),
6246 ];
6247 let d =
6248 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6249 let mut acc = d[0] * act.sx;
6250 for &(j, xv) in &act.outliers {
6251 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6252 acc += w * s * xv;
6253 }
6254 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6256 }
6257 }
6258 r_lo += 2;
6259 }
6260 }
6261 #[cfg(target_arch = "x86_64")]
6262 let row_start = r_lo;
6263 #[cfg(not(target_arch = "x86_64"))]
6264 let row_start = start;
6265 for r in row_start..end {
6266 v.scales_into(r, gpr, &mut sc);
6267 let mut bi = 0usize;
6268 #[cfg(target_arch = "x86_64")]
6269 if blocked_ok {
6270 while bi + 8 <= alen {
6271 let xs = [
6272 acts[abase + bi].xq.as_slice(),
6273 acts[abase + bi + 1].xq.as_slice(),
6274 acts[abase + bi + 2].xq.as_slice(),
6275 acts[abase + bi + 3].xq.as_slice(),
6276 acts[abase + bi + 4].xq.as_slice(),
6277 acts[abase + bi + 5].xq.as_slice(),
6278 acts[abase + bi + 6].xq.as_slice(),
6279 acts[abase + bi + 7].xq.as_slice(),
6280 ];
6281 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6282 for k in 0..8 {
6283 let act = &acts[abase + bi + k];
6284 let mut acc = d[k] * act.sx;
6285 for &(j, xv) in &act.outliers {
6286 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6287 acc += w * s * xv;
6288 }
6289 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6291 }
6292 bi += 8;
6293 }
6294 while bi + 4 <= alen {
6295 let xs = [
6296 acts[abase + bi].xq.as_slice(),
6297 acts[abase + bi + 1].xq.as_slice(),
6298 acts[abase + bi + 2].xq.as_slice(),
6299 acts[abase + bi + 3].xq.as_slice(),
6300 ];
6301 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6302 for k in 0..4 {
6303 let act = &acts[abase + bi + k];
6304 let mut acc = d[k] * act.sx;
6305 for &(j, xv) in &act.outliers {
6306 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6307 acc += w * s * xv;
6308 }
6309 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6311 }
6312 bi += 4;
6313 }
6314 }
6315 #[cfg(target_arch = "aarch64")]
6316 if blocked_ok {
6317 while bi + 4 <= alen {
6318 let xs = [
6319 acts[abase + bi].xq.as_slice(),
6320 acts[abase + bi + 1].xq.as_slice(),
6321 acts[abase + bi + 2].xq.as_slice(),
6322 acts[abase + bi + 3].xq.as_slice(),
6323 ];
6324 let d = unsafe {
6325 if q4tp_v1() {
6326 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6327 } else {
6328 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6329 }
6330 };
6331 for k in 0..4 {
6332 let act = &acts[abase + bi + k];
6333 let mut acc = d[k] * act.sx;
6334 for &(j, xv) in &act.outliers {
6335 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6336 acc += w * s * xv;
6337 }
6338 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6340 }
6341 bi += 4;
6342 }
6343 }
6344 let _ = blocked_ok;
6345 while bi < alen {
6346 let act = &acts[abase + bi];
6347 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6348 for &(j, xv) in &act.outliers {
6349 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6350 acc += w * s * xv;
6351 }
6352 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
6354 bi += 1;
6355 }
6356 }
6357 }
6358 };
6359 dispatch_rows(pool, rows, &run);
6360 return;
6361 }
6362
6363 let run = |start: usize, end: usize| {
6364 let mut sc = vec![0f32; gpr];
6365 for r in start..end {
6366 v.scales_into(r, gpr, &mut sc);
6367 for bi in 0..b {
6368 let x = &xs_all[bi * cols..(bi + 1) * cols];
6369 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
6371 }
6372 }
6373 };
6374 dispatch_rows(pool, rows, &run);
6375}
6376
6377fn q4t_matvec(
6379 bytes: &[u8],
6380 x: &[f32],
6381 rows: usize,
6382 cols: usize,
6383 out: &mut [f32],
6384 pool: Option<&Pool>,
6385) {
6386 debug_assert_eq!(out.len(), rows);
6387 let gpr = cols / GROUP_SIZE;
6388 let out_addr = SendMut(out.as_mut_ptr());
6389 if a8w8_enabled() {
6390 let act = split_act(x);
6391 let run = move |start: usize, end: usize| {
6392 for r in start..end {
6393 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6394 for &(j, xv) in &act.outliers {
6395 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6396 acc += w * s * xv;
6397 }
6398 unsafe { *out_addr.at(r) = acc };
6400 }
6401 };
6402 dispatch_rows(pool, rows, &run);
6403 return;
6404 }
6405 let run = move |start: usize, end: usize| {
6406 for r in start..end {
6407 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
6409 }
6410 };
6411 dispatch_rows(pool, rows, &run);
6412}
6413
6414#[allow(clippy::too_many_arguments)]
6416fn q4t_matvec2(
6417 bytes: &[u8],
6418 x1: &[f32],
6419 x2: &[f32],
6420 rows: usize,
6421 cols: usize,
6422 o1: &mut [f32],
6423 o2: &mut [f32],
6424 pool: Option<&Pool>,
6425) {
6426 let gpr = cols / GROUP_SIZE;
6427 let p1 = SendMut(o1.as_mut_ptr());
6428 let p2 = SendMut(o2.as_mut_ptr());
6429 if a8w8_enabled() {
6430 let a1 = split_act(x1);
6431 let a2 = split_act(x2);
6432 let run = move |start: usize, end: usize| {
6433 for r in start..end {
6434 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
6435 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
6436 for &(j, xv) in &a1.outliers {
6437 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6438 v1 += w * s * xv;
6439 }
6440 for &(j, xv) in &a2.outliers {
6441 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6442 v2 += w * s * xv;
6443 }
6444 unsafe {
6446 *p1.at(r) = v1;
6447 *p2.at(r) = v2;
6448 }
6449 }
6450 };
6451 dispatch_rows(pool, rows, &run);
6452 return;
6453 }
6454 let run = move |start: usize, end: usize| {
6455 for r in start..end {
6456 unsafe {
6458 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
6459 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
6460 }
6461 }
6462 };
6463 dispatch_rows(pool, rows, &run);
6464}
6465
6466#[allow(clippy::too_many_arguments)]
6468#[cfg(target_os = "macos")]
6474fn dequant_matmat_accel(
6475 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
6476 xs_all: &[f32],
6477 b: usize,
6478 rows: usize,
6479 cols: usize,
6480 out: &mut [f32],
6481 pool: Option<&Pool>,
6482) {
6483 const TR: usize = 2048;
6484 thread_local! {
6485 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6486 }
6487 WTILE.with(|wt| {
6488 let mut wtile = wt.borrow_mut();
6489 wtile.resize(TR * cols, 0.0);
6490 let mut r0 = 0usize;
6491 while r0 < rows {
6492 let tr = TR.min(rows - r0);
6493 let wt_addr = SendMut(wtile.as_mut_ptr());
6494 let run = |start: usize, end: usize| {
6495 for r in start..end {
6496 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
6498 dequant_row(r0 + r, dst);
6499 }
6500 };
6501 dispatch_rows(pool, tr, &run);
6502 unsafe {
6503 accel_blas::cblas_sgemm(
6504 101, 111, 112, b as i32,
6508 tr as i32,
6509 cols as i32,
6510 1.0,
6511 xs_all.as_ptr(),
6512 cols as i32,
6513 wtile.as_ptr(),
6514 cols as i32,
6515 0.0,
6516 out.as_mut_ptr().add(r0),
6517 rows as i32,
6518 );
6519 }
6520 r0 += tr;
6521 }
6522 });
6523}
6524
6525fn q4t_matmat(
6526 bytes: &[u8],
6527 xs_all: &[f32],
6528 b: usize,
6529 rows: usize,
6530 cols: usize,
6531 out: &mut [f32],
6532 pool: Option<&Pool>,
6533) {
6534 debug_assert_eq!(out.len(), b * rows);
6535 let gpr = cols / GROUP_SIZE;
6536 #[cfg(target_os = "macos")]
6540 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6541 dequant_matmat_accel(
6542 &|r, dst| {
6543 for gi in 0..gpr {
6544 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6545 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6546 for (k, &bb) in tile[2..].iter().enumerate() {
6547 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
6548 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
6549 }
6550 }
6551 },
6552 xs_all,
6553 b,
6554 rows,
6555 cols,
6556 out,
6557 pool,
6558 );
6559 return;
6560 }
6561 let out_addr = SendMut(out.as_mut_ptr());
6562 if a8w8_enabled() {
6563 let acts: Vec<SplitAct> = (0..b)
6564 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6565 .collect();
6566 let acts = &acts;
6567 #[cfg(target_arch = "x86_64")]
6568 let blocked_ok = avx2_enabled() && blocked_enabled();
6569 #[cfg(target_arch = "aarch64")]
6570 let blocked_ok = sdot_enabled() && blocked_enabled();
6571 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
6572 let blocked_ok = false;
6573 let run = move |start: usize, end: usize| {
6574 for r in start..end {
6575 let mut bi = 0usize;
6576 #[cfg(target_arch = "aarch64")]
6577 if blocked_ok {
6578 while bi + 4 <= acts.len() {
6579 let xs = [
6580 acts[bi].xq.as_slice(),
6581 acts[bi + 1].xq.as_slice(),
6582 acts[bi + 2].xq.as_slice(),
6583 acts[bi + 3].xq.as_slice(),
6584 ];
6585 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
6586 for k in 0..4 {
6587 let act = &acts[bi + k];
6588 let mut acc = d[k] * act.sx;
6589 for &(j, xv) in &act.outliers {
6590 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6591 acc += w * sc * xv;
6592 }
6593 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6595 }
6596 bi += 4;
6597 }
6598 }
6599 #[cfg(target_arch = "x86_64")]
6600 if blocked_ok {
6601 while bi + 4 <= acts.len() {
6602 let xs = [
6603 acts[bi].xq.as_slice(),
6604 acts[bi + 1].xq.as_slice(),
6605 acts[bi + 2].xq.as_slice(),
6606 acts[bi + 3].xq.as_slice(),
6607 ];
6608 let d = unsafe {
6609 if vnni_tiles_enabled() {
6610 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
6611 } else {
6612 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
6613 }
6614 };
6615 for k in 0..4 {
6616 let act = &acts[bi + k];
6617 let mut acc = d[k] * act.sx;
6618 for &(j, xv) in &act.outliers {
6619 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6620 acc += w * sc * xv;
6621 }
6622 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6624 }
6625 bi += 4;
6626 }
6627 }
6628 let _ = blocked_ok;
6629 while bi < acts.len() {
6630 let act = &acts[bi];
6631 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6632 for &(j, xv) in &act.outliers {
6633 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6634 acc += w * s * xv;
6635 }
6636 unsafe { *out_addr.at(bi * rows + r) = acc };
6638 bi += 1;
6639 }
6640 }
6641 };
6642 dispatch_rows(pool, rows, &run);
6643 return;
6644 }
6645 let run = move |start: usize, end: usize| {
6646 for r in start..end {
6647 for bi in 0..b {
6648 let x = &xs_all[bi * cols..(bi + 1) * cols];
6649 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
6651 }
6652 }
6653 };
6654 dispatch_rows(pool, rows, &run);
6655}
6656
6657fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
6666 (0..gpr)
6667 .map(|gi| {
6668 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
6669 .iter()
6670 .map(|&v| v as i32)
6671 .sum()
6672 })
6673 .collect()
6674}
6675
6676#[inline]
6680#[allow(unreachable_code)]
6681#[cfg(target_arch = "x86_64")]
6686#[target_feature(enable = "avx2")]
6687unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6688 unsafe {
6690 use core::arch::x86_64::*;
6691 let expand = _mm256_setr_epi8(
6693 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,
6694 3, 3, 3,
6695 );
6696 let bitsel = _mm256_setr_epi8(
6697 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6698 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6699 );
6700 let ones8 = _mm256_set1_epi8(1);
6701 let ones16 = _mm256_set1_epi16(1);
6702 let mut acc = 0f32;
6703 for gi in 0..gpr {
6704 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6705 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6706 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6707 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6708 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6709 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6710 let sel = _mm256_and_si256(x, mask);
6711 let p16 = _mm256_maddubs_epi16(ones8, sel);
6713 let d32 = _mm256_madd_epi16(p16, ones16);
6714 let hi128 = _mm256_extracti128_si256::<1>(d32);
6715 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6716 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6717 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6718 let msum = _mm_cvtsi128_si32(s32);
6719 let d = 2 * msum - gsum[gi];
6722 acc += d as f32 * s;
6723 }
6724 acc
6725 }
6726}
6727
6728#[cfg(target_arch = "x86_64")]
6731#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6732unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6733 unsafe {
6735 use core::arch::x86_64::*;
6736 let expand = _mm256_setr_epi8(
6737 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,
6738 3, 3, 3,
6739 );
6740 let bitsel = _mm256_setr_epi8(
6741 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6742 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6743 );
6744 let ones8 = _mm256_set1_epi8(1);
6745 let mut acc = 0f32;
6746 for gi in 0..gpr {
6747 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6748 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6749 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6750 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6751 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6752 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6753 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6754 let d = 2 * msum - gsum[gi];
6755 acc += d as f32 * s;
6756 }
6757 acc
6758 }
6759}
6760
6761#[cfg(target_arch = "x86_64")]
6763#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6764unsafe fn dot_q1_row_1x4_vnni(
6765 bytes: &[u8],
6766 r: usize,
6767 gpr: usize,
6768 xs: [&[i8]; 4],
6769 gsums: [&[i32]; 4],
6770) -> [f32; 4] {
6771 unsafe {
6773 use core::arch::x86_64::*;
6774 let expand = _mm256_setr_epi8(
6775 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,
6776 3, 3, 3,
6777 );
6778 let bitsel = _mm256_setr_epi8(
6779 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6780 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6781 );
6782 let ones8 = _mm256_set1_epi8(1);
6783 let mut acc = [0f32; 4];
6784 for gi in 0..gpr {
6785 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6786 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6787 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6788 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6789 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6790 for (k, xq) in xs.iter().enumerate() {
6791 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6792 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6793 let d = 2 * msum - gsums[k][gi];
6794 acc[k] += d as f32 * s;
6795 }
6796 }
6797 acc
6798 }
6799}
6800
6801#[cfg(target_arch = "x86_64")]
6804#[target_feature(enable = "avx2")]
6805unsafe fn dot_q1_row_1x4_avx2(
6806 bytes: &[u8],
6807 r: usize,
6808 gpr: usize,
6809 xs: [&[i8]; 4],
6810 gsums: [&[i32]; 4],
6811) -> [f32; 4] {
6812 unsafe {
6814 use core::arch::x86_64::*;
6815 let expand = _mm256_setr_epi8(
6816 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,
6817 3, 3, 3,
6818 );
6819 let bitsel = _mm256_setr_epi8(
6820 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6821 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6822 );
6823 let ones8 = _mm256_set1_epi8(1);
6824 let ones16 = _mm256_set1_epi16(1);
6825 let mut acc = [0f32; 4];
6826 for gi in 0..gpr {
6827 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6828 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6829 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6830 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6831 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6832 for (k, xq) in xs.iter().enumerate() {
6833 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6834 let sel = _mm256_and_si256(x, mask);
6835 let p16 = _mm256_maddubs_epi16(ones8, sel);
6836 let d32 = _mm256_madd_epi16(p16, ones16);
6837 let hi128 = _mm256_extracti128_si256::<1>(d32);
6838 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6839 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6840 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6841 let msum = _mm_cvtsi128_si32(s32);
6842 let d = 2 * msum - gsums[k][gi];
6843 acc[k] += d as f32 * s;
6844 }
6845 }
6846 acc
6847 }
6848}
6849
6850#[allow(unreachable_code)]
6851fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6852 #[cfg(target_arch = "aarch64")]
6853 unsafe {
6854 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6855 }
6856 #[cfg(target_arch = "x86_64")]
6857 if avx2_enabled() {
6858 unsafe {
6859 if vnni_tiles_enabled() {
6860 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6861 }
6862 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6863 }
6864 }
6865 let _ = gsum;
6866 let mut acc = 0f32;
6867 for gi in 0..gpr {
6868 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6869 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6870 let mut d = 0i32;
6871 for (j, &b) in tile[2..].iter().enumerate() {
6872 for k in 0..8 {
6873 let w = ((b >> k) & 1) as i32 * 2 - 1;
6874 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6875 }
6876 }
6877 acc += d as f32 * s;
6878 }
6879 acc
6880}
6881
6882#[cfg(target_arch = "aarch64")]
6891#[target_feature(enable = "neon,dotprod")]
6892unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6893 unsafe {
6896 use core::arch::aarch64::*;
6897 use core::arch::asm;
6898 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6899 let m = vld1q_u8(MASKS.as_ptr());
6900 macro_rules! tile_dot {
6902 ($t:expr, $x:expr) => {{
6903 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6904 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6905 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6906 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6907 let x0 = vld1q_s8($x);
6908 let x1 = vld1q_s8($x.add(16));
6909 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6910 asm!(
6911 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6912 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6913 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6914 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6915 options(pure, nomem, nostack),
6916 );
6917 vaddq_s32(a0, a1)
6918 }};
6919 }
6920 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6929 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6930 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6931 const IW11: [u8; 16] = [
6932 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6933 ];
6934 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6935 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6936 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6937 let isc = vld1_u8(ISC.as_ptr());
6938 macro_rules! tile_dot_tbl {
6940 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6941 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6942 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6943 let x0 = vld1q_s8($x);
6944 let x1 = vld1q_s8($x.add(16));
6945 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6946 asm!(
6947 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6948 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6949 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6950 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6951 options(pure, nomem, nostack),
6952 );
6953 vaddq_s32(a0, a1)
6954 }};
6955 }
6956 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6957 let row_base = r * gpr * Q1_TILE;
6958 let abs_end = bytes.len();
6959 let xp = xq.as_ptr();
6960 let gp = gsum.as_ptr();
6961 let mut accv = vdupq_n_f32(0.0);
6962 let mut gi = 0;
6963 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6966 let t0 = base.add(gi * Q1_TILE);
6967 let ld_a = vld1q_u8(t0);
6968 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6969 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6970 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6971 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6972 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6973 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6975 let g = vld1q_s32(gp.add(gi));
6976 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6977 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6978 let scf: float32x4_t;
6979 asm!(
6980 "fcvtl {o:v}.4s, {i:v}.4h",
6981 o = out(vreg) scf, i = in(vreg) sc16,
6982 options(pure, nomem, nostack),
6983 );
6984 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6985 gi += 4;
6986 }
6987 let mut acc = vaddvq_f32(accv);
6988 while gi < gpr {
6989 let t = base.add(gi * Q1_TILE);
6990 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6991 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6992 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6993 gi += 1;
6994 }
6995 acc
6996 }
6997}
6998
6999#[cfg(target_arch = "aarch64")]
7004#[target_feature(enable = "neon,dotprod")]
7005unsafe fn dot_q1_row_1x4_sdot(
7006 bytes: &[u8],
7007 r: usize,
7008 gpr: usize,
7009 xs: [&[i8]; 4],
7010 gs: [&[i32]; 4],
7011) -> [f32; 4] {
7012 unsafe {
7014 use core::arch::aarch64::*;
7015 use core::arch::asm;
7016 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7017 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7018 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7019 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7020 const IW11: [u8; 16] = [
7021 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7022 ];
7023 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7024 let m = vld1q_u8(MASKS.as_ptr());
7025 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7026 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7027 let isc = vld1_u8(ISC.as_ptr());
7028 macro_rules! sdot2 {
7029 ($w0:expr, $w1:expr, $x:expr) => {{
7030 let x0 = vld1q_s8($x);
7031 let x1 = vld1q_s8($x.add(16));
7032 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7033 asm!(
7034 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7035 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7036 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7037 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7038 options(pure, nomem, nostack),
7039 );
7040 vaddq_s32(a0, a1)
7041 }};
7042 }
7043 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7044 let row_base = r * gpr * Q1_TILE;
7045 let abs_end = bytes.len();
7046 let mut accv = [vdupq_n_f32(0.0); 4];
7047 let mut gi = 0;
7048 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7049 let t0 = base.add(gi * Q1_TILE);
7050 let ld_a = vld1q_u8(t0);
7051 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7052 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7054 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7055 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7056 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7057 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7058 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7059 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7060 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7061 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7062 let scf: float32x4_t;
7063 asm!(
7064 "fcvtl {o:v}.4s, {i:v}.4h",
7065 o = out(vreg) scf, i = in(vreg) sc16,
7066 options(pure, nomem, nostack),
7067 );
7068 for k in 0..4 {
7069 let xp = xs[k].as_ptr();
7070 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7071 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7072 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7073 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7074 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7075 let g = vld1q_s32(gs[k].as_ptr().add(gi));
7076 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7077 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7078 }
7079 gi += 4;
7080 }
7081 let mut acc = [
7082 vaddvq_f32(accv[0]),
7083 vaddvq_f32(accv[1]),
7084 vaddvq_f32(accv[2]),
7085 vaddvq_f32(accv[3]),
7086 ];
7087 while gi < gpr {
7088 let t = base.add(gi * Q1_TILE);
7089 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7090 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7091 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7092 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7093 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7094 for k in 0..4 {
7095 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7096 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7097 }
7098 gi += 1;
7099 }
7100 acc
7101 }
7102}
7103
7104#[inline]
7106fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7107 let gi = j / GROUP_SIZE;
7108 let k = j % GROUP_SIZE;
7109 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7110 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7111 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7112 ((bit as i32 * 2 - 1) as f32, s)
7113}
7114
7115#[inline]
7117fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7118 let mut acc = 0f32;
7119 for gi in 0..gpr {
7120 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7121 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7122 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7123 let mut ga = 0f32;
7124 for (j, &b) in tile[2..].iter().enumerate() {
7125 for k in 0..8 {
7126 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7127 }
7128 }
7129 acc += ga * s;
7130 }
7131 acc
7132}
7133
7134#[allow(clippy::too_many_arguments)]
7137fn q1_range_a8w8(
7138 bytes: &[u8],
7139 gpr: usize,
7140 act: &SplitAct,
7141 gsum: &[i32],
7142 out: SendMut,
7143 start: usize,
7144 end: usize,
7145) {
7146 for r in start..end {
7147 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7148 for &(j, xv) in &act.outliers {
7149 let (w, s) = q1_outlier(bytes, r, gpr, j);
7150 acc += w * s * xv;
7151 }
7152 unsafe { *out.at(r) = acc };
7154 }
7155}
7156
7157fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7159 for r in start..end {
7160 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
7162 }
7163}
7164
7165fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
7170 let entries = base_len + (rows + 1) * 4;
7171 (base_len, entries, entries <= bytes.len())
7172}
7173
7174#[inline]
7176fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
7177 let o = rp_off + r * 4;
7178 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
7179}
7180
7181const SIGN5: [[f32; 5]; 256] = {
7185 let mut lut = [[0.0f32; 5]; 256];
7186 let pow3 = [1u16, 3, 9, 27, 81];
7187 let mut byte = 0usize;
7188 while byte < 256 {
7189 let mut i = 0usize;
7190 while i < 5 {
7191 let code = (byte as u16 / pow3[i]) % 3;
7192 lut[byte][i] = if code == 1 {
7193 1.0
7194 } else if code == 2 {
7195 -1.0
7196 } else {
7197 0.0
7198 };
7199 i += 1;
7200 }
7201 byte += 1;
7202 }
7203 lut
7204};
7205
7206const SIGN5_I8: [[i8; 5]; 256] = {
7208 let mut lut = [[0i8; 5]; 256];
7209 let pow3 = [1u16, 3, 9, 27, 81];
7210 let mut byte = 0usize;
7211 while byte < 256 {
7212 let mut i = 0usize;
7213 while i < 5 {
7214 let code = (byte as u16 / pow3[i]) % 3;
7215 lut[byte][i] = if code == 1 {
7216 1
7217 } else if code == 2 {
7218 -1
7219 } else {
7220 0
7221 };
7222 i += 1;
7223 }
7224 byte += 1;
7225 }
7226 lut
7227};
7228
7229const SIGN5_U64: [u64; 256] = {
7235 let mut lut = [0u64; 256];
7236 let pow3 = [1u16, 3, 9, 27, 81];
7237 let mut byte = 0usize;
7238 while byte < 256 {
7239 let mut v = 0u64;
7240 let mut i = 0usize;
7241 while i < 5 {
7242 let code = (byte as u16 / pow3[i]) % 3;
7243 let s: u8 = if code == 1 {
7244 1
7245 } else if code == 2 {
7246 0xFF
7247 } else {
7248 0
7249 };
7250 v |= (s as u64) << (i * 8);
7251 i += 1;
7252 }
7253 lut[byte] = v;
7254 byte += 1;
7255 }
7256 lut
7257};
7258
7259#[inline]
7264fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
7265 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7266 let off = (r * gpr + j / GROUP_SIZE) * TILE;
7267 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7268 let within = j % GROUP_SIZE;
7269 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
7270}
7271
7272#[cfg(target_arch = "aarch64")]
7275#[target_feature(enable = "neon,dotprod")]
7276#[inline]
7277unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
7278 unsafe {
7280 use core::arch::aarch64::*;
7281 use core::arch::asm;
7282 let w0 = vld1q_s8(w);
7283 let w1 = vld1q_s8(w.add(16));
7284 let x0 = vld1q_s8(x);
7285 let x1 = vld1q_s8(x.add(16));
7286 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7287 asm!(
7288 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7289 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7290 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7291 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7292 options(pure, nomem, nostack),
7293 );
7294 vaddvq_s32(vaddq_s32(a0, a1))
7295 }
7296}
7297
7298#[cfg(target_arch = "x86_64")]
7301#[target_feature(enable = "avx2")]
7302#[inline]
7303unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
7304 unsafe {
7306 use core::arch::x86_64::*;
7307 let wv = _mm256_loadu_si256(w as *const __m256i);
7308 let xv = _mm256_loadu_si256(x as *const __m256i);
7309 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7310 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
7311 let hi128 = _mm256_extracti128_si256::<1>(d);
7312 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7313 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7314 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7315 _mm_cvtsi128_si32(s32)
7316 }
7317}
7318
7319#[inline]
7324fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
7325 debug_assert!(dst.len() >= 40);
7326 unsafe {
7329 let p = dst.as_mut_ptr();
7330 for bi in 0..7 {
7331 core::ptr::write_unaligned(
7332 p.add(bi * 5) as *mut u64,
7333 SIGN5_U64[*codes.add(bi) as usize],
7334 );
7335 }
7336 }
7337}
7338
7339#[inline]
7344fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
7345 #[cfg(target_arch = "aarch64")]
7346 unsafe {
7347 return sdot32_i8(w, x);
7348 }
7349 #[cfg(target_arch = "x86_64")]
7350 unsafe {
7351 return i8dot32_avx2(w, x);
7352 }
7353 #[allow(unreachable_code)]
7354 unsafe {
7355 let mut s = 0i32;
7356 for k in 0..GROUP_SIZE {
7357 s += *w.add(k) as i32 * *x.add(k) as i32;
7358 }
7359 s
7360 }
7361}
7362
7363#[inline]
7364unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
7365 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
7366 (
7367 SIGN5_U64[*codes as usize],
7368 SIGN5_U64[*codes.add(1) as usize],
7369 SIGN5_U64[*codes.add(2) as usize],
7370 SIGN5_U64[*codes.add(3) as usize],
7371 SIGN5_U64[*codes.add(4) as usize],
7372 SIGN5_U64[*codes.add(5) as usize],
7373 SIGN5_U64[*codes.add(6) as usize],
7374 )
7375 };
7376
7377 let u0 = s0 | (s1 << 40);
7378 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
7379 let u2 = (s3 >> 8) | (s4 << 32);
7380 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
7381
7382 (u0, u1, u2, u3)
7383}
7384
7385#[cfg(target_arch = "aarch64")]
7389#[target_feature(enable = "neon,dotprod")]
7390unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7391 use core::arch::aarch64::*;
7392 use core::arch::asm;
7393 unsafe {
7394 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7395 let mut acc = 0f32;
7396 let bytes_ptr = bytes.as_ptr();
7397 let xq_ptr = xq.as_ptr();
7398 let row_off = r * gpr * TILE;
7399
7400 let gpr2 = gpr & !1;
7401 let mut gi = 0;
7402 while gi < gpr2 {
7403 let off0 = row_off + gi * TILE;
7404 let off1 = off0 + TILE;
7405 let s0 = f16_to_f32(u16::from_le_bytes([
7406 *bytes_ptr.add(off0),
7407 *bytes_ptr.add(off0 + 1),
7408 ]));
7409 let s1 = f16_to_f32(u16::from_le_bytes([
7410 *bytes_ptr.add(off1),
7411 *bytes_ptr.add(off1 + 1),
7412 ]));
7413
7414 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7415 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7416
7417 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7418 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7419 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7420 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7421
7422 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7423 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7424 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
7425 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
7426
7427 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
7428 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7429 asm!(
7430 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
7431 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
7432 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
7433 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
7434 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
7435 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
7436 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
7437 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
7438 options(pure, nomem, nostack),
7439 );
7440 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
7441 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
7442 acc += d0 as f32 * s0 + d1 as f32 * s1;
7443 gi += 2;
7444 }
7445
7446 if gi < gpr {
7447 let off = row_off + gi * TILE;
7448 let s = f16_to_f32(u16::from_le_bytes([
7449 *bytes_ptr.add(off),
7450 *bytes_ptr.add(off + 1),
7451 ]));
7452 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7453 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7454 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7455 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7456 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7457 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7458 asm!(
7459 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7460 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7461 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7462 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7463 options(pure, nomem, nostack),
7464 );
7465 let d = vaddvq_s32(vaddq_s32(a0, a1));
7466 acc += d as f32 * s;
7467 }
7468 acc
7469 }
7470}
7471
7472#[cfg(target_arch = "x86_64")]
7474#[target_feature(enable = "avx2")]
7475unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7476 use core::arch::x86_64::*;
7477 unsafe {
7478 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7479 let mut acc = 0f32;
7480 let bytes_ptr = bytes.as_ptr();
7481 let xq_ptr = xq.as_ptr();
7482 let row_off = r * gpr * TILE;
7483
7484 let ones = _mm256_set1_epi16(1);
7485 for gi in 0..gpr {
7486 let off = row_off + gi * TILE;
7487 let s = f16_to_f32(u16::from_le_bytes([
7488 *bytes_ptr.add(off),
7489 *bytes_ptr.add(off + 1),
7490 ]));
7491 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7492 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7493 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7494 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7495 let d256 = _mm256_madd_epi16(p16, ones);
7496 let d128 = _mm_add_epi32(
7497 _mm256_castsi256_si128(d256),
7498 _mm256_extracti128_si256(d256, 1),
7499 );
7500 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
7501 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
7502 acc += d32 as f32 * s;
7503 }
7504 acc
7505 }
7506}
7507
7508#[cfg(target_arch = "x86_64")]
7510#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7511unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7512 use core::arch::x86_64::*;
7513 unsafe {
7515 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7516 let mut acc = 0f32;
7517 let bytes_ptr = bytes.as_ptr();
7518 let xq_ptr = xq.as_ptr();
7519 let row_off = r * gpr * TILE;
7520 for gi in 0..gpr {
7521 let off = row_off + gi * TILE;
7522 let s = f16_to_f32(u16::from_le_bytes([
7523 *bytes_ptr.add(off),
7524 *bytes_ptr.add(off + 1),
7525 ]));
7526 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7527 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7528 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7529 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7530 acc += d as f32 * s;
7531 }
7532 acc
7533 }
7534}
7535
7536#[inline]
7540fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7541 #[cfg(target_arch = "aarch64")]
7542 unsafe {
7543 return q1t_dot_row_sdot(bytes, r, gpr, xq);
7544 }
7545 #[cfg(target_arch = "x86_64")]
7546 unsafe {
7547 if vnni_tiles_enabled() {
7548 return q1t_dot_row_vnni(bytes, r, gpr, xq);
7549 }
7550 return q1t_dot_row_avx2(bytes, r, gpr, xq);
7551 }
7552 #[allow(unreachable_code)]
7553 {
7554 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7555 let mut acc = 0f32;
7556 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
7558 let off = (r * gpr + gi) * TILE;
7559 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7560 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
7561 let mut d = 0i32;
7562 for k in 0..GROUP_SIZE {
7563 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
7564 }
7565 acc += d as f32 * s;
7566 }
7567 acc
7568 }
7569}
7570
7571fn q1t_row_outlier_correction(
7578 bytes: &[u8],
7579 r: usize,
7580 rp_off: usize,
7581 entries_off: usize,
7582 has_ov: bool,
7583 x: &[f32],
7584) -> f32 {
7585 if !has_ov {
7586 return 0.0;
7587 }
7588 let (c0, c1) = (
7589 q1t_rowptr(bytes, rp_off, r),
7590 q1t_rowptr(bytes, rp_off, r + 1),
7591 );
7592 let mut corr = 0f32;
7593 for p in c0..c1 {
7594 let e = entries_off + p * 4;
7595 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7596 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7597 corr += val * x[col];
7598 }
7599 corr
7600}
7601
7602fn q1t_dequant_row(
7606 bytes: &[u8],
7607 r: usize,
7608 gpr: usize,
7609 rp_off: usize,
7610 entries_off: usize,
7611 has_ov: bool,
7612 buf: &mut [f32],
7613) {
7614 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7615 for g in 0..gpr {
7616 let off = (r * gpr + g) * TILE;
7617 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7618 let codes = &bytes[off + 2..off + TILE];
7619 let bc = g * GROUP_SIZE;
7620 for bi in 0..6 {
7622 let lut = &SIGN5[codes[bi] as usize];
7623 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
7624 for i in 0..5 {
7625 d[i] = lut[i] * s;
7626 }
7627 }
7628 let lut = &SIGN5[codes[6] as usize];
7629 buf[bc + 30] = lut[0] * s;
7630 buf[bc + 31] = lut[1] * s;
7631 }
7632 if !has_ov {
7633 return;
7634 }
7635 let (c0, c1) = (
7636 q1t_rowptr(bytes, rp_off, r),
7637 q1t_rowptr(bytes, rp_off, r + 1),
7638 );
7639 for p in c0..c1 {
7640 let e = entries_off + p * 4;
7641 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7642 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7643 }
7644}
7645
7646fn q1t_add_overlay(
7650 bytes: &[u8],
7651 x: &[f32],
7652 rows: usize,
7653 cols: usize,
7654 out: &mut [f32],
7655 pool: Option<&Pool>,
7656) {
7657 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7658 let gpr = cols / GROUP_SIZE;
7659 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7660 if !has_ov {
7661 return;
7662 }
7663 let out_addr = SendMut(out.as_mut_ptr());
7664 let run = move |start: usize, end: usize| {
7665 for r in start..end {
7666 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7667 unsafe { *out_addr.at(r) += corr };
7669 }
7670 };
7671 dispatch_rows(pool, rows, &run);
7672}
7673
7674#[allow(clippy::too_many_arguments)]
7677fn q1t_range_a8w8(
7678 bytes: &[u8],
7679 gpr: usize,
7680 rp_off: usize,
7681 ent_off: usize,
7682 has_ov: bool,
7683 act: &SplitAct,
7684 x: &[f32],
7685 out: SendMut,
7686 start: usize,
7687 end: usize,
7688) {
7689 for r in start..end {
7690 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7691 for &(j, xv) in &act.outliers {
7692 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7693 }
7694 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7695 unsafe { *out.at(r) = acc };
7697 }
7698}
7699
7700#[allow(clippy::too_many_arguments)]
7703fn q1t_range_f32_batch(
7704 bytes: &[u8],
7705 gpr: usize,
7706 rp_off: usize,
7707 ent_off: usize,
7708 has_ov: bool,
7709 x: &[f32],
7710 out: SendMut,
7711 start: usize,
7712 end: usize,
7713) {
7714 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7715 let mut sg = [0f32; GROUP_SIZE];
7716 for r in start..end {
7717 let mut acc = 0f32;
7718 for g in 0..gpr {
7719 let off = (r * gpr + g) * TILE;
7720 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7721 let codes = &bytes[off + 2..off + TILE];
7722 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7723 for bi in 0..6 {
7724 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7725 }
7726 let lut = &SIGN5[codes[6] as usize];
7727 sg[30] = lut[0];
7728 sg[31] = lut[1];
7729 let mut gsum = 0f32;
7730 for k in 0..GROUP_SIZE {
7731 gsum += sg[k] * xg[k];
7732 }
7733 acc += s * gsum;
7734 }
7735 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7736 unsafe { *out.at(r) = acc };
7738 }
7739}
7740
7741fn q1t_matvec(
7745 bytes: &[u8],
7746 x: &[f32],
7747 rows: usize,
7748 cols: usize,
7749 out: &mut [f32],
7750 pool: Option<&Pool>,
7751) {
7752 debug_assert_eq!(out.len(), rows);
7753 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7754 let gpr = cols / GROUP_SIZE;
7755 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7756 let out_addr = SendMut(out.as_mut_ptr());
7757 if a8w8_enabled() {
7761 let act = split_act(x);
7762 let act = &act;
7763 let run = move |start: usize, end: usize| {
7764 for r in start..end {
7765 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7766 for &(j, xv) in &act.outliers {
7767 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7768 }
7769 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7770 unsafe { *out_addr.at(r) = acc };
7772 }
7773 };
7774 dispatch_rows(pool, rows, &run);
7775 return;
7776 }
7777 let run = move |start: usize, end: usize| {
7778 let mut sg = [0f32; GROUP_SIZE];
7782 for r in start..end {
7783 let mut acc = 0f32;
7784 for g in 0..gpr {
7785 let off = (r * gpr + g) * TILE;
7786 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7787 let codes = &bytes[off + 2..off + TILE];
7788 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7789 for bi in 0..6 {
7790 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7791 }
7792 let lut = &SIGN5[codes[6] as usize];
7793 sg[30] = lut[0];
7794 sg[31] = lut[1];
7795 let mut gsum = 0f32;
7796 for k in 0..GROUP_SIZE {
7797 gsum += sg[k] * xg[k];
7798 }
7799 acc += s * gsum;
7800 }
7801 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7802 unsafe { *out_addr.at(r) = acc };
7803 }
7804 };
7805 dispatch_rows(pool, rows, &run);
7806}
7807
7808#[cfg(target_arch = "aarch64")]
7814#[target_feature(enable = "neon,dotprod")]
7815unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7816 use core::arch::aarch64::*;
7817 use core::arch::asm;
7818 unsafe {
7820 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7821 let bytes_ptr = bytes.as_ptr();
7822 let row_off = r * gpr * TILE;
7823 let xp = [xa.as_ptr(), xb.as_ptr()];
7824 let mut acc = [0f32; 2];
7825 macro_rules! sdot2 {
7826 ($w0:expr, $w1:expr, $x:expr) => {{
7827 let x0 = vld1q_s8($x);
7828 let x1 = vld1q_s8($x.add(16));
7829 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7830 asm!(
7831 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7832 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7833 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7834 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7835 options(pure, nomem, nostack),
7836 );
7837 vaddvq_s32(vaddq_s32(a0, a1))
7838 }};
7839 }
7840 let gpr2 = gpr & !1;
7841 let mut gi = 0;
7842 while gi < gpr2 {
7843 let off0 = row_off + gi * TILE;
7844 let off1 = off0 + TILE;
7845 let s0 = f16_to_f32(u16::from_le_bytes([
7846 *bytes_ptr.add(off0),
7847 *bytes_ptr.add(off0 + 1),
7848 ]));
7849 let s1 = f16_to_f32(u16::from_le_bytes([
7850 *bytes_ptr.add(off1),
7851 *bytes_ptr.add(off1 + 1),
7852 ]));
7853 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7854 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7855 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7856 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7857 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7858 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7859 for k in 0..2 {
7860 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7861 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7862 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7863 }
7864 gi += 2;
7865 }
7866 if gi < gpr {
7867 let off = row_off + gi * TILE;
7868 let s = f16_to_f32(u16::from_le_bytes([
7869 *bytes_ptr.add(off),
7870 *bytes_ptr.add(off + 1),
7871 ]));
7872 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7873 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7874 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7875 for k in 0..2 {
7876 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7877 acc[k] += d as f32 * s;
7878 }
7879 }
7880 acc
7881 }
7882}
7883
7884fn q1t_matvec2(
7890 bytes: &[u8],
7891 x1: &[f32],
7892 x2: &[f32],
7893 rows: usize,
7894 cols: usize,
7895 o1: &mut [f32],
7896 o2: &mut [f32],
7897 pool: Option<&Pool>,
7898) {
7899 debug_assert_eq!(o1.len(), rows);
7900 debug_assert_eq!(o2.len(), rows);
7901 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7902 let gpr = cols / GROUP_SIZE;
7903 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7904 let out1 = SendMut(o1.as_mut_ptr());
7905 let out2 = SendMut(o2.as_mut_ptr());
7906 if a8w8_enabled() {
7907 let a1 = split_act(x1);
7908 let a2 = split_act(x2);
7909 let (a1, a2) = (&a1, &a2);
7910 let run = move |start: usize, end: usize| {
7911 for r in start..end {
7912 #[cfg(target_arch = "aarch64")]
7913 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7916 #[cfg(not(target_arch = "aarch64"))]
7917 let ds = [
7918 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7919 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7920 ];
7921 let mut acc1 = ds[0] * a1.sx;
7922 for &(j, xv) in &a1.outliers {
7923 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7924 }
7925 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7926 let mut acc2 = ds[1] * a2.sx;
7927 for &(j, xv) in &a2.outliers {
7928 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7929 }
7930 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7931 unsafe {
7933 *out1.at(r) = acc1;
7934 *out2.at(r) = acc2;
7935 }
7936 }
7937 };
7938 dispatch_rows(pool, rows, &run);
7939 return;
7940 }
7941 let run = move |start: usize, end: usize| {
7942 let mut sg = [0f32; GROUP_SIZE];
7945 for r in start..end {
7946 let mut acc1 = 0f32;
7947 let mut acc2 = 0f32;
7948 for g in 0..gpr {
7949 let off = (r * gpr + g) * TILE;
7950 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7951 let codes = &bytes[off + 2..off + TILE];
7952 for bi in 0..6 {
7953 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7954 }
7955 let lut = &SIGN5[codes[6] as usize];
7956 sg[30] = lut[0];
7957 sg[31] = lut[1];
7958 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7959 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7960 let mut gsum1 = 0f32;
7961 for k in 0..GROUP_SIZE {
7962 gsum1 += sg[k] * xg1[k];
7963 }
7964 acc1 += s * gsum1;
7965 let mut gsum2 = 0f32;
7966 for k in 0..GROUP_SIZE {
7967 gsum2 += sg[k] * xg2[k];
7968 }
7969 acc2 += s * gsum2;
7970 }
7971 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7972 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7973 unsafe {
7975 *out1.at(r) = acc1;
7976 *out2.at(r) = acc2;
7977 }
7978 }
7979 };
7980 dispatch_rows(pool, rows, &run);
7981}
7982
7983fn q1t_matmat(
7986 bytes: &[u8],
7987 xs: &[f32],
7988 b: usize,
7989 rows: usize,
7990 cols: usize,
7991 out: &mut [f32],
7992 pool: Option<&Pool>,
7993) {
7994 debug_assert_eq!(out.len(), b * rows);
7995 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7996 let gpr = cols / GROUP_SIZE;
7997 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7998 let out_addr = SendMut(out.as_mut_ptr());
7999 if a8w8_enabled() {
8003 let acts: Vec<SplitAct> = (0..b)
8004 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
8005 .collect();
8006 let acts = &acts;
8007 let run = move |start: usize, end: usize| {
8008 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
8012 for g in 0..gpr {
8013 let off = (r * gpr + g) * TILE;
8014 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8015 q1t_unpack_group_i8(
8016 bytes.as_ptr().wrapping_add(off + 2),
8017 &mut sg[g * GROUP_SIZE..],
8018 );
8019 }
8020 for bi in 0..b {
8021 let act = &acts[bi];
8022 let mut isum = 0f32;
8023 for g in 0..gpr {
8024 let d = q1t_i8dot32(
8025 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
8026 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
8027 );
8028 isum += d as f32 * sc[g];
8029 }
8030 let mut acc = isum * act.sx;
8031 for &(j, xv) in &act.outliers {
8032 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8033 }
8034 accs[bi] = acc;
8035 }
8036 if has_ov {
8040 let (c0, c1) = (
8041 q1t_rowptr(bytes, rp_off, r),
8042 q1t_rowptr(bytes, rp_off, r + 1),
8043 );
8044 for p in c0..c1 {
8045 let e = ent_off + p * 4;
8046 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8047 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8048 for bi in 0..b {
8049 accs[bi] += val * xs[bi * cols + col];
8050 }
8051 }
8052 }
8053 for bi in 0..b {
8054 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
8055 }
8056 }
8057 };
8058 dispatch_rows(pool, rows, &run);
8059 return;
8060 }
8061 let run = move |start: usize, end: usize| {
8062 let mut buf = vec![0f32; cols];
8063 for r in start..end {
8064 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
8065 for bi in 0..b {
8066 let xr = &xs[bi * cols..(bi + 1) * cols];
8067 let mut acc = 0f32;
8068 for j in 0..cols {
8069 acc += buf[j] * xr[j];
8070 }
8071 unsafe { *out_addr.at(bi * rows + r) = acc };
8072 }
8073 }
8074 };
8075 dispatch_rows(pool, rows, &run);
8076}
8077
8078fn q1_matvec(
8079 bytes: &[u8],
8080 x: &[f32],
8081 rows: usize,
8082 cols: usize,
8083 out: &mut [f32],
8084 pool: Option<&Pool>,
8085) {
8086 debug_assert_eq!(out.len(), rows);
8087 let gpr = cols / GROUP_SIZE;
8088 let out_addr = SendMut(out.as_mut_ptr());
8089 if a8w8_enabled() {
8090 let act = split_act(x);
8091 let gsum = q1_group_sums(&act.xq, gpr);
8092 let (act, gsum) = (&act, &gsum);
8093 let run = move |start: usize, end: usize| {
8094 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
8095 };
8096 dispatch_rows(pool, rows, &run);
8097 return;
8098 }
8099 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
8100 dispatch_rows(pool, rows, &run);
8101}
8102
8103#[allow(clippy::too_many_arguments)]
8105fn q1_matvec2(
8106 bytes: &[u8],
8107 x1: &[f32],
8108 x2: &[f32],
8109 rows: usize,
8110 cols: usize,
8111 o1: &mut [f32],
8112 o2: &mut [f32],
8113 pool: Option<&Pool>,
8114) {
8115 let gpr = cols / GROUP_SIZE;
8116 let p1 = SendMut(o1.as_mut_ptr());
8117 let p2 = SendMut(o2.as_mut_ptr());
8118 if a8w8_enabled() {
8119 let a1 = split_act(x1);
8120 let a2 = split_act(x2);
8121 let g1 = q1_group_sums(&a1.xq, gpr);
8122 let g2 = q1_group_sums(&a2.xq, gpr);
8123 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
8124 let run = move |start: usize, end: usize| {
8125 for r in start..end {
8126 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
8127 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
8128 for &(j, xv) in &a1.outliers {
8129 let (w, s) = q1_outlier(bytes, r, gpr, j);
8130 v1 += w * s * xv;
8131 }
8132 for &(j, xv) in &a2.outliers {
8133 let (w, s) = q1_outlier(bytes, r, gpr, j);
8134 v2 += w * s * xv;
8135 }
8136 unsafe {
8138 *p1.at(r) = v1;
8139 *p2.at(r) = v2;
8140 }
8141 }
8142 };
8143 dispatch_rows(pool, rows, &run);
8144 return;
8145 }
8146 let run = move |start: usize, end: usize| {
8147 for r in start..end {
8148 unsafe {
8150 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
8151 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
8152 }
8153 }
8154 };
8155 dispatch_rows(pool, rows, &run);
8156}
8157
8158#[allow(clippy::too_many_arguments)]
8160fn q1_matmat(
8161 bytes: &[u8],
8162 xs_all: &[f32],
8163 b: usize,
8164 rows: usize,
8165 cols: usize,
8166 out: &mut [f32],
8167 pool: Option<&Pool>,
8168) {
8169 debug_assert_eq!(out.len(), b * rows);
8170 let gpr = cols / GROUP_SIZE;
8171 let out_addr = SendMut(out.as_mut_ptr());
8172 if a8w8_enabled() {
8173 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
8174 .map(|bi| {
8175 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
8176 let gsum = q1_group_sums(&act.xq, gpr);
8177 (act, gsum)
8178 })
8179 .collect();
8180 let acts = &acts;
8181 #[cfg(target_arch = "x86_64")]
8182 let blocked_ok = avx2_enabled() && blocked_enabled();
8183 #[cfg(target_arch = "aarch64")]
8184 let blocked_ok = sdot_enabled() && blocked_enabled();
8185 let run = move |start: usize, end: usize| {
8186 for r in start..end {
8187 let mut bi = 0usize;
8188 #[cfg(target_arch = "aarch64")]
8191 if blocked_ok {
8192 while bi + 4 <= acts.len() {
8193 let xs = [
8194 acts[bi].0.xq.as_slice(),
8195 acts[bi + 1].0.xq.as_slice(),
8196 acts[bi + 2].0.xq.as_slice(),
8197 acts[bi + 3].0.xq.as_slice(),
8198 ];
8199 let gs = [
8200 acts[bi].1.as_slice(),
8201 acts[bi + 1].1.as_slice(),
8202 acts[bi + 2].1.as_slice(),
8203 acts[bi + 3].1.as_slice(),
8204 ];
8205 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
8206 for k in 0..4 {
8207 let (act, _) = &acts[bi + k];
8208 let mut acc = d[k] * act.sx;
8209 for &(j, xv) in &act.outliers {
8210 let (w, sc) = q1_outlier(bytes, r, gpr, j);
8211 acc += w * sc * xv;
8212 }
8213 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8215 }
8216 bi += 4;
8217 }
8218 }
8219 #[cfg(target_arch = "x86_64")]
8220 if blocked_ok {
8221 while bi + 4 <= acts.len() {
8222 let xs = [
8223 acts[bi].0.xq.as_slice(),
8224 acts[bi + 1].0.xq.as_slice(),
8225 acts[bi + 2].0.xq.as_slice(),
8226 acts[bi + 3].0.xq.as_slice(),
8227 ];
8228 let gs = [
8229 acts[bi].1.as_slice(),
8230 acts[bi + 1].1.as_slice(),
8231 acts[bi + 2].1.as_slice(),
8232 acts[bi + 3].1.as_slice(),
8233 ];
8234 let d = unsafe {
8235 if vnni_tiles_enabled() {
8236 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
8237 } else {
8238 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
8239 }
8240 };
8241 for k in 0..4 {
8242 let (act, _) = &acts[bi + k];
8243 let mut acc = d[k] * act.sx;
8244 for &(j, xv) in &act.outliers {
8245 let (w, sc) = q1_outlier(bytes, r, gpr, j);
8246 acc += w * sc * xv;
8247 }
8248 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8250 }
8251 bi += 4;
8252 }
8253 }
8254 while bi < acts.len() {
8255 let (act, gsum) = &acts[bi];
8256 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
8257 for &(j, xv) in &act.outliers {
8258 let (w, s) = q1_outlier(bytes, r, gpr, j);
8259 acc += w * s * xv;
8260 }
8261 unsafe { *out_addr.at(bi * rows + r) = acc };
8263 bi += 1;
8264 }
8265 }
8266 };
8267 dispatch_rows(pool, rows, &run);
8268 return;
8269 }
8270 let run = move |start: usize, end: usize| {
8271 for r in start..end {
8272 for bi in 0..b {
8273 let x = &xs_all[bi * cols..(bi + 1) * cols];
8274 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
8276 }
8277 }
8278 };
8279 dispatch_rows(pool, rows, &run);
8280}
8281
8282fn q4matvec(
8288 bytes: &[u8],
8289 x: &[f32],
8290 rows: usize,
8291 cols: usize,
8292 out: &mut [f32],
8293 pool: Option<&Pool>,
8294) {
8295 debug_assert_eq!(out.len(), rows);
8296 let (packed, scales) = q4_split(bytes, rows, cols);
8297 let gpr = cols / GROUP_SIZE;
8298 let out_addr = SendMut(out.as_mut_ptr());
8299
8300 if a8w8_enabled() {
8301 let act = split_act(x);
8302 let run = move |start: usize, end: usize| {
8303 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
8304 };
8305 dispatch_rows(pool, rows, &run);
8306 return;
8307 }
8308
8309 let run =
8310 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
8311 dispatch_rows(pool, rows, &run);
8312}
8313
8314#[inline]
8317#[allow(unreachable_code)]
8318#[cfg(target_arch = "x86_64")]
8323#[target_feature(enable = "avx2")]
8324unsafe fn dot_q4b_row_1x4_avx2(
8325 buf: &[u8],
8326 scales: &[u8],
8327 g0: usize,
8328 gpr: usize,
8329 xs: [&[i8]; 4],
8330) -> [f32; 4] {
8331 unsafe {
8333 use core::arch::x86_64::*;
8334 let ones = _mm256_set1_epi16(1);
8335 let mut acc = [0f32; 4];
8336 for gi in 0..gpr {
8337 let s = f16_to_f32(u16::from_le_bytes([
8338 scales[(g0 + gi) * 2],
8339 scales[(g0 + gi) * 2 + 1],
8340 ]));
8341 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8342 let aw = _mm256_abs_epi8(w);
8343 for (k, xq) in xs.iter().enumerate() {
8344 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8345 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8346 let d = _mm256_madd_epi16(p16, ones);
8347 let hi128 = _mm256_extracti128_si256::<1>(d);
8348 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8349 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8350 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8351 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
8352 }
8353 }
8354 acc
8355 }
8356}
8357
8358#[cfg(target_arch = "x86_64")]
8360#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8361unsafe fn dot_q4b_row_1x4_vnni(
8362 buf: &[u8],
8363 scales: &[u8],
8364 g0: usize,
8365 gpr: usize,
8366 xs: [&[i8]; 4],
8367) -> [f32; 4] {
8368 unsafe {
8370 use core::arch::x86_64::*;
8371 let mut acc = [0f32; 4];
8372 for gi in 0..gpr {
8373 let s = f16_to_f32(u16::from_le_bytes([
8374 scales[(g0 + gi) * 2],
8375 scales[(g0 + gi) * 2 + 1],
8376 ]));
8377 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8378 let aw = _mm256_abs_epi8(w);
8379 for (k, xq) in xs.iter().enumerate() {
8380 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8381 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8382 acc[k] += d as f32 * s;
8383 }
8384 }
8385 acc
8386 }
8387}
8388
8389#[cfg(target_arch = "x86_64")]
8395#[target_feature(enable = "avx2")]
8396unsafe fn dot_q4b_row_1x4_sx_avx2(
8397 buf: &[u8],
8398 scales: &[u8],
8399 g0: usize,
8400 gpr: usize,
8401 xs: [&[i8]; 4],
8402 sxs: [f32; 4],
8403) -> [f32; 4] {
8404 unsafe {
8406 use core::arch::x86_64::*;
8407 let ones = _mm256_set1_epi16(1);
8408 let mut acc = [0f32; 4];
8409 for gi in 0..gpr {
8410 let s = f16_to_f32(u16::from_le_bytes([
8411 scales[(g0 + gi) * 2],
8412 scales[(g0 + gi) * 2 + 1],
8413 ]));
8414 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8415 let aw = _mm256_abs_epi8(w);
8416 for (k, xq) in xs.iter().enumerate() {
8417 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8418 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8419 let d = _mm256_madd_epi16(p16, ones);
8420 let hi128 = _mm256_extracti128_si256::<1>(d);
8421 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8422 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8423 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8424 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
8425 }
8426 }
8427 acc
8428 }
8429}
8430
8431#[cfg(target_arch = "x86_64")]
8434#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8435unsafe fn dot_q4b_row_1x4_sx_vnni(
8436 buf: &[u8],
8437 scales: &[u8],
8438 g0: usize,
8439 gpr: usize,
8440 xs: [&[i8]; 4],
8441 sxs: [f32; 4],
8442) -> [f32; 4] {
8443 unsafe {
8445 use core::arch::x86_64::*;
8446 let mut acc = [0f32; 4];
8447 for gi in 0..gpr {
8448 let s = f16_to_f32(u16::from_le_bytes([
8449 scales[(g0 + gi) * 2],
8450 scales[(g0 + gi) * 2 + 1],
8451 ]));
8452 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8453 let aw = _mm256_abs_epi8(w);
8454 for (k, xq) in xs.iter().enumerate() {
8455 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8456 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8457 acc[k] += (d as f32 * sxs[k]) * s;
8458 }
8459 }
8460 acc
8461 }
8462}
8463
8464#[allow(unreachable_code)]
8465fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8466 #[cfg(target_arch = "aarch64")]
8467 unsafe {
8468 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
8469 }
8470 #[cfg(target_arch = "x86_64")]
8471 unsafe {
8472 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
8473 }
8474 let mut acc = 0f32;
8475 for gi in 0..gpr {
8476 let g = g0 + gi;
8477 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8478 let mut d = 0i32;
8479 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8480 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
8481 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
8482 }
8483 acc += d as f32 * s;
8484 }
8485 acc
8486}
8487
8488#[inline]
8490#[allow(unreachable_code)]
8491fn dot_q4_row_i8_2(
8492 packed: &[u8],
8493 scales: &[u8],
8494 g0: usize,
8495 gpr: usize,
8496 xq1: &[i8],
8497 xq2: &[i8],
8498) -> (f32, f32) {
8499 #[cfg(target_arch = "aarch64")]
8500 unsafe {
8501 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
8502 }
8503 #[cfg(target_arch = "x86_64")]
8504 unsafe {
8505 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
8506 }
8507 (
8508 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
8509 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
8510 )
8511}
8512
8513#[allow(clippy::too_many_arguments)]
8516fn q4_range_a8w8(
8517 packed: &[u8],
8518 scales: &[u8],
8519 gpr: usize,
8520 cols: usize,
8521 act: &SplitAct,
8522 out: SendMut,
8523 start: usize,
8524 end: usize,
8525) {
8526 for r in start..end {
8527 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
8528 for &(j, xv) in &act.outliers {
8530 let flat = r * cols + j;
8531 let byte = packed[flat / 2];
8532 let nib = if flat & 1 == 0 {
8533 byte & 0x0F
8534 } else {
8535 byte >> 4
8536 };
8537 let s = f16_to_f32(u16::from_le_bytes([
8538 scales[(flat / GROUP_SIZE) * 2],
8539 scales[(flat / GROUP_SIZE) * 2 + 1],
8540 ]));
8541 acc += ((nib as i32 - 8) as f32) * s * xv;
8542 }
8543 unsafe { *out.at(r) = acc };
8545 }
8546}
8547
8548#[allow(clippy::too_many_arguments)]
8551fn q4_range2_a8w8(
8552 packed: &[u8],
8553 scales: &[u8],
8554 gpr: usize,
8555 cols: usize,
8556 a1: &SplitAct,
8557 a2: &SplitAct,
8558 p1: SendMut,
8559 p2: SendMut,
8560 start: usize,
8561 end: usize,
8562) {
8563 for r in start..end {
8564 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
8565 let mut acc1 = s1 * a1.sx;
8566 let mut acc2 = s2 * a2.sx;
8567 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
8569 for &(j, xv) in outliers {
8570 let flat = r * cols + j;
8571 let byte = packed[flat / 2];
8572 let nib = if flat & 1 == 0 {
8573 byte & 0x0F
8574 } else {
8575 byte >> 4
8576 };
8577 let s = f16_to_f32(u16::from_le_bytes([
8578 scales[(flat / GROUP_SIZE) * 2],
8579 scales[(flat / GROUP_SIZE) * 2 + 1],
8580 ]));
8581 *acc += ((nib as i32 - 8) as f32) * s * xv;
8582 }
8583 };
8584 fix(&a1.outliers, &mut acc1);
8585 fix(&a2.outliers, &mut acc2);
8586 unsafe {
8588 *p1.at(r) = acc1;
8589 *p2.at(r) = acc2;
8590 }
8591 }
8592}
8593
8594fn q4_range_f32(
8596 packed: &[u8],
8597 scales: &[u8],
8598 gpr: usize,
8599 x: &[f32],
8600 out: SendMut,
8601 start: usize,
8602 end: usize,
8603) {
8604 for r in start..end {
8605 let mut acc = 0f32;
8606 for gi in 0..gpr {
8607 let g = r * gpr + gi;
8608 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8609 let pk = &packed[g * 16..(g + 1) * 16];
8610 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8611 let mut ga = 0f32;
8612 for (k, &b) in pk.iter().enumerate() {
8613 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
8614 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
8615 }
8616 acc += ga * s;
8617 }
8618 unsafe { *out.at(r) = acc };
8620 }
8621}
8622
8623#[allow(clippy::too_many_arguments)]
8627fn q4matvec2(
8628 bytes: &[u8],
8629 x1: &[f32],
8630 x2: &[f32],
8631 rows: usize,
8632 cols: usize,
8633 o1: &mut [f32],
8634 o2: &mut [f32],
8635 pool: Option<&Pool>,
8636) {
8637 debug_assert_eq!(o1.len(), rows);
8638 debug_assert_eq!(o2.len(), rows);
8639 let (packed, scales) = q4_split(bytes, rows, cols);
8640 let gpr = cols / GROUP_SIZE;
8641
8642 if a8w8_enabled() {
8643 let a1 = split_act(x1);
8644 let a2 = split_act(x2);
8645 let p1 = SendMut(o1.as_mut_ptr());
8646 let p2 = SendMut(o2.as_mut_ptr());
8647 let run = move |start: usize, end: usize| {
8648 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
8649 };
8650 dispatch_rows(pool, rows, &run);
8651 return;
8652 }
8653
8654 let p1 = SendMut(o1.as_mut_ptr());
8655 let p2 = SendMut(o2.as_mut_ptr());
8656 let run = move |start: usize, end: usize| {
8657 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
8658 };
8659 dispatch_rows(pool, rows, &run);
8660}
8661
8662#[allow(clippy::too_many_arguments)]
8664fn q4_range2_f32(
8665 packed: &[u8],
8666 scales: &[u8],
8667 gpr: usize,
8668 x1: &[f32],
8669 x2: &[f32],
8670 p1: SendMut,
8671 p2: SendMut,
8672 start: usize,
8673 end: usize,
8674) {
8675 for r in start..end {
8676 let (mut acc1, mut acc2) = (0f32, 0f32);
8677 for gi in 0..gpr {
8678 let g = r * gpr + gi;
8679 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8680 let pk = &packed[g * 16..(g + 1) * 16];
8681 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8682 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8683 let (mut g1, mut g2) = (0f32, 0f32);
8684 for (k, &b) in pk.iter().enumerate() {
8685 let wl = (b & 0x0F) as f32 - 8.0;
8686 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
8687 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
8688 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
8689 }
8690 acc1 += g1 * s;
8691 acc2 += g2 * s;
8692 }
8693 unsafe {
8695 *p1.at(r) = acc1;
8696 *p2.at(r) = acc2;
8697 }
8698 }
8699}
8700
8701thread_local! {
8702 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8705 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8706}
8707
8708#[allow(clippy::too_many_arguments)]
8714fn q4matmat(
8715 bytes: &[u8],
8716 xs_all: &[f32],
8717 b: usize,
8718 rows: usize,
8719 cols: usize,
8720 out: &mut [f32],
8721 pool: Option<&Pool>,
8722) {
8723 debug_assert_eq!(xs_all.len(), b * cols);
8724 debug_assert_eq!(out.len(), b * rows);
8725 let (packed, scales) = q4_split(bytes, rows, cols);
8726 let gpr = cols / GROUP_SIZE;
8727 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8728
8729 if a8w8_enabled() {
8730 let acts: Vec<SplitAct> = (0..b)
8731 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8732 .collect();
8733 let acts = &acts;
8734 let out_addr = SendMut(out.as_mut_ptr());
8735 let run = move |start: usize, end: usize| {
8736 ROW_I8.with(|rb| {
8737 let mut buf = rb.borrow_mut();
8738 buf.resize(cols, 0);
8739 for r in start..end {
8740 for gi in 0..gpr {
8744 let g = r * gpr + gi;
8745 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8746 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
8747 buf[gi * GROUP_SIZE + k * 2 + 1] =
8748 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
8749 }
8750 }
8751 let mut bi = 0usize;
8752 #[cfg(target_arch = "x86_64")]
8753 if avx2_enabled() && blocked_enabled() {
8754 while bi + 4 <= acts.len() {
8755 let xs = [
8756 acts[bi].xq.as_slice(),
8757 acts[bi + 1].xq.as_slice(),
8758 acts[bi + 2].xq.as_slice(),
8759 acts[bi + 3].xq.as_slice(),
8760 ];
8761 let d = unsafe {
8762 if vnni_tiles_enabled() {
8763 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
8764 } else {
8765 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
8766 }
8767 };
8768 for k in 0..4 {
8769 let act = &acts[bi + k];
8770 let mut acc = d[k] * act.sx;
8771 for &(j, xv) in &act.outliers {
8772 acc += (buf[j] as i8) as f32
8773 * gscale((r * cols + j) / GROUP_SIZE)
8774 * xv;
8775 }
8776 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8778 }
8779 bi += 4;
8780 }
8781 }
8782 while bi < acts.len() {
8783 let act = &acts[bi];
8784 let mut acc = 0f32;
8785 for gi in 0..gpr {
8786 let d = dot_i8_i8(
8787 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8788 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8789 );
8790 acc += d as f32 * gscale(r * gpr + gi);
8791 }
8792 acc *= act.sx;
8793 for &(j, xv) in &act.outliers {
8795 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
8796 }
8797 unsafe { *out_addr.at(bi * rows + r) = acc };
8799 bi += 1;
8800 }
8801 }
8802 })
8803 };
8804 dispatch_rows(pool, rows, &run);
8805 return;
8806 }
8807
8808 let out_addr = SendMut(out.as_mut_ptr());
8809 let run = move |start: usize, end: usize| {
8810 ROW_F32.with(|rb| {
8811 let mut buf = rb.borrow_mut();
8812 buf.resize(cols, 0.0);
8813 for r in start..end {
8814 for gi in 0..gpr {
8817 let g = r * gpr + gi;
8818 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8819 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8820 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8821 }
8822 }
8823 for bi in 0..b {
8824 let x = &xs_all[bi * cols..(bi + 1) * cols];
8825 let mut acc = 0f32;
8826 for gi in 0..gpr {
8827 let mut ga = 0f32;
8828 for k in 0..GROUP_SIZE / 2 {
8833 let e = gi * GROUP_SIZE + k * 2;
8834 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8835 }
8836 acc += ga * gscale(r * gpr + gi);
8837 }
8838 unsafe { *out_addr.at(bi * rows + r) = acc };
8840 }
8841 }
8842 })
8843 };
8844 dispatch_rows(pool, rows, &run);
8845}
8846
8847#[allow(clippy::too_many_arguments)]
8852fn vbitmatmat(
8853 bytes: &[u8],
8854 offsets: &[usize],
8855 xs_all: &[f32],
8856 b: usize,
8857 rows: usize,
8858 cols: usize,
8859 out: &mut [f32],
8860 pool: Option<&Pool>,
8861) {
8862 debug_assert_eq!(xs_all.len(), b * cols);
8863 debug_assert_eq!(out.len(), b * rows);
8864 debug_assert_eq!(offsets.len(), rows + 1);
8865 let ng = cols / GROUP_SIZE;
8866 let bits = &bytes[..rows];
8867 let sc_off = rows;
8868 let gscale = |r: usize, g: usize| {
8869 let so = (r * ng + g) * 2;
8870 f16_to_f32(u16::from_le_bytes([
8871 bytes[sc_off + so],
8872 bytes[sc_off + so + 1],
8873 ]))
8874 };
8875
8876 let decode_f32 = |r: usize, dst: &mut [f32]| {
8878 let bw = bits[r] as usize;
8879 let l = ((1i32 << (bw - 1)) - 1) as f32;
8880 let data = &bytes[offsets[r]..offsets[r + 1]];
8881 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8882 for d in dst.iter_mut() {
8883 while nbits < bw {
8884 acc = (acc << 8) | data[idx] as u64;
8885 idx += 1;
8886 nbits += 8;
8887 }
8888 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8889 nbits -= bw;
8890 *d = u - l;
8891 }
8892 };
8893
8894 if a8w8_enabled() {
8895 let acts: Vec<SplitAct> = (0..b)
8896 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8897 .collect();
8898 let acts = &acts;
8899 let out_addr = SendMut(out.as_mut_ptr());
8900 let run = move |start: usize, end: usize| {
8901 for r in start..end {
8902 let bw = bits[r] as usize;
8903 if bw == 8 {
8904 ROW_F32.with(|rb| {
8907 let mut buf = rb.borrow_mut();
8908 buf.resize(cols, 0.0);
8909 decode_f32(r, &mut buf);
8910 for bi in 0..b {
8911 let x = &xs_all[bi * cols..(bi + 1) * cols];
8912 let mut dot = 0f32;
8913 for g in 0..ng {
8914 let mut gd = 0f32;
8915 for k in 0..GROUP_SIZE {
8916 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8917 }
8918 dot += gd * gscale(r, g);
8919 }
8920 unsafe { *out_addr.at(bi * rows + r) = dot };
8922 }
8923 });
8924 continue;
8925 }
8926 let l = (1i32 << (bw - 1)) - 1;
8927 let data = &bytes[offsets[r]..offsets[r + 1]];
8928 ROW_I8.with(|rb| {
8929 let mut buf = rb.borrow_mut();
8930 buf.resize(cols, 0);
8931 #[inline(always)]
8932 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8933 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8934 let u = unpack8::<B>(&data[blk * B..]);
8935 for k in 0..8 {
8936 chunk[k] = (u[k] - l) as i8 as u8;
8937 }
8938 }
8939 }
8940 match bw {
8941 3 => fill::<3>(data, l, &mut buf),
8942 4 => vbit_fill4(data, &mut buf),
8943 5 => fill::<5>(data, l, &mut buf),
8944 6 => fill::<6>(data, l, &mut buf),
8945 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8946 }
8947 let mut bi = 0usize;
8948 #[cfg(target_arch = "x86_64")]
8952 if avx2_enabled() && blocked_enabled() {
8953 while bi + 4 <= acts.len() {
8954 let xs = [
8955 acts[bi].xq.as_slice(),
8956 acts[bi + 1].xq.as_slice(),
8957 acts[bi + 2].xq.as_slice(),
8958 acts[bi + 3].xq.as_slice(),
8959 ];
8960 let sxs = [
8961 acts[bi].sx,
8962 acts[bi + 1].sx,
8963 acts[bi + 2].sx,
8964 acts[bi + 3].sx,
8965 ];
8966 let d = unsafe {
8967 if vnni_tiles_enabled() {
8968 dot_q4b_row_1x4_sx_vnni(
8969 &buf,
8970 &bytes[sc_off..],
8971 r * ng,
8972 ng,
8973 xs,
8974 sxs,
8975 )
8976 } else {
8977 dot_q4b_row_1x4_sx_avx2(
8978 &buf,
8979 &bytes[sc_off..],
8980 r * ng,
8981 ng,
8982 xs,
8983 sxs,
8984 )
8985 }
8986 };
8987 for k in 0..4 {
8988 let act = &acts[bi + k];
8989 let mut dot = d[k];
8990 for &(j, xv) in &act.outliers {
8991 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8992 }
8993 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8995 }
8996 bi += 4;
8997 }
8998 }
8999 while bi < acts.len() {
9000 let act = &acts[bi];
9001 let mut dot = 0f32;
9002 for g in 0..ng {
9003 let d = dot_i8_i8(
9004 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9005 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9006 ) as f32
9007 * act.sx;
9008 dot += d * gscale(r, g);
9009 }
9010 for &(j, xv) in &act.outliers {
9011 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9012 }
9013 unsafe { *out_addr.at(bi * rows + r) = dot };
9015 bi += 1;
9016 }
9017 });
9018 }
9019 };
9020 dispatch_rows(pool, rows, &run);
9021 return;
9022 }
9023
9024 let out_addr = SendMut(out.as_mut_ptr());
9025 let run = move |start: usize, end: usize| {
9026 ROW_F32.with(|rb| {
9027 let mut buf = rb.borrow_mut();
9028 buf.resize(cols, 0.0);
9029 for r in start..end {
9030 decode_f32(r, &mut buf);
9031 for bi in 0..b {
9032 let x = &xs_all[bi * cols..(bi + 1) * cols];
9033 let mut dot = 0f32;
9034 for g in 0..ng {
9035 let mut gd = 0f32;
9036 for k in 0..GROUP_SIZE {
9037 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9038 }
9039 dot += gd * gscale(r, g);
9040 }
9041 unsafe { *out_addr.at(bi * rows + r) = dot };
9043 }
9044 }
9045 })
9046 };
9047 dispatch_rows(pool, rows, &run);
9048}
9049
9050pub(crate) fn gpu_batch_job<'a>(
9054 t: &'a QTensor,
9055 x: &[f32],
9056) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
9057 match t {
9058 QTensor::Mapped {
9059 model,
9060 idx,
9061 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
9062 rows,
9063 cols,
9064 row_scale,
9065 col_field,
9066 ..
9067 } => Some((
9068 model.clone(),
9069 crate::gpu::BatchJob {
9070 idx: *idx,
9071 rows: *rows,
9072 cols: *cols,
9073 row_scale,
9074 xs: prescale(x, col_field, *dt).into_owned(),
9075 layout: crate::gpu::BatchLayout::Q8,
9076 },
9077 )),
9078 QTensor::Mapped {
9080 model,
9081 idx,
9082 dtype: TensorDtype::Q1,
9083 rows,
9084 cols,
9085 ..
9086 } => Some((
9087 model.clone(),
9088 crate::gpu::BatchJob {
9089 idx: *idx,
9090 rows: *rows,
9091 cols: *cols,
9092 row_scale: &[],
9093 xs: x.to_vec(),
9094 layout: crate::gpu::BatchLayout::Q1,
9095 },
9096 )),
9097 QTensor::Mapped {
9102 model,
9103 idx,
9104 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
9105 rows,
9106 cols,
9107 ..
9108 } => Some((
9109 model.clone(),
9110 crate::gpu::BatchJob {
9111 idx: *idx,
9112 rows: *rows,
9113 cols: *cols,
9114 row_scale: &[],
9115 xs: x.to_vec(),
9116 layout: if *dt == TensorDtype::Q4Tiled {
9117 crate::gpu::BatchLayout::Q4t
9118 } else {
9119 crate::gpu::BatchLayout::Q4tp
9120 },
9121 },
9122 )),
9123 _ => None,
9124 }
9125}
9126
9127thread_local! {
9128 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9129 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9130}
9131
9132pub(crate) fn prescale<'a>(
9133 x: &'a [f32],
9134 col_field: &[f32],
9135 dtype: TensorDtype,
9136) -> std::borrow::Cow<'a, [f32]> {
9137 if dtype == TensorDtype::Q8_2f {
9138 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
9139 } else {
9140 std::borrow::Cow::Borrowed(x)
9141 }
9142}
9143
9144pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
9147 x: &[f32],
9148 col_field: &[f32],
9149 dtype: TensorDtype,
9150 buf_id: u8,
9151 f: F,
9152) -> R {
9153 if dtype == TensorDtype::Q8_2f {
9154 if buf_id == 1 {
9155 PRESCALE_BUF1.with(|b| {
9156 let mut buf = b.borrow_mut();
9157 buf.clear();
9158 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9159 f(&buf)
9160 })
9161 } else {
9162 PRESCALE_BUF2.with(|b| {
9163 let mut buf = b.borrow_mut();
9164 buf.clear();
9165 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9166 f(&buf)
9167 })
9168 }
9169 } else {
9170 f(x)
9171 }
9172}
9173
9174#[cfg(target_arch = "x86_64")]
9179pub(crate) fn avx2_enabled() -> bool {
9180 use std::sync::OnceLock;
9181 static ON: OnceLock<bool> = OnceLock::new();
9182 *ON.get_or_init(|| {
9183 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
9184 && std::arch::is_x86_feature_detected!("avx2")
9185 && std::arch::is_x86_feature_detected!("fma")
9186 })
9187}
9188
9189#[cfg(target_arch = "x86_64")]
9194fn avx2_a8w8_enabled() -> bool {
9195 use std::sync::OnceLock;
9196 static ON: OnceLock<bool> = OnceLock::new();
9197 *ON.get_or_init(|| {
9198 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
9199 })
9200}
9201
9202#[inline]
9206pub(crate) fn a8w8_enabled() -> bool {
9207 #[cfg(target_arch = "aarch64")]
9208 {
9209 sdot_enabled()
9210 }
9211 #[cfg(target_arch = "x86_64")]
9212 {
9213 avx2_a8w8_enabled()
9214 }
9215 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
9216 {
9217 false
9218 }
9219}
9220
9221#[inline]
9224#[allow(unreachable_code)]
9225fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
9226 #[cfg(target_arch = "aarch64")]
9227 unsafe {
9228 return dot_i8_sdot(w, xq);
9229 }
9230 #[cfg(target_arch = "x86_64")]
9231 unsafe {
9232 if avx512vnni_enabled() {
9233 return dot_i8_i8_vnni(w, xq);
9234 }
9235 return dot_i8_i8_avx2(w, xq);
9236 }
9237 w.iter()
9238 .zip(xq)
9239 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
9240 .sum()
9241}
9242
9243#[cfg(target_arch = "x86_64")]
9247fn avx512vnni_enabled() -> bool {
9248 use std::sync::OnceLock;
9249 static ON: OnceLock<bool> = OnceLock::new();
9250 *ON.get_or_init(|| {
9251 std::env::var("CMF_AVX512")
9252 .map(|v| v != "0")
9253 .unwrap_or(true)
9254 && std::arch::is_x86_feature_detected!("avx512f")
9255 && std::arch::is_x86_feature_detected!("avx512bw")
9256 && std::arch::is_x86_feature_detected!("avx512vl")
9257 && std::arch::is_x86_feature_detected!("avx512vnni")
9258 })
9259}
9260
9261#[cfg(target_arch = "x86_64")]
9269fn vnni_tiles_enabled() -> bool {
9270 use std::sync::OnceLock;
9271 static ON: OnceLock<bool> = OnceLock::new();
9272 *ON.get_or_init(|| {
9273 std::env::var("CMF_VNNI_TILES")
9274 .map(|v| v != "0")
9275 .unwrap_or(true)
9276 && avx512vnni_enabled()
9277 })
9278}
9279
9280#[cfg(target_arch = "x86_64")]
9285#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9286#[inline]
9287unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
9288 unsafe {
9290 use core::arch::x86_64::*;
9291 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
9292 let hi128 = _mm256_extracti128_si256::<1>(d);
9293 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9294 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9295 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9296 _mm_cvtsi128_si32(s32)
9297 }
9298}
9299
9300#[cfg(target_arch = "x86_64")]
9305#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9306unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9307 unsafe {
9309 use core::arch::x86_64::*;
9310 let n = w.len();
9311 let mut j = 0usize;
9312 let mut total: i32;
9313 {
9318 #[inline(always)]
9319 unsafe fn step(
9320 w: *const u8,
9321 x: *const i8,
9322 acc: core::arch::x86_64::__m512i,
9323 ) -> core::arch::x86_64::__m512i {
9324 unsafe {
9325 use core::arch::x86_64::*;
9326 let wv = _mm512_loadu_si512(w as *const _);
9327 let xv = _mm512_loadu_si512(x as *const _);
9328 let aw = _mm512_abs_epi8(wv);
9329 let neg = _mm512_movepi8_mask(wv);
9330 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
9331 _mm512_dpbusd_epi32(acc, aw, sx)
9332 }
9333 }
9334 let (mut a0, mut a1, mut a2, mut a3) = (
9335 _mm512_setzero_si512(),
9336 _mm512_setzero_si512(),
9337 _mm512_setzero_si512(),
9338 _mm512_setzero_si512(),
9339 );
9340 while j + 256 <= n {
9341 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9342 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
9343 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
9344 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
9345 j += 256;
9346 }
9347 while j + 64 <= n {
9348 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9349 j += 64;
9350 }
9351 let s01 = _mm512_add_epi32(a0, a1);
9352 let s23 = _mm512_add_epi32(a2, a3);
9353 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
9354 }
9355 if j + 32 <= n {
9357 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9358 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9359 let d = _mm256_dpbusd_epi32(
9360 _mm256_setzero_si256(),
9361 _mm256_abs_epi8(wv),
9362 _mm256_sign_epi8(xv, wv),
9363 );
9364 let hi128 = _mm256_extracti128_si256::<1>(d);
9365 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9366 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9367 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9368 total += _mm_cvtsi128_si32(s32);
9369 j += 32;
9370 }
9371 while j < n {
9372 total += (w[j] as i8) as i32 * xq[j] as i32;
9373 j += 1;
9374 }
9375 total
9376 }
9377}
9378
9379#[cfg(target_arch = "x86_64")]
9381#[target_feature(enable = "avx2,fma")]
9382unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
9383 unsafe {
9385 use core::arch::x86_64::*;
9386 let n = x.len();
9387 let wp = w.as_ptr();
9388 let xp = x.as_ptr();
9389 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
9390 let mut j = 0usize;
9391 while j + 16 <= n {
9392 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
9393 let lo = _mm256_cvtepi8_epi32(wb);
9394 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
9395 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
9396 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
9397 j += 16;
9398 }
9399 let acc = _mm256_add_ps(a0, a1);
9400 let hi128 = _mm256_extractf128_ps::<1>(acc);
9401 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
9402 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
9403 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
9404 let mut sum = _mm_cvtss_f32(s32);
9405 while j < n {
9406 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
9407 j += 1;
9408 }
9409 sum
9410 }
9411}
9412
9413#[cfg(target_arch = "x86_64")]
9418#[target_feature(enable = "avx2")]
9419unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
9420 unsafe {
9422 use core::arch::x86_64::*;
9423 let n = w.len();
9424 let ones = _mm256_set1_epi16(1);
9425 let mut acc = _mm256_setzero_si256();
9426 let mut j = 0usize;
9427 while j + 32 <= n {
9428 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9429 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9430 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9431 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
9432 j += 32;
9433 }
9434 let hi128 = _mm256_extracti128_si256::<1>(acc);
9435 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
9436 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9437 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9438 let mut s = _mm_cvtsi128_si32(s32);
9439 while j < n {
9440 s += (w[j] as i8) as i32 * xq[j] as i32;
9441 j += 1;
9442 }
9443 s
9444 }
9445}
9446
9447#[cfg(target_arch = "aarch64")]
9451#[target_feature(enable = "neon,i8mm")]
9452unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9453 unsafe {
9455 use core::arch::aarch64::*;
9456 use core::arch::asm;
9457 let n = w0.len();
9458 let w0p = w0.as_ptr() as *const i8;
9459 let w1p = w1.as_ptr() as *const i8;
9460 let mut acc01 = vdupq_n_s32(0);
9463 let mut acc23 = vdupq_n_s32(0);
9464 let mut i = 0usize;
9465 while i + 8 <= n {
9466 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
9467 let xb01 = vcombine_s8(
9468 vld1_s8(xs[0].as_ptr().add(i)),
9469 vld1_s8(xs[1].as_ptr().add(i)),
9470 );
9471 let xb23 = vcombine_s8(
9472 vld1_s8(xs[2].as_ptr().add(i)),
9473 vld1_s8(xs[3].as_ptr().add(i)),
9474 );
9475 asm!(
9476 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
9477 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
9478 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
9479 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
9480 options(pure, nomem, nostack),
9481 );
9482 i += 8;
9483 }
9484 let mut out = [[0i32; 4]; 2];
9485 let a01: [i32; 4] = core::mem::transmute(acc01);
9486 let a23: [i32; 4] = core::mem::transmute(acc23);
9487 out[0][0] = a01[0];
9488 out[0][1] = a01[1];
9489 out[1][0] = a01[2];
9490 out[1][1] = a01[3];
9491 out[0][2] = a23[0];
9492 out[0][3] = a23[1];
9493 out[1][2] = a23[2];
9494 out[1][3] = a23[3];
9495 if i < n {
9496 for (k, x) in xs.iter().enumerate() {
9497 for j in i..n {
9498 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9499 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9500 }
9501 }
9502 }
9503 out
9504 }
9505}
9506
9507#[cfg(target_arch = "aarch64")]
9511#[target_feature(enable = "neon,dotprod")]
9512unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9513 unsafe {
9515 use core::arch::aarch64::*;
9516 use core::arch::asm;
9517 let n = w0.len();
9518 let w0p = w0.as_ptr() as *const i8;
9519 let w1p = w1.as_ptr() as *const i8;
9520 let mut acc = [[vdupq_n_s32(0); 4]; 2];
9521 let mut i = 0usize;
9522 while i + 16 <= n {
9523 let wv0 = vld1q_s8(w0p.add(i));
9524 let wv1 = vld1q_s8(w1p.add(i));
9525 for (k, x) in xs.iter().enumerate() {
9526 let xv = vld1q_s8(x.as_ptr().add(i));
9527 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
9528 asm!(
9529 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
9530 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
9531 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9532 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
9533 options(pure, nomem, nostack),
9534 );
9535 acc[0][k] = a0;
9536 acc[1][k] = a1;
9537 }
9538 i += 16;
9539 }
9540 let mut out = [[0i32; 4]; 2];
9541 for r in 0..2 {
9542 for k in 0..4 {
9543 out[r][k] = vaddvq_s32(acc[r][k]);
9544 }
9545 }
9546 if i < n {
9547 for (k, x) in xs.iter().enumerate() {
9548 for j in i..n {
9549 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9550 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9551 }
9552 }
9553 }
9554 out
9555 }
9556}
9557
9558#[cfg(target_arch = "x86_64")]
9564#[target_feature(enable = "avx2")]
9565unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9566 unsafe {
9568 use core::arch::x86_64::*;
9569 let n = w0.len();
9570 let ones = _mm256_set1_epi16(1);
9571 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
9572 let mut j = 0usize;
9573 while j + 32 <= n {
9574 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
9575 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
9576 let aw0 = _mm256_abs_epi8(wv0);
9577 let aw1 = _mm256_abs_epi8(wv1);
9578 for (k, x) in xs.iter().enumerate() {
9579 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
9580 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
9581 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
9582 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
9583 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
9584 }
9585 j += 32;
9586 }
9587 let mut out = [[0i32; 4]; 2];
9588 for r in 0..2 {
9589 for k in 0..4 {
9590 let a = acc[r][k];
9591 let hi128 = _mm256_extracti128_si256::<1>(a);
9592 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
9593 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9594 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9595 out[r][k] = _mm_cvtsi128_si32(s32);
9596 }
9597 }
9598 if j < n {
9599 for (k, x) in xs.iter().enumerate() {
9600 for i in j..n {
9601 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
9602 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
9603 }
9604 }
9605 }
9606 out
9607 }
9608}
9609
9610#[cfg(target_arch = "x86_64")]
9615#[inline]
9616fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
9617 let dot = if avx512vnni_enabled() && row.len() >= 64 {
9618 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
9619 } else {
9620 unsafe { dot_i8_i8_avx2(row, &act.xq) }
9621 };
9622 let mut acc = dot as f32 * act.sx;
9623 for &(j, xv) in &act.outliers {
9624 acc += (row[j] as i8) as f32 * xv;
9625 }
9626 acc
9627}
9628
9629#[cfg(target_arch = "x86_64")]
9633#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9634unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9635 unsafe {
9637 use core::arch::x86_64::*;
9638 let n = w.len();
9639 let flip = _mm512_set1_epi8(-128); #[inline(always)]
9641 unsafe fn step(
9642 w: *const u8,
9643 x: *const i8,
9644 flip: core::arch::x86_64::__m512i,
9645 acc: core::arch::x86_64::__m512i,
9646 ) -> core::arch::x86_64::__m512i {
9647 unsafe {
9648 use core::arch::x86_64::*;
9649 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
9650 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
9651 }
9652 }
9653 let (mut a0, mut a1, mut a2, mut a3) = (
9654 _mm512_setzero_si512(),
9655 _mm512_setzero_si512(),
9656 _mm512_setzero_si512(),
9657 _mm512_setzero_si512(),
9658 );
9659 let mut j = 0usize;
9660 while j + 256 <= n {
9661 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9662 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
9663 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
9664 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
9665 j += 256;
9666 }
9667 while j + 64 <= n {
9668 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9669 j += 64;
9670 }
9671 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
9672 _mm512_add_epi32(a0, a1),
9673 _mm512_add_epi32(a2, a3),
9674 ));
9675 while j < n {
9677 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
9678 j += 1;
9679 }
9680 total
9681 }
9682}
9683
9684#[cfg(target_arch = "x86_64")]
9690#[target_feature(enable = "avx2")]
9691unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9692 unsafe {
9695 use core::arch::x86_64::*;
9696 let lomask = _mm_set1_epi8(0x0F);
9697 let eight = _mm256_set1_epi8(8);
9698 let ones = _mm256_set1_epi16(1);
9699 let mut acc = 0f32;
9700 for gi in 0..gpr {
9701 let g = g0 + gi;
9702 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9703 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9704 let lo = _mm_and_si128(b, lomask);
9705 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9706 let w = _mm256_sub_epi8(
9707 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9708 eight,
9709 );
9710 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9711 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
9712 let d = _mm256_madd_epi16(p16, ones);
9713 let hi128 = _mm256_extracti128_si256::<1>(d);
9714 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9715 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9716 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9717 acc += _mm_cvtsi128_si32(s32) as f32 * s;
9718 }
9719 acc
9720 }
9721}
9722
9723#[cfg(target_arch = "x86_64")]
9726#[target_feature(enable = "avx2")]
9727unsafe fn dot_q4_row_avx2_2(
9728 packed: &[u8],
9729 scales: &[u8],
9730 g0: usize,
9731 gpr: usize,
9732 xq1: &[i8],
9733 xq2: &[i8],
9734) -> (f32, f32) {
9735 unsafe {
9737 use core::arch::x86_64::*;
9738 let lomask = _mm_set1_epi8(0x0F);
9739 let eight = _mm256_set1_epi8(8);
9740 let ones = _mm256_set1_epi16(1);
9741 let (mut acc1, mut acc2) = (0f32, 0f32);
9742 #[inline(always)]
9743 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
9744 unsafe {
9745 use core::arch::x86_64::*;
9746 let hi128 = _mm256_extracti128_si256::<1>(d);
9747 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9748 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9749 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9750 _mm_cvtsi128_si32(s32)
9751 }
9752 }
9753 for gi in 0..gpr {
9754 let g = g0 + gi;
9755 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9756 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9757 let lo = _mm_and_si128(b, lomask);
9758 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9759 let w = _mm256_sub_epi8(
9760 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9761 eight,
9762 );
9763 let aw = _mm256_abs_epi8(w);
9764 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9765 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9766 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
9767 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
9768 acc1 += hsum(d1) as f32 * s;
9769 acc2 += hsum(d2) as f32 * s;
9770 }
9771 (acc1, acc2)
9772 }
9773}
9774
9775#[cfg(target_arch = "x86_64")]
9777fn q8_range_avx2(
9778 q: &[u8],
9779 row_scale: &[f32],
9780 act: &SplitAct,
9781 cols: usize,
9782 out_addr: SendMut,
9783 start: usize,
9784 end: usize,
9785) {
9786 for o in start..end {
9787 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9788 unsafe { *out_addr.at(o) = v };
9790 }
9791}
9792
9793#[cfg(target_arch = "x86_64")]
9795#[allow(clippy::too_many_arguments)]
9796fn q8_range2_avx2(
9797 q: &[u8],
9798 row_scale: &[f32],
9799 a1: &SplitAct,
9800 a2: &SplitAct,
9801 cols: usize,
9802 p1: SendMut,
9803 p2: SendMut,
9804 start: usize,
9805 end: usize,
9806) {
9807 for o in start..end {
9808 let row = &q[o * cols..(o + 1) * cols];
9809 unsafe {
9811 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9812 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9813 }
9814 }
9815}
9816
9817#[cfg(target_arch = "aarch64")]
9828fn i8mm_enabled() -> bool {
9829 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9830 *ON.get_or_init(|| {
9831 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9832 && std::arch::is_aarch64_feature_detected!("i8mm")
9833 })
9834}
9835
9836#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9840fn sdot_enabled() -> bool {
9841 use std::sync::OnceLock;
9842 static ON: OnceLock<bool> = OnceLock::new();
9843 *ON.get_or_init(|| {
9844 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9845 if !want {
9846 return false;
9847 }
9848
9849 #[cfg(target_arch = "aarch64")]
9850 {
9851 if std::arch::is_aarch64_feature_detected!("dotprod") {
9852 return true;
9853 }
9854 #[cfg(target_os = "android")]
9855 {
9856 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9857 if cpuinfo.lines().any(|l| {
9858 (l.starts_with("Features") || l.starts_with("features"))
9859 && l.contains("asimddp")
9860 }) {
9861 return true;
9862 }
9863 }
9864 }
9865 false
9866 }
9867 #[cfg(not(target_arch = "aarch64"))]
9868 {
9869 false
9870 }
9871 })
9872}
9873
9874struct SplitAct {
9879 xq: Vec<i8>,
9880 sx: f32,
9881 outliers: Vec<(usize, f32)>,
9882 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9885 xsum: i32,
9886}
9887
9888thread_local! {
9889 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9892 const { std::cell::RefCell::new(Vec::new()) };
9893}
9894
9895impl Drop for SplitAct {
9896 fn drop(&mut self) {
9897 let buf = std::mem::take(&mut self.xq);
9898 if buf.capacity() > 0 {
9899 XQ_FREE.with(|f| {
9900 let mut f = f.borrow_mut();
9901 if f.len() < 16 {
9902 f.push(buf);
9903 }
9904 });
9905 }
9906 }
9907}
9908
9909thread_local! {
9910 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9917}
9918
9919#[inline]
9922fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9923 KROW.with(|s| {
9924 let mut b = s.borrow_mut();
9925 if b.len() < n {
9926 b.resize(n, 0.0);
9927 }
9928 f(&mut b[..n])
9929 })
9930}
9931
9932#[inline(always)]
9943fn q8_round(t: f32) -> i8 {
9944 let t = t.clamp(-127.0, 127.0);
9945 let i = t as i32;
9946 let f = t - i as f32;
9947 let r = if f >= 0.5 {
9948 i + 1
9949 } else if f <= -0.5 {
9950 i - 1
9951 } else {
9952 i
9953 };
9954 r as i8
9955}
9956
9957fn split_act(x: &[f32]) -> SplitAct {
9958 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
9959 let n = x.len();
9960 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9961 let thr = 8.0 * rms;
9962 let mut outliers: Vec<(usize, f32)> = Vec::new();
9966 let mut amax = 0f32;
9967 for (j, &v) in x.iter().enumerate() {
9968 let a = v.abs();
9969 if a > thr {
9970 outliers.push((j, v));
9971 } else if a > amax {
9972 amax = a;
9973 }
9974 }
9975 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9976 let inv = 1.0 / sx;
9977 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9978 xq.clear();
9979 xq.reserve(n);
9980 if outliers.is_empty() {
9981 xq.extend(
9982 x.iter()
9983 .map(|&v| q8_round(v * inv)),
9984 );
9985 } else {
9986 xq.extend(x.iter().map(|&v| {
9988 if v.abs() > thr {
9989 0
9990 } else {
9991 q8_round(v * inv)
9992 }
9993 }));
9994 }
9995 let xsum = xq.iter().map(|&v| v as i32).sum();
9996 SplitAct {
9997 xq,
9998 sx,
9999 outliers,
10000 xsum,
10001 }
10002}
10003
10004fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
10005 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10006 let n = x.len();
10007 let rms = (x
10008 .iter()
10009 .zip(col)
10010 .map(|(&a, &c)| {
10011 let v = a * c;
10012 (v * v) as f64
10013 })
10014 .sum::<f64>()
10015 / n.max(1) as f64)
10016 .sqrt() as f32;
10017 let thr = 8.0 * rms;
10018
10019 let mut outliers = Vec::new();
10020 let mut amax = 0f32;
10021 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
10022 let v = a * c;
10023 let s = v.abs();
10024 if s > thr {
10025 outliers.push((j, v));
10026 } else if s > amax {
10027 amax = s;
10028 }
10029 }
10030
10031 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10032 let inv = 1.0 / sx;
10033 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10034 xq.clear();
10035 xq.reserve(n);
10036 if outliers.is_empty() {
10037 xq.extend(
10038 x.iter()
10039 .zip(col)
10040 .map(|(&a, &c)| q8_round((a * c) * inv)),
10041 );
10042 } else {
10043 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
10044 let v = a * c;
10045 if v.abs() > thr {
10046 0
10047 } else {
10048 q8_round(v * inv)
10049 }
10050 }));
10051 }
10052 let xsum = xq.iter().map(|&v| v as i32).sum();
10053 SplitAct {
10054 xq,
10055 sx,
10056 outliers,
10057 xsum,
10058 }
10059}
10060
10061#[cfg(target_arch = "aarch64")]
10064#[target_feature(enable = "neon,dotprod")]
10065unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
10066 unsafe {
10068 use core::arch::aarch64::*;
10069 use core::arch::asm;
10070 let wp = w.as_ptr() as *const i8;
10071 let n = w.len();
10072 let (mut a0, mut a1, mut a2, mut a3) = (
10073 vdupq_n_s32(0),
10074 vdupq_n_s32(0),
10075 vdupq_n_s32(0),
10076 vdupq_n_s32(0),
10077 );
10078 let mut i = 0;
10079 while i + 64 <= n {
10080 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10081 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
10082 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
10083 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
10084 asm!(
10085 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10086 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10087 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
10088 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
10089 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10090 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10091 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
10092 options(pure, nomem, nostack),
10093 );
10094 i += 64;
10095 }
10096 while i + 16 <= n {
10097 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10098 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
10099 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
10100 i += 16;
10101 }
10102 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
10103 while i < n {
10104 s += (*wp.add(i)) as i32 * xq[i] as i32;
10105 i += 1;
10106 }
10107 s
10108 }
10109}
10110
10111#[cfg(target_arch = "aarch64")]
10115#[target_feature(enable = "neon,dotprod")]
10116unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
10117 unsafe {
10119 use core::arch::aarch64::*;
10120 use core::arch::asm;
10121 let n = xq.len();
10122 let px = xq.as_ptr();
10123 let (p0, p1, p2, p3) = (
10124 w0.as_ptr() as *const i8,
10125 w1.as_ptr() as *const i8,
10126 w2.as_ptr() as *const i8,
10127 w3.as_ptr() as *const i8,
10128 );
10129 let (mut a0, mut a1, mut a2, mut a3) = (
10130 vdupq_n_s32(0),
10131 vdupq_n_s32(0),
10132 vdupq_n_s32(0),
10133 vdupq_n_s32(0),
10134 );
10135 let mut i = 0;
10136 while i + 16 <= n {
10137 let x = vld1q_s8(px.add(i));
10138 let v0 = vld1q_s8(p0.add(i));
10139 let v1 = vld1q_s8(p1.add(i));
10140 let v2 = vld1q_s8(p2.add(i));
10141 let v3 = vld1q_s8(p3.add(i));
10142 asm!(
10143 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10144 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10145 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10146 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10147 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10148 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10149 options(pure, nomem, nostack),
10150 );
10151 i += 16;
10152 }
10153 let mut r = [
10154 vaddvq_s32(a0),
10155 vaddvq_s32(a1),
10156 vaddvq_s32(a2),
10157 vaddvq_s32(a3),
10158 ];
10159 while i < n {
10160 let xi = *px.add(i) as i32;
10161 r[0] += (*p0.add(i)) as i32 * xi;
10162 r[1] += (*p1.add(i)) as i32 * xi;
10163 r[2] += (*p2.add(i)) as i32 * xi;
10164 r[3] += (*p3.add(i)) as i32 * xi;
10165 i += 1;
10166 }
10167 r
10168 }
10169}
10170
10171#[cfg(target_arch = "aarch64")]
10178#[target_feature(enable = "neon,dotprod")]
10179unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
10180 unsafe {
10183 use core::arch::aarch64::*;
10184 use core::arch::asm;
10185 let n = xq.len();
10186 let px = xq.as_ptr();
10187 let pg = g.as_ptr() as *const i8;
10188 let (mut a0, mut a1, mut a2, mut a3) = (
10189 vdupq_n_s32(0),
10190 vdupq_n_s32(0),
10191 vdupq_n_s32(0),
10192 vdupq_n_s32(0),
10193 );
10194 let mut i = 0;
10195 while i + 16 <= n {
10196 let x = vld1q_s8(px.add(i));
10197 let base = pg.add(4 * i);
10198 let v0 = vld1q_s8(base);
10199 let v1 = vld1q_s8(base.add(16));
10200 let v2 = vld1q_s8(base.add(32));
10201 let v3 = vld1q_s8(base.add(48));
10202 asm!(
10203 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10204 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10205 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10206 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10207 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10208 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10209 options(pure, nomem, nostack),
10210 );
10211 i += 16;
10212 }
10213 [
10214 vaddvq_s32(a0),
10215 vaddvq_s32(a1),
10216 vaddvq_s32(a2),
10217 vaddvq_s32(a3),
10218 ]
10219 }
10220}
10221
10222#[cfg(target_arch = "aarch64")]
10228fn q8_range_sdot(
10229 q: &[u8],
10230 rep: &[u8],
10231 row_scale: &[f32],
10232 act: &SplitAct,
10233 cols: usize,
10234 out_addr: SendMut,
10235 start: usize,
10236 end: usize,
10237) {
10238 let mut o = start;
10239 if !rep.is_empty() {
10242 while o < end && o % 4 != 0 {
10243 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10244 unsafe { *out_addr.at(o) = v };
10245 o += 1;
10246 }
10247 }
10248 while o + 4 <= end {
10249 let r = if rep.is_empty() {
10250 unsafe {
10251 dot_i8_sdot_4rows(
10252 &q[o * cols..(o + 1) * cols],
10253 &q[(o + 1) * cols..(o + 2) * cols],
10254 &q[(o + 2) * cols..(o + 3) * cols],
10255 &q[(o + 3) * cols..(o + 4) * cols],
10256 &act.xq,
10257 )
10258 }
10259 } else {
10260 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
10261 };
10262 for k in 0..4 {
10263 let mut acc = r[k] as f32 * act.sx;
10264 for &(j, xv) in &act.outliers {
10265 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
10266 }
10267 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
10269 }
10270 o += 4;
10271 }
10272 while o < end {
10273 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10274 unsafe { *out_addr.at(o) = v };
10275 o += 1;
10276 }
10277}
10278
10279#[cfg(target_arch = "aarch64")]
10282#[allow(clippy::too_many_arguments)]
10283fn q8_range2_sdot(
10284 q: &[u8],
10285 row_scale: &[f32],
10286 a1: &SplitAct,
10287 a2: &SplitAct,
10288 cols: usize,
10289 p1: SendMut,
10290 p2: SendMut,
10291 start: usize,
10292 end: usize,
10293) {
10294 for o in start..end {
10295 let row = &q[o * cols..(o + 1) * cols];
10296 unsafe {
10298 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
10299 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
10300 }
10301 }
10302}
10303
10304#[allow(clippy::too_many_arguments)]
10306fn q8_range2_f32(
10307 q: &[u8],
10308 row_scale: &[f32],
10309 x1: &[f32],
10310 x2: &[f32],
10311 cols: usize,
10312 p1: SendMut,
10313 p2: SendMut,
10314 start: usize,
10315 end: usize,
10316) {
10317 for o in start..end {
10318 let row = &q[o * cols..(o + 1) * cols];
10319 unsafe {
10321 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
10322 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
10323 }
10324 }
10325}
10326
10327fn q8_range_f32(
10329 q: &[u8],
10330 row_scale: &[f32],
10331 xs: &[f32],
10332 cols: usize,
10333 out_addr: SendMut,
10334 start: usize,
10335 end: usize,
10336) {
10337 for o in start..end {
10338 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10339 unsafe { *out_addr.at(o) = v };
10341 }
10342}
10343
10344#[inline]
10348fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
10349 #[cfg(target_arch = "aarch64")]
10350 return row_dot_sdot(row, act);
10351 #[cfg(target_arch = "x86_64")]
10352 return row_dot_avx2(row, act);
10353 #[allow(unreachable_code)]
10354 q8_row_dot_scalar(row, act)
10355}
10356
10357#[allow(dead_code)]
10358fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
10359 let mut acc = 0i32;
10360 for (k, &b) in row.iter().enumerate() {
10361 acc += (b as i8) as i32 * act.xq[k] as i32;
10362 }
10363 let mut acc = acc as f32 * act.sx;
10364 for &(j, xv) in &act.outliers {
10365 acc += (row[j] as i8) as f32 * xv;
10366 }
10367 acc
10368}
10369
10370#[cfg(target_arch = "aarch64")]
10373#[inline]
10374fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
10375 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
10376 for &(j, xv) in &act.outliers {
10377 acc += (row[j] as i8) as f32 * xv;
10378 }
10379 acc
10380}
10381
10382#[cfg(target_arch = "aarch64")]
10390#[target_feature(enable = "neon,dotprod")]
10391unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10392 unsafe {
10395 use core::arch::aarch64::*;
10396 use core::arch::asm;
10397 let lomask = vdupq_n_u8(0x0F);
10398 let eight = vdupq_n_s8(8);
10399 let mut acc = 0f32;
10400 for gi in 0..gpr {
10401 let g = g0 + gi;
10402 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10403 let b = vld1q_u8(packed.as_ptr().add(g * 16));
10404 let lo = vandq_u8(b, lomask);
10405 let hi = vshrq_n_u8::<4>(b);
10406 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10407 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10408 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
10409 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
10410 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10411 asm!(
10412 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
10413 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
10414 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10415 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
10416 options(pure, nomem, nostack),
10417 );
10418 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10419 }
10420 acc
10421 }
10422}
10423
10424#[cfg(target_arch = "aarch64")]
10429#[target_feature(enable = "neon,dotprod")]
10430unsafe fn dot_q4_row_sdot2(
10431 packed: &[u8],
10432 scales: &[u8],
10433 g0: usize,
10434 gpr: usize,
10435 xq1: &[i8],
10436 xq2: &[i8],
10437) -> (f32, f32) {
10438 unsafe {
10441 use core::arch::aarch64::*;
10442 use core::arch::asm;
10443 let lomask = vdupq_n_u8(0x0F);
10444 let eight = vdupq_n_s8(8);
10445 let (mut acc1, mut acc2) = (0f32, 0f32);
10446 for gi in 0..gpr {
10447 let g = g0 + gi;
10448 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10449 let b = vld1q_u8(packed.as_ptr().add(g * 16));
10450 let lo = vandq_u8(b, lomask);
10451 let hi = vshrq_n_u8::<4>(b);
10452 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10453 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10454 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
10455 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
10456 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
10457 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
10458 let (mut a0, mut a1, mut b0, mut b1) = (
10459 vdupq_n_s32(0),
10460 vdupq_n_s32(0),
10461 vdupq_n_s32(0),
10462 vdupq_n_s32(0),
10463 );
10464 asm!(
10465 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
10466 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
10467 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
10468 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
10469 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10470 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
10471 e0 = in(vreg) e0, e1 = in(vreg) e1,
10472 x10 = in(vreg) x10, x11 = in(vreg) x11,
10473 x20 = in(vreg) x20, x21 = in(vreg) x21,
10474 options(pure, nomem, nostack),
10475 );
10476 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10477 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
10478 }
10479 (acc1, acc2)
10480 }
10481}
10482
10483#[inline]
10488pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
10489 #[cfg(target_arch = "aarch64")]
10490 unsafe {
10491 return axpy_i8_f32_neon(acc, row, w);
10492 }
10493 #[cfg(target_arch = "x86_64")]
10494 if avx2_enabled() {
10495 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
10496 }
10497 #[allow(unreachable_code)]
10498 {
10499 for (a, &b) in acc.iter_mut().zip(row) {
10500 *a += w * b as f32;
10501 }
10502 }
10503}
10504
10505#[cfg(target_arch = "x86_64")]
10507#[target_feature(enable = "avx2,fma")]
10508unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
10509 unsafe {
10511 use core::arch::x86_64::*;
10512 let n = acc.len().min(row.len());
10513 let ap = acc.as_mut_ptr();
10514 let rp = row.as_ptr();
10515 let wv = _mm256_set1_ps(w);
10516 let mut j = 0usize;
10517 while j + 16 <= n {
10518 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
10519 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
10520 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
10521 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
10522 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
10523 _mm256_storeu_ps(ap.add(j), v0);
10524 _mm256_storeu_ps(ap.add(j + 8), v1);
10525 j += 16;
10526 }
10527 while j < n {
10528 *ap.add(j) += w * (*rp.add(j)) as f32;
10529 j += 1;
10530 }
10531 }
10532}
10533
10534#[cfg(target_arch = "aarch64")]
10535#[target_feature(enable = "neon")]
10536unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
10537 unsafe {
10539 use core::arch::aarch64::*;
10540 let n = acc.len().min(row.len());
10541 let ap = acc.as_mut_ptr();
10542 let rp = row.as_ptr();
10543 let wv = vdupq_n_f32(w);
10544 let mut j = 0usize;
10545 while j + 16 <= n {
10546 let rb = vld1q_s8(rp.add(j));
10547 let lo = vmovl_s8(vget_low_s8(rb));
10548 let hi = vmovl_s8(vget_high_s8(rb));
10549 for (off, half) in [(0, lo), (8, hi)] {
10550 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
10551 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
10552 let o = j + off;
10553 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
10554 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
10555 }
10556 j += 16;
10557 }
10558 while j < n {
10559 *ap.add(j) += w * (*rp.add(j)) as f32;
10560 j += 1;
10561 }
10562 }
10563}
10564
10565#[inline]
10568pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
10569 #[cfg(target_arch = "aarch64")]
10570 unsafe {
10571 return dot_i8_f32_neon(w, x);
10572 }
10573 #[cfg(target_arch = "x86_64")]
10574 if avx2_enabled() {
10575 return unsafe { dot_i8_f32_avx2(w, x) };
10576 }
10577 #[allow(unreachable_code)]
10578 {
10579 let mut sum = 0.0f32;
10580 for (j, &b) in w.iter().enumerate() {
10581 sum += (b as i8) as f32 * x[j];
10582 }
10583 sum
10584 }
10585}
10586
10587#[inline]
10591fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10592 #[cfg(target_arch = "aarch64")]
10593 unsafe {
10594 return dot_i8_col_f32_neon(w, x, col);
10595 }
10596 #[allow(unreachable_code)]
10597 {
10598 let mut sum = 0.0f32;
10599 for (j, &b) in w.iter().enumerate() {
10600 sum += (b as i8) as f32 * x[j] * col[j];
10601 }
10602 sum
10603 }
10604}
10605
10606#[cfg(target_arch = "aarch64")]
10607#[target_feature(enable = "neon")]
10608unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10609 unsafe {
10611 use core::arch::aarch64::*;
10612 let n = x.len();
10613 let wp = w.as_ptr() as *const i8;
10614 let xp = x.as_ptr();
10615 let cp = col.as_ptr();
10616 let (mut a0, mut a1, mut a2, mut a3) = (
10617 vdupq_n_f32(0.0),
10618 vdupq_n_f32(0.0),
10619 vdupq_n_f32(0.0),
10620 vdupq_n_f32(0.0),
10621 );
10622 let mut j = 0usize;
10623 while j + 16 <= n {
10624 let wb = vld1q_s8(wp.add(j));
10625 let lo = vmovl_s8(vget_low_s8(wb));
10626 let hi = vmovl_s8(vget_high_s8(wb));
10627 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10628 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10629 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10630 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10631 a0 = vfmaq_f32(
10632 a0,
10633 w0,
10634 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
10635 );
10636 a1 = vfmaq_f32(
10637 a1,
10638 w1,
10639 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
10640 );
10641 a2 = vfmaq_f32(
10642 a2,
10643 w2,
10644 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
10645 );
10646 a3 = vfmaq_f32(
10647 a3,
10648 w3,
10649 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
10650 );
10651 j += 16;
10652 }
10653 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10654 while j < n {
10655 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
10656 j += 1;
10657 }
10658 sum
10659 }
10660}
10661
10662#[cfg(target_arch = "aarch64")]
10663#[target_feature(enable = "neon")]
10664unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
10665 unsafe {
10667 use core::arch::aarch64::*;
10668 let n = x.len();
10669 let wp = w.as_ptr() as *const i8;
10670 let xp = x.as_ptr();
10671 let (mut a0, mut a1, mut a2, mut a3) = (
10672 vdupq_n_f32(0.0),
10673 vdupq_n_f32(0.0),
10674 vdupq_n_f32(0.0),
10675 vdupq_n_f32(0.0),
10676 );
10677 let mut j = 0usize;
10678 while j + 16 <= n {
10679 let wb = vld1q_s8(wp.add(j));
10680 let lo = vmovl_s8(vget_low_s8(wb));
10681 let hi = vmovl_s8(vget_high_s8(wb));
10682 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10683 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10684 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10685 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10686 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
10687 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
10688 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
10689 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
10690 j += 16;
10691 }
10692 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10693 while j < n {
10694 sum += (*wp.add(j)) as f32 * *xp.add(j);
10695 j += 1;
10696 }
10697 sum
10698 }
10699}
10700
10701#[allow(clippy::too_many_arguments)]
10702fn qmatvec(
10703 q: &[u8],
10704 rep: &[u8],
10705 row_scale: &[f32],
10706 x: &[f32],
10707 col_field: &[f32],
10708 dtype: TensorDtype,
10709 rows: usize,
10710 cols: usize,
10711 out: &mut [f32],
10712 pool: Option<&Pool>,
10713) {
10714 debug_assert_eq!(out.len(), rows);
10715 #[cfg(not(target_arch = "aarch64"))]
10716 let _ = rep;
10717
10718 #[cfg(target_arch = "aarch64")]
10719 if sdot_enabled() {
10720 let act = if dtype == TensorDtype::Q8_2f {
10721 split_act_q8_2f(x, col_field)
10722 } else {
10723 split_act(x)
10724 };
10725 let out_addr = SendMut(out.as_mut_ptr());
10726 let run_range = |start: usize, end: usize| {
10727 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
10728 };
10729 match pool {
10730 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10731 _ => run_range(0, rows),
10732 }
10733 return;
10734 }
10735 #[cfg(target_arch = "x86_64")]
10738 if avx2_a8w8_enabled() {
10739 let act = if dtype == TensorDtype::Q8_2f {
10740 split_act_q8_2f(x, col_field)
10741 } else {
10742 split_act(x)
10743 };
10744 let out_addr = SendMut(out.as_mut_ptr());
10745 let run_range = |start: usize, end: usize| {
10746 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
10747 };
10748 match pool {
10749 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10750 _ => run_range(0, rows),
10751 }
10752 return;
10753 }
10754
10755 prescale_with(x, col_field, dtype, 1, |xs| {
10756 let out_addr = SendMut(out.as_mut_ptr());
10757 let run_range = move |start: usize, end: usize| {
10758 for o in start..end {
10759 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10760 unsafe { *out_addr.at(o) = v };
10762 }
10763 };
10764 match pool {
10765 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10766 _ => run_range(0, rows),
10767 }
10768 });
10769}
10770
10771#[allow(clippy::too_many_arguments)]
10772fn qmatvec2(
10773 q: &[u8],
10774 row_scale: &[f32],
10775 x1: &[f32],
10776 x2: &[f32],
10777 col_field: &[f32],
10778 dtype: TensorDtype,
10779 rows: usize,
10780 cols: usize,
10781 o1: &mut [f32],
10782 o2: &mut [f32],
10783 pool: Option<&Pool>,
10784) {
10785 #[cfg(target_arch = "aarch64")]
10786 if sdot_enabled() {
10787 let a1s = if dtype == TensorDtype::Q8_2f {
10788 split_act_q8_2f(x1, col_field)
10789 } else {
10790 split_act(x1)
10791 };
10792 let a2s = if dtype == TensorDtype::Q8_2f {
10793 split_act_q8_2f(x2, col_field)
10794 } else {
10795 split_act(x2)
10796 };
10797 let p1 = SendMut(o1.as_mut_ptr());
10798 let p2 = SendMut(o2.as_mut_ptr());
10799 let run_range = |start: usize, end: usize| {
10800 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10801 };
10802 match pool {
10803 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10804 _ => run_range(0, rows),
10805 }
10806 return;
10807 }
10808 #[cfg(target_arch = "x86_64")]
10809 if avx2_a8w8_enabled() {
10810 let a1s = if dtype == TensorDtype::Q8_2f {
10811 split_act_q8_2f(x1, col_field)
10812 } else {
10813 split_act(x1)
10814 };
10815 let a2s = if dtype == TensorDtype::Q8_2f {
10816 split_act_q8_2f(x2, col_field)
10817 } else {
10818 split_act(x2)
10819 };
10820 let p1 = SendMut(o1.as_mut_ptr());
10821 let p2 = SendMut(o2.as_mut_ptr());
10822 let run_range = |start: usize, end: usize| {
10823 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10824 };
10825 match pool {
10826 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10827 _ => run_range(0, rows),
10828 }
10829 return;
10830 }
10831
10832 prescale_with(x1, col_field, dtype, 1, |x1s| {
10833 prescale_with(x2, col_field, dtype, 2, |x2s| {
10834 let p1 = SendMut(o1.as_mut_ptr());
10835 let p2 = SendMut(o2.as_mut_ptr());
10836 let run_range = move |start: usize, end: usize| {
10837 for o in start..end {
10838 let row = &q[o * cols..(o + 1) * cols];
10839 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10840 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10841 unsafe {
10843 *p1.at(o) = s1;
10844 *p2.at(o) = s2;
10845 }
10846 }
10847 };
10848 match pool {
10849 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10850 _ => run_range(0, rows),
10851 }
10852 });
10853 });
10854}
10855
10856#[derive(Clone, Copy)]
10857struct SendMut(*mut f32);
10858unsafe impl Send for SendMut {}
10859unsafe impl Sync for SendMut {}
10860
10861impl SendMut {
10862 #[inline]
10863 fn at(self, i: usize) -> *mut f32 {
10864 unsafe { self.0.add(i) }
10865 }
10866}
10867
10868#[cfg(test)]
10869mod tests {
10870 #[test]
10874 fn q8_round_is_round_clamp() {
10875 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
10876 let mut probe = vec![
10877 0.0f32,
10878 -0.0,
10879 f32::NAN,
10880 f32::INFINITY,
10881 f32::NEG_INFINITY,
10882 f32::MAX,
10883 f32::MIN,
10884 1e30,
10885 -1e30,
10886 f32::MIN_POSITIVE,
10887 -f32::MIN_POSITIVE,
10888 ];
10889 for k in -300i32..=300 {
10890 let h = k as f32 * 0.5;
10891 let up = f32::from_bits(h.to_bits() + 1);
10892 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
10893 for t in [h, up, down] {
10894 probe.push(t);
10895 probe.push(-t);
10896 }
10897 }
10898 let mut t = -140.0f32;
10899 while t < 140.0 {
10900 probe.push(t);
10901 t += 0.000_731;
10902 }
10903 for t in probe {
10904 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
10905 }
10906 }
10907
10908 use super::*;
10909
10910 #[test]
10911 fn q2tp_i8_dot_matches_exact_on_grid() {
10912 let (rows, cols) = (5, 64);
10916 let gpr = cols / GROUP_SIZE;
10917 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10921 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10922 .collect();
10923 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10924 let x: Vec<f32> = (0..cols)
10925 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10926 .collect();
10927 let act = split_act(&x);
10928 assert!(
10929 act.outliers.is_empty(),
10930 "on-grid input must have no outliers"
10931 );
10932 let gsum = q1_group_sums(&act.xq, gpr);
10933 for r in 0..rows {
10934 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10935 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10936 assert!(
10937 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10938 "row {r}: exact {exact} vs i8 {fast}"
10939 );
10940 }
10941 }
10942
10943 #[test]
10944 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
10945 let (rows, cols) = (1usize, GROUP_SIZE);
10946 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
10947 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
10951 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
10952 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
10954 let mut raw = vec![0.0f32; rows];
10955 let mut affine = vec![0.0f32; rows];
10956 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
10957 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
10958 assert_eq!(raw, vec![-24.0]);
10959 assert_eq!(affine, vec![-8.0]);
10960 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
10961 }
10962
10963 #[cfg(target_arch = "x86_64")]
10964 #[test]
10965 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
10966 if !std::arch::is_x86_feature_detected!("avx2") {
10971 return;
10972 }
10973 let mut seed = 0x9e3779b9u32;
10974 let mut next = || {
10975 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
10976 seed
10977 };
10978 for _ in 0..20_000 {
10979 let mut ch = [0u8; Q2TP_CHUNK];
10980 let mut x = [0i8; GROUP_SIZE];
10981 for b in &mut ch {
10982 *b = next() as u8;
10983 }
10984 for v in &mut x {
10985 *v = (next() >> 24) as i8;
10986 }
10987 let mut reference = 0i32;
10988 for (k, &b) in ch.iter().enumerate() {
10989 reference += (b & 3) as i32 * x[k * 4] as i32;
10990 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
10991 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
10992 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
10993 }
10994 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
10997 assert_eq!(got, reference, "packed q2 lane mismatch");
10998 }
10999 }
11000
11001 #[test]
11002 fn q8_row_dot_fast_matches_scalar() {
11003 let cols = 96;
11006 let row: Vec<u8> = (0..cols)
11007 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
11008 .collect();
11009 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
11010 let act = split_act(&x);
11011 let fast = q8_row_dot(&row, &act);
11012 let scalar = q8_row_dot_scalar(&row, &act);
11013 assert!(
11014 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
11015 "fast {fast} vs scalar {scalar}"
11016 );
11017 }
11018
11019 #[test]
11020 fn f32_matvec_matches_matvec_rows_bitexact() {
11021 let (rows, cols) = (300, 40);
11022 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
11023 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
11024 let qt = QTensor::from_f32(w.clone(), rows, cols);
11025
11026 let mut a = vec![0.0f32; rows];
11027 matvec_rows(None, &w, &x, &mut a);
11028 let mut b = vec![0.0f32; rows];
11029 qt.matvec(&x, &mut b, None);
11030 assert_eq!(a, b);
11031 }
11032
11033 #[test]
11034 fn sdot_kernel_exact_on_grid() {
11035 eprintln!("sdot_enabled = {}", sdot_enabled());
11040 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
11042 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11043 .collect();
11044 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
11045 let x: Vec<f32> = (0..cols)
11046 .map(|i| match i % 3 {
11047 0 => 1.0,
11048 1 => -1.0,
11049 _ => 0.0,
11050 })
11051 .collect();
11052 let mut a = vec![0.0f32; rows];
11053 qmatvec(
11054 &w,
11055 &[],
11056 &scales,
11057 &x,
11058 &[],
11059 TensorDtype::Q8Row,
11060 rows,
11061 cols,
11062 &mut a,
11063 None,
11064 );
11065 for o in 0..rows {
11066 let mut acc = 0.0f32;
11067 for j in 0..cols {
11068 acc += (w[o * cols + j] as i8) as f32 * x[j];
11069 }
11070 let expect = acc * scales[o];
11071 assert!(
11072 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11073 "row {o}: {} vs {expect}",
11074 a[o]
11075 );
11076 }
11077 }
11078
11079 #[test]
11080 fn q1_tbl_fast_path_matches_reference() {
11081 let (rows, cols) = (5, 256);
11086 let gpr = cols / GROUP_SIZE;
11087 let mut bytes = Vec::new();
11088 for t in 0..rows * gpr {
11089 let s = 0.007 + (t % 11) as f32 * 0.004;
11090 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11091 for j in 0..4 {
11092 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
11093 }
11094 }
11095 let x: Vec<f32> = (0..cols)
11096 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
11097 .collect();
11098 let mut w = vec![0.0f32; rows * cols];
11099 cortiq_core::quant::dequant_q1(&bytes, &mut w);
11100 let mut got = vec![0.0f32; rows];
11101 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11102 for o in 0..rows {
11103 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11104 assert!(
11105 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11106 "row {o}: {} vs {expect}",
11107 got[o]
11108 );
11109 }
11110 let b = 5usize;
11113 let mut xs_all = Vec::new();
11114 for bi in 0..b {
11115 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
11116 }
11117 let mut mm = vec![0.0f32; b * rows];
11118 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
11119 for bi in 0..b {
11120 let mut single = vec![0.0f32; rows];
11121 q1_matvec(
11122 &bytes,
11123 &xs_all[bi * cols..(bi + 1) * cols],
11124 rows,
11125 cols,
11126 &mut single,
11127 None,
11128 );
11129 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
11130 }
11131 }
11132
11133 #[test]
11134 fn q1_kernels_match_exact_reference() {
11135 let (rows, cols) = (7, 96);
11137 let gpr = cols / GROUP_SIZE;
11138 let mut bytes = Vec::new();
11139 for t in 0..rows * gpr {
11140 let s = 0.01 + (t % 13) as f32 * 0.003;
11141 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11142 for j in 0..4 {
11143 bytes.push(((t * 31 + j * 97) % 251) as u8);
11144 }
11145 }
11146 let x: Vec<f32> = (0..cols)
11148 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
11149 .collect();
11150 let mut w = vec![0.0f32; rows * cols];
11152 cortiq_core::quant::dequant_q1(&bytes, &mut w);
11153 let mut expect = vec![0.0f32; rows];
11154 for o in 0..rows {
11155 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11156 }
11157 let mut got = vec![0.0f32; rows];
11158 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11159 for o in 0..rows {
11160 assert!(
11161 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
11162 "row {o}: {} vs {}",
11163 got[o],
11164 expect[o]
11165 );
11166 }
11167 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
11169 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
11170 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
11171 assert_eq!(a1, got);
11172 let mut xs = x.clone();
11173 xs.extend_from_slice(&x2);
11174 let mut mm = vec![0.0f32; 2 * rows];
11175 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
11176 assert_eq!(&mm[..rows], got.as_slice());
11177 assert_eq!(&mm[rows..], a2.as_slice());
11178 }
11179
11180 #[test]
11181 fn repack_is_bit_identical() {
11182 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
11188 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
11189 .collect();
11190 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
11191 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
11192 let rep = q8_repack_layout(&w, rows, cols);
11193 for g in 0..rows / 4 {
11195 for c in 0..cols / 16 {
11196 for lane in 0..4 {
11197 assert_eq!(
11198 &rep[g * 4 * cols + c * 64 + lane * 16
11199 ..g * 4 * cols + c * 64 + lane * 16 + 16],
11200 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
11201 );
11202 }
11203 }
11204 }
11205 let mut a = vec![0.0f32; rows];
11206 qmatvec(
11207 &w,
11208 &[],
11209 &scales,
11210 &x,
11211 &[],
11212 TensorDtype::Q8Row,
11213 rows,
11214 cols,
11215 &mut a,
11216 None,
11217 );
11218 let mut b = vec![0.0f32; rows];
11219 qmatvec(
11220 &w,
11221 &rep,
11222 &scales,
11223 &x,
11224 &[],
11225 TensorDtype::Q8Row,
11226 rows,
11227 cols,
11228 &mut b,
11229 None,
11230 );
11231 assert_eq!(a, b, "full-range repack output diverged");
11232
11233 #[cfg(target_arch = "aarch64")]
11234 if sdot_enabled() {
11235 let act = split_act(&x);
11237 let mut c1 = vec![0.0f32; rows];
11238 let mut c2 = vec![0.0f32; rows];
11239 q8_range_sdot(
11240 &w,
11241 &[],
11242 &scales,
11243 &act,
11244 cols,
11245 SendMut(c1.as_mut_ptr()),
11246 3,
11247 rows - 2,
11248 );
11249 q8_range_sdot(
11250 &w,
11251 &rep,
11252 &scales,
11253 &act,
11254 cols,
11255 SendMut(c2.as_mut_ptr()),
11256 3,
11257 rows - 2,
11258 );
11259 assert_eq!(c1, c2, "unaligned-range repack output diverged");
11260 }
11261 }
11262
11263 #[test]
11264 fn sdot_a8w8_noise_is_bounded() {
11265 let (rows, cols) = (16, 512);
11269 let w: Vec<u8> = (0..rows * cols)
11270 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11271 .collect();
11272 let scales = vec![0.01f32; rows];
11273 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
11274 let mut a = vec![0.0f32; rows];
11275 qmatvec(
11276 &w,
11277 &[],
11278 &scales,
11279 &x,
11280 &[],
11281 TensorDtype::Q8Row,
11282 rows,
11283 cols,
11284 &mut a,
11285 None,
11286 );
11287 let (mut num, mut den) = (0f64, 0f64);
11288 for o in 0..rows {
11289 let mut acc = 0.0f32;
11290 for j in 0..cols {
11291 acc += (w[o * cols + j] as i8) as f32 * x[j];
11292 }
11293 let expect = acc * scales[o];
11294 num += ((a[o] - expect) as f64).powi(2);
11295 den += (expect as f64).powi(2);
11296 }
11297 let rel = (num / den.max(1e-12)).sqrt();
11298 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
11299 }
11300
11301 #[test]
11302 fn i8_dot_neon_matches_scalar() {
11303 let n = 100;
11304 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
11305 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
11306 let mut scalar = 0.0f32;
11307 for j in 0..n {
11308 scalar += (w[j] as i8) as f32 * x[j];
11309 }
11310 let fast = dot_i8_f32(&w, &x);
11311 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
11312 }
11313
11314 #[test]
11316 fn vbitmatvec_matches_full_dequant() {
11317 let (rows, cols) = (6, 64);
11318 let ng = cols / GROUP_SIZE;
11319 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11321 let mut bytes = bits.clone();
11322 for g in 0..rows * ng {
11323 let s = 0.02 + 0.001 * g as f32;
11324 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11325 }
11326 for r in 0..rows {
11327 let b = bits[r] as usize;
11328 let (mut acc, mut nb) = (0u64, 0usize);
11329 let mut rowbytes = Vec::new();
11330 for i in 0..cols {
11331 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11332 acc = (acc << b) | v;
11333 nb += b;
11334 while nb >= 8 {
11335 nb -= 8;
11336 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11337 }
11338 }
11339 if nb > 0 {
11340 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11341 }
11342 bytes.extend_from_slice(&rowbytes);
11343 }
11344 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11345
11346 let mut reference = vec![0f32; rows * cols];
11347 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
11348 let mut expect = vec![0f32; rows];
11349 for r in 0..rows {
11350 expect[r] = reference[r * cols..(r + 1) * cols]
11351 .iter()
11352 .zip(&x)
11353 .map(|(w, xv)| w * xv)
11354 .sum();
11355 }
11356 let mut got = vec![0f32; rows];
11357 let offsets = vbit_row_offsets(&bytes, rows, cols);
11358 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
11359 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11363 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11364 for r in 0..rows {
11365 assert!(
11366 (got[r] - expect[r]).abs() < tol * scale,
11367 "row {r}: {} vs {}",
11368 got[r],
11369 expect[r]
11370 );
11371 }
11372 }
11373
11374 #[test]
11379 #[cfg(target_arch = "x86_64")]
11380 fn vbit_matmat_blocked_matches_per_row() {
11381 let (rows, cols, b) = (64usize, 128usize, 9usize);
11382 let ng = cols / GROUP_SIZE;
11383 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
11384 let mut bytes = bits.clone();
11385 for g in 0..rows * ng {
11386 let sc = 0.02 + 0.0005 * g as f32;
11387 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11388 }
11389 for r in 0..rows {
11390 let bw = bits[r] as usize;
11391 let (mut acc, mut nb) = (0u64, 0usize);
11392 let mut rowbytes = Vec::new();
11393 for i in 0..cols {
11394 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11395 acc = (acc << bw) | v;
11396 nb += bw;
11397 while nb >= 8 {
11398 nb -= 8;
11399 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11400 }
11401 }
11402 if nb > 0 {
11403 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11404 }
11405 bytes.extend_from_slice(&rowbytes);
11406 }
11407 let x: Vec<f32> = (0..b * cols)
11408 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11409 .collect();
11410 let offsets = vbit_row_offsets(&bytes, rows, cols);
11411 let mut y_a = vec![0f32; b * rows];
11412 let mut y_b = vec![0f32; b * rows];
11413 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11414 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
11415 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11416 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
11417 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11418 let max_d = y_a
11419 .iter()
11420 .zip(&y_b)
11421 .map(|(p, q)| (p - q).abs())
11422 .fold(0.0f32, f32::max);
11423 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
11424 }
11425
11426 #[test]
11434 fn q4t_matmat_blocked_matches_per_row() {
11435 let (rows, cols, b) = (16usize, 64usize, 9usize);
11436 let gpr = cols / GROUP_SIZE;
11437 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11438 for r in 0..rows {
11439 for g in 0..gpr {
11440 let t = (r * gpr + g) * Q4_TILE;
11441 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
11442 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11443 for k in 0..16 {
11444 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11445 }
11446 }
11447 }
11448 let x: Vec<f32> = (0..b * cols)
11449 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11450 .collect();
11451 let mut y_blk = vec![0f32; b * rows];
11452 let mut y_row = vec![0f32; b * rows];
11453 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11454 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
11455 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11456 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
11457 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11458 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
11459 }
11460
11461 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
11468 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
11469 let gpr = cols / GROUP_SIZE;
11470 let stride = q4tp_code_stride(gpr);
11471 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
11472 let mut b = vec![0u8; codes_off + rows * stride];
11473 for r in 0..rows {
11474 for g in 0..gpr {
11475 let t = (r * gpr + g) * Q4TP_NIB;
11476 for k in 0..16 {
11477 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11478 }
11479 }
11480 let lo = -6.0 - 0.03 * (r % 17) as f32;
11481 let step = 0.01 + 0.004 * (r % 11) as f32;
11482 let p = params_off + r * 4;
11483 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
11484 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
11485 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
11486 for g in 0..gpr {
11487 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
11488 }
11489 }
11490 b
11491 }
11492
11493 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
11497 let gpr = cols / GROUP_SIZE;
11498 let v = Q4tpView::new(bytes, rows, cols);
11499 let mut out = vec![0u8; rows * gpr * Q4_TILE];
11500 let mut sc = vec![0f32; gpr];
11501 for r in 0..rows {
11502 v.scales_into(r, gpr, &mut sc);
11503 for g in 0..gpr {
11504 let t = (r * gpr + g) * Q4_TILE;
11505 let s = sc[g];
11506 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11507 let src = (r * gpr + g) * Q4TP_NIB;
11508 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
11509 }
11510 }
11511 out
11512 }
11513
11514 #[test]
11520 fn q4tp_exact_path_matches_dequant_reference() {
11521 let (rows, cols) = (256usize, 512usize);
11522 let gpr = cols / GROUP_SIZE;
11523 let bytes = synth_q4tp(rows, cols);
11524 let mut w = vec![0f32; rows * cols];
11525 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11526
11527 let x: Vec<f32> = (0..cols)
11528 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11529 .collect();
11530 let v = Q4tpView::new(&bytes, rows, cols);
11531 let mut sc = vec![0f32; gpr];
11532 for r in 0..rows {
11533 v.scales_into(r, gpr, &mut sc);
11534 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
11535 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
11536 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11540 assert!(
11541 (got - want).abs() <= 1e-5 * mag,
11542 "row {r}: kernel {got} vs dequant {want}"
11543 );
11544 }
11545 }
11546
11547 #[test]
11553 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
11554 let (rows, cols) = (256usize, 512usize);
11555 let bytes = synth_q4tp(rows, cols);
11556 let twin = q4tp_as_q4t(&bytes, rows, cols);
11557 let x: Vec<f32> = (0..cols)
11558 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11559 .collect();
11560
11561 let mut got = vec![0f32; rows];
11562 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
11563 let mut want = vec![0f32; rows];
11564 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
11565
11566 let mut w = vec![0f32; rows * cols];
11569 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11570 for r in 0..rows {
11571 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11572 assert!(
11573 (got[r] - want[r]).abs() <= 1e-3 * mag,
11574 "row {r}: q4tp {} vs q4t {}",
11575 got[r],
11576 want[r]
11577 );
11578 }
11579 }
11580
11581 #[test]
11586 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
11587 let (rows, cols, b) = (256usize, 512usize, 5usize);
11588 let bytes = synth_q4tp(rows, cols);
11589 let twin = q4tp_as_q4t(&bytes, rows, cols);
11590 let xs: Vec<f32> = (0..b * cols)
11591 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11592 .collect();
11593
11594 let mut got = vec![0f32; b * rows];
11595 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
11596 let mut want = vec![0f32; b * rows];
11597 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
11598
11599 let mut w = vec![0f32; rows * cols];
11600 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11601 for t in 0..b {
11602 for r in 0..rows {
11603 let mag: f32 = (0..cols)
11604 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
11605 .sum();
11606 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
11607 assert!(
11608 (g - wa).abs() <= 1e-3 * mag,
11609 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
11610 );
11611 }
11612 }
11613 }
11614
11615 #[test]
11616 fn q4tp_matvec2_matches_the_single_stream_kernel() {
11617 let (rows, cols) = (128usize, 256usize);
11618 let gpr = cols / GROUP_SIZE;
11619 let bytes = synth_q4tp(rows, cols);
11620 let xs: Vec<f32> = (0..2 * cols)
11621 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11622 .collect();
11623
11624 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
11625 q4tp_matvec2(
11626 &bytes,
11627 &xs[..cols],
11628 &xs[cols..],
11629 rows,
11630 cols,
11631 &mut o1,
11632 &mut o2,
11633 None,
11634 );
11635
11636 let v = Q4tpView::new(&bytes, rows, cols);
11639 let mut sc = vec![0f32; gpr];
11640 for r in 0..rows {
11641 v.scales_into(r, gpr, &mut sc);
11642 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
11643 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
11644 }
11645 }
11646
11647 #[test]
11654 fn q4tp_matvec_keeps_pace_with_q4t() {
11655 let (rows, cols) = (4096usize, 3072usize);
11656 let bytes = synth_q4tp(rows, cols);
11657 let twin = q4tp_as_q4t(&bytes, rows, cols);
11658 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
11659 let mut o = vec![0f32; rows];
11660 let n = 12;
11661 let mut best = (f64::MAX, f64::MAX);
11662 for _ in 0..3 {
11665 let t0 = std::time::Instant::now();
11666 for _ in 0..n {
11667 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
11668 }
11669 best.0 = best.0.min(t0.elapsed().as_secs_f64());
11670 let t0 = std::time::Instant::now();
11671 for _ in 0..n {
11672 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
11673 }
11674 best.1 = best.1.min(t0.elapsed().as_secs_f64());
11675 }
11676 let ratio = best.1 / best.0;
11677 println!(
11678 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
11679 best.0 * 1e3 / n as f64,
11680 best.1 * 1e3 / n as f64
11681 );
11682 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
11683 }
11684
11685 #[cfg(target_os = "macos")]
11686 #[test]
11687 fn q4t_matmat_accel_matches_dequant_reference() {
11688 if !accel_gemm_enabled() {
11689 return; }
11691 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
11693 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11694 for r in 0..rows {
11695 for g in 0..gpr {
11696 let t = (r * gpr + g) * Q4_TILE;
11697 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
11698 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11699 for k in 0..16 {
11700 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11701 }
11702 }
11703 }
11704 let x: Vec<f32> = (0..b * cols)
11705 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11706 .collect();
11707 let mut got = vec![0f32; b * rows];
11708 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
11709 let mut w = vec![0f32; rows * cols];
11711 for r in 0..rows {
11712 for g in 0..gpr {
11713 let t = (r * gpr + g) * Q4_TILE;
11714 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
11715 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
11716 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
11717 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
11718 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
11719 }
11720 }
11721 }
11722 for bi in 0..b {
11723 for r in 0..rows {
11724 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
11725 let d = (got[bi * rows + r] - want).abs();
11726 assert!(
11727 d <= want.abs().max(1.0) * 1e-4,
11728 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
11729 got[bi * rows + r]
11730 );
11731 }
11732 }
11733 }
11734
11735 #[test]
11736 fn q4matvec_matches_full_dequant() {
11737 let (rows, cols) = (8, 64);
11738 let groups = rows * cols / GROUP_SIZE;
11739 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11741 for i in 0..groups * 16 {
11742 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11743 }
11744 for g in 0..groups {
11745 let s = 0.01 + 0.003 * g as f32;
11746 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11747 }
11748 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11749
11750 let mut reference = vec![0.0f32; rows * cols];
11751 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11752 let mut expect = vec![0.0f32; rows];
11753 for r in 0..rows {
11754 expect[r] = reference[r * cols..(r + 1) * cols]
11755 .iter()
11756 .zip(&x)
11757 .map(|(w, xv)| w * xv)
11758 .sum();
11759 }
11760
11761 let mut got = vec![0.0f32; rows];
11762 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11763 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11767 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11768 for r in 0..rows {
11769 assert!(
11770 (got[r] - expect[r]).abs() < tol * scale,
11771 "row {r}: {} vs {}",
11772 got[r],
11773 expect[r]
11774 );
11775 }
11776 }
11777
11778 #[test]
11781 fn vbitmatvec2_equals_two_singles() {
11782 let (rows, cols) = (6, 64);
11783 let ng = cols / GROUP_SIZE;
11784 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11785 let mut bytes = bits.clone();
11786 for g in 0..rows * ng {
11787 let s = 0.02 + 0.001 * g as f32;
11788 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11789 }
11790 for r in 0..rows {
11791 let b = bits[r] as usize;
11792 let (mut acc, mut nb) = (0u64, 0usize);
11793 let mut rowbytes = Vec::new();
11794 for i in 0..cols {
11795 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11796 acc = (acc << b) | v;
11797 nb += b;
11798 while nb >= 8 {
11799 nb -= 8;
11800 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11801 }
11802 }
11803 if nb > 0 {
11804 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11805 }
11806 bytes.extend_from_slice(&rowbytes);
11807 }
11808 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11809 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
11810 let offsets = vbit_row_offsets(&bytes, rows, cols);
11811
11812 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11813 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
11814 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
11815 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11816 vbitmatvec2(
11817 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
11818 );
11819 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
11820 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
11821 }
11822
11823 #[test]
11825 fn q4matvec2_equals_two_singles() {
11826 let (rows, cols) = (8, 128);
11827 let groups = rows * cols / GROUP_SIZE;
11828 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11829 for i in 0..groups * 16 {
11830 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11831 }
11832 for g in 0..groups {
11833 let s = 0.01 + 0.003 * g as f32;
11834 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11835 }
11836 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11839 x1[9] = 250.0;
11840 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11841
11842 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11843 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
11844 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
11845 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11846 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
11847 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
11848 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
11849 }
11850
11851 #[test]
11854 fn matvec_many_equals_separate_matvecs() {
11855 use crate::pool::Pool;
11856 let (r1, r2, cols) = (300, 200, 64);
11857 let mk = |salt: usize, rows: usize| {
11858 QTensor::from_f32(
11859 (0..rows * cols)
11860 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
11861 .collect(),
11862 rows,
11863 cols,
11864 )
11865 };
11866 let (a, b) = (mk(1, r1), mk(5, r2));
11867 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
11868 let pool = Pool::new(3);
11869
11870 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
11871 a.matvec(&x, &mut ea, Some(&pool));
11872 b.matvec(&x, &mut eb, Some(&pool));
11873 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
11874 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11875 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
11876 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
11877 }
11878
11879 #[test]
11884 fn q4tp_matvec_many_equals_separate_matvecs() {
11885 use crate::pool::Pool;
11886 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
11887
11888 let (r1, r2, cols) = (300usize, 200usize, 64usize);
11889 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
11890 "arch_name": "tiny-q4tp",
11891 "hidden_size": cols,
11892 "intermediate_size": cols * 2,
11893 "num_layers": 1,
11894 "num_attention_heads": 2,
11895 "num_kv_heads": 1,
11896 "head_dim": 32,
11897 "vocab_size": r1,
11898 "layer_types": ["FullAttention"],
11899 "rms_norm_eps": 1e-6,
11900 "max_position_embeddings": 8,
11901 "linear_conv_kernel_dim": 0,
11902 "linear_num_key_heads": 0,
11903 "linear_num_value_heads": 0
11904 }))
11905 .unwrap();
11906 let header = CmfHeader {
11907 format: "cmf".into(),
11908 version: CMF_VERSION,
11909 arch,
11910 quant_type: QuantType::Q4Block,
11911 provenance: None,
11912 tokenizer_config: None,
11913 section_hashes: None,
11914 skills: Vec::new(),
11915 shard: None,
11916 calibration: None,
11917 routing: None,
11918 };
11919 let specs = [
11920 TensorSpec {
11921 name: "q".into(),
11922 dtype: TensorDtype::Q4TiledP,
11923 shape: vec![r1, cols],
11924 data: synth_q4tp(r1, cols),
11925 },
11926 TensorSpec {
11927 name: "kv".into(),
11928 dtype: TensorDtype::Q4TiledP,
11929 shape: vec![r2, cols],
11930 data: synth_q4tp(r2, cols),
11931 },
11932 ];
11933 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
11934 std::fs::create_dir_all(&dir).unwrap();
11935 let path = dir.join("m.cmf");
11936 CmfModel::write(&path, &header, &specs, None, None).unwrap();
11937 let model = Arc::new(CmfModel::open(&path).unwrap());
11938 let (a, b) = (
11939 QTensor::from_model(&model, "q").unwrap(),
11940 QTensor::from_model(&model, "kv").unwrap(),
11941 );
11942 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
11943 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
11944 let x: Vec<f32> = (0..cols)
11945 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
11946 .collect();
11947 let pool = Pool::new(3);
11948 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
11949 a.matvec(&x, &mut ea, Some(&pool));
11950 b.matvec(&x, &mut eb, Some(&pool));
11951 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
11952 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11953 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
11954 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
11955 let _ = std::fs::remove_dir_all(&dir);
11956 }
11957
11958 #[test]
11961 fn batched_matmat_equals_per_position_matvec() {
11962 let (rows, cols, b) = (8, 64, 5);
11963 let groups = rows * cols / GROUP_SIZE;
11965 let mut q4 = Vec::new();
11966 for i in 0..groups * 16 {
11967 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11968 }
11969 for g in 0..groups {
11970 q4.extend_from_slice(
11971 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11972 );
11973 }
11974 let ng = cols / GROUP_SIZE;
11976 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
11977 let mut vb = bits.clone();
11978 for g in 0..rows * ng {
11979 vb.extend_from_slice(
11980 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
11981 );
11982 }
11983 for r in 0..rows {
11984 let bw = bits[r] as usize;
11985 let (mut acc, mut nb) = (0u64, 0usize);
11986 let mut rowbytes = Vec::new();
11987 for i in 0..cols {
11988 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11989 acc = (acc << bw) | v;
11990 nb += bw;
11991 while nb >= 8 {
11992 nb -= 8;
11993 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11994 }
11995 }
11996 if nb > 0 {
11997 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11998 }
11999 vb.extend_from_slice(&rowbytes);
12000 }
12001 let offsets = vbit_row_offsets(&vb, rows, cols);
12002
12003 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12004
12005 let mut got = vec![0f32; b * rows];
12007 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
12008 for bi in 0..b {
12009 let mut expect = vec![0f32; rows];
12010 q4matvec(
12011 &q4,
12012 &xs[bi * cols..(bi + 1) * cols],
12013 rows,
12014 cols,
12015 &mut expect,
12016 None,
12017 );
12018 assert_eq!(
12019 &got[bi * rows..(bi + 1) * rows],
12020 &expect[..],
12021 "q4 batch pos {bi}"
12022 );
12023 }
12024
12025 let mut got = vec![0f32; b * rows];
12027 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
12028 for bi in 0..b {
12029 let mut expect = vec![0f32; rows];
12030 vbitmatvec(
12031 &vb,
12032 &offsets,
12033 &xs[bi * cols..(bi + 1) * cols],
12034 rows,
12035 cols,
12036 &mut expect,
12037 None,
12038 );
12039 assert_eq!(
12040 &got[bi * rows..(bi + 1) * rows],
12041 &expect[..],
12042 "vbit batch pos {bi}"
12043 );
12044 }
12045 }
12046
12047 #[test]
12051 fn q4_tiled_matches_q4_block_bitexact() {
12052 let (rows, cols, b) = (8usize, 128usize, 3usize);
12053 let groups = rows * cols / GROUP_SIZE;
12054 let mut split = Vec::with_capacity(groups * 18);
12055 for i in 0..groups * 16 {
12056 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12057 }
12058 for g in 0..groups {
12059 split.extend_from_slice(
12060 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
12061 );
12062 }
12063 let (packed, scales) = split.split_at(groups * 16);
12065 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
12066 for g in 0..groups {
12067 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
12068 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
12069 }
12070
12071 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12072 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12074
12075 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
12076 q4matvec(&split, &x1, rows, cols, &mut a, None);
12077 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
12078 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
12079
12080 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12081 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
12082 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
12083 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
12084 assert_eq!(a1, t1);
12085 assert_eq!(a2, t2);
12086
12087 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12088 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12089 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
12090 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
12091 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
12092 }
12093
12094 #[test]
12101 fn q4matvec_sdot_outlier_exact() {
12102 let (rows, cols) = (4, 128);
12103 let groups = rows * cols / GROUP_SIZE;
12104 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12105 for i in 0..groups * 16 {
12106 bytes.push(((i * 11 + 5) % 256) as u8);
12107 }
12108 for g in 0..groups {
12109 let s = 0.02 + 0.002 * g as f32;
12110 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12111 }
12112 let mut x: Vec<f32> = (0..cols)
12113 .map(|i| match i % 3 {
12114 0 => 1.0,
12115 1 => -1.0,
12116 _ => 0.0,
12117 })
12118 .collect();
12119 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
12122 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12123 let mut expect = vec![0.0f32; rows];
12124 for r in 0..rows {
12125 expect[r] = reference[r * cols..(r + 1) * cols]
12126 .iter()
12127 .zip(&x)
12128 .map(|(w, xv)| w * xv)
12129 .sum();
12130 }
12131 let mut got = vec![0.0f32; rows];
12132 q4matvec(&bytes, &x, rows, cols, &mut got, None);
12133 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12134 for r in 0..rows {
12135 assert!(
12136 (got[r] - expect[r]).abs() < 2e-3 * scale,
12137 "row {r}: {} vs {} (outlier term must be exact)",
12138 got[r],
12139 expect[r]
12140 );
12141 }
12142 }
12143
12144 #[test]
12148 fn q1t_matvec_matches_reference() {
12149 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
12150 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
12152 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
12153 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
12155 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
12156 let mut bytes = Vec::new();
12157 for r in 0..rows {
12158 for g in 0..gpr {
12159 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
12160 let mut c = [0u8; 7];
12161 for k in 0..GROUP_SIZE {
12162 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
12164 0
12165 } else {
12166 ((k + r * 3 + g) % 3) as u8 };
12168 cortiq_core::quant::q1t_pack(&mut c, k, code);
12169 }
12170 bytes.extend_from_slice(&c);
12171 }
12172 }
12173 let mut row_ptr = vec![0u32; rows + 1];
12176 for &(idx, _) in &outliers {
12177 row_ptr[idx as usize / cols + 1] += 1;
12178 }
12179 for r in 0..rows {
12180 row_ptr[r + 1] += row_ptr[r];
12181 }
12182 for &p in &row_ptr {
12183 bytes.extend_from_slice(&p.to_le_bytes());
12184 }
12185 for &(idx, v) in &outliers {
12186 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
12187 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
12188 }
12189
12190 let mut refw = vec![0f32; rows * cols];
12191 dequant_q1t(&bytes, rows, cols, &mut refw);
12192 let x: Vec<f32> = (0..cols)
12195 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12196 .collect();
12197 let mut expect = vec![0f32; rows];
12198 for r in 0..rows {
12199 let mut a = 0.0f32;
12200 for j in 0..cols {
12201 a += refw[r * cols + j] * x[j];
12202 }
12203 expect[r] = a;
12204 }
12205 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
12206 let mut got = vec![0f32; rows];
12207 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
12208 for r in 0..rows {
12209 assert!(
12210 (got[r] - expect[r]).abs() < tol(expect[r]),
12211 "row {r}: {} vs {}",
12212 got[r],
12213 expect[r]
12214 );
12215 }
12216 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
12218 let mut gm = vec![0f32; 2 * rows];
12219 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
12220 for r in 0..rows {
12221 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
12222 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
12223 }
12224 let xb: Vec<f32> = (0..cols)
12228 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
12229 .collect();
12230 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12231 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
12232 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
12233 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12234 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
12235 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
12236 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
12237 }
12238
12239 #[test]
12242 fn q1t_matvec2_odd_gpr_matches_singles() {
12243 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12244 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
12246 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12247 for r in 0..rows {
12248 for g in 0..gpr {
12249 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
12250 let mut c = [0u8; 7];
12251 for k in 0..GROUP_SIZE {
12252 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
12253 }
12254 bytes.extend_from_slice(&c);
12255 }
12256 }
12257 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12258 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12259 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12260 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12261 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12262 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12263 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12264 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
12265 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
12266 }
12267
12268 #[test]
12272 #[ignore]
12273 fn q1t_matvec2_speed() {
12274 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12275 use std::time::Instant;
12276 let (rows, cols) = (8192usize, 4096usize);
12277 let gpr = cols / GROUP_SIZE;
12278 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12279 for r in 0..rows {
12280 for g in 0..gpr {
12281 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12282 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12283 let mut c = [0u8; 7];
12284 for k in 0..GROUP_SIZE {
12285 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12286 }
12287 bytes.extend_from_slice(&c);
12288 }
12289 }
12290 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12291 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12292 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12293 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12294 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12296 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12297 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
12298 for _ in 0..8 {
12299 let t0 = Instant::now();
12300 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12301 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
12302 let t1 = Instant::now();
12303 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12304 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12305 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
12306 }
12307 assert_eq!(p1, s1);
12308 assert_eq!(p2, s2);
12309 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
12310 }
12311
12312 #[test]
12316 #[ignore]
12317 fn q1t_matvec_speed() {
12318 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
12319 use std::time::Instant;
12320 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
12322 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
12323 for r in 0..rows {
12324 for g in 0..gpr {
12325 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12326 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12327 let mut c = [0u8; 7];
12328 for k in 0..GROUP_SIZE {
12329 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12330 }
12331 bytes.extend_from_slice(&c);
12332 }
12333 }
12334 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
12336 let mut idx = 0usize;
12337 while idx < n {
12338 row_ptr[idx / cols + 1] += 1;
12339 idx += stride;
12340 }
12341 for r in 0..rows {
12342 row_ptr[r + 1] += row_ptr[r];
12343 }
12344 for &p in &row_ptr {
12345 bytes.extend_from_slice(&p.to_le_bytes());
12346 }
12347 let mut idx = 0usize;
12348 while idx < n {
12349 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
12350 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
12351 idx += stride;
12352 }
12353 let x: Vec<f32> = (0..cols)
12356 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12357 .collect();
12358 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
12359
12360 let slow = |out: &mut [f32]| {
12362 let mut buf = vec![0f32; cols];
12363 for r in 0..rows {
12364 for g in 0..gpr {
12365 let off = (r * gpr + g) * Q1T_TILE;
12366 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
12367 let codes = &bytes[off + 2..off + Q1T_TILE];
12368 for k in 0..GROUP_SIZE {
12369 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
12370 1 => s,
12371 2 => -s,
12372 _ => 0.0,
12373 };
12374 }
12375 }
12376 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
12377 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
12378 }
12379 };
12380 let iters = 5;
12381 let mut a = vec![0f32; rows];
12382 slow(&mut a); let t = Instant::now();
12384 for _ in 0..iters {
12385 slow(&mut a);
12386 }
12387 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
12388
12389 let mut b = vec![0f32; rows];
12390 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
12392 for _ in 0..iters {
12393 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
12394 }
12395 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
12396
12397 for r in 0..rows {
12398 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
12399 }
12400 println!(
12401 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
12402 slow_ms / fast_ms
12403 );
12404 }
12405}
12406
12407#[cfg(test)]
12408mod gemm_bench {
12409 #[test]
12420 #[ignore]
12421 fn q4tp_matmat_throughput() {
12422 let b: usize = std::env::var("CMF_BENCH_B")
12426 .ok()
12427 .and_then(|v| v.parse().ok())
12428 .unwrap_or(296);
12429 let (rows, cols) = (9216usize, 2304usize);
12430 let (_, _, _) = (rows, cols, b);
12431 let total =
12432 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
12433 .unwrap();
12434 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
12439 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
12440 let lo = cortiq_core::quant::f32_to_f16(-4.0);
12441 let step = cortiq_core::quant::f32_to_f16(0.1);
12442 for r in 0..rows {
12443 let o = params_off + r * 4;
12444 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
12445 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
12446 }
12447 let _ = codes_off;
12448 let xs: Vec<f32> = (0..b * cols)
12449 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
12450 .collect();
12451 let mut out = vec![0f32; b * rows];
12452 let pool = crate::pool::Pool::from_env();
12453 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12459 let reps: usize = std::env::var("CMF_BENCH_REPS")
12460 .ok()
12461 .and_then(|v| v.parse().ok())
12462 .unwrap_or(10);
12463 let mut best = [f64::MAX; 2];
12464 let mut sums = [0f32; 2];
12465 for _ in 0..reps {
12466 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
12467 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
12468 let t = std::time::Instant::now();
12469 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12470 best[k] = best[k].min(t.elapsed().as_secs_f64());
12471 sums[k] = out.iter().take(64).sum::<f32>();
12472 }
12473 }
12474 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
12475 for (k, name) in ["previous", "tuned "].iter().enumerate() {
12476 println!(
12477 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
12478 best[k] * 1e3,
12479 flops / best[k] / 1e9,
12480 sums[k]
12481 );
12482 }
12483 assert!(
12484 (sums[0] - sums[1]).abs() < 1e-2,
12485 "the tuned kernel changed the result: {} vs {}",
12486 sums[0],
12487 sums[1]
12488 );
12489 }
12490
12491 #[test]
12497 fn q4tp_matmat_blocked_matches_scalar() {
12498 use std::sync::atomic::Ordering::Relaxed;
12499 for &(rows, cols, b) in &[
12507 (64usize, 128usize, 7usize),
12508 (33, 96, 4),
12509 (16, 256, 9),
12510 (192, 2304, 37),
12511 ] {
12512 let total = cortiq_core::quant::expected_nbytes(
12513 cortiq_core::TensorDtype::Q4TiledP,
12514 &[rows, cols],
12515 )
12516 .unwrap();
12517 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
12518 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
12519 let lo = cortiq_core::quant::f32_to_f16(-4.0);
12520 let step = cortiq_core::quant::f32_to_f16(0.1);
12521 for r in 0..rows {
12522 let o = params_off + r * 4;
12523 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
12524 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
12525 }
12526 let xs: Vec<f32> = (0..b * cols)
12527 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
12528 .collect();
12529 let mut got = vec![0f32; b * rows];
12530 let mut want = vec![0f32; b * rows];
12531 let gpr = cols / 32;
12532 let view = super::Q4tpView::new(&bytes, rows, cols);
12533 let pool = crate::pool::Pool::from_env();
12534 super::Q4TP_ALT.store(2, Relaxed);
12535 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
12536 super::Q4TP_ALT.store(1, Relaxed);
12537 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
12538 super::Q4TP_ALT.store(0, Relaxed);
12539 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
12546 let (mut worst, mut at) = (0f32, 0usize);
12547 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
12548 if (g - w).abs() > worst {
12549 worst = (g - w).abs();
12550 at = i;
12551 }
12552 }
12553 assert!(
12554 worst <= 1e-4 * scale,
12555 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
12556 (scale {scale:.3e}) at cell {at}: {} vs {}",
12557 got[at],
12558 want[at]
12559 );
12560
12561 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
12568 for bi in 0..b {
12569 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
12570 for r in 0..rows {
12571 let mut sc = vec![0f32; gpr];
12572 view.scales_into(r, gpr, &mut sc);
12573 let mut exact = 0f64;
12574 for j in 0..cols {
12575 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
12576 exact += w as f64 * sq as f64 * act.xq[j] as f64;
12577 }
12578 exact *= act.sx as f64;
12579 for &(j, xv) in &act.outliers {
12580 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
12581 exact += w as f64 * sq as f64 * xv as f64;
12582 }
12583 let i = bi * rows + r;
12584 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
12585 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
12586 }
12587 }
12588 println!(
12589 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
12590 per-column {e_scalar:.3e}"
12591 );
12592 assert!(
12597 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
12598 "{rows}x{cols} b={b}: error against f64 too large — blocked \
12599 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
12600 );
12601 }
12602 }
12603
12604 #[test]
12608 #[ignore]
12609 fn q4t_matmat_throughput() {
12610 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
12611 let total =
12612 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
12613 .unwrap();
12614 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
12617 let sc = cortiq_core::quant::f32_to_f16(0.02);
12618 for t in bytes.chunks_mut(super::Q4_TILE) {
12619 t[..2].copy_from_slice(&sc.to_le_bytes());
12620 }
12621 let xs: Vec<f32> = (0..b * cols)
12622 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
12623 .collect();
12624 let mut out = vec![0f32; b * rows];
12625 let pool = crate::pool::Pool::from_env();
12626 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12627 let reps: usize = std::env::var("CMF_BENCH_REPS")
12628 .ok()
12629 .and_then(|v| v.parse().ok())
12630 .unwrap_or(10);
12631 let mut best = f64::MAX;
12632 for _ in 0..reps {
12633 let t = std::time::Instant::now();
12634 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12635 best = best.min(t.elapsed().as_secs_f64());
12636 }
12637 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
12638 println!(
12639 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
12640 best * 1e3,
12641 flops / best / 1e9,
12642 out.iter().take(64).sum::<f32>()
12643 );
12644 }
12645}