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 if crate::gptq_capture::capturing() {
1772 if let Self::Mapped { model, idx, .. } = self {
1773 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1774 }
1775 }
1776 match self {
1777 Self::F32 { data, .. } => {
1778 let out_addr = SendMut(out.as_mut_ptr());
1779 let run = |start: usize, end: usize| {
1780 for o in start..end {
1781 let row = &data[o * cols..(o + 1) * cols];
1782 for bi in 0..b {
1783 let x = &xs_all[bi * cols..(bi + 1) * cols];
1784 let mut acc = 0f32;
1785 for j in 0..cols {
1786 acc += row[j] * x[j];
1787 }
1788 unsafe { *out_addr.at(bi * rows + o) = acc };
1789 }
1790 }
1791 };
1792 dispatch_rows(pool, rows, &run);
1793 }
1794 Self::Mapped {
1795 model,
1796 idx,
1797 dtype,
1798 row_scale,
1799 col_field,
1800 vbit_offsets,
1801 ..
1802 } => {
1803 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1804 let mut transformed = Vec::with_capacity(xs_all.len());
1805 for bi in 0..b {
1806 transformed.extend_from_slice(&crate::prism::forward(
1807 model,
1808 &xs_all[bi * cols..(bi + 1) * cols],
1809 ));
1810 }
1811 match dtype {
1812 TensorDtype::Q4Block => {
1813 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1814 }
1815 TensorDtype::Q4Tiled => {
1816 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1817 }
1818 TensorDtype::Q4TiledP => {
1819 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1820 }
1821 TensorDtype::Q2TiledP => {
1822 let affine =
1823 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1824 let gpu_batch_ok = if affine {
1830 b >= 2
1831 } else {
1832 b >= 32 && b * rows * cols >= 128_000_000
1833 };
1834 if gpu_batch_ok
1835 && cols % 32 == 0
1836 && crate::gpu::enabled_here()
1837 && crate::gpu::q2tp_gpu_opt_in()
1838 {
1839 let gpu_ok = if affine {
1840 crate::gpu::q2tp_affine_matmat(
1841 model,
1842 *idx,
1843 &transformed,
1844 b,
1845 rows,
1846 cols,
1847 out,
1848 )
1849 } else {
1850 crate::gpu::q2tp_matmat(
1851 model,
1852 *idx,
1853 &transformed,
1854 b,
1855 rows,
1856 cols,
1857 out,
1858 )
1859 };
1860 if gpu_ok {
1861 return;
1862 }
1863 }
1864 if affine
1868 && b == 1
1869 && cols % 32 == 0
1870 && crate::gpu::enabled_here()
1871 && crate::gpu::q2tp_gpu_opt_in()
1872 && crate::gpu::q2tp_affine_matvec(
1873 model,
1874 *idx,
1875 &transformed[..cols],
1876 rows,
1877 cols,
1878 &mut out[..rows],
1879 )
1880 {
1881 return;
1882 }
1883 if affine {
1884 q2tp_affine_matmat(
1885 self.quant_bytes(),
1886 &transformed,
1887 b,
1888 rows,
1889 cols,
1890 out,
1891 pool,
1892 )
1893 } else {
1894 q2tp_matmat(
1895 self.quant_bytes(),
1896 &transformed,
1897 b,
1898 rows,
1899 cols,
1900 out,
1901 pool,
1902 )
1903 }
1904 }
1905 TensorDtype::Q1 => {
1906 q1_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1907 }
1908 TensorDtype::Q1T => {
1909 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1910 }
1911 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
1912 self.quant_bytes(),
1913 vbit_offsets,
1914 &transformed,
1915 b,
1916 rows,
1917 cols,
1918 out,
1919 pool,
1920 ),
1921 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1922 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1923 .map(|bi| {
1924 prescale(
1925 &transformed[bi * cols..(bi + 1) * cols],
1926 col_field,
1927 *dtype,
1928 )
1929 })
1930 .collect();
1931 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
1932 }
1933 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1934 }
1935 return;
1936 }
1937 if *dtype == TensorDtype::Q4Block {
1938 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1939 return;
1940 }
1941 if *dtype == TensorDtype::Q4TiledP {
1942 if b >= 32
1954 && b * rows * cols >= 128_000_000
1955 && cols % 32 == 0
1956 && !crate::gpu::mm_killed()
1957 && crate::gpu::enabled_here()
1958 {
1959 let class = if b >= 128 {
1960 crate::gpu::OpClass::MatmatWide
1961 } else {
1962 crate::gpu::OpClass::Matmat
1963 };
1964 if let Self::Mapped { model, idx, .. } = self {
1965 if crate::mm_ab::on() {
1977 let mut g = vec![0f32; b * rows];
1978 let t = std::time::Instant::now();
1979 let took = crate::gpu::q4tp_matmat(
1980 model, *idx, xs_all, b, rows, cols, &mut g,
1981 );
1982 let dg = t.elapsed();
1983 let t = std::time::Instant::now();
1984 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1985 let dc = t.elapsed();
1986 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1987 return;
1988 }
1989 let t0 = std::time::Instant::now();
1990 let resident = crate::gpu::weight_is_resident(model, *idx);
1994 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1995 crate::gpu::ProbeArm::Gpu => {
1996 if crate::gpu::q4tp_matmat(
1997 model, *idx, xs_all, b, rows, cols, out,
1998 ) {
1999 let el = t0.elapsed();
2000 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2010 let budget = std::time::Duration::from_secs_f64(
2011 flops / 1.5e12 * 8.0 + 0.020,
2012 );
2013 crate::gpu::mm_budget_check(
2014 "q4tp matmat",
2015 el,
2016 budget,
2017 crate::gpu::probe_was_cold() || !resident,
2018 );
2019 crate::gpu::probe_record(class, true, el);
2020 return;
2021 }
2022 }
2023 crate::gpu::ProbeArm::CpuTimed => {
2024 q4tp_matmat(
2025 self.quant_bytes(),
2026 xs_all,
2027 b,
2028 rows,
2029 cols,
2030 out,
2031 pool,
2032 );
2033 crate::gpu::probe_record(class, false, t0.elapsed());
2034 return;
2035 }
2036 crate::gpu::ProbeArm::Cpu => {}
2037 }
2038 }
2039 }
2040 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2041 return;
2042 }
2043 if *dtype == TensorDtype::Q2TiledP {
2044 if b >= 32
2050 && b * rows * cols >= 128_000_000
2051 && cols % 32 == 0
2052 && !crate::gpu::mm_killed()
2053 && crate::gpu::enabled_here()
2054 {
2055 let class = if b >= 128 {
2056 crate::gpu::OpClass::MatmatWide
2057 } else {
2058 crate::gpu::OpClass::Matmat
2059 };
2060 if let Self::Mapped { model, idx, .. } = self {
2061 let t0 = std::time::Instant::now();
2062 match crate::gpu::probe_arm(class) {
2063 crate::gpu::ProbeArm::Gpu => {
2064 if crate::gpu::q2tp_matmat(
2065 model, *idx, xs_all, b, rows, cols, out,
2066 ) {
2067 crate::gpu::probe_record(class, true, t0.elapsed());
2068 return;
2069 }
2070 }
2071 crate::gpu::ProbeArm::CpuTimed => {
2072 q2tp_matmat(
2073 self.quant_bytes(),
2074 xs_all,
2075 b,
2076 rows,
2077 cols,
2078 out,
2079 pool,
2080 );
2081 crate::gpu::probe_record(class, false, t0.elapsed());
2082 return;
2083 }
2084 crate::gpu::ProbeArm::Cpu => {}
2085 }
2086 }
2087 }
2088 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2093 return;
2094 }
2095 if *dtype == TensorDtype::Q4Tiled {
2096 if b >= 32
2108 && b * rows * cols >= 128_000_000
2109 && cols % 32 == 0
2110 && !crate::gpu::mm_killed()
2111 && crate::gpu::enabled_here()
2112 {
2113 let class = if b >= 128 {
2114 crate::gpu::OpClass::MatmatWide
2115 } else {
2116 crate::gpu::OpClass::Matmat
2117 };
2118 if let Self::Mapped { model, idx, .. } = self {
2119 let t0 = std::time::Instant::now();
2120 match crate::gpu::probe_arm(class) {
2121 crate::gpu::ProbeArm::Gpu => {
2122 if crate::gpu::q4t_matmat(
2123 model, *idx, xs_all, b, rows, cols, out,
2124 ) {
2125 let el = t0.elapsed();
2126 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2136 let budget = std::time::Duration::from_secs_f64(
2137 flops / 1.5e12 * 8.0 + 0.020,
2138 );
2139 crate::gpu::mm_budget_check(
2140 "q4t matmat",
2141 el,
2142 budget,
2143 crate::gpu::probe_was_cold(),
2144 );
2145 crate::gpu::probe_record(class, true, el);
2146 return;
2147 }
2148 }
2149 crate::gpu::ProbeArm::CpuTimed => {
2150 q4t_matmat(
2151 self.quant_bytes(),
2152 xs_all,
2153 b,
2154 rows,
2155 cols,
2156 out,
2157 pool,
2158 );
2159 crate::gpu::probe_record(class, false, t0.elapsed());
2160 return;
2161 }
2162 crate::gpu::ProbeArm::Cpu => {}
2163 }
2164 }
2165 }
2166 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2167 return;
2168 }
2169 if *dtype == TensorDtype::Q1 {
2170 if b >= 32
2173 && b * rows * cols >= 128_000_000
2174 && cols % 64 == 0
2175 && crate::gpu::enabled_here()
2176 {
2177 if let Self::Mapped { model, idx, .. } = self {
2178 let t0 = std::time::Instant::now();
2179 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2180 crate::gpu::ProbeArm::Gpu => {
2181 if crate::gpu::q1_matmat(
2182 model, *idx, xs_all, b, rows, cols, out,
2183 ) {
2184 crate::gpu::probe_record(
2185 crate::gpu::OpClass::Matmat,
2186 true,
2187 t0.elapsed(),
2188 );
2189 return;
2190 }
2191 }
2192 crate::gpu::ProbeArm::CpuTimed => {
2193 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2194 crate::gpu::probe_record(
2195 crate::gpu::OpClass::Matmat,
2196 false,
2197 t0.elapsed(),
2198 );
2199 return;
2200 }
2201 crate::gpu::ProbeArm::Cpu => {}
2202 }
2203 }
2204 }
2205 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2206 return;
2207 }
2208 if *dtype == TensorDtype::Q1T {
2209 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2212 if let Self::Mapped { model, idx, .. } = self {
2213 let t0 = std::time::Instant::now();
2214 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2215 crate::gpu::ProbeArm::Gpu => {
2216 if crate::gpu::q1t_matmat(
2217 model, *idx, xs_all, b, rows, cols, out,
2218 ) {
2219 crate::gpu::probe_record(
2220 crate::gpu::OpClass::Matmat,
2221 true,
2222 t0.elapsed(),
2223 );
2224 return;
2225 }
2226 }
2227 crate::gpu::ProbeArm::CpuTimed => {
2228 q1t_matmat(
2229 self.quant_bytes(),
2230 xs_all,
2231 b,
2232 rows,
2233 cols,
2234 out,
2235 pool,
2236 );
2237 crate::gpu::probe_record(
2238 crate::gpu::OpClass::Matmat,
2239 false,
2240 t0.elapsed(),
2241 );
2242 return;
2243 }
2244 crate::gpu::ProbeArm::Cpu => {}
2245 }
2246 }
2247 }
2248 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2249 return;
2250 }
2251 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2252 vbitmatmat(
2253 self.quant_bytes(),
2254 vbit_offsets,
2255 xs_all,
2256 b,
2257 rows,
2258 cols,
2259 out,
2260 pool,
2261 );
2262 return;
2263 }
2264 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2265 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2266 .collect();
2267 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2273 if let Self::Mapped { model, idx, .. } = self {
2274 let t0 = std::time::Instant::now();
2275 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2276 crate::gpu::ProbeArm::Gpu
2277 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2278 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2279 {
2280 let q = self.quant_bytes();
2284 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2285 return;
2286 }
2287 crate::gpu::ProbeArm::Gpu => {
2288 let flat: Vec<f32> =
2289 pre.iter().flat_map(|v| v.iter().copied()).collect();
2290 if crate::gpu::q8_matmat(
2291 model, *idx, row_scale, &flat, b, rows, cols, out,
2292 ) {
2293 crate::gpu::probe_record(
2294 crate::gpu::OpClass::Matmat,
2295 true,
2296 t0.elapsed(),
2297 );
2298 return;
2299 }
2300 }
2301 crate::gpu::ProbeArm::CpuTimed => {
2302 let q = self.quant_bytes();
2303 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2304 crate::gpu::probe_record(
2305 crate::gpu::OpClass::Matmat,
2306 false,
2307 t0.elapsed(),
2308 );
2309 return;
2310 }
2311 crate::gpu::ProbeArm::Cpu => {}
2312 }
2313 }
2314 }
2315 let q = self.quant_bytes();
2316 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2317 }
2318 }
2319 }
2320}
2321
2322impl QTensor {
2323 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2333 let (rows, cols) = (self.rows(), self.cols());
2334 let Self::Mapped {
2335 model,
2336 idx,
2337 dtype,
2338 row_scale,
2339 col_field,
2340 ..
2341 } = self
2342 else {
2343 return false;
2344 };
2345 if crate::prism::has_contract(model) {
2346 return false;
2347 }
2348 match *dtype {
2349 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2350 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2354 if *dtype == TensorDtype::Q8_2f
2357 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2358 && crate::gpu::q8_matmat_2f(
2359 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2360 )
2361 {
2362 return true;
2363 }
2364 let flat: Vec<f32> = (0..b)
2365 .flat_map(|bi| {
2366 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2367 })
2368 .collect();
2369 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2370 }
2371 _ => false,
2372 }
2373 }
2374
2375 pub fn matvec_many<const N: usize>(
2382 ts: [&QTensor; N],
2383 x: &[f32],
2384 mut outs: [&mut [f32]; N],
2385 pool: Option<&Pool>,
2386 ) {
2387 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2388 if ts.iter().any(|t| t.has_prism_contract()) {
2389 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2393 t.matvec(x, o, pool);
2394 }
2395 return;
2396 }
2397 let uniform_q8 = ts.iter().all(|t| {
2398 matches!(
2399 t,
2400 Self::Mapped {
2401 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2402 ..
2403 }
2404 )
2405 });
2406 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2407 let uniform_q4 = ts.iter().all(|t| {
2408 matches!(
2409 t,
2410 Self::Mapped {
2411 dtype: TensorDtype::Q4Block,
2412 ..
2413 }
2414 )
2415 });
2416 let uniform_vbit = ts.iter().all(|t| {
2417 matches!(
2418 t,
2419 Self::Mapped {
2420 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2421 ..
2422 }
2423 )
2424 });
2425 let uniform_q1 = ts.iter().all(|t| {
2426 matches!(
2427 t,
2428 Self::Mapped {
2429 dtype: TensorDtype::Q1,
2430 ..
2431 }
2432 )
2433 });
2434 let uniform_q1t = ts.iter().all(|t| {
2435 matches!(
2436 t,
2437 Self::Mapped {
2438 dtype: TensorDtype::Q1T,
2439 ..
2440 }
2441 )
2442 });
2443 let uniform_q4tp = ts.iter().all(|t| {
2448 matches!(
2449 t,
2450 Self::Mapped {
2451 dtype: TensorDtype::Q4TiledP,
2452 ..
2453 }
2454 )
2455 }) && ts
2456 .iter()
2457 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2458 let Some(pool) = pool else {
2459 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2460 t.matvec(x, o, None);
2461 }
2462 return;
2463 };
2464 if total_rows < 256
2465 || !(uniform_q8
2466 || uniform_f32
2467 || uniform_q4
2468 || uniform_vbit
2469 || uniform_q1
2470 || uniform_q1t
2471 || uniform_q4tp)
2472 {
2473 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2474 t.matvec(x, o, Some(pool));
2475 }
2476 return;
2477 }
2478
2479 if uniform_q4tp {
2480 let cols = ts[0].cols();
2486 let gpr = cols / GROUP_SIZE;
2487 let views: Vec<Q4tpView> = ts
2488 .iter()
2489 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2490 .collect();
2491 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2492 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2493 let locate = |flat: usize| -> (usize, usize) {
2495 let mut acc = 0;
2496 for (i, &r) in rows_of.iter().enumerate() {
2497 if flat < acc + r {
2498 return (i, flat - acc);
2499 }
2500 acc += r;
2501 }
2502 (rows_of.len() - 1, 0)
2503 };
2504 let (views, outs_addr) = (&views, &outs_addr);
2505 if a8w8_enabled() {
2506 let act = split_act(x);
2507 let act = &act;
2508 let run = |start: usize, end: usize| {
2509 let mut sc = vec![0f32; gpr];
2510 for flat in start..end {
2511 let (t, r) = locate(flat);
2512 let v = &views[t];
2513 v.scales_into(r, gpr, &mut sc);
2514 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2515 for &(j, xv) in &act.outliers {
2516 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2517 acc += w * s * xv;
2518 }
2519 unsafe { *outs_addr[t].at(r) = acc };
2521 }
2522 };
2523 pool.run_rows(total_rows, &run);
2524 } else {
2525 let run = |start: usize, end: usize| {
2526 let mut sc = vec![0f32; gpr];
2527 for flat in start..end {
2528 let (t, r) = locate(flat);
2529 let v = &views[t];
2530 v.scales_into(r, gpr, &mut sc);
2531 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2533 }
2534 };
2535 pool.run_rows(total_rows, &run);
2536 }
2537 return;
2538 }
2539
2540 if uniform_q1 {
2541 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2544 if a8w8_enabled() {
2545 let act = split_act(x);
2546 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2547 let (act, gsum) = (&act, &gsum);
2548 let closures: [_; N] = std::array::from_fn(|i| {
2549 let (bytes, gpr, out) =
2550 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2551 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2552 });
2553 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2554 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2555 pool.run_many(&parts);
2556 } else {
2557 let closures: [_; N] = std::array::from_fn(|i| {
2558 let (bytes, gpr, out) =
2559 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2560 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2561 });
2562 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2563 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2564 pool.run_many(&parts);
2565 }
2566 return;
2567 }
2568
2569 if uniform_q1t {
2570 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2574 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2575 if a8w8_enabled() {
2576 let act = split_act(x);
2577 let act = &act;
2578 let x_ref = x;
2579 let closures: [_; N] = std::array::from_fn(|i| {
2580 let bytes = ts[i].quant_bytes();
2581 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2582 let gpr = cols / GROUP_SIZE;
2583 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2584 let out = outs_addr[i];
2585 move |s: usize, e: usize| {
2586 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2587 }
2588 });
2589 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2590 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2591 pool.run_many(&parts);
2592 } else {
2593 let x_ref = x;
2594 let closures: [_; N] = std::array::from_fn(|i| {
2595 let bytes = ts[i].quant_bytes();
2596 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2597 let gpr = cols / GROUP_SIZE;
2598 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2599 let out = outs_addr[i];
2600 move |s: usize, e: usize| {
2601 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2602 }
2603 });
2604 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2605 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2606 pool.run_many(&parts);
2607 }
2608 return;
2609 }
2610
2611 if uniform_q4 || uniform_vbit {
2612 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2613 if a8w8_enabled() {
2615 let act = split_act(x);
2616 let act = &act;
2617 if uniform_q4 {
2618 let closures: [_; N] = std::array::from_fn(|i| {
2619 let (packed, scales) =
2620 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2621 let (gpr, cols, out) =
2622 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2623 move |s: usize, e: usize| {
2624 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2625 }
2626 });
2627 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2628 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2629 pool.run_many(&parts);
2630 } else {
2631 let closures: [_; N] = std::array::from_fn(|i| {
2632 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2633 unreachable!()
2634 };
2635 let (bytes, rows, cols, out) = (
2636 ts[i].quant_bytes(),
2637 ts[i].rows(),
2638 ts[i].cols(),
2639 outs_addr[i],
2640 );
2641 move |s: usize, e: usize| {
2642 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2643 }
2644 });
2645 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2646 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2647 pool.run_many(&parts);
2648 }
2649 return;
2650 }
2651 if uniform_q4 {
2652 let closures: [_; N] = std::array::from_fn(|i| {
2653 let (packed, scales) =
2654 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2655 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2656 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2657 });
2658 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2659 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2660 pool.run_many(&parts);
2661 } else {
2662 let closures: [_; N] = std::array::from_fn(|i| {
2663 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2664 unreachable!()
2665 };
2666 let (bytes, rows, cols, out) = (
2667 ts[i].quant_bytes(),
2668 ts[i].rows(),
2669 ts[i].cols(),
2670 outs_addr[i],
2671 );
2672 move |s: usize, e: usize| {
2673 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2674 }
2675 });
2676 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2677 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2678 pool.run_many(&parts);
2679 }
2680 return;
2681 }
2682
2683 if uniform_f32 {
2684 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2685 let closures: [_; N] = std::array::from_fn(|i| {
2686 let Self::F32 { data, cols, .. } = ts[i] else {
2687 unreachable!()
2688 };
2689 let out = outs_addr[i];
2690 move |start: usize, end: usize| {
2691 for o in start..end {
2692 let row = &data[o * cols..(o + 1) * cols];
2693 let mut sum = 0.0f32;
2694 for j in 0..*cols {
2695 sum += row[j] * x[j];
2696 }
2697 unsafe { *out.at(o) = sum };
2699 }
2700 }
2701 });
2702 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2703 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2704 pool.run_many(&parts);
2705 return;
2706 }
2707
2708 struct Ctx<'a> {
2711 bytes: &'a [u8],
2712 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2713 rep: &'a [u8],
2714 row_scale: &'a [f32],
2715 cols: usize,
2716 xs: std::borrow::Cow<'a, [f32]>,
2717 }
2718 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2719 let Self::Mapped {
2720 dtype,
2721 cols,
2722 row_scale,
2723 col_field,
2724 repack,
2725 ..
2726 } = ts[i]
2727 else {
2728 unreachable!()
2729 };
2730 Ctx {
2731 bytes: ts[i].quant_bytes(),
2732 rep: repack,
2733 row_scale,
2734 cols: *cols,
2735 xs: prescale(x, col_field, *dtype),
2736 }
2737 });
2738 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2739 #[cfg(target_arch = "aarch64")]
2740 if sdot_enabled() {
2741 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2742 let closures: [_; N] = std::array::from_fn(|i| {
2743 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2744 move |start: usize, end: usize| {
2745 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2746 }
2747 });
2748 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2749 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2750 pool.run_many(&parts);
2751 return;
2752 }
2753 #[cfg(target_arch = "x86_64")]
2754 if avx2_a8w8_enabled() {
2755 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2756 let closures: [_; N] = std::array::from_fn(|i| {
2757 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2758 move |start: usize, end: usize| {
2759 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2760 }
2761 });
2762 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2763 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2764 pool.run_many(&parts);
2765 return;
2766 }
2767 let closures: [_; N] = std::array::from_fn(|i| {
2768 let (c, out) = (&ctxs[i], outs_addr[i]);
2769 move |start: usize, end: usize| {
2770 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2771 }
2772 });
2773 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2774 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2775 pool.run_many(&parts);
2776 }
2777}
2778
2779impl QTensor {
2780 #[allow(clippy::needless_range_loop)]
2785 pub fn matvec2_many<const N: usize>(
2786 ts: [&QTensor; N],
2787 x1: &[f32],
2788 x2: &[f32],
2789 mut o1s: [&mut [f32]; N],
2790 mut o2s: [&mut [f32]; N],
2791 pool: Option<&Pool>,
2792 ) {
2793 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2794 if ts.iter().any(|t| t.has_prism_contract()) {
2795 for i in 0..N {
2796 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2797 }
2798 return;
2799 }
2800 let uniform_q8 = ts.iter().all(|t| {
2801 matches!(
2802 t,
2803 Self::Mapped {
2804 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2805 ..
2806 }
2807 )
2808 });
2809 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2810 let uniform_q4 = ts.iter().all(|t| {
2811 matches!(
2812 t,
2813 Self::Mapped {
2814 dtype: TensorDtype::Q4Block,
2815 ..
2816 }
2817 )
2818 });
2819 let uniform_vbit = ts.iter().all(|t| {
2820 matches!(
2821 t,
2822 Self::Mapped {
2823 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2824 ..
2825 }
2826 )
2827 });
2828 let fusable = pool.is_some()
2829 && total_rows >= 256
2830 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2831 if !fusable {
2832 for i in 0..N {
2833 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2834 }
2835 return;
2836 }
2837 let pool = pool.unwrap();
2838
2839 if uniform_q4 || uniform_vbit {
2840 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2841 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2842 if a8w8_enabled() {
2844 let a1 = split_act(x1);
2845 let a2 = split_act(x2);
2846 let (a1, a2) = (&a1, &a2);
2847 if uniform_q4 {
2848 let closures: [_; N] = std::array::from_fn(|i| {
2849 let (packed, scales) =
2850 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2851 let (gpr, cols, o1, o2) =
2852 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2853 move |s: usize, e: usize| {
2854 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2855 }
2856 });
2857 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2858 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2859 pool.run_many(&parts);
2860 } else {
2861 let closures: [_; N] = std::array::from_fn(|i| {
2862 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2863 unreachable!()
2864 };
2865 let (bytes, rows, cols, o1, o2) = (
2866 ts[i].quant_bytes(),
2867 ts[i].rows(),
2868 ts[i].cols(),
2869 p1[i],
2870 p2[i],
2871 );
2872 move |s: usize, e: usize| {
2873 vbit_range2_a8w8(
2874 bytes,
2875 vbit_offsets,
2876 x1,
2877 x2,
2878 a1,
2879 a2,
2880 rows,
2881 cols,
2882 o1,
2883 o2,
2884 s,
2885 e,
2886 )
2887 }
2888 });
2889 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2890 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2891 pool.run_many(&parts);
2892 }
2893 return;
2894 }
2895 if uniform_q4 {
2896 let closures: [_; N] = std::array::from_fn(|i| {
2897 let (packed, scales) =
2898 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2899 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2900 move |s: usize, e: usize| {
2901 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2902 }
2903 });
2904 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2905 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2906 pool.run_many(&parts);
2907 } else {
2908 let closures: [_; N] = std::array::from_fn(|i| {
2909 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2910 unreachable!()
2911 };
2912 let (bytes, rows, cols, o1, o2) = (
2913 ts[i].quant_bytes(),
2914 ts[i].rows(),
2915 ts[i].cols(),
2916 p1[i],
2917 p2[i],
2918 );
2919 move |s: usize, e: usize| {
2920 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2921 }
2922 });
2923 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2924 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2925 pool.run_many(&parts);
2926 }
2927 return;
2928 }
2929
2930 if uniform_f32 {
2931 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2932 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2933 let closures: [_; N] = std::array::from_fn(|i| {
2934 let Self::F32 { data, cols, .. } = ts[i] else {
2935 unreachable!()
2936 };
2937 let (o1, o2) = (p1[i], p2[i]);
2938 move |start: usize, end: usize| {
2939 for o in start..end {
2940 let row = &data[o * cols..(o + 1) * cols];
2941 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2942 for j in 0..*cols {
2943 s1 += row[j] * x1[j];
2944 s2 += row[j] * x2[j];
2945 }
2946 unsafe {
2948 *o1.at(o) = s1;
2949 *o2.at(o) = s2;
2950 }
2951 }
2952 }
2953 });
2954 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2955 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2956 pool.run_many(&parts);
2957 return;
2958 }
2959
2960 struct Ctx<'a> {
2961 bytes: &'a [u8],
2962 row_scale: &'a [f32],
2963 cols: usize,
2964 xs1: std::borrow::Cow<'a, [f32]>,
2965 xs2: std::borrow::Cow<'a, [f32]>,
2966 }
2967 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2968 let Self::Mapped {
2969 dtype,
2970 cols,
2971 row_scale,
2972 col_field,
2973 ..
2974 } = ts[i]
2975 else {
2976 unreachable!()
2977 };
2978 Ctx {
2979 bytes: ts[i].quant_bytes(),
2980 row_scale,
2981 cols: *cols,
2982 xs1: prescale(x1, col_field, *dtype),
2983 xs2: prescale(x2, col_field, *dtype),
2984 }
2985 });
2986 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2987 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2988 #[cfg(target_arch = "aarch64")]
2989 if sdot_enabled() {
2990 let acts: [(SplitAct, SplitAct); N] =
2991 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2992 let closures: [_; N] = std::array::from_fn(|i| {
2993 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2994 move |start: usize, end: usize| {
2995 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2996 }
2997 });
2998 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2999 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3000 pool.run_many(&parts);
3001 return;
3002 }
3003 #[cfg(target_arch = "x86_64")]
3004 if avx2_a8w8_enabled() {
3005 let acts: [(SplitAct, SplitAct); N] =
3006 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3007 let closures: [_; N] = std::array::from_fn(|i| {
3008 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3009 move |start: usize, end: usize| {
3010 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3011 }
3012 });
3013 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3014 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3015 pool.run_many(&parts);
3016 return;
3017 }
3018 let closures: [_; N] = std::array::from_fn(|i| {
3019 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3020 move |start: usize, end: usize| {
3021 q8_range2_f32(
3022 c.bytes,
3023 c.row_scale,
3024 &c.xs1,
3025 &c.xs2,
3026 c.cols,
3027 o1,
3028 o2,
3029 start,
3030 end,
3031 )
3032 }
3033 });
3034 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3035 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3036 pool.run_many(&parts);
3037 }
3038
3039 pub fn matvec_silu_mul(
3044 gate: &QTensor,
3045 up: &QTensor,
3046 x: &[f32],
3047 out: &mut [f32],
3048 pool: Option<&Pool>,
3049 ) -> bool {
3050 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3051 }
3052
3053 pub fn matvec_silu_mul_limited(
3059 gate: &QTensor,
3060 up: &QTensor,
3061 x: &[f32],
3062 out: &mut [f32],
3063 limit: f32,
3064 pool: Option<&Pool>,
3065 ) -> bool {
3066 if gate.has_prism_contract() || up.has_prism_contract() {
3067 return false;
3071 }
3072 let inter = gate.rows();
3073 debug_assert_eq!(up.rows(), inter);
3074 debug_assert_eq!(out.len(), inter);
3075 debug_assert_eq!(gate.cols(), up.cols());
3076 if !a8w8_enabled() {
3077 return false;
3078 }
3079 let act = split_act(x);
3080 let act = &act;
3081 let x_ref = x;
3082 let out_addr = SendMut(out.as_mut_ptr());
3083
3084 match (gate, up) {
3085 (
3087 Self::Mapped {
3088 dtype: TensorDtype::Q4Block,
3089 ..
3090 },
3091 Self::Mapped {
3092 dtype: TensorDtype::Q4Block,
3093 ..
3094 },
3095 ) => {
3096 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3097 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3098 let gpr = gate.cols() / GROUP_SIZE;
3099 let cols = gate.cols();
3100 let run = move |start: usize, end: usize| {
3101 for r in start..end {
3102 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3103 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3104 for &(j, xv) in &act.outliers {
3105 let flat = r * cols + j;
3106 let gb = gp[flat / 2];
3107 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3108 let gsc = f16_to_f32(u16::from_le_bytes([
3109 gs[(flat / GROUP_SIZE) * 2],
3110 gs[(flat / GROUP_SIZE) * 2 + 1],
3111 ]));
3112 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3113 let ub = up_p[flat / 2];
3114 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3115 let usc = f16_to_f32(u16::from_le_bytes([
3116 up_s[(flat / GROUP_SIZE) * 2],
3117 up_s[(flat / GROUP_SIZE) * 2 + 1],
3118 ]));
3119 uv += ((un as i32 - 8) as f32) * usc * xv;
3120 }
3121 let silu_g = gv / (1.0 + (-gv).exp());
3122 unsafe { *out_addr.at(r) = silu_g * uv };
3124 }
3125 };
3126 dispatch_rows(pool, inter, &run);
3127 true
3128 }
3129 (
3133 Self::Mapped {
3134 dtype: TensorDtype::Q4Tiled,
3135 ..
3136 },
3137 Self::Mapped {
3138 dtype: TensorDtype::Q4Tiled,
3139 ..
3140 },
3141 ) => {
3142 let g_bytes = gate.quant_bytes();
3143 let u_bytes = up.quant_bytes();
3144 let gpr = gate.cols() / GROUP_SIZE;
3145 let run = move |start: usize, end: usize| {
3146 for r in start..end {
3147 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3148 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3149 for &(j, xv) in &act.outliers {
3150 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3151 gv += w * s * xv;
3152 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3153 uv += w * s * xv;
3154 }
3155 let silu_g = gv / (1.0 + (-gv).exp());
3156 unsafe { *out_addr.at(r) = silu_g * uv };
3158 }
3159 };
3160 dispatch_rows(pool, inter, &run);
3161 true
3162 }
3163 (
3166 Self::Mapped {
3167 dtype: TensorDtype::Q4TiledP,
3168 ..
3169 },
3170 Self::Mapped {
3171 dtype: TensorDtype::Q4TiledP,
3172 ..
3173 },
3174 ) => {
3175 let cols = gate.cols();
3176 let gpr = cols / GROUP_SIZE;
3177 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3178 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3179 let run = |start: usize, end: usize| {
3180 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3181 for r in start..end {
3182 gv_view.scales_into(r, gpr, &mut gsc);
3183 uv_view.scales_into(r, gpr, &mut usc);
3184 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3185 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3186 for &(j, xv) in &act.outliers {
3187 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3188 gv += w * s * xv;
3189 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3190 uv += w * s * xv;
3191 }
3192 let silu_g = gv / (1.0 + (-gv).exp());
3193 unsafe { *out_addr.at(r) = silu_g * uv };
3195 }
3196 };
3197 dispatch_rows(pool, inter, &run);
3198 true
3199 }
3200 (
3206 Self::Mapped {
3207 dtype: TensorDtype::Q1,
3208 ..
3209 },
3210 Self::Mapped {
3211 dtype: TensorDtype::Q1,
3212 ..
3213 },
3214 ) => {
3215 let g_bytes = gate.quant_bytes();
3216 let u_bytes = up.quant_bytes();
3217 let gpr = gate.cols() / GROUP_SIZE;
3218 let gsum = q1_group_sums(&act.xq, gpr);
3219 let gsum = &gsum;
3220 let run = move |start: usize, end: usize| {
3221 for r in start..end {
3222 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
3223 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
3224 for &(j, xv) in &act.outliers {
3225 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
3226 gv += w * s * xv;
3227 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
3228 uv += w * s * xv;
3229 }
3230 let silu_g = gv / (1.0 + (-gv).exp());
3231 unsafe { *out_addr.at(r) = silu_g * uv };
3233 }
3234 };
3235 dispatch_rows(pool, inter, &run);
3236 true
3237 }
3238 (
3242 Self::Mapped {
3243 dtype: TensorDtype::Q2TiledP,
3244 ..
3245 },
3246 Self::Mapped {
3247 dtype: TensorDtype::Q2TiledP,
3248 ..
3249 },
3250 ) => {
3251 let cols = gate.cols();
3252 let gpr = cols / GROUP_SIZE;
3253 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3254 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3255 let gsum = q1_group_sums(&act.xq, gpr);
3256 let gsum = &gsum;
3257 let run = move |start: usize, end: usize| {
3258 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3259 for r in start..end {
3260 gv_view.scales_into(r, gpr, &mut gsc);
3261 uv_view.scales_into(r, gpr, &mut usc);
3262 let mut gv =
3263 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
3264 let mut uv =
3265 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
3266 for &(j, xv) in &act.outliers {
3267 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3268 gv += w * s * xv;
3269 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
3270 uv += w * s * xv;
3271 }
3272 let silu_g = gv / (1.0 + (-gv).exp());
3273 unsafe { *out_addr.at(r) = silu_g * uv };
3275 }
3276 };
3277 dispatch_rows(pool, inter, &run);
3278 true
3279 }
3280 (
3285 Self::Mapped {
3286 dtype: TensorDtype::Q8Row,
3287 row_scale: g_rs,
3288 ..
3289 },
3290 Self::Mapped {
3291 dtype: TensorDtype::Q8Row,
3292 row_scale: u_rs,
3293 ..
3294 },
3295 ) => {
3296 let g_bytes = gate.quant_bytes();
3297 let u_bytes = up.quant_bytes();
3298 let cols = gate.cols();
3299 let run = move |start: usize, end: usize| {
3300 for r in start..end {
3301 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3302 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3303 let silu_g = gv / (1.0 + (-gv).exp());
3304 unsafe { *out_addr.at(r) = silu_g * uv };
3306 }
3307 };
3308 dispatch_rows(pool, inter, &run);
3309 true
3310 }
3311 (
3313 Self::Mapped {
3314 dtype: TensorDtype::Q1T,
3315 ..
3316 },
3317 Self::Mapped {
3318 dtype: TensorDtype::Q1T,
3319 ..
3320 },
3321 ) => {
3322 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3323 let g_bytes = gate.quant_bytes();
3324 let u_bytes = up.quant_bytes();
3325 let gpr = gate.cols() / GROUP_SIZE;
3326 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3327 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3328 let run = move |start: usize, end: usize| {
3329 for r in start..end {
3330 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3331 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3332 for &(j, xv) in &act.outliers {
3333 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3334 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3335 }
3336 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3337 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3338 let silu_g = gv / (1.0 + (-gv).exp());
3339 unsafe { *out_addr.at(r) = silu_g * uv };
3341 }
3342 };
3343 dispatch_rows(pool, inter, &run);
3344 true
3345 }
3346 _ => false,
3347 }
3348 }
3349
3350 pub fn moe_gate_up_many(
3364 pairs: &[(&QTensor, &QTensor)],
3365 x: &[f32],
3366 outs: &mut [Vec<f32>],
3367 pool: Option<&Pool>,
3368 ) -> bool {
3369 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
3370 return false;
3371 }
3372 let inter = pairs[0].0.rows();
3373 let cols = pairs[0].0.cols();
3374 if cols % GROUP_SIZE != 0 {
3375 return false;
3376 }
3377 let gpr = cols / GROUP_SIZE;
3378 let q2 = matches!(
3381 pairs[0].0,
3382 Self::Mapped {
3383 dtype: TensorDtype::Q2TiledP,
3384 ..
3385 }
3386 );
3387 let want = if q2 {
3388 TensorDtype::Q2TiledP
3389 } else {
3390 TensorDtype::Q4TiledP
3391 };
3392 let mut views = Vec::with_capacity(pairs.len() * 2);
3393 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3394 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3395 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3396 if !both
3397 || g.rows() != inter
3398 || u.rows() != inter
3399 || g.cols() != cols
3400 || u.cols() != cols
3401 || o.len() != inter
3402 {
3403 return false;
3404 }
3405 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3406 views.push(mk(g.quant_bytes(), inter, cols));
3407 views.push(mk(u.quant_bytes(), inter, cols));
3408 }
3409 let act = split_act(x);
3410 let gsum = if q2 {
3411 q1_group_sums(&act.xq, gpr)
3412 } else {
3413 Vec::new()
3414 };
3415 let (act, gsum) = (&act, &gsum);
3416 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3417 let (views, ptrs) = (&views, &ptrs);
3418 let run = |start: usize, end: usize| {
3419 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3420 for flat in start..end {
3421 let (e, r) = (flat / inter, flat % inter);
3422 let gv_view = &views[e * 2];
3423 let uv_view = &views[e * 2 + 1];
3424 gv_view.scales_into(r, gpr, &mut gsc);
3425 uv_view.scales_into(r, gpr, &mut usc);
3426 let (mut gv, mut uv) = if q2 {
3427 (
3428 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
3429 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
3430 )
3431 } else {
3432 (
3433 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
3434 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
3435 )
3436 };
3437 for &(j, xv) in &act.outliers {
3438 let (og, ou) = if q2 {
3439 (
3440 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3441 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
3442 )
3443 } else {
3444 (
3445 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3446 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
3447 )
3448 };
3449 gv += og.0 * og.1 * xv;
3450 uv += ou.0 * ou.1 * xv;
3451 }
3452 let silu_g = gv / (1.0 + (-gv).exp());
3453 unsafe { *ptrs[e].at(r) = silu_g * uv };
3455 }
3456 };
3457 dispatch_rows(pool, pairs.len() * inter, &run);
3458 true
3459 }
3460
3461 pub fn moe_down_many(
3470 downs: &[&QTensor],
3471 gs: &[Vec<f32>],
3472 weights: &[f32],
3473 out: &mut [f32],
3474 pool: Option<&Pool>,
3475 ) -> bool {
3476 if downs.is_empty()
3477 || downs.len() != gs.len()
3478 || downs.len() != weights.len()
3479 || !a8w8_enabled()
3480 {
3481 return false;
3482 }
3483 let rows = out.len();
3484 let cols = downs[0].cols();
3485 if cols % GROUP_SIZE != 0 {
3486 return false;
3487 }
3488 let gpr = cols / GROUP_SIZE;
3489 let mut views = Vec::with_capacity(downs.len());
3490 for (d, g) in downs.iter().zip(gs.iter()) {
3491 if !matches!(
3492 d,
3493 Self::Mapped {
3494 dtype: TensorDtype::Q4TiledP,
3495 ..
3496 }
3497 ) || d.rows() != rows
3498 || d.cols() != cols
3499 || g.len() != cols
3500 {
3501 return false;
3502 }
3503 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3504 }
3505 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3507 let out_addr = SendMut(out.as_mut_ptr());
3515 let (views, acts, weights) = (&views, &acts, &weights);
3516 let run = |start: usize, end: usize| {
3517 let mut sc = vec![0f32; gpr];
3518 for r in start..end {
3519 let mut acc = 0f32;
3520 for (e, v) in views.iter().enumerate() {
3521 v.scales_into(r, gpr, &mut sc);
3522 let a = &acts[e];
3523 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3524 for &(j, xv) in &a.outliers {
3525 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3526 d += w * s * xv;
3527 }
3528 acc += weights[e] * d;
3529 }
3530 unsafe { *out_addr.at(r) = acc };
3532 }
3533 };
3534 dispatch_rows(pool, rows, &run);
3535 true
3536 }
3537}
3538
3539#[cfg(target_os = "macos")]
3544mod accel_blas {
3545 #[link(name = "Accelerate", kind = "framework")]
3546 unsafe extern "C" {
3547 pub fn cblas_sgemm(
3548 order: i32,
3549 trans_a: i32,
3550 trans_b: i32,
3551 m: i32,
3552 n: i32,
3553 k: i32,
3554 alpha: f32,
3555 a: *const f32,
3556 lda: i32,
3557 b: *const f32,
3558 ldb: i32,
3559 beta: f32,
3560 c: *mut f32,
3561 ldc: i32,
3562 );
3563 }
3564}
3565
3566#[cfg(target_os = "macos")]
3567pub(crate) fn accel_gemm_enabled() -> bool {
3568 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3569 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3570}
3571
3572#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3575pub(crate) fn accel_gemm_enabled() -> bool {
3576 true
3577}
3578
3579#[cfg(target_arch = "aarch64")]
3586#[allow(clippy::too_many_arguments)]
3587pub(crate) fn neon_gemm_rm(
3588 m: usize,
3589 n: usize,
3590 k: usize,
3591 alpha: f32,
3592 a: &[f32],
3593 lda: usize,
3594 b_mat: &[f32],
3595 ldb: usize,
3596 b_rows_are_n: bool,
3597 c: &mut [f32],
3598 ldc: usize,
3599) {
3600 debug_assert!(a.len() >= (m - 1) * lda + k);
3601 debug_assert!(c.len() >= (m - 1) * ldc + n);
3602 unsafe {
3604 use core::arch::aarch64::*;
3605 let mut i = 0usize;
3606 while i < m {
3607 let mi = (m - i).min(4);
3608 let mut j = 0usize;
3609 while j < n {
3610 let nj = (n - j).min(8);
3611 if mi == 4 && nj == 8 {
3612 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3613 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3614 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3615 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3616 for p in 0..k {
3617 let (b0, b1) = if b_rows_are_n {
3618 let base = b_mat.as_ptr().add(j * ldb + p);
3621 let g = |o: usize| *base.add(o * ldb);
3622 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3623 } else {
3624 let base = b_mat.as_ptr().add(p * ldb + j);
3625 (
3626 [*base, *base.add(1), *base.add(2), *base.add(3)],
3627 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3628 )
3629 };
3630 let bv0 = vld1q_f32(b0.as_ptr());
3631 let bv1 = vld1q_f32(b1.as_ptr());
3632 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3633 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3634 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3635 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3636 c0a = vfmaq_f32(c0a, a0, bv0);
3637 c0b = vfmaq_f32(c0b, a0, bv1);
3638 c1a = vfmaq_f32(c1a, a1, bv0);
3639 c1b = vfmaq_f32(c1b, a1, bv1);
3640 c2a = vfmaq_f32(c2a, a2, bv0);
3641 c2b = vfmaq_f32(c2b, a2, bv1);
3642 c3a = vfmaq_f32(c3a, a3, bv0);
3643 c3b = vfmaq_f32(c3b, a3, bv1);
3644 }
3645 let al = vdupq_n_f32(alpha);
3646 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3647 .iter()
3648 .enumerate()
3649 {
3650 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3651 vst1q_f32(dst, vmulq_f32(*ca, al));
3652 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3653 }
3654 } else {
3655 for r in 0..mi {
3656 for q in 0..nj {
3657 let mut acc = 0f32;
3658 for p in 0..k {
3659 let bv = if b_rows_are_n {
3660 b_mat[(j + q) * ldb + p]
3661 } else {
3662 b_mat[p * ldb + j + q]
3663 };
3664 acc += a[(i + r) * lda + p] * bv;
3665 }
3666 c[(i + r) * ldc + j + q] = acc * alpha;
3667 }
3668 }
3669 }
3670 j += nj;
3671 }
3672 i += mi;
3673 }
3674 }
3675}
3676
3677#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3679#[allow(clippy::too_many_arguments)]
3680pub(crate) fn sgemm_rm(
3681 m: usize,
3682 n: usize,
3683 k: usize,
3684 alpha: f32,
3685 a: &[f32],
3686 lda: usize,
3687 b_mat: &[f32],
3688 ldb: usize,
3689 b_rows_are_n: bool,
3690 c: &mut [f32],
3691 ldc: usize,
3692) {
3693 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3694}
3695
3696#[allow(clippy::too_many_arguments)]
3700pub fn sgemm_public(
3701 m: usize,
3702 n: usize,
3703 k: usize,
3704 alpha: f32,
3705 a: &[f32],
3706 lda: usize,
3707 b_mat: &[f32],
3708 ldb: usize,
3709 b_rows_are_n: bool,
3710 c: &mut [f32],
3711 ldc: usize,
3712) {
3713 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3714 {
3715 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3716 }
3717 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3722 {
3723 for i in 0..m {
3724 for j in 0..n {
3725 let mut acc = 0f32;
3726 for p in 0..k {
3727 let bv = if b_rows_are_n {
3728 b_mat[j * ldb + p]
3729 } else {
3730 b_mat[p * ldb + j]
3731 };
3732 acc += a[i * lda + p] * bv;
3733 }
3734 c[i * ldc + j] = alpha * acc;
3735 }
3736 }
3737 }
3738}
3739
3740#[cfg(target_os = "macos")]
3743#[allow(clippy::too_many_arguments)]
3744pub(crate) fn sgemm_rm(
3745 m: usize,
3746 n: usize,
3747 k: usize,
3748 alpha: f32,
3749 a: &[f32],
3750 lda: usize,
3751 b_mat: &[f32],
3752 ldb: usize,
3753 b_rows_are_n: bool,
3754 c: &mut [f32],
3755 ldc: usize,
3756) {
3757 debug_assert!(a.len() >= (m - 1) * lda + k);
3758 debug_assert!(c.len() >= (m - 1) * ldc + n);
3759 #[cfg(target_arch = "aarch64")]
3764 if std::env::var("CMF_FORCE_NEON_GEMM")
3765 .map(|v| v == "1")
3766 .unwrap_or(false)
3767 {
3768 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3769 }
3770 unsafe {
3771 accel_blas::cblas_sgemm(
3772 101, 111, if b_rows_are_n { 112 } else { 111 },
3775 m as i32,
3776 n as i32,
3777 k as i32,
3778 alpha,
3779 a.as_ptr(),
3780 lda as i32,
3781 b_mat.as_ptr(),
3782 ldb as i32,
3783 0.0,
3784 c.as_mut_ptr(),
3785 ldc as i32,
3786 );
3787 }
3788}
3789
3790#[cfg(target_os = "macos")]
3797fn qmatmat_accel(
3798 q: &[u8],
3799 row_scale: &[f32],
3800 pre: &[std::borrow::Cow<'_, [f32]>],
3801 rows: usize,
3802 cols: usize,
3803 out: &mut [f32],
3804 pool: Option<&Pool>,
3805) {
3806 const TR: usize = 2048;
3811 let b = pre.len();
3812 thread_local! {
3813 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3814 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3815 }
3816 XPANEL.with(|xp| {
3817 WTILE.with(|wt| {
3818 let mut xpanel = xp.borrow_mut();
3819 xpanel.clear();
3820 for x in pre {
3821 xpanel.extend_from_slice(x);
3822 }
3823 let mut wtile = wt.borrow_mut();
3824 wtile.resize(TR * cols, 0.0);
3825 let mut r0 = 0usize;
3826 while r0 < rows {
3827 let tr = TR.min(rows - r0);
3828 let wt_addr = SendMut(wtile.as_mut_ptr());
3830 let run = |start: usize, end: usize| {
3831 for r in start..end {
3832 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3833 let s = row_scale[r0 + r];
3834 let dst =
3836 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3837 for (d, &v) in dst.iter_mut().zip(row) {
3838 *d = (v as i8) as f32 * s;
3839 }
3840 }
3841 };
3842 dispatch_rows(pool, tr, &run);
3843 unsafe {
3845 accel_blas::cblas_sgemm(
3846 101, 111, 112, b as i32,
3850 tr as i32,
3851 cols as i32,
3852 1.0,
3853 xpanel.as_ptr(),
3854 cols as i32,
3855 wtile.as_ptr(),
3856 cols as i32,
3857 0.0,
3858 out.as_mut_ptr().add(r0),
3859 rows as i32,
3860 );
3861 }
3862 r0 += tr;
3863 }
3864 })
3865 });
3866}
3867
3868fn qmatmat(
3869 q: &[u8],
3870 row_scale: &[f32],
3871 pre: &[std::borrow::Cow<'_, [f32]>],
3872 rows: usize,
3873 cols: usize,
3874 out: &mut [f32],
3875 pool: Option<&Pool>,
3876) {
3877 let b = pre.len();
3878 debug_assert_eq!(out.len(), b * rows);
3879 #[cfg(target_os = "macos")]
3884 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3885 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3886 return;
3887 }
3888 #[cfg(target_arch = "aarch64")]
3889 if sdot_enabled() {
3890 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3891 let out_addr = SendMut(out.as_mut_ptr());
3892 let blocked_ok = blocked_enabled();
3895 let use_i8mm = i8mm_enabled();
3896 if blocked_ok {
3897 let run = |start: usize, end: usize| {
3898 let mut o = start;
3899 while o < end {
3900 if o + 2 <= end {
3901 let r0 = &q[o * cols..(o + 1) * cols];
3902 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3903 let mut bi = 0usize;
3904 while bi + 4 <= acts.len() {
3905 let xs = [
3906 acts[bi].xq.as_slice(),
3907 acts[bi + 1].xq.as_slice(),
3908 acts[bi + 2].xq.as_slice(),
3909 acts[bi + 3].xq.as_slice(),
3910 ];
3911 let d = if use_i8mm {
3912 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3913 } else {
3914 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3915 };
3916 for (r, row) in [r0, r1].into_iter().enumerate() {
3917 for k in 0..4 {
3918 let act = &acts[bi + k];
3919 let mut v = d[r][k] as f32 * act.sx;
3920 for &(j, xv) in &act.outliers {
3921 v += (row[j] as i8) as f32 * xv;
3922 }
3923 unsafe {
3924 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3925 };
3926 }
3927 }
3928 bi += 4;
3929 }
3930 while bi < acts.len() {
3931 for (r, row) in [r0, r1].into_iter().enumerate() {
3932 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3933 unsafe { *out_addr.at(bi * rows + o + r) = v };
3934 }
3935 bi += 1;
3936 }
3937 o += 2;
3938 } else {
3939 let row = &q[o * cols..(o + 1) * cols];
3940 for (bi, act) in acts.iter().enumerate() {
3941 let v = row_dot_sdot(row, act) * row_scale[o];
3942 unsafe { *out_addr.at(bi * rows + o) = v };
3943 }
3944 o += 1;
3945 }
3946 }
3947 };
3948 dispatch_rows(pool, rows, &run);
3949 return;
3950 }
3951 let run = |start: usize, end: usize| {
3952 for o in start..end {
3953 let row = &q[o * cols..(o + 1) * cols];
3954 for (bi, act) in acts.iter().enumerate() {
3955 let v = row_dot_sdot(row, act) * row_scale[o];
3956 unsafe { *out_addr.at(bi * rows + o) = v };
3957 }
3958 }
3959 };
3960 dispatch_rows(pool, rows, &run);
3961 return;
3962 }
3963 #[cfg(target_arch = "x86_64")]
3968 if avx2_a8w8_enabled() {
3969 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3970 let out_addr = SendMut(out.as_mut_ptr());
3971 let blocked_ok = blocked_enabled();
3974 if !avx512vnni_enabled() && blocked_ok {
3975 let run = |start: usize, end: usize| {
3976 let mut o = start;
3977 while o < end {
3978 if o + 2 <= end {
3979 let r0 = &q[o * cols..(o + 1) * cols];
3980 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3981 let mut bi = 0usize;
3982 while bi + 4 <= acts.len() {
3983 let xs = [
3984 acts[bi].xq.as_slice(),
3985 acts[bi + 1].xq.as_slice(),
3986 acts[bi + 2].xq.as_slice(),
3987 acts[bi + 3].xq.as_slice(),
3988 ];
3989 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3990 for (r, row) in [r0, r1].into_iter().enumerate() {
3991 for k in 0..4 {
3992 let act = &acts[bi + k];
3993 let mut v = d[r][k] as f32 * act.sx;
3994 for &(j, xv) in &act.outliers {
3995 v += (row[j] as i8) as f32 * xv;
3996 }
3997 unsafe {
3998 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3999 };
4000 }
4001 }
4002 bi += 4;
4003 }
4004 while bi < acts.len() {
4005 for (r, row) in [r0, r1].into_iter().enumerate() {
4006 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4007 unsafe { *out_addr.at(bi * rows + o + r) = v };
4008 }
4009 bi += 1;
4010 }
4011 o += 2;
4012 } else {
4013 let row = &q[o * cols..(o + 1) * cols];
4014 for (bi, act) in acts.iter().enumerate() {
4015 let v = row_dot_avx2(row, act) * row_scale[o];
4016 unsafe { *out_addr.at(bi * rows + o) = v };
4017 }
4018 o += 1;
4019 }
4020 }
4021 };
4022 dispatch_rows(pool, rows, &run);
4023 return;
4024 }
4025 let run = |start: usize, end: usize| {
4026 for o in start..end {
4027 let row = &q[o * cols..(o + 1) * cols];
4028 for (bi, act) in acts.iter().enumerate() {
4029 let v = row_dot_avx2(row, act) * row_scale[o];
4030 unsafe { *out_addr.at(bi * rows + o) = v };
4031 }
4032 }
4033 };
4034 dispatch_rows(pool, rows, &run);
4035 return;
4036 }
4037 let out_addr = SendMut(out.as_mut_ptr());
4038 let run = |start: usize, end: usize| {
4039 for o in start..end {
4040 let row = &q[o * cols..(o + 1) * cols];
4041 for (bi, x) in pre.iter().enumerate() {
4042 let mut acc = 0f32;
4043 for j in 0..cols {
4044 acc += (row[j] as i8) as f32 * x[j];
4045 }
4046 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4047 }
4048 }
4049 };
4050 dispatch_rows(pool, rows, &run);
4051}
4052
4053fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4056 match pool {
4057 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4058 _ => run(0, rows),
4059 }
4060}
4061
4062fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4064 let groups = rows * cols / GROUP_SIZE;
4065 bytes.split_at(groups * 16)
4066}
4067
4068#[inline]
4073fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4074 #[cfg(target_arch = "aarch64")]
4075 unsafe {
4076 return vbit_fill4_neon(data, buf);
4077 }
4078 #[cfg(target_arch = "x86_64")]
4079 if avx2_enabled() {
4080 return unsafe { vbit_fill4_avx2(data, buf) };
4081 }
4082 #[allow(unreachable_code)]
4083 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4084 let u = unpack8::<4>(&data[blk * 4..]);
4085 for k in 0..8 {
4086 chunk[k] = (u[k] - 7) as i8 as u8;
4087 }
4088 }
4089}
4090
4091#[cfg(target_arch = "aarch64")]
4092#[target_feature(enable = "neon")]
4093unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4094 unsafe {
4097 use core::arch::aarch64::*;
4098 let n = buf.len();
4099 let mask = vdupq_n_u8(0x0F);
4100 let seven = vdupq_n_s8(7);
4101 let mut g = 0usize;
4102 while g * 32 + 32 <= n {
4103 let b = vld1q_u8(data.as_ptr().add(g * 16));
4104 let hi = vshrq_n_u8::<4>(b);
4105 let lo = vandq_u8(b, mask);
4106 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4107 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4108 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4109 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4110 g += 1;
4111 }
4112 }
4113}
4114
4115#[cfg(target_arch = "x86_64")]
4116#[target_feature(enable = "avx2")]
4117unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4118 unsafe {
4120 use core::arch::x86_64::*;
4121 let n = buf.len();
4122 let mask = _mm_set1_epi8(0x0F);
4123 let seven = _mm256_set1_epi8(7);
4124 let mut g = 0usize;
4125 while g * 32 + 32 <= n {
4126 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4127 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4128 let lo = _mm_and_si128(b, mask);
4129 let z = _mm256_sub_epi8(
4130 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4131 seven,
4132 );
4133 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4134 g += 1;
4135 }
4136 }
4137}
4138
4139#[inline(always)]
4144fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4145 let mut acc = 0u64;
4146 for i in 0..B {
4147 acc = (acc << 8) | data[i] as u64;
4148 }
4149 let mask = (1u64 << B) - 1;
4150 let mut out = [0i32; 8];
4151 for (k, o) in out.iter_mut().enumerate() {
4152 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4153 }
4154 out
4155}
4156
4157#[allow(clippy::too_many_arguments)]
4163fn vbitmatvec(
4164 bytes: &[u8],
4165 offsets: &[usize],
4166 x: &[f32],
4167 rows: usize,
4168 cols: usize,
4169 out: &mut [f32],
4170 pool: Option<&Pool>,
4171) {
4172 debug_assert_eq!(out.len(), rows);
4173 debug_assert_eq!(offsets.len(), rows + 1);
4174
4175 if a8w8_enabled() {
4179 let act = split_act(x);
4180 let out_addr = SendMut(out.as_mut_ptr());
4181 let run = move |start: usize, end: usize| {
4182 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4183 };
4184 dispatch_rows(pool, rows, &run);
4185 return;
4186 }
4187
4188 let out_addr = SendMut(out.as_mut_ptr());
4189 let run = move |start: usize, end: usize| {
4190 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4191 };
4192 dispatch_rows(pool, rows, &run);
4193}
4194
4195#[allow(clippy::too_many_arguments)]
4199fn vbit_range_a8w8(
4200 bytes: &[u8],
4201 offsets: &[usize],
4202 x: &[f32],
4203 act: &SplitAct,
4204 rows: usize,
4205 cols: usize,
4206 out: SendMut,
4207 start: usize,
4208 end: usize,
4209) {
4210 let ng = cols / GROUP_SIZE;
4211 let bits = &bytes[..rows];
4212 let sc_off = rows;
4213 let row_dot = |r: usize| -> f32 {
4214 let b = bits[r] as usize;
4215 let l = (1i32 << (b - 1)) - 1;
4216 let mask = (1u64 << b) - 1;
4217 let data = &bytes[offsets[r]..offsets[r + 1]];
4218 if b == 8 {
4219 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4221 let mut dot = 0f32;
4222 for g in 0..ng {
4223 let so = (r * ng + g) * 2;
4224 let sgf = f16_to_f32(u16::from_le_bytes([
4225 bytes[sc_off + so],
4226 bytes[sc_off + so + 1],
4227 ]));
4228 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4229 let mut gd = 0f32;
4230 for &xv in xg.iter() {
4231 if nbits < 8 {
4232 acc = (acc << 8) | data[idx] as u64;
4233 idx += 1;
4234 nbits += 8;
4235 }
4236 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4237 nbits -= 8;
4238 gd += (u - l) as f32 * xv;
4239 }
4240 dot += gd * sgf;
4241 }
4242 return dot;
4243 }
4244 thread_local! {
4248 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4249 const { std::cell::RefCell::new(Vec::new()) };
4250 }
4251 #[inline(always)]
4252 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4253 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4254 let u = unpack8::<B>(&data[blk * B..]);
4255 for k in 0..8 {
4256 chunk[k] = (u[k] - l) as i8 as u8;
4257 }
4258 }
4259 }
4260 let _ = mask;
4261 VBIT_SCRATCH.with(|scratch| {
4262 let mut buf = scratch.borrow_mut();
4263 buf.resize(cols, 0);
4264 match b {
4265 3 => fill::<3>(data, l, &mut buf),
4266 4 => vbit_fill4(data, &mut buf),
4267 5 => fill::<5>(data, l, &mut buf),
4268 6 => fill::<6>(data, l, &mut buf),
4269 _ => unreachable!(),
4270 }
4271 let mut dot = 0f32;
4272 for g in 0..ng {
4273 let so = (r * ng + g) * 2;
4274 let s = f16_to_f32(u16::from_le_bytes([
4275 bytes[sc_off + so],
4276 bytes[sc_off + so + 1],
4277 ]));
4278 let d = dot_i8_i8(
4279 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4280 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4281 ) as f32
4282 * act.sx;
4283 dot += d * s;
4284 }
4285 for &(j, xv) in &act.outliers {
4286 let so = (r * ng + j / GROUP_SIZE) * 2;
4287 let s = f16_to_f32(u16::from_le_bytes([
4288 bytes[sc_off + so],
4289 bytes[sc_off + so + 1],
4290 ]));
4291 dot += (buf[j] as i8) as f32 * s * xv;
4293 }
4294 dot
4295 })
4296 };
4297 for r in start..end {
4298 unsafe { *out.at(r) = row_dot(r) };
4300 }
4301}
4302
4303#[allow(clippy::too_many_arguments)]
4305fn vbit_range_f32(
4306 bytes: &[u8],
4307 offsets: &[usize],
4308 x: &[f32],
4309 rows: usize,
4310 cols: usize,
4311 out: SendMut,
4312 start: usize,
4313 end: usize,
4314) {
4315 let ng = cols / GROUP_SIZE;
4316 let bits = &bytes[..rows];
4317 let sc_off = rows;
4318 #[inline(always)]
4322 fn dot_row<const B: usize>(
4323 data: &[u8],
4324 bytes: &[u8],
4325 sc_off: usize,
4326 r: usize,
4327 ng: usize,
4328 x: &[f32],
4329 ) -> f32 {
4330 let l = ((1i32 << (B - 1)) - 1) as f32;
4331 let gbytes = GROUP_SIZE * B / 8;
4332 let mut dot = 0f32;
4333 for g in 0..ng {
4334 let so = (r * ng + g) * 2;
4335 let s = f16_to_f32(u16::from_le_bytes([
4336 bytes[sc_off + so],
4337 bytes[sc_off + so + 1],
4338 ]));
4339 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4340 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4341 let mut gd = 0f32;
4342 for blk in 0..GROUP_SIZE / 8 {
4343 let u = unpack8::<B>(&gd0[blk * B..]);
4344 let xb = &xg[blk * 8..blk * 8 + 8];
4345 for k in 0..8 {
4346 gd += (u[k] as f32 - l) * xb[k];
4347 }
4348 }
4349 dot += gd * s;
4350 }
4351 dot
4352 }
4353 for r in start..end {
4354 let data = &bytes[offsets[r]..offsets[r + 1]];
4355 let v = match bits[r] {
4356 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4357 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4358 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4359 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4360 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4361 b => unreachable!("vbit bit-width {b} (validated at load)"),
4362 };
4363 unsafe { *out.at(r) = v };
4365 }
4366}
4367
4368#[allow(clippy::too_many_arguments)]
4373fn vbitmatvec2(
4374 bytes: &[u8],
4375 offsets: &[usize],
4376 x1: &[f32],
4377 x2: &[f32],
4378 rows: usize,
4379 cols: usize,
4380 o1: &mut [f32],
4381 o2: &mut [f32],
4382 pool: Option<&Pool>,
4383) {
4384 debug_assert_eq!(o1.len(), rows);
4385 debug_assert_eq!(o2.len(), rows);
4386
4387 if a8w8_enabled() {
4388 let a1 = split_act(x1);
4389 let a2 = split_act(x2);
4390 let p1 = SendMut(o1.as_mut_ptr());
4391 let p2 = SendMut(o2.as_mut_ptr());
4392 let run = move |start: usize, end: usize| {
4393 vbit_range2_a8w8(
4394 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4395 )
4396 };
4397 dispatch_rows(pool, rows, &run);
4398 return;
4399 }
4400
4401 let p1 = SendMut(o1.as_mut_ptr());
4402 let p2 = SendMut(o2.as_mut_ptr());
4403 let run = move |start: usize, end: usize| {
4404 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4405 };
4406 dispatch_rows(pool, rows, &run);
4407}
4408
4409#[allow(clippy::too_many_arguments)]
4413fn vbit_range2_a8w8(
4414 bytes: &[u8],
4415 offsets: &[usize],
4416 x1: &[f32],
4417 x2: &[f32],
4418 a1: &SplitAct,
4419 a2: &SplitAct,
4420 rows: usize,
4421 cols: usize,
4422 p1: SendMut,
4423 p2: SendMut,
4424 start: usize,
4425 end: usize,
4426) {
4427 let ng = cols / GROUP_SIZE;
4428 let bits = &bytes[..rows];
4429 let sc_off = rows;
4430 let row_dots = |r: usize| -> (f32, f32) {
4431 let b = bits[r] as usize;
4432 let l = (1i32 << (b - 1)) - 1;
4433 let data = &bytes[offsets[r]..offsets[r + 1]];
4434 if b == 8 {
4435 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4438 let (mut d1, mut d2) = (0f32, 0f32);
4439 for g in 0..ng {
4440 let so = (r * ng + g) * 2;
4441 let sgf = f16_to_f32(u16::from_le_bytes([
4442 bytes[sc_off + so],
4443 bytes[sc_off + so + 1],
4444 ]));
4445 let (mut g1, mut g2) = (0f32, 0f32);
4446 for k in 0..GROUP_SIZE {
4447 if nbits < 8 {
4448 acc = (acc << 8) | data[idx] as u64;
4449 idx += 1;
4450 nbits += 8;
4451 }
4452 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4453 nbits -= 8;
4454 let w = (u - l) as f32;
4455 g1 += w * x1[g * GROUP_SIZE + k];
4456 g2 += w * x2[g * GROUP_SIZE + k];
4457 }
4458 d1 += g1 * sgf;
4459 d2 += g2 * sgf;
4460 }
4461 return (d1, d2);
4462 }
4463 thread_local! {
4464 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4465 const { std::cell::RefCell::new(Vec::new()) };
4466 }
4467 #[inline(always)]
4468 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4469 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4470 let u = unpack8::<B>(&data[blk * B..]);
4471 for k in 0..8 {
4472 chunk[k] = (u[k] - l) as i8 as u8;
4473 }
4474 }
4475 }
4476 VBIT_SCRATCH2.with(|scratch| {
4477 let mut buf = scratch.borrow_mut();
4478 buf.resize(cols, 0);
4479 match b {
4480 3 => fill::<3>(data, l, &mut buf),
4481 4 => vbit_fill4(data, &mut buf),
4482 5 => fill::<5>(data, l, &mut buf),
4483 6 => fill::<6>(data, l, &mut buf),
4484 _ => unreachable!(),
4485 }
4486 let (mut d1, mut d2) = (0f32, 0f32);
4487 for g in 0..ng {
4488 let so = (r * ng + g) * 2;
4489 let s = f16_to_f32(u16::from_le_bytes([
4490 bytes[sc_off + so],
4491 bytes[sc_off + so + 1],
4492 ]));
4493 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4494 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
4495 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
4496 d1 += v1 * s;
4497 d2 += v2 * s;
4498 }
4499 for &(j, xv) in &a1.outliers {
4500 let so = (r * ng + j / GROUP_SIZE) * 2;
4501 let s = f16_to_f32(u16::from_le_bytes([
4502 bytes[sc_off + so],
4503 bytes[sc_off + so + 1],
4504 ]));
4505 d1 += (buf[j] as i8) as f32 * s * xv;
4506 }
4507 for &(j, xv) in &a2.outliers {
4508 let so = (r * ng + j / GROUP_SIZE) * 2;
4509 let s = f16_to_f32(u16::from_le_bytes([
4510 bytes[sc_off + so],
4511 bytes[sc_off + so + 1],
4512 ]));
4513 d2 += (buf[j] as i8) as f32 * s * xv;
4514 }
4515 (d1, d2)
4516 })
4517 };
4518 for r in start..end {
4519 let (v1, v2) = row_dots(r);
4520 unsafe {
4522 *p1.at(r) = v1;
4523 *p2.at(r) = v2;
4524 }
4525 }
4526}
4527
4528#[allow(clippy::too_many_arguments)]
4532fn vbit_range2_f32(
4533 bytes: &[u8],
4534 offsets: &[usize],
4535 x1: &[f32],
4536 x2: &[f32],
4537 rows: usize,
4538 cols: usize,
4539 p1: SendMut,
4540 p2: SendMut,
4541 start: usize,
4542 end: usize,
4543) {
4544 let ng = cols / GROUP_SIZE;
4545 let bits = &bytes[..rows];
4546 let sc_off = rows;
4547 #[inline(always)]
4548 #[allow(clippy::too_many_arguments)]
4549 fn dot_row2<const B: usize>(
4550 data: &[u8],
4551 bytes: &[u8],
4552 sc_off: usize,
4553 r: usize,
4554 ng: usize,
4555 x1: &[f32],
4556 x2: &[f32],
4557 ) -> (f32, f32) {
4558 let l = ((1i32 << (B - 1)) - 1) as f32;
4559 let gbytes = GROUP_SIZE * B / 8;
4560 let (mut d1, mut d2) = (0f32, 0f32);
4561 for g in 0..ng {
4562 let so = (r * ng + g) * 2;
4563 let s = f16_to_f32(u16::from_le_bytes([
4564 bytes[sc_off + so],
4565 bytes[sc_off + so + 1],
4566 ]));
4567 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4568 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4569 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4570 let (mut g1, mut g2) = (0f32, 0f32);
4571 for blk in 0..GROUP_SIZE / 8 {
4572 let u = unpack8::<B>(&gd0[blk * B..]);
4573 for k in 0..8 {
4574 let w = u[k] as f32 - l;
4575 g1 += w * x1g[blk * 8 + k];
4576 g2 += w * x2g[blk * 8 + k];
4577 }
4578 }
4579 d1 += g1 * s;
4580 d2 += g2 * s;
4581 }
4582 (d1, d2)
4583 }
4584 for r in start..end {
4585 let data = &bytes[offsets[r]..offsets[r + 1]];
4586 let (v1, v2) = match bits[r] {
4587 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4588 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4589 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4590 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4591 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4592 b => unreachable!("vbit bit-width {b} (validated at load)"),
4593 };
4594 unsafe {
4596 *p1.at(r) = v1;
4597 *p2.at(r) = v2;
4598 }
4599 }
4600}
4601
4602#[inline]
4609#[allow(unreachable_code)]
4610fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4611 #[cfg(target_arch = "aarch64")]
4612 unsafe {
4613 return dot_q4t_row_sdot(bytes, r, gpr, xq);
4614 }
4615 #[cfg(target_arch = "x86_64")]
4616 unsafe {
4617 if vnni_tiles_enabled() {
4618 return dot_q4t_row_vnni(bytes, r, gpr, xq);
4619 }
4620 return dot_q4t_row_avx2(bytes, r, gpr, xq);
4621 }
4622 let mut acc = 0f32;
4623 for gi in 0..gpr {
4624 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4625 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4626 let mut d = 0i32;
4627 for (k, &b) in tile[2..].iter().enumerate() {
4628 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4629 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4630 }
4631 acc += d as f32 * s;
4632 }
4633 acc
4634}
4635
4636#[cfg(target_arch = "aarch64")]
4637#[target_feature(enable = "neon,dotprod")]
4638unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4639 unsafe {
4642 use core::arch::aarch64::*;
4643 use core::arch::asm;
4644 let lomask = vdupq_n_u8(0x0F);
4645 let eight = vdupq_n_s8(8);
4646 let mut acc = 0f32;
4647 for gi in 0..gpr {
4648 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4649 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4650 let b = vld1q_u8(t.add(2));
4651 let lo = vandq_u8(b, lomask);
4652 let hi = vshrq_n_u8::<4>(b);
4653 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4654 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4655 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4656 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4657 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4658 asm!(
4659 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4660 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4661 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4662 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4663 options(pure, nomem, nostack),
4664 );
4665 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4666 }
4667 acc
4668 }
4669}
4670
4671#[cfg(target_arch = "x86_64")]
4672#[target_feature(enable = "avx2")]
4673unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4674 unsafe {
4676 use core::arch::x86_64::*;
4677 let lomask = _mm_set1_epi8(0x0F);
4678 let eight = _mm256_set1_epi8(8);
4679 let ones = _mm256_set1_epi16(1);
4680 let mut acc = 0f32;
4681 for gi in 0..gpr {
4682 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4683 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4684 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4685 let lo = _mm_and_si128(b, lomask);
4686 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4687 let w = _mm256_sub_epi8(
4688 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4689 eight,
4690 );
4691 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4692 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4693 let d = _mm256_madd_epi16(p16, ones);
4694 let hi128 = _mm256_extracti128_si256::<1>(d);
4695 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4696 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4697 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4698 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4699 }
4700 acc
4701 }
4702}
4703
4704#[cfg(target_arch = "x86_64")]
4708#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4709unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4710 unsafe {
4712 use core::arch::x86_64::*;
4713 let lomask = _mm_set1_epi8(0x0F);
4714 let eight = _mm256_set1_epi8(8);
4715 let mut acc = 0f32;
4716 for gi in 0..gpr {
4717 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4718 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4719 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4720 let lo = _mm_and_si128(b, lomask);
4721 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4722 let w = _mm256_sub_epi8(
4723 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4724 eight,
4725 );
4726 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4727 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4728 acc += d as f32 * s;
4729 }
4730 acc
4731 }
4732}
4733
4734#[cfg(target_arch = "x86_64")]
4739#[target_feature(enable = "avx2,fma")]
4744unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4745 unsafe {
4747 use core::arch::x86_64::*;
4748 let lomask = _mm_set1_epi8(0x0F);
4749 let eight = _mm256_set1_epi8(8);
4750 let ones = _mm256_set1_epi16(1);
4751 let mut f0 = _mm256_setzero_ps();
4764 let mut f1 = _mm256_setzero_ps();
4765 let mut f2 = _mm256_setzero_ps();
4766 let mut f3 = _mm256_setzero_ps();
4767 for gi in 0..gpr {
4768 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4769 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4770 let sv = _mm256_set1_ps(s);
4771 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4772 let lo = _mm_and_si128(bb, lomask);
4773 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4774 let w = _mm256_sub_epi8(
4775 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4776 eight,
4777 );
4778 let aw = _mm256_abs_epi8(w);
4779 let off = gi * GROUP_SIZE;
4780 let dot = |xq: &[i8]| {
4781 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4782 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4783 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4784 };
4785 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4786 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4787 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4788 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4789 }
4790 [
4791 hsum256_ps(f0),
4792 hsum256_ps(f1),
4793 hsum256_ps(f2),
4794 hsum256_ps(f3),
4795 ]
4796 }
4797}
4798
4799#[cfg(target_arch = "x86_64")]
4802#[target_feature(enable = "avx2")]
4803#[inline]
4804unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4805 unsafe {
4807 use core::arch::x86_64::*;
4808 let hi = _mm256_extractf128_ps::<1>(v);
4809 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4810 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4811 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4812 _mm_cvtss_f32(s)
4813 }
4814}
4815
4816#[cfg(target_arch = "x86_64")]
4818#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4819unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4820 unsafe {
4822 use core::arch::x86_64::*;
4823 let lomask = _mm_set1_epi8(0x0F);
4824 let eight = _mm256_set1_epi8(8);
4825 let mut f0 = _mm256_setzero_ps();
4828 let mut f1 = _mm256_setzero_ps();
4829 let mut f2 = _mm256_setzero_ps();
4830 let mut f3 = _mm256_setzero_ps();
4831 for gi in 0..gpr {
4832 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4833 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4834 let sv = _mm256_set1_ps(s);
4835 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4836 let lo = _mm_and_si128(bb, lomask);
4837 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4838 let w = _mm256_sub_epi8(
4839 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4840 eight,
4841 );
4842 let aw = _mm256_abs_epi8(w);
4843 let off = gi * GROUP_SIZE;
4844 let dot = |xq: &[i8]| {
4845 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4846 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4847 _mm256_setzero_si256(),
4848 aw,
4849 _mm256_sign_epi8(x, w),
4850 ))
4851 };
4852 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4853 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4854 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4855 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4856 }
4857 let acc = [
4858 hsum256_ps(f0),
4859 hsum256_ps(f1),
4860 hsum256_ps(f2),
4861 hsum256_ps(f3),
4862 ];
4863 acc
4864 }
4865}
4866
4867#[cfg(target_arch = "aarch64")]
4872#[target_feature(enable = "neon,dotprod")]
4873unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4874 unsafe {
4876 use core::arch::aarch64::*;
4877 use core::arch::asm;
4878 let lomask = vdupq_n_u8(0x0F);
4879 let eight = vdupq_n_s8(8);
4880 let mut acc = [0f32; 4];
4881 for gi in 0..gpr {
4882 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4883 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4884 let b = vld1q_u8(t.add(2));
4885 let lo = vandq_u8(b, lomask);
4886 let hi = vshrq_n_u8::<4>(b);
4887 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4888 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4889 for (k, xq) in xs.iter().enumerate() {
4890 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4891 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4892 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4893 asm!(
4894 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4895 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4896 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4897 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4898 options(pure, nomem, nostack),
4899 );
4900 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4901 }
4902 }
4903 acc
4904 }
4905}
4906
4907#[inline]
4909fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4910 let gi = j / GROUP_SIZE;
4911 let k = j % GROUP_SIZE;
4912 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4913 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4914 let byte = tile[2 + k / 2];
4915 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4916 ((nib as i32 - 8) as f32, s)
4917}
4918
4919#[inline]
4922fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4923 let mut acc = 0f32;
4924 for gi in 0..gpr {
4925 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4926 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4927 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4928 let mut ga = 0f32;
4929 for (k, &b) in tile[2..].iter().enumerate() {
4930 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4931 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4932 }
4933 acc += ga * s;
4934 }
4935 acc
4936}
4937
4938struct Q4tpView<'a> {
4942 nib: &'a [u8],
4943 params: &'a [u8],
4944 codes: &'a [u8],
4945 stride: usize,
4946 zero_rung: bool,
4948}
4949
4950impl<'a> Q4tpView<'a> {
4951 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4952 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4953 Self {
4954 nib: &bytes[..params_off],
4955 params: &bytes[params_off..codes_off],
4956 codes: &bytes[codes_off..],
4957 stride,
4958 zero_rung: false,
4959 }
4960 }
4961
4962 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4964 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4965 Self {
4966 nib: &bytes[..params_off],
4967 params: &bytes[params_off..codes_off],
4968 codes: &bytes[codes_off..],
4969 stride,
4970 zero_rung: true,
4971 }
4972 }
4973
4974 #[inline]
4990 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4991 let tab = if self.zero_rung {
4992 q2tp_ladder(self.params, r)
4993 } else {
4994 q4tp_ladder(self.params, r)
4995 };
4996 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4997 let out = &mut out[..gpr];
4998 let mut chunks = out.chunks_exact_mut(8);
4999 let mut ci = 0usize;
5000 for c in &mut chunks {
5001 let w = u64::from(codes[ci])
5002 | u64::from(codes[ci + 1]) << 8
5003 | u64::from(codes[ci + 2]) << 16
5004 | u64::from(codes[ci + 3]) << 24
5005 | u64::from(codes[ci + 4]) << 32;
5006 for (k, o) in c.iter_mut().enumerate() {
5007 *o = tab[((w >> (5 * k)) & 31) as usize];
5008 }
5009 ci += 5;
5010 }
5011 let tail = &codes[ci..];
5014 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5015 *o = tab[q4tp_code(tail, k)];
5016 }
5017 }
5018}
5019
5020#[inline]
5021fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5022 #[cfg(target_arch = "aarch64")]
5023 unsafe {
5024 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5025 }
5026 #[cfg(target_arch = "x86_64")]
5027 unsafe {
5028 if vnni_tiles_enabled() {
5029 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5030 }
5031 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5032 }
5033 #[allow(unreachable_code)]
5034 {
5035 let mut acc = 0f32;
5036 for gi in 0..gpr {
5037 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5038 let s = scales[gi];
5039 let mut d = 0i32;
5040 for (k, &b) in tile.iter().enumerate() {
5041 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5042 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5043 }
5044 acc += d as f32 * s;
5045 }
5046 acc
5047 }
5048}
5049
5050#[cfg(target_arch = "aarch64")]
5053#[target_feature(enable = "neon,dotprod")]
5054unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5055 unsafe {
5058 use core::arch::aarch64::*;
5059 use core::arch::asm;
5060 let lomask = vdupq_n_u8(0x0F);
5061 let eight = vdupq_n_s8(8);
5062 let mut acc = 0f32;
5063 for gi in 0..gpr {
5064 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5065 let s = *scales.get_unchecked(gi);
5066 let b = vld1q_u8(t);
5067 let lo = vandq_u8(b, lomask);
5068 let hi = vshrq_n_u8::<4>(b);
5069 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5070 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5071 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5072 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5073 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5074 asm!(
5075 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5076 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5077 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5078 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5079 options(pure, nomem, nostack),
5080 );
5081 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5082 }
5083 acc
5084 }
5085}
5086
5087#[cfg(target_arch = "x86_64")]
5088#[target_feature(enable = "avx2")]
5089unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5090 unsafe {
5092 use core::arch::x86_64::*;
5093 let lomask = _mm_set1_epi8(0x0F);
5094 let eight = _mm256_set1_epi8(8);
5095 let ones = _mm256_set1_epi16(1);
5096 let mut acc = 0f32;
5097 for gi in 0..gpr {
5098 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5099 let s = *scales.get_unchecked(gi);
5100 let b = _mm_loadu_si128(t as *const __m128i);
5101 let lo = _mm_and_si128(b, lomask);
5102 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5103 let w = _mm256_sub_epi8(
5104 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5105 eight,
5106 );
5107 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5108 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5109 let d = _mm256_madd_epi16(p16, ones);
5110 let hi128 = _mm256_extracti128_si256::<1>(d);
5111 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5112 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5113 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5114 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5115 }
5116 acc
5117 }
5118}
5119
5120#[cfg(target_arch = "x86_64")]
5123#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5124unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5125 unsafe {
5127 use core::arch::x86_64::*;
5128 let lomask = _mm_set1_epi8(0x0F);
5129 let eight = _mm256_set1_epi8(8);
5130 let mut acc = 0f32;
5131 for gi in 0..gpr {
5132 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5133 let s = *scales.get_unchecked(gi);
5134 let b = _mm_loadu_si128(t as *const __m128i);
5135 let lo = _mm_and_si128(b, lomask);
5136 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5137 let w = _mm256_sub_epi8(
5138 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5139 eight,
5140 );
5141 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5142 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5143 }
5144 acc
5145 }
5146}
5147
5148#[inline]
5151fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5152 let mut acc = 0f32;
5153 for gi in 0..gpr {
5154 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5155 let s = scales[gi];
5156 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5157 let mut ga = 0f32;
5158 for (k, &b) in tile.iter().enumerate() {
5159 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5160 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5161 }
5162 acc += ga * s;
5163 }
5164 acc
5165}
5166
5167#[inline]
5170fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5171 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5172 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5173 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5174 ((n as i32 - 8) as f32, scales[gi])
5175}
5176
5177fn q4tp_matvec(
5179 bytes: &[u8],
5180 x: &[f32],
5181 rows: usize,
5182 cols: usize,
5183 out: &mut [f32],
5184 pool: Option<&Pool>,
5185) {
5186 debug_assert_eq!(out.len(), rows);
5187 let gpr = cols / GROUP_SIZE;
5188 let v = Q4tpView::new(bytes, rows, cols);
5189 let out_addr = SendMut(out.as_mut_ptr());
5190 if a8w8_enabled() {
5191 let act = split_act(x);
5192 let run = |start: usize, end: usize| {
5193 with_krow(gpr, |sc| {
5195 for r in start..end {
5196 v.scales_into(r, gpr, sc);
5197 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5198 for &(j, xv) in &act.outliers {
5199 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5200 acc += w * s * xv;
5201 }
5202 unsafe { *out_addr.at(r) = acc };
5204 }
5205 })
5206 };
5207 dispatch_rows(pool, rows, &run);
5208 return;
5209 }
5210 let run = |start: usize, end: usize| {
5211 with_krow(gpr, |sc| {
5212 for r in start..end {
5213 v.scales_into(r, gpr, sc);
5214 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5216 }
5217 })
5218 };
5219 dispatch_rows(pool, rows, &run);
5220}
5221
5222#[allow(clippy::too_many_arguments)]
5225fn q4tp_matvec2(
5226 bytes: &[u8],
5227 x1: &[f32],
5228 x2: &[f32],
5229 rows: usize,
5230 cols: usize,
5231 o1: &mut [f32],
5232 o2: &mut [f32],
5233 pool: Option<&Pool>,
5234) {
5235 let gpr = cols / GROUP_SIZE;
5236 let v = Q4tpView::new(bytes, rows, cols);
5237 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5238 let run = |start: usize, end: usize| {
5239 let mut sc = vec![0f32; gpr];
5240 for r in start..end {
5241 v.scales_into(r, gpr, &mut sc);
5242 unsafe {
5244 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5245 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5246 }
5247 }
5248 };
5249 dispatch_rows(pool, rows, &run);
5250}
5251
5252#[inline]
5255fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5256 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5257 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5258 let c = (byte >> (2 * (k % 4))) & 3;
5259 (c as f32 - 1.5, scales[gi])
5260}
5261
5262#[cfg(target_arch = "x86_64")]
5263const Q2TP_DECODE_U32: [u32; 256] = {
5264 let mut tab = [0u32; 256];
5265 let mut b = 0usize;
5266 while b < 256 {
5267 tab[b] = ((b as u32) & 3)
5268 | ((((b as u32) >> 2) & 3) << 8)
5269 | ((((b as u32) >> 4) & 3) << 16)
5270 | ((((b as u32) >> 6) & 3) << 24);
5271 b += 1;
5272 }
5273 tab
5274};
5275
5276#[cfg(target_arch = "x86_64")]
5280#[target_feature(enable = "avx2")]
5281unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5282 use core::arch::x86_64::*;
5283 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5284 let codes = _mm256_setr_epi32(
5285 Q2TP_DECODE_U32[ch[0] as usize] as i32,
5286 Q2TP_DECODE_U32[ch[1] as usize] as i32,
5287 Q2TP_DECODE_U32[ch[2] as usize] as i32,
5288 Q2TP_DECODE_U32[ch[3] as usize] as i32,
5289 Q2TP_DECODE_U32[ch[4] as usize] as i32,
5290 Q2TP_DECODE_U32[ch[5] as usize] as i32,
5291 Q2TP_DECODE_U32[ch[6] as usize] as i32,
5292 Q2TP_DECODE_U32[ch[7] as usize] as i32,
5293 );
5294 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5295 let pair = _mm256_maddubs_epi16(codes, xv);
5296 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5297 let sum128 = _mm_add_epi32(
5298 _mm256_castsi256_si128(quad),
5299 _mm256_extracti128_si256(quad, 1),
5300 );
5301 let sum64 = _mm_hadd_epi32(sum128, sum128);
5302 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5303}
5304
5305#[inline]
5312fn dot_q2tp_row_i8(
5313 chunks: &[u8],
5314 r: usize,
5315 gpr: usize,
5316 xq: &[i8],
5317 gsum: &[i32],
5318 scales: &[f32],
5319) -> f32 {
5320 let mut acc = 0f32;
5321 let base = r * gpr * Q2TP_CHUNK;
5322 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5323 let mut codes = [0i8; GROUP_SIZE];
5324 #[cfg(target_arch = "x86_64")]
5325 let avx2 = std::arch::is_x86_feature_detected!("avx2");
5326 for gi in 0..gpr {
5327 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5328 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5329 #[cfg(target_arch = "aarch64")]
5330 let dot = unsafe {
5336 use core::arch::aarch64::*;
5337 let b = vld1_u8(ch.as_ptr());
5338 let three = vdup_n_u8(3);
5339 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5340 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5341 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5342 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5343 let x4 = vld4_s8(xg.as_ptr());
5344 let mut acc4 = vdupq_n_s32(0);
5345 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5346 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5347 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5348 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5349 vaddvq_s32(acc4)
5350 };
5351 #[cfg(target_arch = "x86_64")]
5352 let dot: i32 = if avx2 {
5353 unsafe { q2tp_code_dot_avx2(ch, xg) }
5356 } else {
5357 ch.iter()
5358 .enumerate()
5359 .map(|(k, &b)| {
5360 ((b & 3) as i32) * xg[k * 4] as i32
5361 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5362 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5363 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5364 })
5365 .sum()
5366 };
5367 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5368 let dot: i32 = {
5369 for (k, &b) in ch.iter().enumerate() {
5370 codes[k * 4] = (b & 3) as i8;
5371 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5372 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5373 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5374 }
5375 codes
5376 .iter()
5377 .zip(xg)
5378 .map(|(&c, &x)| c as i32 * x as i32)
5379 .sum()
5380 };
5381 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5382 }
5383 acc
5384}
5385
5386fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5390 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5391}
5392
5393#[inline]
5397fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5398 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
5399}
5400
5401#[inline]
5402fn q2tp_row_exact_center(
5403 chunks: &[u8],
5404 r: usize,
5405 gpr: usize,
5406 x: &[f32],
5407 scales: &[f32],
5408 center: f32,
5409) -> f32 {
5410 let mut acc = 0f32;
5411 for gi in 0..gpr {
5412 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
5413 let s = scales[gi];
5414 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5415 let mut g = 0f32;
5416 for (k, &b) in ch.iter().enumerate() {
5417 g += ((b & 3) as f32 - center) * xb[k * 4]
5418 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
5419 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
5420 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
5421 }
5422 acc += s * g;
5423 }
5424 acc
5425}
5426
5427fn q2tp_matvec(
5428 bytes: &[u8],
5429 x: &[f32],
5430 rows: usize,
5431 cols: usize,
5432 out: &mut [f32],
5433 pool: Option<&Pool>,
5434) {
5435 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
5436}
5437
5438fn q2tp_affine_matvec(
5439 bytes: &[u8],
5440 x: &[f32],
5441 rows: usize,
5442 cols: usize,
5443 out: &mut [f32],
5444 pool: Option<&Pool>,
5445) {
5446 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
5447}
5448
5449fn q2tp_matvec_mode(
5450 bytes: &[u8],
5451 x: &[f32],
5452 rows: usize,
5453 cols: usize,
5454 out: &mut [f32],
5455 pool: Option<&Pool>,
5456 affine: bool,
5457) {
5458 debug_assert_eq!(out.len(), rows);
5459 let gpr = cols / GROUP_SIZE;
5460 let v = Q4tpView::new_q2(bytes, rows, cols);
5461 let out_addr = SendMut(out.as_mut_ptr());
5462 if !affine && a8w8_enabled() {
5467 let act = split_act(x);
5468 let gsum = q1_group_sums(&act.xq, gpr);
5469 let (act, gsum) = (&act, &gsum);
5470 let run = move |start: usize, end: usize| {
5471 with_krow(gpr, |sc| {
5472 for r in start..end {
5473 v.scales_into(r, gpr, sc);
5474 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
5475 for &(j, xv) in &act.outliers {
5476 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
5477 acc += w * s * xv;
5478 }
5479 unsafe { *out_addr.at(r) = acc };
5481 }
5482 })
5483 };
5484 dispatch_rows(pool, rows, &run);
5485 return;
5486 }
5487 let run = |start: usize, end: usize| {
5488 with_krow(gpr, |sc| {
5489 for r in start..end {
5490 v.scales_into(r, gpr, sc);
5491 unsafe {
5493 *out_addr.at(r) = if affine {
5494 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
5495 } else {
5496 q2tp_row_exact(v.nib, r, gpr, x, sc)
5497 }
5498 };
5499 }
5500 })
5501 };
5502 dispatch_rows(pool, rows, &run);
5503}
5504
5505#[allow(clippy::too_many_arguments)]
5507fn q2tp_matvec2(
5508 bytes: &[u8],
5509 x1: &[f32],
5510 x2: &[f32],
5511 rows: usize,
5512 cols: usize,
5513 o1: &mut [f32],
5514 o2: &mut [f32],
5515 pool: Option<&Pool>,
5516) {
5517 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
5518}
5519
5520#[allow(clippy::too_many_arguments)]
5521fn q2tp_affine_matvec2(
5522 bytes: &[u8],
5523 x1: &[f32],
5524 x2: &[f32],
5525 rows: usize,
5526 cols: usize,
5527 o1: &mut [f32],
5528 o2: &mut [f32],
5529 pool: Option<&Pool>,
5530) {
5531 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
5532}
5533
5534#[allow(clippy::too_many_arguments)]
5535fn q2tp_matvec2_mode(
5536 bytes: &[u8],
5537 x1: &[f32],
5538 x2: &[f32],
5539 rows: usize,
5540 cols: usize,
5541 o1: &mut [f32],
5542 o2: &mut [f32],
5543 pool: Option<&Pool>,
5544 affine: bool,
5545) {
5546 let gpr = cols / GROUP_SIZE;
5547 let v = Q4tpView::new_q2(bytes, rows, cols);
5548 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5549 let run = |start: usize, end: usize| {
5550 let mut sc = vec![0f32; gpr];
5551 for r in start..end {
5552 v.scales_into(r, gpr, &mut sc);
5553 unsafe {
5555 *p1.at(r) = if affine {
5556 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
5557 } else {
5558 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
5559 };
5560 *p2.at(r) = if affine {
5561 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
5562 } else {
5563 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
5564 };
5565 }
5566 }
5567 };
5568 dispatch_rows(pool, rows, &run);
5569}
5570
5571pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
5578 let gpr = cols / GROUP_SIZE;
5583 let v = Q4tpView::new_q2(bytes, rows, cols);
5584 with_krow(gpr, |sc| {
5585 for r in 0..rows {
5586 v.scales_into(r, gpr, sc);
5587 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
5588 }
5589 });
5590}
5591
5592pub fn q2tp_affine_matvec_for_test(
5595 bytes: &[u8],
5596 x: &[f32],
5597 rows: usize,
5598 cols: usize,
5599 out: &mut [f32],
5600) {
5601 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
5602}
5603
5604pub fn q2tp_matmat_for_test(
5605 bytes: &[u8],
5606 xs_all: &[f32],
5607 b: usize,
5608 rows: usize,
5609 cols: usize,
5610 out: &mut [f32],
5611) {
5612 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
5613}
5614
5615fn q2tp_matmat(
5616 bytes: &[u8],
5617 xs_all: &[f32],
5618 b: usize,
5619 rows: usize,
5620 cols: usize,
5621 out: &mut [f32],
5622 pool: Option<&Pool>,
5623) {
5624 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
5625}
5626
5627fn q2tp_affine_matmat(
5628 bytes: &[u8],
5629 xs_all: &[f32],
5630 b: usize,
5631 rows: usize,
5632 cols: usize,
5633 out: &mut [f32],
5634 pool: Option<&Pool>,
5635) {
5636 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
5637}
5638
5639fn q2tp_matmat_mode(
5640 bytes: &[u8],
5641 xs_all: &[f32],
5642 b: usize,
5643 rows: usize,
5644 cols: usize,
5645 out: &mut [f32],
5646 pool: Option<&Pool>,
5647 affine: bool,
5648) {
5649 debug_assert_eq!(out.len(), b * rows);
5650 let gpr = cols / GROUP_SIZE;
5651 let v = Q4tpView::new_q2(bytes, rows, cols);
5652 let out_addr = SendMut(out.as_mut_ptr());
5653 let run = |start: usize, end: usize| {
5654 let mut sc = vec![0f32; gpr];
5655 for r in start..end {
5656 v.scales_into(r, gpr, &mut sc);
5657 for bi in 0..b {
5658 let x = &xs_all[bi * cols..(bi + 1) * cols];
5659 unsafe {
5661 *out_addr.at(bi * rows + r) = if affine {
5662 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
5663 } else {
5664 q2tp_row_exact(v.nib, r, gpr, x, &sc)
5665 }
5666 };
5667 }
5668 }
5669 };
5670 dispatch_rows(pool, rows, &run);
5671}
5672
5673#[cfg(target_arch = "aarch64")]
5677#[target_feature(enable = "neon,dotprod")]
5678unsafe fn dot_q4tp_row_1x4_sdot_v1(
5679 nib: &[u8],
5680 r: usize,
5681 gpr: usize,
5682 xs: [&[i8]; 4],
5683 scales: &[f32],
5684) -> [f32; 4] {
5685 unsafe {
5686 use core::arch::aarch64::*;
5687 use core::arch::asm;
5688 let lomask = vdupq_n_u8(0x0F);
5689 let eight = vdupq_n_s8(8);
5690 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
5691 for gi in 0..gpr {
5692 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5693 let s = *scales.get_unchecked(gi);
5694 let bb = vld1q_u8(t);
5695 let lo = vandq_u8(bb, lomask);
5696 let hi = vshrq_n_u8::<4>(bb);
5697 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5698 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5699 let mut d = [0f32; 4];
5700 for (k, dk) in d.iter_mut().enumerate() {
5701 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5702 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5703 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5704 asm!(
5705 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5706 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5707 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5708 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5709 options(pure, nomem, nostack),
5710 );
5711 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5712 }
5713 f0 += d[0];
5714 f1 += d[1];
5715 f2 += d[2];
5716 f3 += d[3];
5717 }
5718 [f0, f1, f2, f3]
5719 }
5720}
5721
5722#[allow(dead_code)]
5729static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
5730
5731#[cfg(target_arch = "x86_64")]
5737fn q4tp_blocked_x86() -> bool {
5738 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5739 1 => false,
5740 2 => avx512vnni_enabled(),
5745 _ => blocked_enabled() && avx512vnni_enabled(),
5749 }
5750}
5751
5752#[cfg(target_arch = "aarch64")]
5754#[allow(dead_code)]
5755fn q4tp_v1() -> bool {
5756 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5757 1 => true,
5758 2 => false,
5759 _ => {
5760 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5761 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
5762 }
5763 }
5764}
5765
5766#[cfg(target_arch = "x86_64")]
5778#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5779unsafe fn dot_q4tp_2x8_avx512(
5780 nib: &[u8],
5781 r0: usize,
5782 gpr: usize,
5783 xs: [&[i8]; 8],
5784 sc0: &[f32],
5785 sc1: &[f32],
5786) -> [[f32; 8]; 2] {
5787 unsafe {
5790 use core::arch::x86_64::*;
5791 let lomask = _mm256_set1_epi8(0x0F);
5792 let eight = _mm256_set1_epi8(8);
5793 let zero = _mm512_setzero_si512();
5794 let mut v0 = [_mm512_setzero_ps(); 8];
5795 let mut v1 = [_mm512_setzero_ps(); 8];
5796 let pairs = gpr / 2;
5797 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
5798 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5799 let bb = _mm256_loadu_si256(t as *const __m256i);
5800 let lo = _mm256_and_si256(bb, lomask);
5801 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5802 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5803 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5804 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5805 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5806 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5807 };
5808 for gp in 0..pairs {
5809 let gi = gp * 2;
5810 let (wa0, neg0) = unpack(r0, gi);
5811 let (wa1, neg1) = unpack(r0 + 1, gi);
5812 let off = gi * GROUP_SIZE;
5813 let sv = |sc: &[f32]| {
5814 _mm512_insertf32x8::<1>(
5815 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5816 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5817 )
5818 };
5819 let s0 = sv(sc0);
5820 let s1 = sv(sc1);
5821 for k in 0..8 {
5822 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5823 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5824 zero,
5825 wa0,
5826 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5827 ));
5828 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5829 zero,
5830 wa1,
5831 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5832 ));
5833 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5834 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5835 }
5836 }
5837 let mut acc = [[0f32; 8]; 2];
5838 for k in 0..8 {
5839 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5840 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5841 }
5842 if gpr % 2 == 1 {
5843 let off = (gpr - 1) * GROUP_SIZE;
5844 for j in off..off + GROUP_SIZE {
5845 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5846 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5847 for k in 0..8 {
5848 let x = *xs[k].get_unchecked(j) as f32;
5849 acc[0][k] += w0 * sa * x;
5850 acc[1][k] += w1 * sb * x;
5851 }
5852 }
5853 }
5854 acc
5855 }
5856}
5857
5858#[cfg(target_arch = "x86_64")]
5863#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5864unsafe fn dot_q4tp_row_1x8_avx512(
5865 nib: &[u8],
5866 r: usize,
5867 gpr: usize,
5868 xs: [&[i8]; 8],
5869 scales: &[f32],
5870) -> [f32; 8] {
5871 unsafe {
5873 use core::arch::x86_64::*;
5874 let lomask = _mm256_set1_epi8(0x0F);
5875 let eight = _mm256_set1_epi8(8);
5876 let zero = _mm512_setzero_si512();
5877 let (mut v0, mut v1, mut v2, mut v3) = (
5878 _mm512_setzero_ps(),
5879 _mm512_setzero_ps(),
5880 _mm512_setzero_ps(),
5881 _mm512_setzero_ps(),
5882 );
5883 let (mut v4, mut v5, mut v6, mut v7) = (
5884 _mm512_setzero_ps(),
5885 _mm512_setzero_ps(),
5886 _mm512_setzero_ps(),
5887 _mm512_setzero_ps(),
5888 );
5889 let pairs = gpr / 2;
5890 for gp in 0..pairs {
5891 let gi = gp * 2;
5892 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5893 let bb = _mm256_loadu_si256(t as *const __m256i);
5894 let lo = _mm256_and_si256(bb, lomask);
5895 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5896 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5901 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5902 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5903 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5904 let wabs = _mm512_abs_epi8(w);
5905 let neg = _mm512_movepi8_mask(w);
5906 let off = gi * GROUP_SIZE;
5907 let sv = _mm512_insertf32x8::<1>(
5908 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5909 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5910 );
5911 let dot = |x: &[i8]| -> __m512 {
5912 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5913 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5914 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5915 };
5916 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5917 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5918 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5919 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5920 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5921 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5922 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5923 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5924 }
5925 let mut acc = [
5926 _mm512_reduce_add_ps(v0),
5927 _mm512_reduce_add_ps(v1),
5928 _mm512_reduce_add_ps(v2),
5929 _mm512_reduce_add_ps(v3),
5930 _mm512_reduce_add_ps(v4),
5931 _mm512_reduce_add_ps(v5),
5932 _mm512_reduce_add_ps(v6),
5933 _mm512_reduce_add_ps(v7),
5934 ];
5935 if gpr % 2 == 1 {
5938 let off = (gpr - 1) * GROUP_SIZE;
5939 for j in off..off + GROUP_SIZE {
5940 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5941 let ws = w * s;
5942 for k in 0..8 {
5943 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5944 }
5945 }
5946 }
5947 acc
5948 }
5949}
5950
5951#[cfg(target_arch = "x86_64")]
5964#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5965unsafe fn dot_q4tp_row_1x4_avx512(
5966 nib: &[u8],
5967 r: usize,
5968 gpr: usize,
5969 xs: [&[i8]; 4],
5970 scales: &[f32],
5971) -> [f32; 4] {
5972 unsafe {
5974 use core::arch::x86_64::*;
5975 let lomask = _mm256_set1_epi8(0x0F);
5976 let eight = _mm256_set1_epi8(8);
5977 let zero = _mm512_setzero_si512();
5978 let (mut v0, mut v1, mut v2, mut v3) = (
5979 _mm512_setzero_ps(),
5980 _mm512_setzero_ps(),
5981 _mm512_setzero_ps(),
5982 _mm512_setzero_ps(),
5983 );
5984 let pairs = gpr / 2;
5985 for gp in 0..pairs {
5986 let gi = gp * 2;
5987 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5988 let bb = _mm256_loadu_si256(t as *const __m256i);
5989 let lo = _mm256_and_si256(bb, lomask);
5990 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5991 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5996 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5997 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5998 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5999 let wabs = _mm512_abs_epi8(w);
6000 let neg = _mm512_movepi8_mask(w);
6001 let off = gi * GROUP_SIZE;
6002 let sv = _mm512_insertf32x8::<1>(
6003 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6004 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6005 );
6006 let dot = |x: &[i8]| -> __m512 {
6007 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6008 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6009 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6010 };
6011 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6012 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6013 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6014 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6015 }
6016 let mut acc = [
6017 _mm512_reduce_add_ps(v0),
6018 _mm512_reduce_add_ps(v1),
6019 _mm512_reduce_add_ps(v2),
6020 _mm512_reduce_add_ps(v3),
6021 ];
6022 if gpr % 2 == 1 {
6025 let off = (gpr - 1) * GROUP_SIZE;
6026 for j in off..off + GROUP_SIZE {
6027 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6028 let ws = w * s;
6029 for k in 0..4 {
6030 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6031 }
6032 }
6033 }
6034 acc
6035 }
6036}
6037
6038#[cfg(target_arch = "aarch64")]
6042#[target_feature(enable = "neon,dotprod")]
6043unsafe fn dot_q4tp_row_1x4_sdot(
6044 nib: &[u8],
6045 r: usize,
6046 gpr: usize,
6047 xs: [&[i8]; 4],
6048 scales: &[f32],
6049) -> [f32; 4] {
6050 unsafe {
6052 use core::arch::aarch64::*;
6053 use core::arch::asm;
6054 let lomask = vdupq_n_u8(0x0F);
6055 let eight = vdupq_n_s8(8);
6056 let (mut v0, mut v1, mut v2, mut v3) = (
6072 vdupq_n_f32(0.0),
6073 vdupq_n_f32(0.0),
6074 vdupq_n_f32(0.0),
6075 vdupq_n_f32(0.0),
6076 );
6077 for gi in 0..gpr {
6078 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6079 let s = *scales.get_unchecked(gi);
6080 let bb = vld1q_u8(t);
6081 let lo = vandq_u8(bb, lomask);
6082 let hi = vshrq_n_u8::<4>(bb);
6083 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6084 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6085 let off = gi * GROUP_SIZE;
6086 let dot4 = |x: &[i8]| -> int32x4_t {
6087 let x0 = vld1q_s8(x.as_ptr().add(off));
6088 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6089 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6090 asm!(
6091 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6092 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6093 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6094 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6095 options(pure, nomem, nostack),
6096 );
6097 vaddq_s32(a0, a1)
6098 };
6099 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6100 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6101 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6102 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6103 }
6104 [
6105 vaddvq_f32(v0),
6106 vaddvq_f32(v1),
6107 vaddvq_f32(v2),
6108 vaddvq_f32(v3),
6109 ]
6110 }
6111}
6112
6113fn q4tp_matmat(
6117 bytes: &[u8],
6118 xs_all: &[f32],
6119 b: usize,
6120 rows: usize,
6121 cols: usize,
6122 out: &mut [f32],
6123 pool: Option<&Pool>,
6124) {
6125 debug_assert_eq!(out.len(), b * rows);
6126 let gpr = cols / GROUP_SIZE;
6127 let v = Q4tpView::new(bytes, rows, cols);
6128
6129 #[cfg(target_os = "macos")]
6131 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6132 dequant_matmat_accel(
6133 &|r, dst| {
6134 let mut sc = [0f32; 32];
6135 let mut scv;
6136 let s: &[f32] = if gpr <= 32 {
6137 v.scales_into(r, gpr, &mut sc);
6138 &sc[..gpr]
6139 } else {
6140 scv = vec![0f32; gpr];
6141 v.scales_into(r, gpr, &mut scv);
6142 &scv
6143 };
6144 for gi in 0..gpr {
6145 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6146 for (k, &bb) in tile.iter().enumerate() {
6147 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6148 dst[gi * GROUP_SIZE + k * 2 + 1] =
6149 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6150 }
6151 }
6152 },
6153 xs_all,
6154 b,
6155 rows,
6156 cols,
6157 out,
6158 pool,
6159 );
6160 return;
6161 }
6162
6163 let out_addr = SendMut(out.as_mut_ptr());
6164 if a8w8_enabled() {
6165 let acts: Vec<SplitAct> = (0..b)
6166 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6167 .collect();
6168 let acts = &acts;
6169 #[cfg(target_arch = "aarch64")]
6170 let blocked_ok = sdot_enabled() && blocked_enabled();
6171 #[cfg(target_arch = "x86_64")]
6178 let blocked_ok = q4tp_blocked_x86();
6179 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6180 let blocked_ok = false;
6181 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6191 .ok()
6192 .and_then(|v| v.parse().ok())
6193 .filter(|v| *v > 0)
6194 .unwrap_or(256);
6195 let run = |start: usize, end: usize| {
6196 for abase in (0..acts.len()).step_by(panel_cols) {
6197 let alen = (acts.len() - abase).min(panel_cols);
6198 let mut sc = vec![0f32; gpr];
6199 #[cfg(target_arch = "x86_64")]
6200 let mut r_lo = start;
6201 #[cfg(target_arch = "x86_64")]
6202 if blocked_ok && alen >= 8 {
6203 let mut sc1 = vec![0f32; gpr];
6204 while r_lo + 2 <= end {
6205 v.scales_into(r_lo, gpr, &mut sc);
6206 v.scales_into(r_lo + 1, gpr, &mut sc1);
6207 let mut bi = 0usize;
6208 while bi + 8 <= alen {
6209 let xs = [
6210 acts[abase + bi].xq.as_slice(),
6211 acts[abase + bi + 1].xq.as_slice(),
6212 acts[abase + bi + 2].xq.as_slice(),
6213 acts[abase + bi + 3].xq.as_slice(),
6214 acts[abase + bi + 4].xq.as_slice(),
6215 acts[abase + bi + 5].xq.as_slice(),
6216 acts[abase + bi + 6].xq.as_slice(),
6217 acts[abase + bi + 7].xq.as_slice(),
6218 ];
6219 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6220 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6221 for k in 0..8 {
6222 let act = &acts[abase + bi + k];
6223 let mut acc = dr[k] * act.sx;
6224 for &(j, xv) in &act.outliers {
6225 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6226 acc += w * s * xv;
6227 }
6228 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6230 }
6231 }
6232 bi += 8;
6233 }
6234 for row in [r_lo, r_lo + 1] {
6237 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6238 for b2 in bi..alen {
6239 let act = &acts[abase + b2];
6240 let xs4 = [
6241 act.xq.as_slice(),
6242 act.xq.as_slice(),
6243 act.xq.as_slice(),
6244 act.xq.as_slice(),
6245 ];
6246 let d =
6247 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6248 let mut acc = d[0] * act.sx;
6249 for &(j, xv) in &act.outliers {
6250 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6251 acc += w * s * xv;
6252 }
6253 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6255 }
6256 }
6257 r_lo += 2;
6258 }
6259 }
6260 #[cfg(target_arch = "x86_64")]
6261 let row_start = r_lo;
6262 #[cfg(not(target_arch = "x86_64"))]
6263 let row_start = start;
6264 for r in row_start..end {
6265 v.scales_into(r, gpr, &mut sc);
6266 let mut bi = 0usize;
6267 #[cfg(target_arch = "x86_64")]
6268 if blocked_ok {
6269 while bi + 8 <= alen {
6270 let xs = [
6271 acts[abase + bi].xq.as_slice(),
6272 acts[abase + bi + 1].xq.as_slice(),
6273 acts[abase + bi + 2].xq.as_slice(),
6274 acts[abase + bi + 3].xq.as_slice(),
6275 acts[abase + bi + 4].xq.as_slice(),
6276 acts[abase + bi + 5].xq.as_slice(),
6277 acts[abase + bi + 6].xq.as_slice(),
6278 acts[abase + bi + 7].xq.as_slice(),
6279 ];
6280 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6281 for k in 0..8 {
6282 let act = &acts[abase + bi + k];
6283 let mut acc = d[k] * act.sx;
6284 for &(j, xv) in &act.outliers {
6285 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6286 acc += w * s * xv;
6287 }
6288 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6290 }
6291 bi += 8;
6292 }
6293 while bi + 4 <= alen {
6294 let xs = [
6295 acts[abase + bi].xq.as_slice(),
6296 acts[abase + bi + 1].xq.as_slice(),
6297 acts[abase + bi + 2].xq.as_slice(),
6298 acts[abase + bi + 3].xq.as_slice(),
6299 ];
6300 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6301 for k in 0..4 {
6302 let act = &acts[abase + bi + k];
6303 let mut acc = d[k] * act.sx;
6304 for &(j, xv) in &act.outliers {
6305 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6306 acc += w * s * xv;
6307 }
6308 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6310 }
6311 bi += 4;
6312 }
6313 }
6314 #[cfg(target_arch = "aarch64")]
6315 if blocked_ok {
6316 while bi + 4 <= alen {
6317 let xs = [
6318 acts[abase + bi].xq.as_slice(),
6319 acts[abase + bi + 1].xq.as_slice(),
6320 acts[abase + bi + 2].xq.as_slice(),
6321 acts[abase + bi + 3].xq.as_slice(),
6322 ];
6323 let d = unsafe {
6324 if q4tp_v1() {
6325 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6326 } else {
6327 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6328 }
6329 };
6330 for k in 0..4 {
6331 let act = &acts[abase + bi + k];
6332 let mut acc = d[k] * act.sx;
6333 for &(j, xv) in &act.outliers {
6334 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6335 acc += w * s * xv;
6336 }
6337 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6339 }
6340 bi += 4;
6341 }
6342 }
6343 let _ = blocked_ok;
6344 while bi < alen {
6345 let act = &acts[abase + bi];
6346 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6347 for &(j, xv) in &act.outliers {
6348 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6349 acc += w * s * xv;
6350 }
6351 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
6353 bi += 1;
6354 }
6355 }
6356 }
6357 };
6358 dispatch_rows(pool, rows, &run);
6359 return;
6360 }
6361
6362 let run = |start: usize, end: usize| {
6363 let mut sc = vec![0f32; gpr];
6364 for r in start..end {
6365 v.scales_into(r, gpr, &mut sc);
6366 for bi in 0..b {
6367 let x = &xs_all[bi * cols..(bi + 1) * cols];
6368 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
6370 }
6371 }
6372 };
6373 dispatch_rows(pool, rows, &run);
6374}
6375
6376fn q4t_matvec(
6378 bytes: &[u8],
6379 x: &[f32],
6380 rows: usize,
6381 cols: usize,
6382 out: &mut [f32],
6383 pool: Option<&Pool>,
6384) {
6385 debug_assert_eq!(out.len(), rows);
6386 let gpr = cols / GROUP_SIZE;
6387 let out_addr = SendMut(out.as_mut_ptr());
6388 if a8w8_enabled() {
6389 let act = split_act(x);
6390 let run = move |start: usize, end: usize| {
6391 for r in start..end {
6392 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6393 for &(j, xv) in &act.outliers {
6394 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6395 acc += w * s * xv;
6396 }
6397 unsafe { *out_addr.at(r) = acc };
6399 }
6400 };
6401 dispatch_rows(pool, rows, &run);
6402 return;
6403 }
6404 let run = move |start: usize, end: usize| {
6405 for r in start..end {
6406 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
6408 }
6409 };
6410 dispatch_rows(pool, rows, &run);
6411}
6412
6413#[allow(clippy::too_many_arguments)]
6415fn q4t_matvec2(
6416 bytes: &[u8],
6417 x1: &[f32],
6418 x2: &[f32],
6419 rows: usize,
6420 cols: usize,
6421 o1: &mut [f32],
6422 o2: &mut [f32],
6423 pool: Option<&Pool>,
6424) {
6425 let gpr = cols / GROUP_SIZE;
6426 let p1 = SendMut(o1.as_mut_ptr());
6427 let p2 = SendMut(o2.as_mut_ptr());
6428 if a8w8_enabled() {
6429 let a1 = split_act(x1);
6430 let a2 = split_act(x2);
6431 let run = move |start: usize, end: usize| {
6432 for r in start..end {
6433 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
6434 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
6435 for &(j, xv) in &a1.outliers {
6436 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6437 v1 += w * s * xv;
6438 }
6439 for &(j, xv) in &a2.outliers {
6440 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6441 v2 += w * s * xv;
6442 }
6443 unsafe {
6445 *p1.at(r) = v1;
6446 *p2.at(r) = v2;
6447 }
6448 }
6449 };
6450 dispatch_rows(pool, rows, &run);
6451 return;
6452 }
6453 let run = move |start: usize, end: usize| {
6454 for r in start..end {
6455 unsafe {
6457 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
6458 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
6459 }
6460 }
6461 };
6462 dispatch_rows(pool, rows, &run);
6463}
6464
6465#[allow(clippy::too_many_arguments)]
6467#[cfg(target_os = "macos")]
6473fn dequant_matmat_accel(
6474 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
6475 xs_all: &[f32],
6476 b: usize,
6477 rows: usize,
6478 cols: usize,
6479 out: &mut [f32],
6480 pool: Option<&Pool>,
6481) {
6482 const TR: usize = 2048;
6483 thread_local! {
6484 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6485 }
6486 WTILE.with(|wt| {
6487 let mut wtile = wt.borrow_mut();
6488 wtile.resize(TR * cols, 0.0);
6489 let mut r0 = 0usize;
6490 while r0 < rows {
6491 let tr = TR.min(rows - r0);
6492 let wt_addr = SendMut(wtile.as_mut_ptr());
6493 let run = |start: usize, end: usize| {
6494 for r in start..end {
6495 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
6497 dequant_row(r0 + r, dst);
6498 }
6499 };
6500 dispatch_rows(pool, tr, &run);
6501 unsafe {
6502 accel_blas::cblas_sgemm(
6503 101, 111, 112, b as i32,
6507 tr as i32,
6508 cols as i32,
6509 1.0,
6510 xs_all.as_ptr(),
6511 cols as i32,
6512 wtile.as_ptr(),
6513 cols as i32,
6514 0.0,
6515 out.as_mut_ptr().add(r0),
6516 rows as i32,
6517 );
6518 }
6519 r0 += tr;
6520 }
6521 });
6522}
6523
6524fn q4t_matmat(
6525 bytes: &[u8],
6526 xs_all: &[f32],
6527 b: usize,
6528 rows: usize,
6529 cols: usize,
6530 out: &mut [f32],
6531 pool: Option<&Pool>,
6532) {
6533 debug_assert_eq!(out.len(), b * rows);
6534 let gpr = cols / GROUP_SIZE;
6535 #[cfg(target_os = "macos")]
6539 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6540 dequant_matmat_accel(
6541 &|r, dst| {
6542 for gi in 0..gpr {
6543 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6544 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6545 for (k, &bb) in tile[2..].iter().enumerate() {
6546 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
6547 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
6548 }
6549 }
6550 },
6551 xs_all,
6552 b,
6553 rows,
6554 cols,
6555 out,
6556 pool,
6557 );
6558 return;
6559 }
6560 let out_addr = SendMut(out.as_mut_ptr());
6561 if a8w8_enabled() {
6562 let acts: Vec<SplitAct> = (0..b)
6563 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6564 .collect();
6565 let acts = &acts;
6566 #[cfg(target_arch = "x86_64")]
6567 let blocked_ok = avx2_enabled() && blocked_enabled();
6568 #[cfg(target_arch = "aarch64")]
6569 let blocked_ok = sdot_enabled() && blocked_enabled();
6570 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
6571 let blocked_ok = false;
6572 let run = move |start: usize, end: usize| {
6573 for r in start..end {
6574 let mut bi = 0usize;
6575 #[cfg(target_arch = "aarch64")]
6576 if blocked_ok {
6577 while bi + 4 <= acts.len() {
6578 let xs = [
6579 acts[bi].xq.as_slice(),
6580 acts[bi + 1].xq.as_slice(),
6581 acts[bi + 2].xq.as_slice(),
6582 acts[bi + 3].xq.as_slice(),
6583 ];
6584 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
6585 for k in 0..4 {
6586 let act = &acts[bi + k];
6587 let mut acc = d[k] * act.sx;
6588 for &(j, xv) in &act.outliers {
6589 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6590 acc += w * sc * xv;
6591 }
6592 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6594 }
6595 bi += 4;
6596 }
6597 }
6598 #[cfg(target_arch = "x86_64")]
6599 if blocked_ok {
6600 while bi + 4 <= acts.len() {
6601 let xs = [
6602 acts[bi].xq.as_slice(),
6603 acts[bi + 1].xq.as_slice(),
6604 acts[bi + 2].xq.as_slice(),
6605 acts[bi + 3].xq.as_slice(),
6606 ];
6607 let d = unsafe {
6608 if vnni_tiles_enabled() {
6609 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
6610 } else {
6611 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
6612 }
6613 };
6614 for k in 0..4 {
6615 let act = &acts[bi + k];
6616 let mut acc = d[k] * act.sx;
6617 for &(j, xv) in &act.outliers {
6618 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6619 acc += w * sc * xv;
6620 }
6621 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6623 }
6624 bi += 4;
6625 }
6626 }
6627 let _ = blocked_ok;
6628 while bi < acts.len() {
6629 let act = &acts[bi];
6630 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6631 for &(j, xv) in &act.outliers {
6632 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6633 acc += w * s * xv;
6634 }
6635 unsafe { *out_addr.at(bi * rows + r) = acc };
6637 bi += 1;
6638 }
6639 }
6640 };
6641 dispatch_rows(pool, rows, &run);
6642 return;
6643 }
6644 let run = move |start: usize, end: usize| {
6645 for r in start..end {
6646 for bi in 0..b {
6647 let x = &xs_all[bi * cols..(bi + 1) * cols];
6648 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
6650 }
6651 }
6652 };
6653 dispatch_rows(pool, rows, &run);
6654}
6655
6656fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
6665 (0..gpr)
6666 .map(|gi| {
6667 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
6668 .iter()
6669 .map(|&v| v as i32)
6670 .sum()
6671 })
6672 .collect()
6673}
6674
6675#[inline]
6679#[allow(unreachable_code)]
6680#[cfg(target_arch = "x86_64")]
6685#[target_feature(enable = "avx2")]
6686unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6687 unsafe {
6689 use core::arch::x86_64::*;
6690 let expand = _mm256_setr_epi8(
6692 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,
6693 3, 3, 3,
6694 );
6695 let bitsel = _mm256_setr_epi8(
6696 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6697 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6698 );
6699 let ones8 = _mm256_set1_epi8(1);
6700 let ones16 = _mm256_set1_epi16(1);
6701 let mut acc = 0f32;
6702 for gi in 0..gpr {
6703 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6704 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6705 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6706 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6707 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6708 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6709 let sel = _mm256_and_si256(x, mask);
6710 let p16 = _mm256_maddubs_epi16(ones8, sel);
6712 let d32 = _mm256_madd_epi16(p16, ones16);
6713 let hi128 = _mm256_extracti128_si256::<1>(d32);
6714 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6715 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6716 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6717 let msum = _mm_cvtsi128_si32(s32);
6718 let d = 2 * msum - gsum[gi];
6721 acc += d as f32 * s;
6722 }
6723 acc
6724 }
6725}
6726
6727#[cfg(target_arch = "x86_64")]
6730#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6731unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6732 unsafe {
6734 use core::arch::x86_64::*;
6735 let expand = _mm256_setr_epi8(
6736 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,
6737 3, 3, 3,
6738 );
6739 let bitsel = _mm256_setr_epi8(
6740 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6741 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6742 );
6743 let ones8 = _mm256_set1_epi8(1);
6744 let mut acc = 0f32;
6745 for gi in 0..gpr {
6746 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6747 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6748 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6749 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6750 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6751 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6752 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6753 let d = 2 * msum - gsum[gi];
6754 acc += d as f32 * s;
6755 }
6756 acc
6757 }
6758}
6759
6760#[cfg(target_arch = "x86_64")]
6762#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6763unsafe fn dot_q1_row_1x4_vnni(
6764 bytes: &[u8],
6765 r: usize,
6766 gpr: usize,
6767 xs: [&[i8]; 4],
6768 gsums: [&[i32]; 4],
6769) -> [f32; 4] {
6770 unsafe {
6772 use core::arch::x86_64::*;
6773 let expand = _mm256_setr_epi8(
6774 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,
6775 3, 3, 3,
6776 );
6777 let bitsel = _mm256_setr_epi8(
6778 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6779 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6780 );
6781 let ones8 = _mm256_set1_epi8(1);
6782 let mut acc = [0f32; 4];
6783 for gi in 0..gpr {
6784 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6785 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6786 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6787 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6788 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6789 for (k, xq) in xs.iter().enumerate() {
6790 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6791 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6792 let d = 2 * msum - gsums[k][gi];
6793 acc[k] += d as f32 * s;
6794 }
6795 }
6796 acc
6797 }
6798}
6799
6800#[cfg(target_arch = "x86_64")]
6803#[target_feature(enable = "avx2")]
6804unsafe fn dot_q1_row_1x4_avx2(
6805 bytes: &[u8],
6806 r: usize,
6807 gpr: usize,
6808 xs: [&[i8]; 4],
6809 gsums: [&[i32]; 4],
6810) -> [f32; 4] {
6811 unsafe {
6813 use core::arch::x86_64::*;
6814 let expand = _mm256_setr_epi8(
6815 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,
6816 3, 3, 3,
6817 );
6818 let bitsel = _mm256_setr_epi8(
6819 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6820 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6821 );
6822 let ones8 = _mm256_set1_epi8(1);
6823 let ones16 = _mm256_set1_epi16(1);
6824 let mut acc = [0f32; 4];
6825 for gi in 0..gpr {
6826 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6827 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6828 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6829 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6830 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6831 for (k, xq) in xs.iter().enumerate() {
6832 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6833 let sel = _mm256_and_si256(x, mask);
6834 let p16 = _mm256_maddubs_epi16(ones8, sel);
6835 let d32 = _mm256_madd_epi16(p16, ones16);
6836 let hi128 = _mm256_extracti128_si256::<1>(d32);
6837 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6838 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6839 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6840 let msum = _mm_cvtsi128_si32(s32);
6841 let d = 2 * msum - gsums[k][gi];
6842 acc[k] += d as f32 * s;
6843 }
6844 }
6845 acc
6846 }
6847}
6848
6849#[allow(unreachable_code)]
6850fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6851 #[cfg(target_arch = "aarch64")]
6852 unsafe {
6853 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6854 }
6855 #[cfg(target_arch = "x86_64")]
6856 if avx2_enabled() {
6857 unsafe {
6858 if vnni_tiles_enabled() {
6859 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6860 }
6861 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6862 }
6863 }
6864 let _ = gsum;
6865 let mut acc = 0f32;
6866 for gi in 0..gpr {
6867 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6868 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6869 let mut d = 0i32;
6870 for (j, &b) in tile[2..].iter().enumerate() {
6871 for k in 0..8 {
6872 let w = ((b >> k) & 1) as i32 * 2 - 1;
6873 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6874 }
6875 }
6876 acc += d as f32 * s;
6877 }
6878 acc
6879}
6880
6881#[cfg(target_arch = "aarch64")]
6890#[target_feature(enable = "neon,dotprod")]
6891unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6892 unsafe {
6895 use core::arch::aarch64::*;
6896 use core::arch::asm;
6897 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6898 let m = vld1q_u8(MASKS.as_ptr());
6899 macro_rules! tile_dot {
6901 ($t:expr, $x:expr) => {{
6902 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6903 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6904 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6905 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6906 let x0 = vld1q_s8($x);
6907 let x1 = vld1q_s8($x.add(16));
6908 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6909 asm!(
6910 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6911 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6912 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6913 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6914 options(pure, nomem, nostack),
6915 );
6916 vaddq_s32(a0, a1)
6917 }};
6918 }
6919 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6928 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6929 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6930 const IW11: [u8; 16] = [
6931 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6932 ];
6933 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6934 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6935 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6936 let isc = vld1_u8(ISC.as_ptr());
6937 macro_rules! tile_dot_tbl {
6939 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6940 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6941 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6942 let x0 = vld1q_s8($x);
6943 let x1 = vld1q_s8($x.add(16));
6944 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6945 asm!(
6946 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6947 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6948 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6949 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6950 options(pure, nomem, nostack),
6951 );
6952 vaddq_s32(a0, a1)
6953 }};
6954 }
6955 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6956 let row_base = r * gpr * Q1_TILE;
6957 let abs_end = bytes.len();
6958 let xp = xq.as_ptr();
6959 let gp = gsum.as_ptr();
6960 let mut accv = vdupq_n_f32(0.0);
6961 let mut gi = 0;
6962 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6965 let t0 = base.add(gi * Q1_TILE);
6966 let ld_a = vld1q_u8(t0);
6967 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6968 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6969 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6970 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6971 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6972 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6974 let g = vld1q_s32(gp.add(gi));
6975 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6976 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6977 let scf: float32x4_t;
6978 asm!(
6979 "fcvtl {o:v}.4s, {i:v}.4h",
6980 o = out(vreg) scf, i = in(vreg) sc16,
6981 options(pure, nomem, nostack),
6982 );
6983 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6984 gi += 4;
6985 }
6986 let mut acc = vaddvq_f32(accv);
6987 while gi < gpr {
6988 let t = base.add(gi * Q1_TILE);
6989 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6990 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6991 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6992 gi += 1;
6993 }
6994 acc
6995 }
6996}
6997
6998#[cfg(target_arch = "aarch64")]
7003#[target_feature(enable = "neon,dotprod")]
7004unsafe fn dot_q1_row_1x4_sdot(
7005 bytes: &[u8],
7006 r: usize,
7007 gpr: usize,
7008 xs: [&[i8]; 4],
7009 gs: [&[i32]; 4],
7010) -> [f32; 4] {
7011 unsafe {
7013 use core::arch::aarch64::*;
7014 use core::arch::asm;
7015 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7016 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7017 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7018 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7019 const IW11: [u8; 16] = [
7020 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7021 ];
7022 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7023 let m = vld1q_u8(MASKS.as_ptr());
7024 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7025 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7026 let isc = vld1_u8(ISC.as_ptr());
7027 macro_rules! sdot2 {
7028 ($w0:expr, $w1:expr, $x:expr) => {{
7029 let x0 = vld1q_s8($x);
7030 let x1 = vld1q_s8($x.add(16));
7031 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7032 asm!(
7033 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7034 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7035 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7036 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7037 options(pure, nomem, nostack),
7038 );
7039 vaddq_s32(a0, a1)
7040 }};
7041 }
7042 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7043 let row_base = r * gpr * Q1_TILE;
7044 let abs_end = bytes.len();
7045 let mut accv = [vdupq_n_f32(0.0); 4];
7046 let mut gi = 0;
7047 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7048 let t0 = base.add(gi * Q1_TILE);
7049 let ld_a = vld1q_u8(t0);
7050 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7051 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7053 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7054 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7055 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7056 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7057 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7058 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7059 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7060 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7061 let scf: float32x4_t;
7062 asm!(
7063 "fcvtl {o:v}.4s, {i:v}.4h",
7064 o = out(vreg) scf, i = in(vreg) sc16,
7065 options(pure, nomem, nostack),
7066 );
7067 for k in 0..4 {
7068 let xp = xs[k].as_ptr();
7069 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7070 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7071 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7072 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7073 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7074 let g = vld1q_s32(gs[k].as_ptr().add(gi));
7075 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7076 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7077 }
7078 gi += 4;
7079 }
7080 let mut acc = [
7081 vaddvq_f32(accv[0]),
7082 vaddvq_f32(accv[1]),
7083 vaddvq_f32(accv[2]),
7084 vaddvq_f32(accv[3]),
7085 ];
7086 while gi < gpr {
7087 let t = base.add(gi * Q1_TILE);
7088 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7089 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7090 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7091 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7092 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7093 for k in 0..4 {
7094 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7095 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7096 }
7097 gi += 1;
7098 }
7099 acc
7100 }
7101}
7102
7103#[inline]
7105fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7106 let gi = j / GROUP_SIZE;
7107 let k = j % GROUP_SIZE;
7108 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7109 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7110 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7111 ((bit as i32 * 2 - 1) as f32, s)
7112}
7113
7114#[inline]
7116fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7117 let mut acc = 0f32;
7118 for gi in 0..gpr {
7119 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7120 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7121 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7122 let mut ga = 0f32;
7123 for (j, &b) in tile[2..].iter().enumerate() {
7124 for k in 0..8 {
7125 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7126 }
7127 }
7128 acc += ga * s;
7129 }
7130 acc
7131}
7132
7133#[allow(clippy::too_many_arguments)]
7136fn q1_range_a8w8(
7137 bytes: &[u8],
7138 gpr: usize,
7139 act: &SplitAct,
7140 gsum: &[i32],
7141 out: SendMut,
7142 start: usize,
7143 end: usize,
7144) {
7145 for r in start..end {
7146 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7147 for &(j, xv) in &act.outliers {
7148 let (w, s) = q1_outlier(bytes, r, gpr, j);
7149 acc += w * s * xv;
7150 }
7151 unsafe { *out.at(r) = acc };
7153 }
7154}
7155
7156fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7158 for r in start..end {
7159 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
7161 }
7162}
7163
7164fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
7169 let entries = base_len + (rows + 1) * 4;
7170 (base_len, entries, entries <= bytes.len())
7171}
7172
7173#[inline]
7175fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
7176 let o = rp_off + r * 4;
7177 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
7178}
7179
7180const SIGN5: [[f32; 5]; 256] = {
7184 let mut lut = [[0.0f32; 5]; 256];
7185 let pow3 = [1u16, 3, 9, 27, 81];
7186 let mut byte = 0usize;
7187 while byte < 256 {
7188 let mut i = 0usize;
7189 while i < 5 {
7190 let code = (byte as u16 / pow3[i]) % 3;
7191 lut[byte][i] = if code == 1 {
7192 1.0
7193 } else if code == 2 {
7194 -1.0
7195 } else {
7196 0.0
7197 };
7198 i += 1;
7199 }
7200 byte += 1;
7201 }
7202 lut
7203};
7204
7205const SIGN5_I8: [[i8; 5]; 256] = {
7207 let mut lut = [[0i8; 5]; 256];
7208 let pow3 = [1u16, 3, 9, 27, 81];
7209 let mut byte = 0usize;
7210 while byte < 256 {
7211 let mut i = 0usize;
7212 while i < 5 {
7213 let code = (byte as u16 / pow3[i]) % 3;
7214 lut[byte][i] = if code == 1 {
7215 1
7216 } else if code == 2 {
7217 -1
7218 } else {
7219 0
7220 };
7221 i += 1;
7222 }
7223 byte += 1;
7224 }
7225 lut
7226};
7227
7228const SIGN5_U64: [u64; 256] = {
7234 let mut lut = [0u64; 256];
7235 let pow3 = [1u16, 3, 9, 27, 81];
7236 let mut byte = 0usize;
7237 while byte < 256 {
7238 let mut v = 0u64;
7239 let mut i = 0usize;
7240 while i < 5 {
7241 let code = (byte as u16 / pow3[i]) % 3;
7242 let s: u8 = if code == 1 {
7243 1
7244 } else if code == 2 {
7245 0xFF
7246 } else {
7247 0
7248 };
7249 v |= (s as u64) << (i * 8);
7250 i += 1;
7251 }
7252 lut[byte] = v;
7253 byte += 1;
7254 }
7255 lut
7256};
7257
7258#[inline]
7263fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
7264 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7265 let off = (r * gpr + j / GROUP_SIZE) * TILE;
7266 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7267 let within = j % GROUP_SIZE;
7268 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
7269}
7270
7271#[cfg(target_arch = "aarch64")]
7274#[target_feature(enable = "neon,dotprod")]
7275#[inline]
7276unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
7277 unsafe {
7279 use core::arch::aarch64::*;
7280 use core::arch::asm;
7281 let w0 = vld1q_s8(w);
7282 let w1 = vld1q_s8(w.add(16));
7283 let x0 = vld1q_s8(x);
7284 let x1 = vld1q_s8(x.add(16));
7285 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7286 asm!(
7287 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7288 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7289 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7290 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7291 options(pure, nomem, nostack),
7292 );
7293 vaddvq_s32(vaddq_s32(a0, a1))
7294 }
7295}
7296
7297#[cfg(target_arch = "x86_64")]
7300#[target_feature(enable = "avx2")]
7301#[inline]
7302unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
7303 unsafe {
7305 use core::arch::x86_64::*;
7306 let wv = _mm256_loadu_si256(w as *const __m256i);
7307 let xv = _mm256_loadu_si256(x as *const __m256i);
7308 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7309 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
7310 let hi128 = _mm256_extracti128_si256::<1>(d);
7311 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7312 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7313 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7314 _mm_cvtsi128_si32(s32)
7315 }
7316}
7317
7318#[inline]
7323fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
7324 debug_assert!(dst.len() >= 40);
7325 unsafe {
7328 let p = dst.as_mut_ptr();
7329 for bi in 0..7 {
7330 core::ptr::write_unaligned(
7331 p.add(bi * 5) as *mut u64,
7332 SIGN5_U64[*codes.add(bi) as usize],
7333 );
7334 }
7335 }
7336}
7337
7338#[inline]
7343fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
7344 #[cfg(target_arch = "aarch64")]
7345 unsafe {
7346 return sdot32_i8(w, x);
7347 }
7348 #[cfg(target_arch = "x86_64")]
7349 unsafe {
7350 return i8dot32_avx2(w, x);
7351 }
7352 #[allow(unreachable_code)]
7353 unsafe {
7354 let mut s = 0i32;
7355 for k in 0..GROUP_SIZE {
7356 s += *w.add(k) as i32 * *x.add(k) as i32;
7357 }
7358 s
7359 }
7360}
7361
7362#[inline]
7363unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
7364 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
7365 (
7366 SIGN5_U64[*codes as usize],
7367 SIGN5_U64[*codes.add(1) as usize],
7368 SIGN5_U64[*codes.add(2) as usize],
7369 SIGN5_U64[*codes.add(3) as usize],
7370 SIGN5_U64[*codes.add(4) as usize],
7371 SIGN5_U64[*codes.add(5) as usize],
7372 SIGN5_U64[*codes.add(6) as usize],
7373 )
7374 };
7375
7376 let u0 = s0 | (s1 << 40);
7377 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
7378 let u2 = (s3 >> 8) | (s4 << 32);
7379 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
7380
7381 (u0, u1, u2, u3)
7382}
7383
7384#[cfg(target_arch = "aarch64")]
7388#[target_feature(enable = "neon,dotprod")]
7389unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7390 use core::arch::aarch64::*;
7391 use core::arch::asm;
7392 unsafe {
7393 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7394 let mut acc = 0f32;
7395 let bytes_ptr = bytes.as_ptr();
7396 let xq_ptr = xq.as_ptr();
7397 let row_off = r * gpr * TILE;
7398
7399 let gpr2 = gpr & !1;
7400 let mut gi = 0;
7401 while gi < gpr2 {
7402 let off0 = row_off + gi * TILE;
7403 let off1 = off0 + TILE;
7404 let s0 = f16_to_f32(u16::from_le_bytes([
7405 *bytes_ptr.add(off0),
7406 *bytes_ptr.add(off0 + 1),
7407 ]));
7408 let s1 = f16_to_f32(u16::from_le_bytes([
7409 *bytes_ptr.add(off1),
7410 *bytes_ptr.add(off1 + 1),
7411 ]));
7412
7413 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7414 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7415
7416 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7417 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7418 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7419 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7420
7421 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7422 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7423 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
7424 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
7425
7426 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
7427 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7428 asm!(
7429 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
7430 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
7431 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
7432 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
7433 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
7434 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
7435 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
7436 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
7437 options(pure, nomem, nostack),
7438 );
7439 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
7440 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
7441 acc += d0 as f32 * s0 + d1 as f32 * s1;
7442 gi += 2;
7443 }
7444
7445 if gi < gpr {
7446 let off = row_off + gi * TILE;
7447 let s = f16_to_f32(u16::from_le_bytes([
7448 *bytes_ptr.add(off),
7449 *bytes_ptr.add(off + 1),
7450 ]));
7451 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7452 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7453 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7454 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7455 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7456 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7457 asm!(
7458 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7459 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7460 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7461 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7462 options(pure, nomem, nostack),
7463 );
7464 let d = vaddvq_s32(vaddq_s32(a0, a1));
7465 acc += d as f32 * s;
7466 }
7467 acc
7468 }
7469}
7470
7471#[cfg(target_arch = "x86_64")]
7473#[target_feature(enable = "avx2")]
7474unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7475 use core::arch::x86_64::*;
7476 unsafe {
7477 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7478 let mut acc = 0f32;
7479 let bytes_ptr = bytes.as_ptr();
7480 let xq_ptr = xq.as_ptr();
7481 let row_off = r * gpr * TILE;
7482
7483 let ones = _mm256_set1_epi16(1);
7484 for gi in 0..gpr {
7485 let off = row_off + gi * TILE;
7486 let s = f16_to_f32(u16::from_le_bytes([
7487 *bytes_ptr.add(off),
7488 *bytes_ptr.add(off + 1),
7489 ]));
7490 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7491 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7492 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7493 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7494 let d256 = _mm256_madd_epi16(p16, ones);
7495 let d128 = _mm_add_epi32(
7496 _mm256_castsi256_si128(d256),
7497 _mm256_extracti128_si256(d256, 1),
7498 );
7499 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
7500 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
7501 acc += d32 as f32 * s;
7502 }
7503 acc
7504 }
7505}
7506
7507#[cfg(target_arch = "x86_64")]
7509#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7510unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7511 use core::arch::x86_64::*;
7512 unsafe {
7514 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7515 let mut acc = 0f32;
7516 let bytes_ptr = bytes.as_ptr();
7517 let xq_ptr = xq.as_ptr();
7518 let row_off = r * gpr * TILE;
7519 for gi in 0..gpr {
7520 let off = row_off + gi * TILE;
7521 let s = f16_to_f32(u16::from_le_bytes([
7522 *bytes_ptr.add(off),
7523 *bytes_ptr.add(off + 1),
7524 ]));
7525 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7526 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7527 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7528 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7529 acc += d as f32 * s;
7530 }
7531 acc
7532 }
7533}
7534
7535#[inline]
7539fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7540 #[cfg(target_arch = "aarch64")]
7541 unsafe {
7542 return q1t_dot_row_sdot(bytes, r, gpr, xq);
7543 }
7544 #[cfg(target_arch = "x86_64")]
7545 unsafe {
7546 if vnni_tiles_enabled() {
7547 return q1t_dot_row_vnni(bytes, r, gpr, xq);
7548 }
7549 return q1t_dot_row_avx2(bytes, r, gpr, xq);
7550 }
7551 #[allow(unreachable_code)]
7552 {
7553 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7554 let mut acc = 0f32;
7555 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
7557 let off = (r * gpr + gi) * TILE;
7558 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7559 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
7560 let mut d = 0i32;
7561 for k in 0..GROUP_SIZE {
7562 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
7563 }
7564 acc += d as f32 * s;
7565 }
7566 acc
7567 }
7568}
7569
7570fn q1t_row_outlier_correction(
7577 bytes: &[u8],
7578 r: usize,
7579 rp_off: usize,
7580 entries_off: usize,
7581 has_ov: bool,
7582 x: &[f32],
7583) -> f32 {
7584 if !has_ov {
7585 return 0.0;
7586 }
7587 let (c0, c1) = (
7588 q1t_rowptr(bytes, rp_off, r),
7589 q1t_rowptr(bytes, rp_off, r + 1),
7590 );
7591 let mut corr = 0f32;
7592 for p in c0..c1 {
7593 let e = entries_off + p * 4;
7594 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7595 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7596 corr += val * x[col];
7597 }
7598 corr
7599}
7600
7601fn q1t_dequant_row(
7605 bytes: &[u8],
7606 r: usize,
7607 gpr: usize,
7608 rp_off: usize,
7609 entries_off: usize,
7610 has_ov: bool,
7611 buf: &mut [f32],
7612) {
7613 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7614 for g in 0..gpr {
7615 let off = (r * gpr + g) * TILE;
7616 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7617 let codes = &bytes[off + 2..off + TILE];
7618 let bc = g * GROUP_SIZE;
7619 for bi in 0..6 {
7621 let lut = &SIGN5[codes[bi] as usize];
7622 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
7623 for i in 0..5 {
7624 d[i] = lut[i] * s;
7625 }
7626 }
7627 let lut = &SIGN5[codes[6] as usize];
7628 buf[bc + 30] = lut[0] * s;
7629 buf[bc + 31] = lut[1] * s;
7630 }
7631 if !has_ov {
7632 return;
7633 }
7634 let (c0, c1) = (
7635 q1t_rowptr(bytes, rp_off, r),
7636 q1t_rowptr(bytes, rp_off, r + 1),
7637 );
7638 for p in c0..c1 {
7639 let e = entries_off + p * 4;
7640 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7641 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7642 }
7643}
7644
7645fn q1t_add_overlay(
7649 bytes: &[u8],
7650 x: &[f32],
7651 rows: usize,
7652 cols: usize,
7653 out: &mut [f32],
7654 pool: Option<&Pool>,
7655) {
7656 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7657 let gpr = cols / GROUP_SIZE;
7658 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7659 if !has_ov {
7660 return;
7661 }
7662 let out_addr = SendMut(out.as_mut_ptr());
7663 let run = move |start: usize, end: usize| {
7664 for r in start..end {
7665 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7666 unsafe { *out_addr.at(r) += corr };
7668 }
7669 };
7670 dispatch_rows(pool, rows, &run);
7671}
7672
7673#[allow(clippy::too_many_arguments)]
7676fn q1t_range_a8w8(
7677 bytes: &[u8],
7678 gpr: usize,
7679 rp_off: usize,
7680 ent_off: usize,
7681 has_ov: bool,
7682 act: &SplitAct,
7683 x: &[f32],
7684 out: SendMut,
7685 start: usize,
7686 end: usize,
7687) {
7688 for r in start..end {
7689 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7690 for &(j, xv) in &act.outliers {
7691 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7692 }
7693 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7694 unsafe { *out.at(r) = acc };
7696 }
7697}
7698
7699#[allow(clippy::too_many_arguments)]
7702fn q1t_range_f32_batch(
7703 bytes: &[u8],
7704 gpr: usize,
7705 rp_off: usize,
7706 ent_off: usize,
7707 has_ov: bool,
7708 x: &[f32],
7709 out: SendMut,
7710 start: usize,
7711 end: usize,
7712) {
7713 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7714 let mut sg = [0f32; GROUP_SIZE];
7715 for r in start..end {
7716 let mut acc = 0f32;
7717 for g in 0..gpr {
7718 let off = (r * gpr + g) * TILE;
7719 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7720 let codes = &bytes[off + 2..off + TILE];
7721 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7722 for bi in 0..6 {
7723 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7724 }
7725 let lut = &SIGN5[codes[6] as usize];
7726 sg[30] = lut[0];
7727 sg[31] = lut[1];
7728 let mut gsum = 0f32;
7729 for k in 0..GROUP_SIZE {
7730 gsum += sg[k] * xg[k];
7731 }
7732 acc += s * gsum;
7733 }
7734 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7735 unsafe { *out.at(r) = acc };
7737 }
7738}
7739
7740fn q1t_matvec(
7744 bytes: &[u8],
7745 x: &[f32],
7746 rows: usize,
7747 cols: usize,
7748 out: &mut [f32],
7749 pool: Option<&Pool>,
7750) {
7751 debug_assert_eq!(out.len(), rows);
7752 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7753 let gpr = cols / GROUP_SIZE;
7754 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7755 let out_addr = SendMut(out.as_mut_ptr());
7756 if a8w8_enabled() {
7760 let act = split_act(x);
7761 let act = &act;
7762 let run = move |start: usize, end: usize| {
7763 for r in start..end {
7764 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7765 for &(j, xv) in &act.outliers {
7766 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7767 }
7768 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7769 unsafe { *out_addr.at(r) = acc };
7771 }
7772 };
7773 dispatch_rows(pool, rows, &run);
7774 return;
7775 }
7776 let run = move |start: usize, end: usize| {
7777 let mut sg = [0f32; GROUP_SIZE];
7781 for r in start..end {
7782 let mut acc = 0f32;
7783 for g in 0..gpr {
7784 let off = (r * gpr + g) * TILE;
7785 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7786 let codes = &bytes[off + 2..off + TILE];
7787 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7788 for bi in 0..6 {
7789 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7790 }
7791 let lut = &SIGN5[codes[6] as usize];
7792 sg[30] = lut[0];
7793 sg[31] = lut[1];
7794 let mut gsum = 0f32;
7795 for k in 0..GROUP_SIZE {
7796 gsum += sg[k] * xg[k];
7797 }
7798 acc += s * gsum;
7799 }
7800 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7801 unsafe { *out_addr.at(r) = acc };
7802 }
7803 };
7804 dispatch_rows(pool, rows, &run);
7805}
7806
7807#[cfg(target_arch = "aarch64")]
7813#[target_feature(enable = "neon,dotprod")]
7814unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7815 use core::arch::aarch64::*;
7816 use core::arch::asm;
7817 unsafe {
7819 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7820 let bytes_ptr = bytes.as_ptr();
7821 let row_off = r * gpr * TILE;
7822 let xp = [xa.as_ptr(), xb.as_ptr()];
7823 let mut acc = [0f32; 2];
7824 macro_rules! sdot2 {
7825 ($w0:expr, $w1:expr, $x:expr) => {{
7826 let x0 = vld1q_s8($x);
7827 let x1 = vld1q_s8($x.add(16));
7828 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7829 asm!(
7830 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7831 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7832 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7833 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7834 options(pure, nomem, nostack),
7835 );
7836 vaddvq_s32(vaddq_s32(a0, a1))
7837 }};
7838 }
7839 let gpr2 = gpr & !1;
7840 let mut gi = 0;
7841 while gi < gpr2 {
7842 let off0 = row_off + gi * TILE;
7843 let off1 = off0 + TILE;
7844 let s0 = f16_to_f32(u16::from_le_bytes([
7845 *bytes_ptr.add(off0),
7846 *bytes_ptr.add(off0 + 1),
7847 ]));
7848 let s1 = f16_to_f32(u16::from_le_bytes([
7849 *bytes_ptr.add(off1),
7850 *bytes_ptr.add(off1 + 1),
7851 ]));
7852 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7853 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7854 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7855 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7856 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7857 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7858 for k in 0..2 {
7859 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7860 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7861 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7862 }
7863 gi += 2;
7864 }
7865 if gi < gpr {
7866 let off = row_off + gi * TILE;
7867 let s = f16_to_f32(u16::from_le_bytes([
7868 *bytes_ptr.add(off),
7869 *bytes_ptr.add(off + 1),
7870 ]));
7871 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7872 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7873 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7874 for k in 0..2 {
7875 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7876 acc[k] += d as f32 * s;
7877 }
7878 }
7879 acc
7880 }
7881}
7882
7883fn q1t_matvec2(
7889 bytes: &[u8],
7890 x1: &[f32],
7891 x2: &[f32],
7892 rows: usize,
7893 cols: usize,
7894 o1: &mut [f32],
7895 o2: &mut [f32],
7896 pool: Option<&Pool>,
7897) {
7898 debug_assert_eq!(o1.len(), rows);
7899 debug_assert_eq!(o2.len(), rows);
7900 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7901 let gpr = cols / GROUP_SIZE;
7902 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7903 let out1 = SendMut(o1.as_mut_ptr());
7904 let out2 = SendMut(o2.as_mut_ptr());
7905 if a8w8_enabled() {
7906 let a1 = split_act(x1);
7907 let a2 = split_act(x2);
7908 let (a1, a2) = (&a1, &a2);
7909 let run = move |start: usize, end: usize| {
7910 for r in start..end {
7911 #[cfg(target_arch = "aarch64")]
7912 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7915 #[cfg(not(target_arch = "aarch64"))]
7916 let ds = [
7917 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7918 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7919 ];
7920 let mut acc1 = ds[0] * a1.sx;
7921 for &(j, xv) in &a1.outliers {
7922 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7923 }
7924 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7925 let mut acc2 = ds[1] * a2.sx;
7926 for &(j, xv) in &a2.outliers {
7927 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7928 }
7929 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7930 unsafe {
7932 *out1.at(r) = acc1;
7933 *out2.at(r) = acc2;
7934 }
7935 }
7936 };
7937 dispatch_rows(pool, rows, &run);
7938 return;
7939 }
7940 let run = move |start: usize, end: usize| {
7941 let mut sg = [0f32; GROUP_SIZE];
7944 for r in start..end {
7945 let mut acc1 = 0f32;
7946 let mut acc2 = 0f32;
7947 for g in 0..gpr {
7948 let off = (r * gpr + g) * TILE;
7949 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7950 let codes = &bytes[off + 2..off + TILE];
7951 for bi in 0..6 {
7952 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7953 }
7954 let lut = &SIGN5[codes[6] as usize];
7955 sg[30] = lut[0];
7956 sg[31] = lut[1];
7957 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7958 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7959 let mut gsum1 = 0f32;
7960 for k in 0..GROUP_SIZE {
7961 gsum1 += sg[k] * xg1[k];
7962 }
7963 acc1 += s * gsum1;
7964 let mut gsum2 = 0f32;
7965 for k in 0..GROUP_SIZE {
7966 gsum2 += sg[k] * xg2[k];
7967 }
7968 acc2 += s * gsum2;
7969 }
7970 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7971 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7972 unsafe {
7974 *out1.at(r) = acc1;
7975 *out2.at(r) = acc2;
7976 }
7977 }
7978 };
7979 dispatch_rows(pool, rows, &run);
7980}
7981
7982fn q1t_matmat(
7985 bytes: &[u8],
7986 xs: &[f32],
7987 b: usize,
7988 rows: usize,
7989 cols: usize,
7990 out: &mut [f32],
7991 pool: Option<&Pool>,
7992) {
7993 debug_assert_eq!(out.len(), b * rows);
7994 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7995 let gpr = cols / GROUP_SIZE;
7996 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7997 let out_addr = SendMut(out.as_mut_ptr());
7998 if a8w8_enabled() {
8002 let acts: Vec<SplitAct> = (0..b)
8003 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
8004 .collect();
8005 let acts = &acts;
8006 let run = move |start: usize, end: usize| {
8007 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
8011 for g in 0..gpr {
8012 let off = (r * gpr + g) * TILE;
8013 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8014 q1t_unpack_group_i8(
8015 bytes.as_ptr().wrapping_add(off + 2),
8016 &mut sg[g * GROUP_SIZE..],
8017 );
8018 }
8019 for bi in 0..b {
8020 let act = &acts[bi];
8021 let mut isum = 0f32;
8022 for g in 0..gpr {
8023 let d = q1t_i8dot32(
8024 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
8025 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
8026 );
8027 isum += d as f32 * sc[g];
8028 }
8029 let mut acc = isum * act.sx;
8030 for &(j, xv) in &act.outliers {
8031 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8032 }
8033 accs[bi] = acc;
8034 }
8035 if has_ov {
8039 let (c0, c1) = (
8040 q1t_rowptr(bytes, rp_off, r),
8041 q1t_rowptr(bytes, rp_off, r + 1),
8042 );
8043 for p in c0..c1 {
8044 let e = ent_off + p * 4;
8045 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8046 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8047 for bi in 0..b {
8048 accs[bi] += val * xs[bi * cols + col];
8049 }
8050 }
8051 }
8052 for bi in 0..b {
8053 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
8054 }
8055 }
8056 };
8057 dispatch_rows(pool, rows, &run);
8058 return;
8059 }
8060 let run = move |start: usize, end: usize| {
8061 let mut buf = vec![0f32; cols];
8062 for r in start..end {
8063 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
8064 for bi in 0..b {
8065 let xr = &xs[bi * cols..(bi + 1) * cols];
8066 let mut acc = 0f32;
8067 for j in 0..cols {
8068 acc += buf[j] * xr[j];
8069 }
8070 unsafe { *out_addr.at(bi * rows + r) = acc };
8071 }
8072 }
8073 };
8074 dispatch_rows(pool, rows, &run);
8075}
8076
8077fn q1_matvec(
8078 bytes: &[u8],
8079 x: &[f32],
8080 rows: usize,
8081 cols: usize,
8082 out: &mut [f32],
8083 pool: Option<&Pool>,
8084) {
8085 debug_assert_eq!(out.len(), rows);
8086 let gpr = cols / GROUP_SIZE;
8087 let out_addr = SendMut(out.as_mut_ptr());
8088 if a8w8_enabled() {
8089 let act = split_act(x);
8090 let gsum = q1_group_sums(&act.xq, gpr);
8091 let (act, gsum) = (&act, &gsum);
8092 let run = move |start: usize, end: usize| {
8093 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
8094 };
8095 dispatch_rows(pool, rows, &run);
8096 return;
8097 }
8098 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
8099 dispatch_rows(pool, rows, &run);
8100}
8101
8102#[allow(clippy::too_many_arguments)]
8104fn q1_matvec2(
8105 bytes: &[u8],
8106 x1: &[f32],
8107 x2: &[f32],
8108 rows: usize,
8109 cols: usize,
8110 o1: &mut [f32],
8111 o2: &mut [f32],
8112 pool: Option<&Pool>,
8113) {
8114 let gpr = cols / GROUP_SIZE;
8115 let p1 = SendMut(o1.as_mut_ptr());
8116 let p2 = SendMut(o2.as_mut_ptr());
8117 if a8w8_enabled() {
8118 let a1 = split_act(x1);
8119 let a2 = split_act(x2);
8120 let g1 = q1_group_sums(&a1.xq, gpr);
8121 let g2 = q1_group_sums(&a2.xq, gpr);
8122 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
8123 let run = move |start: usize, end: usize| {
8124 for r in start..end {
8125 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
8126 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
8127 for &(j, xv) in &a1.outliers {
8128 let (w, s) = q1_outlier(bytes, r, gpr, j);
8129 v1 += w * s * xv;
8130 }
8131 for &(j, xv) in &a2.outliers {
8132 let (w, s) = q1_outlier(bytes, r, gpr, j);
8133 v2 += w * s * xv;
8134 }
8135 unsafe {
8137 *p1.at(r) = v1;
8138 *p2.at(r) = v2;
8139 }
8140 }
8141 };
8142 dispatch_rows(pool, rows, &run);
8143 return;
8144 }
8145 let run = move |start: usize, end: usize| {
8146 for r in start..end {
8147 unsafe {
8149 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
8150 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
8151 }
8152 }
8153 };
8154 dispatch_rows(pool, rows, &run);
8155}
8156
8157#[allow(clippy::too_many_arguments)]
8159fn q1_matmat(
8160 bytes: &[u8],
8161 xs_all: &[f32],
8162 b: usize,
8163 rows: usize,
8164 cols: usize,
8165 out: &mut [f32],
8166 pool: Option<&Pool>,
8167) {
8168 debug_assert_eq!(out.len(), b * rows);
8169 let gpr = cols / GROUP_SIZE;
8170 let out_addr = SendMut(out.as_mut_ptr());
8171 if a8w8_enabled() {
8172 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
8173 .map(|bi| {
8174 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
8175 let gsum = q1_group_sums(&act.xq, gpr);
8176 (act, gsum)
8177 })
8178 .collect();
8179 let acts = &acts;
8180 #[cfg(target_arch = "x86_64")]
8181 let blocked_ok = avx2_enabled() && blocked_enabled();
8182 #[cfg(target_arch = "aarch64")]
8183 let blocked_ok = sdot_enabled() && blocked_enabled();
8184 let run = move |start: usize, end: usize| {
8185 for r in start..end {
8186 let mut bi = 0usize;
8187 #[cfg(target_arch = "aarch64")]
8190 if blocked_ok {
8191 while bi + 4 <= acts.len() {
8192 let xs = [
8193 acts[bi].0.xq.as_slice(),
8194 acts[bi + 1].0.xq.as_slice(),
8195 acts[bi + 2].0.xq.as_slice(),
8196 acts[bi + 3].0.xq.as_slice(),
8197 ];
8198 let gs = [
8199 acts[bi].1.as_slice(),
8200 acts[bi + 1].1.as_slice(),
8201 acts[bi + 2].1.as_slice(),
8202 acts[bi + 3].1.as_slice(),
8203 ];
8204 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
8205 for k in 0..4 {
8206 let (act, _) = &acts[bi + k];
8207 let mut acc = d[k] * act.sx;
8208 for &(j, xv) in &act.outliers {
8209 let (w, sc) = q1_outlier(bytes, r, gpr, j);
8210 acc += w * sc * xv;
8211 }
8212 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8214 }
8215 bi += 4;
8216 }
8217 }
8218 #[cfg(target_arch = "x86_64")]
8219 if blocked_ok {
8220 while bi + 4 <= acts.len() {
8221 let xs = [
8222 acts[bi].0.xq.as_slice(),
8223 acts[bi + 1].0.xq.as_slice(),
8224 acts[bi + 2].0.xq.as_slice(),
8225 acts[bi + 3].0.xq.as_slice(),
8226 ];
8227 let gs = [
8228 acts[bi].1.as_slice(),
8229 acts[bi + 1].1.as_slice(),
8230 acts[bi + 2].1.as_slice(),
8231 acts[bi + 3].1.as_slice(),
8232 ];
8233 let d = unsafe {
8234 if vnni_tiles_enabled() {
8235 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
8236 } else {
8237 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
8238 }
8239 };
8240 for k in 0..4 {
8241 let (act, _) = &acts[bi + k];
8242 let mut acc = d[k] * act.sx;
8243 for &(j, xv) in &act.outliers {
8244 let (w, sc) = q1_outlier(bytes, r, gpr, j);
8245 acc += w * sc * xv;
8246 }
8247 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8249 }
8250 bi += 4;
8251 }
8252 }
8253 while bi < acts.len() {
8254 let (act, gsum) = &acts[bi];
8255 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
8256 for &(j, xv) in &act.outliers {
8257 let (w, s) = q1_outlier(bytes, r, gpr, j);
8258 acc += w * s * xv;
8259 }
8260 unsafe { *out_addr.at(bi * rows + r) = acc };
8262 bi += 1;
8263 }
8264 }
8265 };
8266 dispatch_rows(pool, rows, &run);
8267 return;
8268 }
8269 let run = move |start: usize, end: usize| {
8270 for r in start..end {
8271 for bi in 0..b {
8272 let x = &xs_all[bi * cols..(bi + 1) * cols];
8273 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
8275 }
8276 }
8277 };
8278 dispatch_rows(pool, rows, &run);
8279}
8280
8281fn q4matvec(
8287 bytes: &[u8],
8288 x: &[f32],
8289 rows: usize,
8290 cols: usize,
8291 out: &mut [f32],
8292 pool: Option<&Pool>,
8293) {
8294 debug_assert_eq!(out.len(), rows);
8295 let (packed, scales) = q4_split(bytes, rows, cols);
8296 let gpr = cols / GROUP_SIZE;
8297 let out_addr = SendMut(out.as_mut_ptr());
8298
8299 if a8w8_enabled() {
8300 let act = split_act(x);
8301 let run = move |start: usize, end: usize| {
8302 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
8303 };
8304 dispatch_rows(pool, rows, &run);
8305 return;
8306 }
8307
8308 let run =
8309 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
8310 dispatch_rows(pool, rows, &run);
8311}
8312
8313#[inline]
8316#[allow(unreachable_code)]
8317#[cfg(target_arch = "x86_64")]
8322#[target_feature(enable = "avx2")]
8323unsafe fn dot_q4b_row_1x4_avx2(
8324 buf: &[u8],
8325 scales: &[u8],
8326 g0: usize,
8327 gpr: usize,
8328 xs: [&[i8]; 4],
8329) -> [f32; 4] {
8330 unsafe {
8332 use core::arch::x86_64::*;
8333 let ones = _mm256_set1_epi16(1);
8334 let mut acc = [0f32; 4];
8335 for gi in 0..gpr {
8336 let s = f16_to_f32(u16::from_le_bytes([
8337 scales[(g0 + gi) * 2],
8338 scales[(g0 + gi) * 2 + 1],
8339 ]));
8340 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8341 let aw = _mm256_abs_epi8(w);
8342 for (k, xq) in xs.iter().enumerate() {
8343 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8344 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8345 let d = _mm256_madd_epi16(p16, ones);
8346 let hi128 = _mm256_extracti128_si256::<1>(d);
8347 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8348 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8349 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8350 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
8351 }
8352 }
8353 acc
8354 }
8355}
8356
8357#[cfg(target_arch = "x86_64")]
8359#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8360unsafe fn dot_q4b_row_1x4_vnni(
8361 buf: &[u8],
8362 scales: &[u8],
8363 g0: usize,
8364 gpr: usize,
8365 xs: [&[i8]; 4],
8366) -> [f32; 4] {
8367 unsafe {
8369 use core::arch::x86_64::*;
8370 let mut acc = [0f32; 4];
8371 for gi in 0..gpr {
8372 let s = f16_to_f32(u16::from_le_bytes([
8373 scales[(g0 + gi) * 2],
8374 scales[(g0 + gi) * 2 + 1],
8375 ]));
8376 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8377 let aw = _mm256_abs_epi8(w);
8378 for (k, xq) in xs.iter().enumerate() {
8379 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8380 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8381 acc[k] += d as f32 * s;
8382 }
8383 }
8384 acc
8385 }
8386}
8387
8388#[cfg(target_arch = "x86_64")]
8394#[target_feature(enable = "avx2")]
8395unsafe fn dot_q4b_row_1x4_sx_avx2(
8396 buf: &[u8],
8397 scales: &[u8],
8398 g0: usize,
8399 gpr: usize,
8400 xs: [&[i8]; 4],
8401 sxs: [f32; 4],
8402) -> [f32; 4] {
8403 unsafe {
8405 use core::arch::x86_64::*;
8406 let ones = _mm256_set1_epi16(1);
8407 let mut acc = [0f32; 4];
8408 for gi in 0..gpr {
8409 let s = f16_to_f32(u16::from_le_bytes([
8410 scales[(g0 + gi) * 2],
8411 scales[(g0 + gi) * 2 + 1],
8412 ]));
8413 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8414 let aw = _mm256_abs_epi8(w);
8415 for (k, xq) in xs.iter().enumerate() {
8416 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8417 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8418 let d = _mm256_madd_epi16(p16, ones);
8419 let hi128 = _mm256_extracti128_si256::<1>(d);
8420 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8421 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8422 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8423 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
8424 }
8425 }
8426 acc
8427 }
8428}
8429
8430#[cfg(target_arch = "x86_64")]
8433#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8434unsafe fn dot_q4b_row_1x4_sx_vnni(
8435 buf: &[u8],
8436 scales: &[u8],
8437 g0: usize,
8438 gpr: usize,
8439 xs: [&[i8]; 4],
8440 sxs: [f32; 4],
8441) -> [f32; 4] {
8442 unsafe {
8444 use core::arch::x86_64::*;
8445 let mut acc = [0f32; 4];
8446 for gi in 0..gpr {
8447 let s = f16_to_f32(u16::from_le_bytes([
8448 scales[(g0 + gi) * 2],
8449 scales[(g0 + gi) * 2 + 1],
8450 ]));
8451 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8452 let aw = _mm256_abs_epi8(w);
8453 for (k, xq) in xs.iter().enumerate() {
8454 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8455 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8456 acc[k] += (d as f32 * sxs[k]) * s;
8457 }
8458 }
8459 acc
8460 }
8461}
8462
8463#[allow(unreachable_code)]
8464fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8465 #[cfg(target_arch = "aarch64")]
8466 unsafe {
8467 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
8468 }
8469 #[cfg(target_arch = "x86_64")]
8470 unsafe {
8471 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
8472 }
8473 let mut acc = 0f32;
8474 for gi in 0..gpr {
8475 let g = g0 + gi;
8476 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8477 let mut d = 0i32;
8478 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8479 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
8480 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
8481 }
8482 acc += d as f32 * s;
8483 }
8484 acc
8485}
8486
8487#[inline]
8489#[allow(unreachable_code)]
8490fn dot_q4_row_i8_2(
8491 packed: &[u8],
8492 scales: &[u8],
8493 g0: usize,
8494 gpr: usize,
8495 xq1: &[i8],
8496 xq2: &[i8],
8497) -> (f32, f32) {
8498 #[cfg(target_arch = "aarch64")]
8499 unsafe {
8500 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
8501 }
8502 #[cfg(target_arch = "x86_64")]
8503 unsafe {
8504 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
8505 }
8506 (
8507 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
8508 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
8509 )
8510}
8511
8512#[allow(clippy::too_many_arguments)]
8515fn q4_range_a8w8(
8516 packed: &[u8],
8517 scales: &[u8],
8518 gpr: usize,
8519 cols: usize,
8520 act: &SplitAct,
8521 out: SendMut,
8522 start: usize,
8523 end: usize,
8524) {
8525 for r in start..end {
8526 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
8527 for &(j, xv) in &act.outliers {
8529 let flat = r * cols + j;
8530 let byte = packed[flat / 2];
8531 let nib = if flat & 1 == 0 {
8532 byte & 0x0F
8533 } else {
8534 byte >> 4
8535 };
8536 let s = f16_to_f32(u16::from_le_bytes([
8537 scales[(flat / GROUP_SIZE) * 2],
8538 scales[(flat / GROUP_SIZE) * 2 + 1],
8539 ]));
8540 acc += ((nib as i32 - 8) as f32) * s * xv;
8541 }
8542 unsafe { *out.at(r) = acc };
8544 }
8545}
8546
8547#[allow(clippy::too_many_arguments)]
8550fn q4_range2_a8w8(
8551 packed: &[u8],
8552 scales: &[u8],
8553 gpr: usize,
8554 cols: usize,
8555 a1: &SplitAct,
8556 a2: &SplitAct,
8557 p1: SendMut,
8558 p2: SendMut,
8559 start: usize,
8560 end: usize,
8561) {
8562 for r in start..end {
8563 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
8564 let mut acc1 = s1 * a1.sx;
8565 let mut acc2 = s2 * a2.sx;
8566 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
8568 for &(j, xv) in outliers {
8569 let flat = r * cols + j;
8570 let byte = packed[flat / 2];
8571 let nib = if flat & 1 == 0 {
8572 byte & 0x0F
8573 } else {
8574 byte >> 4
8575 };
8576 let s = f16_to_f32(u16::from_le_bytes([
8577 scales[(flat / GROUP_SIZE) * 2],
8578 scales[(flat / GROUP_SIZE) * 2 + 1],
8579 ]));
8580 *acc += ((nib as i32 - 8) as f32) * s * xv;
8581 }
8582 };
8583 fix(&a1.outliers, &mut acc1);
8584 fix(&a2.outliers, &mut acc2);
8585 unsafe {
8587 *p1.at(r) = acc1;
8588 *p2.at(r) = acc2;
8589 }
8590 }
8591}
8592
8593fn q4_range_f32(
8595 packed: &[u8],
8596 scales: &[u8],
8597 gpr: usize,
8598 x: &[f32],
8599 out: SendMut,
8600 start: usize,
8601 end: usize,
8602) {
8603 for r in start..end {
8604 let mut acc = 0f32;
8605 for gi in 0..gpr {
8606 let g = r * gpr + gi;
8607 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8608 let pk = &packed[g * 16..(g + 1) * 16];
8609 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8610 let mut ga = 0f32;
8611 for (k, &b) in pk.iter().enumerate() {
8612 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
8613 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
8614 }
8615 acc += ga * s;
8616 }
8617 unsafe { *out.at(r) = acc };
8619 }
8620}
8621
8622#[allow(clippy::too_many_arguments)]
8626fn q4matvec2(
8627 bytes: &[u8],
8628 x1: &[f32],
8629 x2: &[f32],
8630 rows: usize,
8631 cols: usize,
8632 o1: &mut [f32],
8633 o2: &mut [f32],
8634 pool: Option<&Pool>,
8635) {
8636 debug_assert_eq!(o1.len(), rows);
8637 debug_assert_eq!(o2.len(), rows);
8638 let (packed, scales) = q4_split(bytes, rows, cols);
8639 let gpr = cols / GROUP_SIZE;
8640
8641 if a8w8_enabled() {
8642 let a1 = split_act(x1);
8643 let a2 = split_act(x2);
8644 let p1 = SendMut(o1.as_mut_ptr());
8645 let p2 = SendMut(o2.as_mut_ptr());
8646 let run = move |start: usize, end: usize| {
8647 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
8648 };
8649 dispatch_rows(pool, rows, &run);
8650 return;
8651 }
8652
8653 let p1 = SendMut(o1.as_mut_ptr());
8654 let p2 = SendMut(o2.as_mut_ptr());
8655 let run = move |start: usize, end: usize| {
8656 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
8657 };
8658 dispatch_rows(pool, rows, &run);
8659}
8660
8661#[allow(clippy::too_many_arguments)]
8663fn q4_range2_f32(
8664 packed: &[u8],
8665 scales: &[u8],
8666 gpr: usize,
8667 x1: &[f32],
8668 x2: &[f32],
8669 p1: SendMut,
8670 p2: SendMut,
8671 start: usize,
8672 end: usize,
8673) {
8674 for r in start..end {
8675 let (mut acc1, mut acc2) = (0f32, 0f32);
8676 for gi in 0..gpr {
8677 let g = r * gpr + gi;
8678 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8679 let pk = &packed[g * 16..(g + 1) * 16];
8680 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8681 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8682 let (mut g1, mut g2) = (0f32, 0f32);
8683 for (k, &b) in pk.iter().enumerate() {
8684 let wl = (b & 0x0F) as f32 - 8.0;
8685 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
8686 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
8687 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
8688 }
8689 acc1 += g1 * s;
8690 acc2 += g2 * s;
8691 }
8692 unsafe {
8694 *p1.at(r) = acc1;
8695 *p2.at(r) = acc2;
8696 }
8697 }
8698}
8699
8700thread_local! {
8701 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8704 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8705}
8706
8707#[allow(clippy::too_many_arguments)]
8713fn q4matmat(
8714 bytes: &[u8],
8715 xs_all: &[f32],
8716 b: usize,
8717 rows: usize,
8718 cols: usize,
8719 out: &mut [f32],
8720 pool: Option<&Pool>,
8721) {
8722 debug_assert_eq!(xs_all.len(), b * cols);
8723 debug_assert_eq!(out.len(), b * rows);
8724 let (packed, scales) = q4_split(bytes, rows, cols);
8725 let gpr = cols / GROUP_SIZE;
8726 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8727
8728 if a8w8_enabled() {
8729 let acts: Vec<SplitAct> = (0..b)
8730 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8731 .collect();
8732 let acts = &acts;
8733 let out_addr = SendMut(out.as_mut_ptr());
8734 let run = move |start: usize, end: usize| {
8735 ROW_I8.with(|rb| {
8736 let mut buf = rb.borrow_mut();
8737 buf.resize(cols, 0);
8738 for r in start..end {
8739 for gi in 0..gpr {
8743 let g = r * gpr + gi;
8744 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8745 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
8746 buf[gi * GROUP_SIZE + k * 2 + 1] =
8747 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
8748 }
8749 }
8750 let mut bi = 0usize;
8751 #[cfg(target_arch = "x86_64")]
8752 if avx2_enabled() && blocked_enabled() {
8753 while bi + 4 <= acts.len() {
8754 let xs = [
8755 acts[bi].xq.as_slice(),
8756 acts[bi + 1].xq.as_slice(),
8757 acts[bi + 2].xq.as_slice(),
8758 acts[bi + 3].xq.as_slice(),
8759 ];
8760 let d = unsafe {
8761 if vnni_tiles_enabled() {
8762 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
8763 } else {
8764 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
8765 }
8766 };
8767 for k in 0..4 {
8768 let act = &acts[bi + k];
8769 let mut acc = d[k] * act.sx;
8770 for &(j, xv) in &act.outliers {
8771 acc += (buf[j] as i8) as f32
8772 * gscale((r * cols + j) / GROUP_SIZE)
8773 * xv;
8774 }
8775 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8777 }
8778 bi += 4;
8779 }
8780 }
8781 while bi < acts.len() {
8782 let act = &acts[bi];
8783 let mut acc = 0f32;
8784 for gi in 0..gpr {
8785 let d = dot_i8_i8(
8786 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8787 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8788 );
8789 acc += d as f32 * gscale(r * gpr + gi);
8790 }
8791 acc *= act.sx;
8792 for &(j, xv) in &act.outliers {
8794 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
8795 }
8796 unsafe { *out_addr.at(bi * rows + r) = acc };
8798 bi += 1;
8799 }
8800 }
8801 })
8802 };
8803 dispatch_rows(pool, rows, &run);
8804 return;
8805 }
8806
8807 let out_addr = SendMut(out.as_mut_ptr());
8808 let run = move |start: usize, end: usize| {
8809 ROW_F32.with(|rb| {
8810 let mut buf = rb.borrow_mut();
8811 buf.resize(cols, 0.0);
8812 for r in start..end {
8813 for gi in 0..gpr {
8816 let g = r * gpr + gi;
8817 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8818 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8819 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8820 }
8821 }
8822 for bi in 0..b {
8823 let x = &xs_all[bi * cols..(bi + 1) * cols];
8824 let mut acc = 0f32;
8825 for gi in 0..gpr {
8826 let mut ga = 0f32;
8827 for k in 0..GROUP_SIZE / 2 {
8832 let e = gi * GROUP_SIZE + k * 2;
8833 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8834 }
8835 acc += ga * gscale(r * gpr + gi);
8836 }
8837 unsafe { *out_addr.at(bi * rows + r) = acc };
8839 }
8840 }
8841 })
8842 };
8843 dispatch_rows(pool, rows, &run);
8844}
8845
8846#[allow(clippy::too_many_arguments)]
8851fn vbitmatmat(
8852 bytes: &[u8],
8853 offsets: &[usize],
8854 xs_all: &[f32],
8855 b: usize,
8856 rows: usize,
8857 cols: usize,
8858 out: &mut [f32],
8859 pool: Option<&Pool>,
8860) {
8861 debug_assert_eq!(xs_all.len(), b * cols);
8862 debug_assert_eq!(out.len(), b * rows);
8863 debug_assert_eq!(offsets.len(), rows + 1);
8864 let ng = cols / GROUP_SIZE;
8865 let bits = &bytes[..rows];
8866 let sc_off = rows;
8867 let gscale = |r: usize, g: usize| {
8868 let so = (r * ng + g) * 2;
8869 f16_to_f32(u16::from_le_bytes([
8870 bytes[sc_off + so],
8871 bytes[sc_off + so + 1],
8872 ]))
8873 };
8874
8875 let decode_f32 = |r: usize, dst: &mut [f32]| {
8877 let bw = bits[r] as usize;
8878 let l = ((1i32 << (bw - 1)) - 1) as f32;
8879 let data = &bytes[offsets[r]..offsets[r + 1]];
8880 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8881 for d in dst.iter_mut() {
8882 while nbits < bw {
8883 acc = (acc << 8) | data[idx] as u64;
8884 idx += 1;
8885 nbits += 8;
8886 }
8887 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8888 nbits -= bw;
8889 *d = u - l;
8890 }
8891 };
8892
8893 if a8w8_enabled() {
8894 let acts: Vec<SplitAct> = (0..b)
8895 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8896 .collect();
8897 let acts = &acts;
8898 let out_addr = SendMut(out.as_mut_ptr());
8899 let run = move |start: usize, end: usize| {
8900 for r in start..end {
8901 let bw = bits[r] as usize;
8902 if bw == 8 {
8903 ROW_F32.with(|rb| {
8906 let mut buf = rb.borrow_mut();
8907 buf.resize(cols, 0.0);
8908 decode_f32(r, &mut buf);
8909 for bi in 0..b {
8910 let x = &xs_all[bi * cols..(bi + 1) * cols];
8911 let mut dot = 0f32;
8912 for g in 0..ng {
8913 let mut gd = 0f32;
8914 for k in 0..GROUP_SIZE {
8915 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8916 }
8917 dot += gd * gscale(r, g);
8918 }
8919 unsafe { *out_addr.at(bi * rows + r) = dot };
8921 }
8922 });
8923 continue;
8924 }
8925 let l = (1i32 << (bw - 1)) - 1;
8926 let data = &bytes[offsets[r]..offsets[r + 1]];
8927 ROW_I8.with(|rb| {
8928 let mut buf = rb.borrow_mut();
8929 buf.resize(cols, 0);
8930 #[inline(always)]
8931 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8932 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8933 let u = unpack8::<B>(&data[blk * B..]);
8934 for k in 0..8 {
8935 chunk[k] = (u[k] - l) as i8 as u8;
8936 }
8937 }
8938 }
8939 match bw {
8940 3 => fill::<3>(data, l, &mut buf),
8941 4 => vbit_fill4(data, &mut buf),
8942 5 => fill::<5>(data, l, &mut buf),
8943 6 => fill::<6>(data, l, &mut buf),
8944 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8945 }
8946 let mut bi = 0usize;
8947 #[cfg(target_arch = "x86_64")]
8951 if avx2_enabled() && blocked_enabled() {
8952 while bi + 4 <= acts.len() {
8953 let xs = [
8954 acts[bi].xq.as_slice(),
8955 acts[bi + 1].xq.as_slice(),
8956 acts[bi + 2].xq.as_slice(),
8957 acts[bi + 3].xq.as_slice(),
8958 ];
8959 let sxs = [
8960 acts[bi].sx,
8961 acts[bi + 1].sx,
8962 acts[bi + 2].sx,
8963 acts[bi + 3].sx,
8964 ];
8965 let d = unsafe {
8966 if vnni_tiles_enabled() {
8967 dot_q4b_row_1x4_sx_vnni(
8968 &buf,
8969 &bytes[sc_off..],
8970 r * ng,
8971 ng,
8972 xs,
8973 sxs,
8974 )
8975 } else {
8976 dot_q4b_row_1x4_sx_avx2(
8977 &buf,
8978 &bytes[sc_off..],
8979 r * ng,
8980 ng,
8981 xs,
8982 sxs,
8983 )
8984 }
8985 };
8986 for k in 0..4 {
8987 let act = &acts[bi + k];
8988 let mut dot = d[k];
8989 for &(j, xv) in &act.outliers {
8990 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8991 }
8992 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8994 }
8995 bi += 4;
8996 }
8997 }
8998 while bi < acts.len() {
8999 let act = &acts[bi];
9000 let mut dot = 0f32;
9001 for g in 0..ng {
9002 let d = dot_i8_i8(
9003 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9004 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9005 ) as f32
9006 * act.sx;
9007 dot += d * gscale(r, g);
9008 }
9009 for &(j, xv) in &act.outliers {
9010 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9011 }
9012 unsafe { *out_addr.at(bi * rows + r) = dot };
9014 bi += 1;
9015 }
9016 });
9017 }
9018 };
9019 dispatch_rows(pool, rows, &run);
9020 return;
9021 }
9022
9023 let out_addr = SendMut(out.as_mut_ptr());
9024 let run = move |start: usize, end: usize| {
9025 ROW_F32.with(|rb| {
9026 let mut buf = rb.borrow_mut();
9027 buf.resize(cols, 0.0);
9028 for r in start..end {
9029 decode_f32(r, &mut buf);
9030 for bi in 0..b {
9031 let x = &xs_all[bi * cols..(bi + 1) * cols];
9032 let mut dot = 0f32;
9033 for g in 0..ng {
9034 let mut gd = 0f32;
9035 for k in 0..GROUP_SIZE {
9036 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9037 }
9038 dot += gd * gscale(r, g);
9039 }
9040 unsafe { *out_addr.at(bi * rows + r) = dot };
9042 }
9043 }
9044 })
9045 };
9046 dispatch_rows(pool, rows, &run);
9047}
9048
9049pub(crate) fn gpu_batch_job<'a>(
9053 t: &'a QTensor,
9054 x: &[f32],
9055) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
9056 match t {
9057 QTensor::Mapped {
9058 model,
9059 idx,
9060 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
9061 rows,
9062 cols,
9063 row_scale,
9064 col_field,
9065 ..
9066 } => Some((
9067 model.clone(),
9068 crate::gpu::BatchJob {
9069 idx: *idx,
9070 rows: *rows,
9071 cols: *cols,
9072 row_scale,
9073 xs: prescale(x, col_field, *dt).into_owned(),
9074 layout: crate::gpu::BatchLayout::Q8,
9075 },
9076 )),
9077 QTensor::Mapped {
9079 model,
9080 idx,
9081 dtype: TensorDtype::Q1,
9082 rows,
9083 cols,
9084 ..
9085 } => Some((
9086 model.clone(),
9087 crate::gpu::BatchJob {
9088 idx: *idx,
9089 rows: *rows,
9090 cols: *cols,
9091 row_scale: &[],
9092 xs: x.to_vec(),
9093 layout: crate::gpu::BatchLayout::Q1,
9094 },
9095 )),
9096 QTensor::Mapped {
9101 model,
9102 idx,
9103 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
9104 rows,
9105 cols,
9106 ..
9107 } => Some((
9108 model.clone(),
9109 crate::gpu::BatchJob {
9110 idx: *idx,
9111 rows: *rows,
9112 cols: *cols,
9113 row_scale: &[],
9114 xs: x.to_vec(),
9115 layout: if *dt == TensorDtype::Q4Tiled {
9116 crate::gpu::BatchLayout::Q4t
9117 } else {
9118 crate::gpu::BatchLayout::Q4tp
9119 },
9120 },
9121 )),
9122 _ => None,
9123 }
9124}
9125
9126thread_local! {
9127 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9128 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9129}
9130
9131pub(crate) fn prescale<'a>(
9132 x: &'a [f32],
9133 col_field: &[f32],
9134 dtype: TensorDtype,
9135) -> std::borrow::Cow<'a, [f32]> {
9136 if dtype == TensorDtype::Q8_2f {
9137 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
9138 } else {
9139 std::borrow::Cow::Borrowed(x)
9140 }
9141}
9142
9143pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
9146 x: &[f32],
9147 col_field: &[f32],
9148 dtype: TensorDtype,
9149 buf_id: u8,
9150 f: F,
9151) -> R {
9152 if dtype == TensorDtype::Q8_2f {
9153 if buf_id == 1 {
9154 PRESCALE_BUF1.with(|b| {
9155 let mut buf = b.borrow_mut();
9156 buf.clear();
9157 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9158 f(&buf)
9159 })
9160 } else {
9161 PRESCALE_BUF2.with(|b| {
9162 let mut buf = b.borrow_mut();
9163 buf.clear();
9164 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9165 f(&buf)
9166 })
9167 }
9168 } else {
9169 f(x)
9170 }
9171}
9172
9173#[cfg(target_arch = "x86_64")]
9178pub(crate) fn avx2_enabled() -> bool {
9179 use std::sync::OnceLock;
9180 static ON: OnceLock<bool> = OnceLock::new();
9181 *ON.get_or_init(|| {
9182 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
9183 && std::arch::is_x86_feature_detected!("avx2")
9184 && std::arch::is_x86_feature_detected!("fma")
9185 })
9186}
9187
9188#[cfg(target_arch = "x86_64")]
9193fn avx2_a8w8_enabled() -> bool {
9194 use std::sync::OnceLock;
9195 static ON: OnceLock<bool> = OnceLock::new();
9196 *ON.get_or_init(|| {
9197 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
9198 })
9199}
9200
9201#[inline]
9205pub(crate) fn a8w8_enabled() -> bool {
9206 #[cfg(target_arch = "aarch64")]
9207 {
9208 sdot_enabled()
9209 }
9210 #[cfg(target_arch = "x86_64")]
9211 {
9212 avx2_a8w8_enabled()
9213 }
9214 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
9215 {
9216 false
9217 }
9218}
9219
9220#[inline]
9223#[allow(unreachable_code)]
9224fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
9225 #[cfg(target_arch = "aarch64")]
9226 unsafe {
9227 return dot_i8_sdot(w, xq);
9228 }
9229 #[cfg(target_arch = "x86_64")]
9230 unsafe {
9231 if avx512vnni_enabled() {
9232 return dot_i8_i8_vnni(w, xq);
9233 }
9234 return dot_i8_i8_avx2(w, xq);
9235 }
9236 w.iter()
9237 .zip(xq)
9238 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
9239 .sum()
9240}
9241
9242#[cfg(target_arch = "x86_64")]
9246fn avx512vnni_enabled() -> bool {
9247 use std::sync::OnceLock;
9248 static ON: OnceLock<bool> = OnceLock::new();
9249 *ON.get_or_init(|| {
9250 std::env::var("CMF_AVX512")
9251 .map(|v| v != "0")
9252 .unwrap_or(true)
9253 && std::arch::is_x86_feature_detected!("avx512f")
9254 && std::arch::is_x86_feature_detected!("avx512bw")
9255 && std::arch::is_x86_feature_detected!("avx512vl")
9256 && std::arch::is_x86_feature_detected!("avx512vnni")
9257 })
9258}
9259
9260#[cfg(target_arch = "x86_64")]
9268fn vnni_tiles_enabled() -> bool {
9269 use std::sync::OnceLock;
9270 static ON: OnceLock<bool> = OnceLock::new();
9271 *ON.get_or_init(|| {
9272 std::env::var("CMF_VNNI_TILES")
9273 .map(|v| v != "0")
9274 .unwrap_or(true)
9275 && avx512vnni_enabled()
9276 })
9277}
9278
9279#[cfg(target_arch = "x86_64")]
9284#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9285#[inline]
9286unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
9287 unsafe {
9289 use core::arch::x86_64::*;
9290 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
9291 let hi128 = _mm256_extracti128_si256::<1>(d);
9292 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9293 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9294 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9295 _mm_cvtsi128_si32(s32)
9296 }
9297}
9298
9299#[cfg(target_arch = "x86_64")]
9304#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9305unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9306 unsafe {
9308 use core::arch::x86_64::*;
9309 let n = w.len();
9310 let mut j = 0usize;
9311 let mut total: i32;
9312 {
9317 #[inline(always)]
9318 unsafe fn step(
9319 w: *const u8,
9320 x: *const i8,
9321 acc: core::arch::x86_64::__m512i,
9322 ) -> core::arch::x86_64::__m512i {
9323 unsafe {
9324 use core::arch::x86_64::*;
9325 let wv = _mm512_loadu_si512(w as *const _);
9326 let xv = _mm512_loadu_si512(x as *const _);
9327 let aw = _mm512_abs_epi8(wv);
9328 let neg = _mm512_movepi8_mask(wv);
9329 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
9330 _mm512_dpbusd_epi32(acc, aw, sx)
9331 }
9332 }
9333 let (mut a0, mut a1, mut a2, mut a3) = (
9334 _mm512_setzero_si512(),
9335 _mm512_setzero_si512(),
9336 _mm512_setzero_si512(),
9337 _mm512_setzero_si512(),
9338 );
9339 while j + 256 <= n {
9340 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9341 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
9342 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
9343 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
9344 j += 256;
9345 }
9346 while j + 64 <= n {
9347 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9348 j += 64;
9349 }
9350 let s01 = _mm512_add_epi32(a0, a1);
9351 let s23 = _mm512_add_epi32(a2, a3);
9352 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
9353 }
9354 if j + 32 <= n {
9356 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9357 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9358 let d = _mm256_dpbusd_epi32(
9359 _mm256_setzero_si256(),
9360 _mm256_abs_epi8(wv),
9361 _mm256_sign_epi8(xv, wv),
9362 );
9363 let hi128 = _mm256_extracti128_si256::<1>(d);
9364 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9365 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9366 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9367 total += _mm_cvtsi128_si32(s32);
9368 j += 32;
9369 }
9370 while j < n {
9371 total += (w[j] as i8) as i32 * xq[j] as i32;
9372 j += 1;
9373 }
9374 total
9375 }
9376}
9377
9378#[cfg(target_arch = "x86_64")]
9380#[target_feature(enable = "avx2,fma")]
9381unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
9382 unsafe {
9384 use core::arch::x86_64::*;
9385 let n = x.len();
9386 let wp = w.as_ptr();
9387 let xp = x.as_ptr();
9388 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
9389 let mut j = 0usize;
9390 while j + 16 <= n {
9391 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
9392 let lo = _mm256_cvtepi8_epi32(wb);
9393 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
9394 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
9395 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
9396 j += 16;
9397 }
9398 let acc = _mm256_add_ps(a0, a1);
9399 let hi128 = _mm256_extractf128_ps::<1>(acc);
9400 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
9401 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
9402 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
9403 let mut sum = _mm_cvtss_f32(s32);
9404 while j < n {
9405 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
9406 j += 1;
9407 }
9408 sum
9409 }
9410}
9411
9412#[cfg(target_arch = "x86_64")]
9417#[target_feature(enable = "avx2")]
9418unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
9419 unsafe {
9421 use core::arch::x86_64::*;
9422 let n = w.len();
9423 let ones = _mm256_set1_epi16(1);
9424 let mut acc = _mm256_setzero_si256();
9425 let mut j = 0usize;
9426 while j + 32 <= n {
9427 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9428 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9429 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9430 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
9431 j += 32;
9432 }
9433 let hi128 = _mm256_extracti128_si256::<1>(acc);
9434 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
9435 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9436 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9437 let mut s = _mm_cvtsi128_si32(s32);
9438 while j < n {
9439 s += (w[j] as i8) as i32 * xq[j] as i32;
9440 j += 1;
9441 }
9442 s
9443 }
9444}
9445
9446#[cfg(target_arch = "aarch64")]
9450#[target_feature(enable = "neon,i8mm")]
9451unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9452 unsafe {
9454 use core::arch::aarch64::*;
9455 use core::arch::asm;
9456 let n = w0.len();
9457 let w0p = w0.as_ptr() as *const i8;
9458 let w1p = w1.as_ptr() as *const i8;
9459 let mut acc01 = vdupq_n_s32(0);
9462 let mut acc23 = vdupq_n_s32(0);
9463 let mut i = 0usize;
9464 while i + 8 <= n {
9465 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
9466 let xb01 = vcombine_s8(
9467 vld1_s8(xs[0].as_ptr().add(i)),
9468 vld1_s8(xs[1].as_ptr().add(i)),
9469 );
9470 let xb23 = vcombine_s8(
9471 vld1_s8(xs[2].as_ptr().add(i)),
9472 vld1_s8(xs[3].as_ptr().add(i)),
9473 );
9474 asm!(
9475 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
9476 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
9477 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
9478 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
9479 options(pure, nomem, nostack),
9480 );
9481 i += 8;
9482 }
9483 let mut out = [[0i32; 4]; 2];
9484 let a01: [i32; 4] = core::mem::transmute(acc01);
9485 let a23: [i32; 4] = core::mem::transmute(acc23);
9486 out[0][0] = a01[0];
9487 out[0][1] = a01[1];
9488 out[1][0] = a01[2];
9489 out[1][1] = a01[3];
9490 out[0][2] = a23[0];
9491 out[0][3] = a23[1];
9492 out[1][2] = a23[2];
9493 out[1][3] = a23[3];
9494 if i < n {
9495 for (k, x) in xs.iter().enumerate() {
9496 for j in i..n {
9497 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9498 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9499 }
9500 }
9501 }
9502 out
9503 }
9504}
9505
9506#[cfg(target_arch = "aarch64")]
9510#[target_feature(enable = "neon,dotprod")]
9511unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9512 unsafe {
9514 use core::arch::aarch64::*;
9515 use core::arch::asm;
9516 let n = w0.len();
9517 let w0p = w0.as_ptr() as *const i8;
9518 let w1p = w1.as_ptr() as *const i8;
9519 let mut acc = [[vdupq_n_s32(0); 4]; 2];
9520 let mut i = 0usize;
9521 while i + 16 <= n {
9522 let wv0 = vld1q_s8(w0p.add(i));
9523 let wv1 = vld1q_s8(w1p.add(i));
9524 for (k, x) in xs.iter().enumerate() {
9525 let xv = vld1q_s8(x.as_ptr().add(i));
9526 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
9527 asm!(
9528 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
9529 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
9530 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9531 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
9532 options(pure, nomem, nostack),
9533 );
9534 acc[0][k] = a0;
9535 acc[1][k] = a1;
9536 }
9537 i += 16;
9538 }
9539 let mut out = [[0i32; 4]; 2];
9540 for r in 0..2 {
9541 for k in 0..4 {
9542 out[r][k] = vaddvq_s32(acc[r][k]);
9543 }
9544 }
9545 if i < n {
9546 for (k, x) in xs.iter().enumerate() {
9547 for j in i..n {
9548 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9549 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9550 }
9551 }
9552 }
9553 out
9554 }
9555}
9556
9557#[cfg(target_arch = "x86_64")]
9563#[target_feature(enable = "avx2")]
9564unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9565 unsafe {
9567 use core::arch::x86_64::*;
9568 let n = w0.len();
9569 let ones = _mm256_set1_epi16(1);
9570 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
9571 let mut j = 0usize;
9572 while j + 32 <= n {
9573 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
9574 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
9575 let aw0 = _mm256_abs_epi8(wv0);
9576 let aw1 = _mm256_abs_epi8(wv1);
9577 for (k, x) in xs.iter().enumerate() {
9578 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
9579 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
9580 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
9581 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
9582 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
9583 }
9584 j += 32;
9585 }
9586 let mut out = [[0i32; 4]; 2];
9587 for r in 0..2 {
9588 for k in 0..4 {
9589 let a = acc[r][k];
9590 let hi128 = _mm256_extracti128_si256::<1>(a);
9591 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
9592 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9593 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9594 out[r][k] = _mm_cvtsi128_si32(s32);
9595 }
9596 }
9597 if j < n {
9598 for (k, x) in xs.iter().enumerate() {
9599 for i in j..n {
9600 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
9601 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
9602 }
9603 }
9604 }
9605 out
9606 }
9607}
9608
9609#[cfg(target_arch = "x86_64")]
9614#[inline]
9615fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
9616 let dot = if avx512vnni_enabled() && row.len() >= 64 {
9617 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
9618 } else {
9619 unsafe { dot_i8_i8_avx2(row, &act.xq) }
9620 };
9621 let mut acc = dot as f32 * act.sx;
9622 for &(j, xv) in &act.outliers {
9623 acc += (row[j] as i8) as f32 * xv;
9624 }
9625 acc
9626}
9627
9628#[cfg(target_arch = "x86_64")]
9632#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9633unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9634 unsafe {
9636 use core::arch::x86_64::*;
9637 let n = w.len();
9638 let flip = _mm512_set1_epi8(-128); #[inline(always)]
9640 unsafe fn step(
9641 w: *const u8,
9642 x: *const i8,
9643 flip: core::arch::x86_64::__m512i,
9644 acc: core::arch::x86_64::__m512i,
9645 ) -> core::arch::x86_64::__m512i {
9646 unsafe {
9647 use core::arch::x86_64::*;
9648 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
9649 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
9650 }
9651 }
9652 let (mut a0, mut a1, mut a2, mut a3) = (
9653 _mm512_setzero_si512(),
9654 _mm512_setzero_si512(),
9655 _mm512_setzero_si512(),
9656 _mm512_setzero_si512(),
9657 );
9658 let mut j = 0usize;
9659 while j + 256 <= n {
9660 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9661 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
9662 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
9663 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
9664 j += 256;
9665 }
9666 while j + 64 <= n {
9667 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9668 j += 64;
9669 }
9670 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
9671 _mm512_add_epi32(a0, a1),
9672 _mm512_add_epi32(a2, a3),
9673 ));
9674 while j < n {
9676 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
9677 j += 1;
9678 }
9679 total
9680 }
9681}
9682
9683#[cfg(target_arch = "x86_64")]
9689#[target_feature(enable = "avx2")]
9690unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9691 unsafe {
9694 use core::arch::x86_64::*;
9695 let lomask = _mm_set1_epi8(0x0F);
9696 let eight = _mm256_set1_epi8(8);
9697 let ones = _mm256_set1_epi16(1);
9698 let mut acc = 0f32;
9699 for gi in 0..gpr {
9700 let g = g0 + gi;
9701 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9702 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9703 let lo = _mm_and_si128(b, lomask);
9704 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9705 let w = _mm256_sub_epi8(
9706 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9707 eight,
9708 );
9709 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9710 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
9711 let d = _mm256_madd_epi16(p16, ones);
9712 let hi128 = _mm256_extracti128_si256::<1>(d);
9713 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9714 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9715 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9716 acc += _mm_cvtsi128_si32(s32) as f32 * s;
9717 }
9718 acc
9719 }
9720}
9721
9722#[cfg(target_arch = "x86_64")]
9725#[target_feature(enable = "avx2")]
9726unsafe fn dot_q4_row_avx2_2(
9727 packed: &[u8],
9728 scales: &[u8],
9729 g0: usize,
9730 gpr: usize,
9731 xq1: &[i8],
9732 xq2: &[i8],
9733) -> (f32, f32) {
9734 unsafe {
9736 use core::arch::x86_64::*;
9737 let lomask = _mm_set1_epi8(0x0F);
9738 let eight = _mm256_set1_epi8(8);
9739 let ones = _mm256_set1_epi16(1);
9740 let (mut acc1, mut acc2) = (0f32, 0f32);
9741 #[inline(always)]
9742 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
9743 unsafe {
9744 use core::arch::x86_64::*;
9745 let hi128 = _mm256_extracti128_si256::<1>(d);
9746 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9747 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9748 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9749 _mm_cvtsi128_si32(s32)
9750 }
9751 }
9752 for gi in 0..gpr {
9753 let g = g0 + gi;
9754 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9755 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9756 let lo = _mm_and_si128(b, lomask);
9757 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9758 let w = _mm256_sub_epi8(
9759 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9760 eight,
9761 );
9762 let aw = _mm256_abs_epi8(w);
9763 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9764 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9765 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
9766 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
9767 acc1 += hsum(d1) as f32 * s;
9768 acc2 += hsum(d2) as f32 * s;
9769 }
9770 (acc1, acc2)
9771 }
9772}
9773
9774#[cfg(target_arch = "x86_64")]
9776fn q8_range_avx2(
9777 q: &[u8],
9778 row_scale: &[f32],
9779 act: &SplitAct,
9780 cols: usize,
9781 out_addr: SendMut,
9782 start: usize,
9783 end: usize,
9784) {
9785 for o in start..end {
9786 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9787 unsafe { *out_addr.at(o) = v };
9789 }
9790}
9791
9792#[cfg(target_arch = "x86_64")]
9794#[allow(clippy::too_many_arguments)]
9795fn q8_range2_avx2(
9796 q: &[u8],
9797 row_scale: &[f32],
9798 a1: &SplitAct,
9799 a2: &SplitAct,
9800 cols: usize,
9801 p1: SendMut,
9802 p2: SendMut,
9803 start: usize,
9804 end: usize,
9805) {
9806 for o in start..end {
9807 let row = &q[o * cols..(o + 1) * cols];
9808 unsafe {
9810 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9811 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9812 }
9813 }
9814}
9815
9816#[cfg(target_arch = "aarch64")]
9827fn i8mm_enabled() -> bool {
9828 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9829 *ON.get_or_init(|| {
9830 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9831 && std::arch::is_aarch64_feature_detected!("i8mm")
9832 })
9833}
9834
9835#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9839fn sdot_enabled() -> bool {
9840 use std::sync::OnceLock;
9841 static ON: OnceLock<bool> = OnceLock::new();
9842 *ON.get_or_init(|| {
9843 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9844 if !want {
9845 return false;
9846 }
9847
9848 #[cfg(target_arch = "aarch64")]
9849 {
9850 if std::arch::is_aarch64_feature_detected!("dotprod") {
9851 return true;
9852 }
9853 #[cfg(target_os = "android")]
9854 {
9855 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9856 if cpuinfo.lines().any(|l| {
9857 (l.starts_with("Features") || l.starts_with("features"))
9858 && l.contains("asimddp")
9859 }) {
9860 return true;
9861 }
9862 }
9863 }
9864 false
9865 }
9866 #[cfg(not(target_arch = "aarch64"))]
9867 {
9868 false
9869 }
9870 })
9871}
9872
9873struct SplitAct {
9878 xq: Vec<i8>,
9879 sx: f32,
9880 outliers: Vec<(usize, f32)>,
9881 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9884 xsum: i32,
9885}
9886
9887thread_local! {
9888 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9891 const { std::cell::RefCell::new(Vec::new()) };
9892}
9893
9894impl Drop for SplitAct {
9895 fn drop(&mut self) {
9896 let buf = std::mem::take(&mut self.xq);
9897 if buf.capacity() > 0 {
9898 XQ_FREE.with(|f| {
9899 let mut f = f.borrow_mut();
9900 if f.len() < 16 {
9901 f.push(buf);
9902 }
9903 });
9904 }
9905 }
9906}
9907
9908thread_local! {
9909 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9916}
9917
9918#[inline]
9921fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9922 KROW.with(|s| {
9923 let mut b = s.borrow_mut();
9924 if b.len() < n {
9925 b.resize(n, 0.0);
9926 }
9927 f(&mut b[..n])
9928 })
9929}
9930
9931fn split_act(x: &[f32]) -> SplitAct {
9932 let n = x.len();
9933 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9934 let thr = 8.0 * rms;
9935 let mut outliers: Vec<(usize, f32)> = Vec::new();
9939 let mut amax = 0f32;
9940 for (j, &v) in x.iter().enumerate() {
9941 let a = v.abs();
9942 if a > thr {
9943 outliers.push((j, v));
9944 } else if a > amax {
9945 amax = a;
9946 }
9947 }
9948 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9949 let inv = 1.0 / sx;
9950 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9951 xq.clear();
9952 xq.reserve(n);
9953 if outliers.is_empty() {
9954 xq.extend(
9955 x.iter()
9956 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9957 );
9958 } else {
9959 xq.extend(x.iter().map(|&v| {
9961 if v.abs() > thr {
9962 0
9963 } else {
9964 (v * inv).round().clamp(-127.0, 127.0) as i8
9965 }
9966 }));
9967 }
9968 let xsum = xq.iter().map(|&v| v as i32).sum();
9969 SplitAct {
9970 xq,
9971 sx,
9972 outliers,
9973 xsum,
9974 }
9975}
9976
9977fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9978 let n = x.len();
9979 let rms = (x
9980 .iter()
9981 .zip(col)
9982 .map(|(&a, &c)| {
9983 let v = a * c;
9984 (v * v) as f64
9985 })
9986 .sum::<f64>()
9987 / n.max(1) as f64)
9988 .sqrt() as f32;
9989 let thr = 8.0 * rms;
9990
9991 let mut outliers = Vec::new();
9992 let mut amax = 0f32;
9993 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9994 let v = a * c;
9995 let s = v.abs();
9996 if s > thr {
9997 outliers.push((j, v));
9998 } else if s > amax {
9999 amax = s;
10000 }
10001 }
10002
10003 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10004 let inv = 1.0 / sx;
10005 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10006 xq.clear();
10007 xq.reserve(n);
10008 if outliers.is_empty() {
10009 xq.extend(
10010 x.iter()
10011 .zip(col)
10012 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
10013 );
10014 } else {
10015 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
10016 let v = a * c;
10017 if v.abs() > thr {
10018 0
10019 } else {
10020 (v * inv).round().clamp(-127.0, 127.0) as i8
10021 }
10022 }));
10023 }
10024 let xsum = xq.iter().map(|&v| v as i32).sum();
10025 SplitAct {
10026 xq,
10027 sx,
10028 outliers,
10029 xsum,
10030 }
10031}
10032
10033#[cfg(target_arch = "aarch64")]
10036#[target_feature(enable = "neon,dotprod")]
10037unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
10038 unsafe {
10040 use core::arch::aarch64::*;
10041 use core::arch::asm;
10042 let wp = w.as_ptr() as *const i8;
10043 let n = w.len();
10044 let (mut a0, mut a1, mut a2, mut a3) = (
10045 vdupq_n_s32(0),
10046 vdupq_n_s32(0),
10047 vdupq_n_s32(0),
10048 vdupq_n_s32(0),
10049 );
10050 let mut i = 0;
10051 while i + 64 <= n {
10052 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10053 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
10054 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
10055 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
10056 asm!(
10057 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10058 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10059 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
10060 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
10061 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10062 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10063 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
10064 options(pure, nomem, nostack),
10065 );
10066 i += 64;
10067 }
10068 while i + 16 <= n {
10069 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10070 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
10071 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
10072 i += 16;
10073 }
10074 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
10075 while i < n {
10076 s += (*wp.add(i)) as i32 * xq[i] as i32;
10077 i += 1;
10078 }
10079 s
10080 }
10081}
10082
10083#[cfg(target_arch = "aarch64")]
10087#[target_feature(enable = "neon,dotprod")]
10088unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
10089 unsafe {
10091 use core::arch::aarch64::*;
10092 use core::arch::asm;
10093 let n = xq.len();
10094 let px = xq.as_ptr();
10095 let (p0, p1, p2, p3) = (
10096 w0.as_ptr() as *const i8,
10097 w1.as_ptr() as *const i8,
10098 w2.as_ptr() as *const i8,
10099 w3.as_ptr() as *const i8,
10100 );
10101 let (mut a0, mut a1, mut a2, mut a3) = (
10102 vdupq_n_s32(0),
10103 vdupq_n_s32(0),
10104 vdupq_n_s32(0),
10105 vdupq_n_s32(0),
10106 );
10107 let mut i = 0;
10108 while i + 16 <= n {
10109 let x = vld1q_s8(px.add(i));
10110 let v0 = vld1q_s8(p0.add(i));
10111 let v1 = vld1q_s8(p1.add(i));
10112 let v2 = vld1q_s8(p2.add(i));
10113 let v3 = vld1q_s8(p3.add(i));
10114 asm!(
10115 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10116 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10117 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10118 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10119 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10120 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10121 options(pure, nomem, nostack),
10122 );
10123 i += 16;
10124 }
10125 let mut r = [
10126 vaddvq_s32(a0),
10127 vaddvq_s32(a1),
10128 vaddvq_s32(a2),
10129 vaddvq_s32(a3),
10130 ];
10131 while i < n {
10132 let xi = *px.add(i) as i32;
10133 r[0] += (*p0.add(i)) as i32 * xi;
10134 r[1] += (*p1.add(i)) as i32 * xi;
10135 r[2] += (*p2.add(i)) as i32 * xi;
10136 r[3] += (*p3.add(i)) as i32 * xi;
10137 i += 1;
10138 }
10139 r
10140 }
10141}
10142
10143#[cfg(target_arch = "aarch64")]
10150#[target_feature(enable = "neon,dotprod")]
10151unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
10152 unsafe {
10155 use core::arch::aarch64::*;
10156 use core::arch::asm;
10157 let n = xq.len();
10158 let px = xq.as_ptr();
10159 let pg = g.as_ptr() as *const i8;
10160 let (mut a0, mut a1, mut a2, mut a3) = (
10161 vdupq_n_s32(0),
10162 vdupq_n_s32(0),
10163 vdupq_n_s32(0),
10164 vdupq_n_s32(0),
10165 );
10166 let mut i = 0;
10167 while i + 16 <= n {
10168 let x = vld1q_s8(px.add(i));
10169 let base = pg.add(4 * i);
10170 let v0 = vld1q_s8(base);
10171 let v1 = vld1q_s8(base.add(16));
10172 let v2 = vld1q_s8(base.add(32));
10173 let v3 = vld1q_s8(base.add(48));
10174 asm!(
10175 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10176 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10177 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10178 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10179 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10180 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10181 options(pure, nomem, nostack),
10182 );
10183 i += 16;
10184 }
10185 [
10186 vaddvq_s32(a0),
10187 vaddvq_s32(a1),
10188 vaddvq_s32(a2),
10189 vaddvq_s32(a3),
10190 ]
10191 }
10192}
10193
10194#[cfg(target_arch = "aarch64")]
10200fn q8_range_sdot(
10201 q: &[u8],
10202 rep: &[u8],
10203 row_scale: &[f32],
10204 act: &SplitAct,
10205 cols: usize,
10206 out_addr: SendMut,
10207 start: usize,
10208 end: usize,
10209) {
10210 let mut o = start;
10211 if !rep.is_empty() {
10214 while o < end && o % 4 != 0 {
10215 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10216 unsafe { *out_addr.at(o) = v };
10217 o += 1;
10218 }
10219 }
10220 while o + 4 <= end {
10221 let r = if rep.is_empty() {
10222 unsafe {
10223 dot_i8_sdot_4rows(
10224 &q[o * cols..(o + 1) * cols],
10225 &q[(o + 1) * cols..(o + 2) * cols],
10226 &q[(o + 2) * cols..(o + 3) * cols],
10227 &q[(o + 3) * cols..(o + 4) * cols],
10228 &act.xq,
10229 )
10230 }
10231 } else {
10232 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
10233 };
10234 for k in 0..4 {
10235 let mut acc = r[k] as f32 * act.sx;
10236 for &(j, xv) in &act.outliers {
10237 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
10238 }
10239 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
10241 }
10242 o += 4;
10243 }
10244 while o < end {
10245 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10246 unsafe { *out_addr.at(o) = v };
10247 o += 1;
10248 }
10249}
10250
10251#[cfg(target_arch = "aarch64")]
10254#[allow(clippy::too_many_arguments)]
10255fn q8_range2_sdot(
10256 q: &[u8],
10257 row_scale: &[f32],
10258 a1: &SplitAct,
10259 a2: &SplitAct,
10260 cols: usize,
10261 p1: SendMut,
10262 p2: SendMut,
10263 start: usize,
10264 end: usize,
10265) {
10266 for o in start..end {
10267 let row = &q[o * cols..(o + 1) * cols];
10268 unsafe {
10270 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
10271 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
10272 }
10273 }
10274}
10275
10276#[allow(clippy::too_many_arguments)]
10278fn q8_range2_f32(
10279 q: &[u8],
10280 row_scale: &[f32],
10281 x1: &[f32],
10282 x2: &[f32],
10283 cols: usize,
10284 p1: SendMut,
10285 p2: SendMut,
10286 start: usize,
10287 end: usize,
10288) {
10289 for o in start..end {
10290 let row = &q[o * cols..(o + 1) * cols];
10291 unsafe {
10293 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
10294 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
10295 }
10296 }
10297}
10298
10299fn q8_range_f32(
10301 q: &[u8],
10302 row_scale: &[f32],
10303 xs: &[f32],
10304 cols: usize,
10305 out_addr: SendMut,
10306 start: usize,
10307 end: usize,
10308) {
10309 for o in start..end {
10310 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10311 unsafe { *out_addr.at(o) = v };
10313 }
10314}
10315
10316#[inline]
10320fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
10321 #[cfg(target_arch = "aarch64")]
10322 return row_dot_sdot(row, act);
10323 #[cfg(target_arch = "x86_64")]
10324 return row_dot_avx2(row, act);
10325 #[allow(unreachable_code)]
10326 q8_row_dot_scalar(row, act)
10327}
10328
10329#[allow(dead_code)]
10330fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
10331 let mut acc = 0i32;
10332 for (k, &b) in row.iter().enumerate() {
10333 acc += (b as i8) as i32 * act.xq[k] as i32;
10334 }
10335 let mut acc = acc as f32 * act.sx;
10336 for &(j, xv) in &act.outliers {
10337 acc += (row[j] as i8) as f32 * xv;
10338 }
10339 acc
10340}
10341
10342#[cfg(target_arch = "aarch64")]
10345#[inline]
10346fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
10347 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
10348 for &(j, xv) in &act.outliers {
10349 acc += (row[j] as i8) as f32 * xv;
10350 }
10351 acc
10352}
10353
10354#[cfg(target_arch = "aarch64")]
10362#[target_feature(enable = "neon,dotprod")]
10363unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10364 unsafe {
10367 use core::arch::aarch64::*;
10368 use core::arch::asm;
10369 let lomask = vdupq_n_u8(0x0F);
10370 let eight = vdupq_n_s8(8);
10371 let mut acc = 0f32;
10372 for gi in 0..gpr {
10373 let g = g0 + gi;
10374 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10375 let b = vld1q_u8(packed.as_ptr().add(g * 16));
10376 let lo = vandq_u8(b, lomask);
10377 let hi = vshrq_n_u8::<4>(b);
10378 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10379 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10380 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
10381 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
10382 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10383 asm!(
10384 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
10385 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
10386 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10387 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
10388 options(pure, nomem, nostack),
10389 );
10390 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10391 }
10392 acc
10393 }
10394}
10395
10396#[cfg(target_arch = "aarch64")]
10401#[target_feature(enable = "neon,dotprod")]
10402unsafe fn dot_q4_row_sdot2(
10403 packed: &[u8],
10404 scales: &[u8],
10405 g0: usize,
10406 gpr: usize,
10407 xq1: &[i8],
10408 xq2: &[i8],
10409) -> (f32, f32) {
10410 unsafe {
10413 use core::arch::aarch64::*;
10414 use core::arch::asm;
10415 let lomask = vdupq_n_u8(0x0F);
10416 let eight = vdupq_n_s8(8);
10417 let (mut acc1, mut acc2) = (0f32, 0f32);
10418 for gi in 0..gpr {
10419 let g = g0 + gi;
10420 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10421 let b = vld1q_u8(packed.as_ptr().add(g * 16));
10422 let lo = vandq_u8(b, lomask);
10423 let hi = vshrq_n_u8::<4>(b);
10424 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10425 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10426 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
10427 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
10428 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
10429 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
10430 let (mut a0, mut a1, mut b0, mut b1) = (
10431 vdupq_n_s32(0),
10432 vdupq_n_s32(0),
10433 vdupq_n_s32(0),
10434 vdupq_n_s32(0),
10435 );
10436 asm!(
10437 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
10438 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
10439 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
10440 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
10441 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10442 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
10443 e0 = in(vreg) e0, e1 = in(vreg) e1,
10444 x10 = in(vreg) x10, x11 = in(vreg) x11,
10445 x20 = in(vreg) x20, x21 = in(vreg) x21,
10446 options(pure, nomem, nostack),
10447 );
10448 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10449 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
10450 }
10451 (acc1, acc2)
10452 }
10453}
10454
10455#[inline]
10460pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
10461 #[cfg(target_arch = "aarch64")]
10462 unsafe {
10463 return axpy_i8_f32_neon(acc, row, w);
10464 }
10465 #[cfg(target_arch = "x86_64")]
10466 if avx2_enabled() {
10467 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
10468 }
10469 #[allow(unreachable_code)]
10470 {
10471 for (a, &b) in acc.iter_mut().zip(row) {
10472 *a += w * b as f32;
10473 }
10474 }
10475}
10476
10477#[cfg(target_arch = "x86_64")]
10479#[target_feature(enable = "avx2,fma")]
10480unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
10481 unsafe {
10483 use core::arch::x86_64::*;
10484 let n = acc.len().min(row.len());
10485 let ap = acc.as_mut_ptr();
10486 let rp = row.as_ptr();
10487 let wv = _mm256_set1_ps(w);
10488 let mut j = 0usize;
10489 while j + 16 <= n {
10490 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
10491 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
10492 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
10493 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
10494 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
10495 _mm256_storeu_ps(ap.add(j), v0);
10496 _mm256_storeu_ps(ap.add(j + 8), v1);
10497 j += 16;
10498 }
10499 while j < n {
10500 *ap.add(j) += w * (*rp.add(j)) as f32;
10501 j += 1;
10502 }
10503 }
10504}
10505
10506#[cfg(target_arch = "aarch64")]
10507#[target_feature(enable = "neon")]
10508unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
10509 unsafe {
10511 use core::arch::aarch64::*;
10512 let n = acc.len().min(row.len());
10513 let ap = acc.as_mut_ptr();
10514 let rp = row.as_ptr();
10515 let wv = vdupq_n_f32(w);
10516 let mut j = 0usize;
10517 while j + 16 <= n {
10518 let rb = vld1q_s8(rp.add(j));
10519 let lo = vmovl_s8(vget_low_s8(rb));
10520 let hi = vmovl_s8(vget_high_s8(rb));
10521 for (off, half) in [(0, lo), (8, hi)] {
10522 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
10523 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
10524 let o = j + off;
10525 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
10526 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
10527 }
10528 j += 16;
10529 }
10530 while j < n {
10531 *ap.add(j) += w * (*rp.add(j)) as f32;
10532 j += 1;
10533 }
10534 }
10535}
10536
10537#[inline]
10540pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
10541 #[cfg(target_arch = "aarch64")]
10542 unsafe {
10543 return dot_i8_f32_neon(w, x);
10544 }
10545 #[cfg(target_arch = "x86_64")]
10546 if avx2_enabled() {
10547 return unsafe { dot_i8_f32_avx2(w, x) };
10548 }
10549 #[allow(unreachable_code)]
10550 {
10551 let mut sum = 0.0f32;
10552 for (j, &b) in w.iter().enumerate() {
10553 sum += (b as i8) as f32 * x[j];
10554 }
10555 sum
10556 }
10557}
10558
10559#[inline]
10563fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10564 #[cfg(target_arch = "aarch64")]
10565 unsafe {
10566 return dot_i8_col_f32_neon(w, x, col);
10567 }
10568 #[allow(unreachable_code)]
10569 {
10570 let mut sum = 0.0f32;
10571 for (j, &b) in w.iter().enumerate() {
10572 sum += (b as i8) as f32 * x[j] * col[j];
10573 }
10574 sum
10575 }
10576}
10577
10578#[cfg(target_arch = "aarch64")]
10579#[target_feature(enable = "neon")]
10580unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10581 unsafe {
10583 use core::arch::aarch64::*;
10584 let n = x.len();
10585 let wp = w.as_ptr() as *const i8;
10586 let xp = x.as_ptr();
10587 let cp = col.as_ptr();
10588 let (mut a0, mut a1, mut a2, mut a3) = (
10589 vdupq_n_f32(0.0),
10590 vdupq_n_f32(0.0),
10591 vdupq_n_f32(0.0),
10592 vdupq_n_f32(0.0),
10593 );
10594 let mut j = 0usize;
10595 while j + 16 <= n {
10596 let wb = vld1q_s8(wp.add(j));
10597 let lo = vmovl_s8(vget_low_s8(wb));
10598 let hi = vmovl_s8(vget_high_s8(wb));
10599 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10600 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10601 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10602 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10603 a0 = vfmaq_f32(
10604 a0,
10605 w0,
10606 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
10607 );
10608 a1 = vfmaq_f32(
10609 a1,
10610 w1,
10611 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
10612 );
10613 a2 = vfmaq_f32(
10614 a2,
10615 w2,
10616 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
10617 );
10618 a3 = vfmaq_f32(
10619 a3,
10620 w3,
10621 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
10622 );
10623 j += 16;
10624 }
10625 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10626 while j < n {
10627 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
10628 j += 1;
10629 }
10630 sum
10631 }
10632}
10633
10634#[cfg(target_arch = "aarch64")]
10635#[target_feature(enable = "neon")]
10636unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
10637 unsafe {
10639 use core::arch::aarch64::*;
10640 let n = x.len();
10641 let wp = w.as_ptr() as *const i8;
10642 let xp = x.as_ptr();
10643 let (mut a0, mut a1, mut a2, mut a3) = (
10644 vdupq_n_f32(0.0),
10645 vdupq_n_f32(0.0),
10646 vdupq_n_f32(0.0),
10647 vdupq_n_f32(0.0),
10648 );
10649 let mut j = 0usize;
10650 while j + 16 <= n {
10651 let wb = vld1q_s8(wp.add(j));
10652 let lo = vmovl_s8(vget_low_s8(wb));
10653 let hi = vmovl_s8(vget_high_s8(wb));
10654 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10655 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10656 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10657 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10658 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
10659 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
10660 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
10661 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
10662 j += 16;
10663 }
10664 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10665 while j < n {
10666 sum += (*wp.add(j)) as f32 * *xp.add(j);
10667 j += 1;
10668 }
10669 sum
10670 }
10671}
10672
10673#[allow(clippy::too_many_arguments)]
10674fn qmatvec(
10675 q: &[u8],
10676 rep: &[u8],
10677 row_scale: &[f32],
10678 x: &[f32],
10679 col_field: &[f32],
10680 dtype: TensorDtype,
10681 rows: usize,
10682 cols: usize,
10683 out: &mut [f32],
10684 pool: Option<&Pool>,
10685) {
10686 debug_assert_eq!(out.len(), rows);
10687 #[cfg(not(target_arch = "aarch64"))]
10688 let _ = rep;
10689
10690 #[cfg(target_arch = "aarch64")]
10691 if sdot_enabled() {
10692 let act = if dtype == TensorDtype::Q8_2f {
10693 split_act_q8_2f(x, col_field)
10694 } else {
10695 split_act(x)
10696 };
10697 let out_addr = SendMut(out.as_mut_ptr());
10698 let run_range = |start: usize, end: usize| {
10699 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
10700 };
10701 match pool {
10702 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10703 _ => run_range(0, rows),
10704 }
10705 return;
10706 }
10707 #[cfg(target_arch = "x86_64")]
10710 if avx2_a8w8_enabled() {
10711 let act = if dtype == TensorDtype::Q8_2f {
10712 split_act_q8_2f(x, col_field)
10713 } else {
10714 split_act(x)
10715 };
10716 let out_addr = SendMut(out.as_mut_ptr());
10717 let run_range = |start: usize, end: usize| {
10718 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
10719 };
10720 match pool {
10721 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10722 _ => run_range(0, rows),
10723 }
10724 return;
10725 }
10726
10727 prescale_with(x, col_field, dtype, 1, |xs| {
10728 let out_addr = SendMut(out.as_mut_ptr());
10729 let run_range = move |start: usize, end: usize| {
10730 for o in start..end {
10731 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10732 unsafe { *out_addr.at(o) = v };
10734 }
10735 };
10736 match pool {
10737 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10738 _ => run_range(0, rows),
10739 }
10740 });
10741}
10742
10743#[allow(clippy::too_many_arguments)]
10744fn qmatvec2(
10745 q: &[u8],
10746 row_scale: &[f32],
10747 x1: &[f32],
10748 x2: &[f32],
10749 col_field: &[f32],
10750 dtype: TensorDtype,
10751 rows: usize,
10752 cols: usize,
10753 o1: &mut [f32],
10754 o2: &mut [f32],
10755 pool: Option<&Pool>,
10756) {
10757 #[cfg(target_arch = "aarch64")]
10758 if sdot_enabled() {
10759 let a1s = if dtype == TensorDtype::Q8_2f {
10760 split_act_q8_2f(x1, col_field)
10761 } else {
10762 split_act(x1)
10763 };
10764 let a2s = if dtype == TensorDtype::Q8_2f {
10765 split_act_q8_2f(x2, col_field)
10766 } else {
10767 split_act(x2)
10768 };
10769 let p1 = SendMut(o1.as_mut_ptr());
10770 let p2 = SendMut(o2.as_mut_ptr());
10771 let run_range = |start: usize, end: usize| {
10772 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10773 };
10774 match pool {
10775 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10776 _ => run_range(0, rows),
10777 }
10778 return;
10779 }
10780 #[cfg(target_arch = "x86_64")]
10781 if avx2_a8w8_enabled() {
10782 let a1s = if dtype == TensorDtype::Q8_2f {
10783 split_act_q8_2f(x1, col_field)
10784 } else {
10785 split_act(x1)
10786 };
10787 let a2s = if dtype == TensorDtype::Q8_2f {
10788 split_act_q8_2f(x2, col_field)
10789 } else {
10790 split_act(x2)
10791 };
10792 let p1 = SendMut(o1.as_mut_ptr());
10793 let p2 = SendMut(o2.as_mut_ptr());
10794 let run_range = |start: usize, end: usize| {
10795 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10796 };
10797 match pool {
10798 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10799 _ => run_range(0, rows),
10800 }
10801 return;
10802 }
10803
10804 prescale_with(x1, col_field, dtype, 1, |x1s| {
10805 prescale_with(x2, col_field, dtype, 2, |x2s| {
10806 let p1 = SendMut(o1.as_mut_ptr());
10807 let p2 = SendMut(o2.as_mut_ptr());
10808 let run_range = move |start: usize, end: usize| {
10809 for o in start..end {
10810 let row = &q[o * cols..(o + 1) * cols];
10811 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10812 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10813 unsafe {
10815 *p1.at(o) = s1;
10816 *p2.at(o) = s2;
10817 }
10818 }
10819 };
10820 match pool {
10821 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10822 _ => run_range(0, rows),
10823 }
10824 });
10825 });
10826}
10827
10828#[derive(Clone, Copy)]
10829struct SendMut(*mut f32);
10830unsafe impl Send for SendMut {}
10831unsafe impl Sync for SendMut {}
10832
10833impl SendMut {
10834 #[inline]
10835 fn at(self, i: usize) -> *mut f32 {
10836 unsafe { self.0.add(i) }
10837 }
10838}
10839
10840#[cfg(test)]
10841mod tests {
10842 use super::*;
10843
10844 #[test]
10845 fn q2tp_i8_dot_matches_exact_on_grid() {
10846 let (rows, cols) = (5, 64);
10850 let gpr = cols / GROUP_SIZE;
10851 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10855 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10856 .collect();
10857 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10858 let x: Vec<f32> = (0..cols)
10859 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10860 .collect();
10861 let act = split_act(&x);
10862 assert!(
10863 act.outliers.is_empty(),
10864 "on-grid input must have no outliers"
10865 );
10866 let gsum = q1_group_sums(&act.xq, gpr);
10867 for r in 0..rows {
10868 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10869 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10870 assert!(
10871 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10872 "row {r}: exact {exact} vs i8 {fast}"
10873 );
10874 }
10875 }
10876
10877 #[test]
10878 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
10879 let (rows, cols) = (1usize, GROUP_SIZE);
10880 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
10881 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
10885 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
10886 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
10888 let mut raw = vec![0.0f32; rows];
10889 let mut affine = vec![0.0f32; rows];
10890 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
10891 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
10892 assert_eq!(raw, vec![-24.0]);
10893 assert_eq!(affine, vec![-8.0]);
10894 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
10895 }
10896
10897 #[cfg(target_arch = "x86_64")]
10898 #[test]
10899 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
10900 if !std::arch::is_x86_feature_detected!("avx2") {
10905 return;
10906 }
10907 let mut seed = 0x9e3779b9u32;
10908 let mut next = || {
10909 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
10910 seed
10911 };
10912 for _ in 0..20_000 {
10913 let mut ch = [0u8; Q2TP_CHUNK];
10914 let mut x = [0i8; GROUP_SIZE];
10915 for b in &mut ch {
10916 *b = next() as u8;
10917 }
10918 for v in &mut x {
10919 *v = (next() >> 24) as i8;
10920 }
10921 let mut reference = 0i32;
10922 for (k, &b) in ch.iter().enumerate() {
10923 reference += (b & 3) as i32 * x[k * 4] as i32;
10924 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
10925 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
10926 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
10927 }
10928 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
10931 assert_eq!(got, reference, "packed q2 lane mismatch");
10932 }
10933 }
10934
10935 #[test]
10936 fn q8_row_dot_fast_matches_scalar() {
10937 let cols = 96;
10940 let row: Vec<u8> = (0..cols)
10941 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10942 .collect();
10943 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10944 let act = split_act(&x);
10945 let fast = q8_row_dot(&row, &act);
10946 let scalar = q8_row_dot_scalar(&row, &act);
10947 assert!(
10948 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10949 "fast {fast} vs scalar {scalar}"
10950 );
10951 }
10952
10953 #[test]
10954 fn f32_matvec_matches_matvec_rows_bitexact() {
10955 let (rows, cols) = (300, 40);
10956 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10957 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10958 let qt = QTensor::from_f32(w.clone(), rows, cols);
10959
10960 let mut a = vec![0.0f32; rows];
10961 matvec_rows(None, &w, &x, &mut a);
10962 let mut b = vec![0.0f32; rows];
10963 qt.matvec(&x, &mut b, None);
10964 assert_eq!(a, b);
10965 }
10966
10967 #[test]
10968 fn sdot_kernel_exact_on_grid() {
10969 eprintln!("sdot_enabled = {}", sdot_enabled());
10974 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
10976 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10977 .collect();
10978 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10979 let x: Vec<f32> = (0..cols)
10980 .map(|i| match i % 3 {
10981 0 => 1.0,
10982 1 => -1.0,
10983 _ => 0.0,
10984 })
10985 .collect();
10986 let mut a = vec![0.0f32; rows];
10987 qmatvec(
10988 &w,
10989 &[],
10990 &scales,
10991 &x,
10992 &[],
10993 TensorDtype::Q8Row,
10994 rows,
10995 cols,
10996 &mut a,
10997 None,
10998 );
10999 for o in 0..rows {
11000 let mut acc = 0.0f32;
11001 for j in 0..cols {
11002 acc += (w[o * cols + j] as i8) as f32 * x[j];
11003 }
11004 let expect = acc * scales[o];
11005 assert!(
11006 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11007 "row {o}: {} vs {expect}",
11008 a[o]
11009 );
11010 }
11011 }
11012
11013 #[test]
11014 fn q1_tbl_fast_path_matches_reference() {
11015 let (rows, cols) = (5, 256);
11020 let gpr = cols / GROUP_SIZE;
11021 let mut bytes = Vec::new();
11022 for t in 0..rows * gpr {
11023 let s = 0.007 + (t % 11) as f32 * 0.004;
11024 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11025 for j in 0..4 {
11026 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
11027 }
11028 }
11029 let x: Vec<f32> = (0..cols)
11030 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
11031 .collect();
11032 let mut w = vec![0.0f32; rows * cols];
11033 cortiq_core::quant::dequant_q1(&bytes, &mut w);
11034 let mut got = vec![0.0f32; rows];
11035 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11036 for o in 0..rows {
11037 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11038 assert!(
11039 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11040 "row {o}: {} vs {expect}",
11041 got[o]
11042 );
11043 }
11044 let b = 5usize;
11047 let mut xs_all = Vec::new();
11048 for bi in 0..b {
11049 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
11050 }
11051 let mut mm = vec![0.0f32; b * rows];
11052 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
11053 for bi in 0..b {
11054 let mut single = vec![0.0f32; rows];
11055 q1_matvec(
11056 &bytes,
11057 &xs_all[bi * cols..(bi + 1) * cols],
11058 rows,
11059 cols,
11060 &mut single,
11061 None,
11062 );
11063 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
11064 }
11065 }
11066
11067 #[test]
11068 fn q1_kernels_match_exact_reference() {
11069 let (rows, cols) = (7, 96);
11071 let gpr = cols / GROUP_SIZE;
11072 let mut bytes = Vec::new();
11073 for t in 0..rows * gpr {
11074 let s = 0.01 + (t % 13) as f32 * 0.003;
11075 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11076 for j in 0..4 {
11077 bytes.push(((t * 31 + j * 97) % 251) as u8);
11078 }
11079 }
11080 let x: Vec<f32> = (0..cols)
11082 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
11083 .collect();
11084 let mut w = vec![0.0f32; rows * cols];
11086 cortiq_core::quant::dequant_q1(&bytes, &mut w);
11087 let mut expect = vec![0.0f32; rows];
11088 for o in 0..rows {
11089 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11090 }
11091 let mut got = vec![0.0f32; rows];
11092 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11093 for o in 0..rows {
11094 assert!(
11095 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
11096 "row {o}: {} vs {}",
11097 got[o],
11098 expect[o]
11099 );
11100 }
11101 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
11103 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
11104 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
11105 assert_eq!(a1, got);
11106 let mut xs = x.clone();
11107 xs.extend_from_slice(&x2);
11108 let mut mm = vec![0.0f32; 2 * rows];
11109 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
11110 assert_eq!(&mm[..rows], got.as_slice());
11111 assert_eq!(&mm[rows..], a2.as_slice());
11112 }
11113
11114 #[test]
11115 fn repack_is_bit_identical() {
11116 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
11122 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
11123 .collect();
11124 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
11125 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
11126 let rep = q8_repack_layout(&w, rows, cols);
11127 for g in 0..rows / 4 {
11129 for c in 0..cols / 16 {
11130 for lane in 0..4 {
11131 assert_eq!(
11132 &rep[g * 4 * cols + c * 64 + lane * 16
11133 ..g * 4 * cols + c * 64 + lane * 16 + 16],
11134 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
11135 );
11136 }
11137 }
11138 }
11139 let mut a = vec![0.0f32; rows];
11140 qmatvec(
11141 &w,
11142 &[],
11143 &scales,
11144 &x,
11145 &[],
11146 TensorDtype::Q8Row,
11147 rows,
11148 cols,
11149 &mut a,
11150 None,
11151 );
11152 let mut b = vec![0.0f32; rows];
11153 qmatvec(
11154 &w,
11155 &rep,
11156 &scales,
11157 &x,
11158 &[],
11159 TensorDtype::Q8Row,
11160 rows,
11161 cols,
11162 &mut b,
11163 None,
11164 );
11165 assert_eq!(a, b, "full-range repack output diverged");
11166
11167 #[cfg(target_arch = "aarch64")]
11168 if sdot_enabled() {
11169 let act = split_act(&x);
11171 let mut c1 = vec![0.0f32; rows];
11172 let mut c2 = vec![0.0f32; rows];
11173 q8_range_sdot(
11174 &w,
11175 &[],
11176 &scales,
11177 &act,
11178 cols,
11179 SendMut(c1.as_mut_ptr()),
11180 3,
11181 rows - 2,
11182 );
11183 q8_range_sdot(
11184 &w,
11185 &rep,
11186 &scales,
11187 &act,
11188 cols,
11189 SendMut(c2.as_mut_ptr()),
11190 3,
11191 rows - 2,
11192 );
11193 assert_eq!(c1, c2, "unaligned-range repack output diverged");
11194 }
11195 }
11196
11197 #[test]
11198 fn sdot_a8w8_noise_is_bounded() {
11199 let (rows, cols) = (16, 512);
11203 let w: Vec<u8> = (0..rows * cols)
11204 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11205 .collect();
11206 let scales = vec![0.01f32; rows];
11207 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
11208 let mut a = vec![0.0f32; rows];
11209 qmatvec(
11210 &w,
11211 &[],
11212 &scales,
11213 &x,
11214 &[],
11215 TensorDtype::Q8Row,
11216 rows,
11217 cols,
11218 &mut a,
11219 None,
11220 );
11221 let (mut num, mut den) = (0f64, 0f64);
11222 for o in 0..rows {
11223 let mut acc = 0.0f32;
11224 for j in 0..cols {
11225 acc += (w[o * cols + j] as i8) as f32 * x[j];
11226 }
11227 let expect = acc * scales[o];
11228 num += ((a[o] - expect) as f64).powi(2);
11229 den += (expect as f64).powi(2);
11230 }
11231 let rel = (num / den.max(1e-12)).sqrt();
11232 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
11233 }
11234
11235 #[test]
11236 fn i8_dot_neon_matches_scalar() {
11237 let n = 100;
11238 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
11239 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
11240 let mut scalar = 0.0f32;
11241 for j in 0..n {
11242 scalar += (w[j] as i8) as f32 * x[j];
11243 }
11244 let fast = dot_i8_f32(&w, &x);
11245 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
11246 }
11247
11248 #[test]
11250 fn vbitmatvec_matches_full_dequant() {
11251 let (rows, cols) = (6, 64);
11252 let ng = cols / GROUP_SIZE;
11253 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11255 let mut bytes = bits.clone();
11256 for g in 0..rows * ng {
11257 let s = 0.02 + 0.001 * g as f32;
11258 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11259 }
11260 for r in 0..rows {
11261 let b = bits[r] as usize;
11262 let (mut acc, mut nb) = (0u64, 0usize);
11263 let mut rowbytes = Vec::new();
11264 for i in 0..cols {
11265 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11266 acc = (acc << b) | v;
11267 nb += b;
11268 while nb >= 8 {
11269 nb -= 8;
11270 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11271 }
11272 }
11273 if nb > 0 {
11274 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11275 }
11276 bytes.extend_from_slice(&rowbytes);
11277 }
11278 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11279
11280 let mut reference = vec![0f32; rows * cols];
11281 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
11282 let mut expect = vec![0f32; rows];
11283 for r in 0..rows {
11284 expect[r] = reference[r * cols..(r + 1) * cols]
11285 .iter()
11286 .zip(&x)
11287 .map(|(w, xv)| w * xv)
11288 .sum();
11289 }
11290 let mut got = vec![0f32; rows];
11291 let offsets = vbit_row_offsets(&bytes, rows, cols);
11292 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
11293 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11297 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11298 for r in 0..rows {
11299 assert!(
11300 (got[r] - expect[r]).abs() < tol * scale,
11301 "row {r}: {} vs {}",
11302 got[r],
11303 expect[r]
11304 );
11305 }
11306 }
11307
11308 #[test]
11313 #[cfg(target_arch = "x86_64")]
11314 fn vbit_matmat_blocked_matches_per_row() {
11315 let (rows, cols, b) = (64usize, 128usize, 9usize);
11316 let ng = cols / GROUP_SIZE;
11317 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
11318 let mut bytes = bits.clone();
11319 for g in 0..rows * ng {
11320 let sc = 0.02 + 0.0005 * g as f32;
11321 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11322 }
11323 for r in 0..rows {
11324 let bw = bits[r] as usize;
11325 let (mut acc, mut nb) = (0u64, 0usize);
11326 let mut rowbytes = Vec::new();
11327 for i in 0..cols {
11328 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11329 acc = (acc << bw) | v;
11330 nb += bw;
11331 while nb >= 8 {
11332 nb -= 8;
11333 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11334 }
11335 }
11336 if nb > 0 {
11337 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11338 }
11339 bytes.extend_from_slice(&rowbytes);
11340 }
11341 let x: Vec<f32> = (0..b * cols)
11342 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11343 .collect();
11344 let offsets = vbit_row_offsets(&bytes, rows, cols);
11345 let mut y_a = vec![0f32; b * rows];
11346 let mut y_b = vec![0f32; b * rows];
11347 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11348 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
11349 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11350 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
11351 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11352 let max_d = y_a
11353 .iter()
11354 .zip(&y_b)
11355 .map(|(p, q)| (p - q).abs())
11356 .fold(0.0f32, f32::max);
11357 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
11358 }
11359
11360 #[test]
11368 fn q4t_matmat_blocked_matches_per_row() {
11369 let (rows, cols, b) = (16usize, 64usize, 9usize);
11370 let gpr = cols / GROUP_SIZE;
11371 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11372 for r in 0..rows {
11373 for g in 0..gpr {
11374 let t = (r * gpr + g) * Q4_TILE;
11375 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
11376 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11377 for k in 0..16 {
11378 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11379 }
11380 }
11381 }
11382 let x: Vec<f32> = (0..b * cols)
11383 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11384 .collect();
11385 let mut y_blk = vec![0f32; b * rows];
11386 let mut y_row = vec![0f32; b * rows];
11387 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11388 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
11389 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11390 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
11391 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11392 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
11393 }
11394
11395 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
11402 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
11403 let gpr = cols / GROUP_SIZE;
11404 let stride = q4tp_code_stride(gpr);
11405 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
11406 let mut b = vec![0u8; codes_off + rows * stride];
11407 for r in 0..rows {
11408 for g in 0..gpr {
11409 let t = (r * gpr + g) * Q4TP_NIB;
11410 for k in 0..16 {
11411 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11412 }
11413 }
11414 let lo = -6.0 - 0.03 * (r % 17) as f32;
11415 let step = 0.01 + 0.004 * (r % 11) as f32;
11416 let p = params_off + r * 4;
11417 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
11418 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
11419 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
11420 for g in 0..gpr {
11421 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
11422 }
11423 }
11424 b
11425 }
11426
11427 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
11431 let gpr = cols / GROUP_SIZE;
11432 let v = Q4tpView::new(bytes, rows, cols);
11433 let mut out = vec![0u8; rows * gpr * Q4_TILE];
11434 let mut sc = vec![0f32; gpr];
11435 for r in 0..rows {
11436 v.scales_into(r, gpr, &mut sc);
11437 for g in 0..gpr {
11438 let t = (r * gpr + g) * Q4_TILE;
11439 let s = sc[g];
11440 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11441 let src = (r * gpr + g) * Q4TP_NIB;
11442 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
11443 }
11444 }
11445 out
11446 }
11447
11448 #[test]
11454 fn q4tp_exact_path_matches_dequant_reference() {
11455 let (rows, cols) = (256usize, 512usize);
11456 let gpr = cols / GROUP_SIZE;
11457 let bytes = synth_q4tp(rows, cols);
11458 let mut w = vec![0f32; rows * cols];
11459 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11460
11461 let x: Vec<f32> = (0..cols)
11462 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11463 .collect();
11464 let v = Q4tpView::new(&bytes, rows, cols);
11465 let mut sc = vec![0f32; gpr];
11466 for r in 0..rows {
11467 v.scales_into(r, gpr, &mut sc);
11468 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
11469 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
11470 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11474 assert!(
11475 (got - want).abs() <= 1e-5 * mag,
11476 "row {r}: kernel {got} vs dequant {want}"
11477 );
11478 }
11479 }
11480
11481 #[test]
11487 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
11488 let (rows, cols) = (256usize, 512usize);
11489 let bytes = synth_q4tp(rows, cols);
11490 let twin = q4tp_as_q4t(&bytes, rows, cols);
11491 let x: Vec<f32> = (0..cols)
11492 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11493 .collect();
11494
11495 let mut got = vec![0f32; rows];
11496 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
11497 let mut want = vec![0f32; rows];
11498 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
11499
11500 let mut w = vec![0f32; rows * cols];
11503 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11504 for r in 0..rows {
11505 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11506 assert!(
11507 (got[r] - want[r]).abs() <= 1e-3 * mag,
11508 "row {r}: q4tp {} vs q4t {}",
11509 got[r],
11510 want[r]
11511 );
11512 }
11513 }
11514
11515 #[test]
11520 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
11521 let (rows, cols, b) = (256usize, 512usize, 5usize);
11522 let bytes = synth_q4tp(rows, cols);
11523 let twin = q4tp_as_q4t(&bytes, rows, cols);
11524 let xs: Vec<f32> = (0..b * cols)
11525 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11526 .collect();
11527
11528 let mut got = vec![0f32; b * rows];
11529 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
11530 let mut want = vec![0f32; b * rows];
11531 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
11532
11533 let mut w = vec![0f32; rows * cols];
11534 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11535 for t in 0..b {
11536 for r in 0..rows {
11537 let mag: f32 = (0..cols)
11538 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
11539 .sum();
11540 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
11541 assert!(
11542 (g - wa).abs() <= 1e-3 * mag,
11543 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
11544 );
11545 }
11546 }
11547 }
11548
11549 #[test]
11550 fn q4tp_matvec2_matches_the_single_stream_kernel() {
11551 let (rows, cols) = (128usize, 256usize);
11552 let gpr = cols / GROUP_SIZE;
11553 let bytes = synth_q4tp(rows, cols);
11554 let xs: Vec<f32> = (0..2 * cols)
11555 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11556 .collect();
11557
11558 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
11559 q4tp_matvec2(
11560 &bytes,
11561 &xs[..cols],
11562 &xs[cols..],
11563 rows,
11564 cols,
11565 &mut o1,
11566 &mut o2,
11567 None,
11568 );
11569
11570 let v = Q4tpView::new(&bytes, rows, cols);
11573 let mut sc = vec![0f32; gpr];
11574 for r in 0..rows {
11575 v.scales_into(r, gpr, &mut sc);
11576 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
11577 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
11578 }
11579 }
11580
11581 #[test]
11588 fn q4tp_matvec_keeps_pace_with_q4t() {
11589 let (rows, cols) = (4096usize, 3072usize);
11590 let bytes = synth_q4tp(rows, cols);
11591 let twin = q4tp_as_q4t(&bytes, rows, cols);
11592 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
11593 let mut o = vec![0f32; rows];
11594 let n = 12;
11595 let mut best = (f64::MAX, f64::MAX);
11596 for _ in 0..3 {
11599 let t0 = std::time::Instant::now();
11600 for _ in 0..n {
11601 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
11602 }
11603 best.0 = best.0.min(t0.elapsed().as_secs_f64());
11604 let t0 = std::time::Instant::now();
11605 for _ in 0..n {
11606 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
11607 }
11608 best.1 = best.1.min(t0.elapsed().as_secs_f64());
11609 }
11610 let ratio = best.1 / best.0;
11611 println!(
11612 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
11613 best.0 * 1e3 / n as f64,
11614 best.1 * 1e3 / n as f64
11615 );
11616 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
11617 }
11618
11619 #[cfg(target_os = "macos")]
11620 #[test]
11621 fn q4t_matmat_accel_matches_dequant_reference() {
11622 if !accel_gemm_enabled() {
11623 return; }
11625 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
11627 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11628 for r in 0..rows {
11629 for g in 0..gpr {
11630 let t = (r * gpr + g) * Q4_TILE;
11631 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
11632 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11633 for k in 0..16 {
11634 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11635 }
11636 }
11637 }
11638 let x: Vec<f32> = (0..b * cols)
11639 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11640 .collect();
11641 let mut got = vec![0f32; b * rows];
11642 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
11643 let mut w = vec![0f32; rows * cols];
11645 for r in 0..rows {
11646 for g in 0..gpr {
11647 let t = (r * gpr + g) * Q4_TILE;
11648 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
11649 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
11650 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
11651 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
11652 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
11653 }
11654 }
11655 }
11656 for bi in 0..b {
11657 for r in 0..rows {
11658 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
11659 let d = (got[bi * rows + r] - want).abs();
11660 assert!(
11661 d <= want.abs().max(1.0) * 1e-4,
11662 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
11663 got[bi * rows + r]
11664 );
11665 }
11666 }
11667 }
11668
11669 #[test]
11670 fn q4matvec_matches_full_dequant() {
11671 let (rows, cols) = (8, 64);
11672 let groups = rows * cols / GROUP_SIZE;
11673 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11675 for i in 0..groups * 16 {
11676 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11677 }
11678 for g in 0..groups {
11679 let s = 0.01 + 0.003 * g as f32;
11680 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11681 }
11682 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11683
11684 let mut reference = vec![0.0f32; rows * cols];
11685 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11686 let mut expect = vec![0.0f32; rows];
11687 for r in 0..rows {
11688 expect[r] = reference[r * cols..(r + 1) * cols]
11689 .iter()
11690 .zip(&x)
11691 .map(|(w, xv)| w * xv)
11692 .sum();
11693 }
11694
11695 let mut got = vec![0.0f32; rows];
11696 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11697 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11701 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11702 for r in 0..rows {
11703 assert!(
11704 (got[r] - expect[r]).abs() < tol * scale,
11705 "row {r}: {} vs {}",
11706 got[r],
11707 expect[r]
11708 );
11709 }
11710 }
11711
11712 #[test]
11715 fn vbitmatvec2_equals_two_singles() {
11716 let (rows, cols) = (6, 64);
11717 let ng = cols / GROUP_SIZE;
11718 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11719 let mut bytes = bits.clone();
11720 for g in 0..rows * ng {
11721 let s = 0.02 + 0.001 * g as f32;
11722 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11723 }
11724 for r in 0..rows {
11725 let b = bits[r] as usize;
11726 let (mut acc, mut nb) = (0u64, 0usize);
11727 let mut rowbytes = Vec::new();
11728 for i in 0..cols {
11729 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11730 acc = (acc << b) | v;
11731 nb += b;
11732 while nb >= 8 {
11733 nb -= 8;
11734 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11735 }
11736 }
11737 if nb > 0 {
11738 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11739 }
11740 bytes.extend_from_slice(&rowbytes);
11741 }
11742 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11743 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
11744 let offsets = vbit_row_offsets(&bytes, rows, cols);
11745
11746 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11747 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
11748 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
11749 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11750 vbitmatvec2(
11751 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
11752 );
11753 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
11754 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
11755 }
11756
11757 #[test]
11759 fn q4matvec2_equals_two_singles() {
11760 let (rows, cols) = (8, 128);
11761 let groups = rows * cols / GROUP_SIZE;
11762 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11763 for i in 0..groups * 16 {
11764 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11765 }
11766 for g in 0..groups {
11767 let s = 0.01 + 0.003 * g as f32;
11768 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11769 }
11770 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11773 x1[9] = 250.0;
11774 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11775
11776 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11777 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
11778 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
11779 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11780 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
11781 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
11782 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
11783 }
11784
11785 #[test]
11788 fn matvec_many_equals_separate_matvecs() {
11789 use crate::pool::Pool;
11790 let (r1, r2, cols) = (300, 200, 64);
11791 let mk = |salt: usize, rows: usize| {
11792 QTensor::from_f32(
11793 (0..rows * cols)
11794 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
11795 .collect(),
11796 rows,
11797 cols,
11798 )
11799 };
11800 let (a, b) = (mk(1, r1), mk(5, r2));
11801 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
11802 let pool = Pool::new(3);
11803
11804 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
11805 a.matvec(&x, &mut ea, Some(&pool));
11806 b.matvec(&x, &mut eb, Some(&pool));
11807 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
11808 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11809 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
11810 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
11811 }
11812
11813 #[test]
11818 fn q4tp_matvec_many_equals_separate_matvecs() {
11819 use crate::pool::Pool;
11820 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
11821
11822 let (r1, r2, cols) = (300usize, 200usize, 64usize);
11823 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
11824 "arch_name": "tiny-q4tp",
11825 "hidden_size": cols,
11826 "intermediate_size": cols * 2,
11827 "num_layers": 1,
11828 "num_attention_heads": 2,
11829 "num_kv_heads": 1,
11830 "head_dim": 32,
11831 "vocab_size": r1,
11832 "layer_types": ["FullAttention"],
11833 "rms_norm_eps": 1e-6,
11834 "max_position_embeddings": 8,
11835 "linear_conv_kernel_dim": 0,
11836 "linear_num_key_heads": 0,
11837 "linear_num_value_heads": 0
11838 }))
11839 .unwrap();
11840 let header = CmfHeader {
11841 format: "cmf".into(),
11842 version: CMF_VERSION,
11843 arch,
11844 quant_type: QuantType::Q4Block,
11845 provenance: None,
11846 tokenizer_config: None,
11847 section_hashes: None,
11848 skills: Vec::new(),
11849 shard: None,
11850 calibration: None,
11851 routing: None,
11852 };
11853 let specs = [
11854 TensorSpec {
11855 name: "q".into(),
11856 dtype: TensorDtype::Q4TiledP,
11857 shape: vec![r1, cols],
11858 data: synth_q4tp(r1, cols),
11859 },
11860 TensorSpec {
11861 name: "kv".into(),
11862 dtype: TensorDtype::Q4TiledP,
11863 shape: vec![r2, cols],
11864 data: synth_q4tp(r2, cols),
11865 },
11866 ];
11867 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
11868 std::fs::create_dir_all(&dir).unwrap();
11869 let path = dir.join("m.cmf");
11870 CmfModel::write(&path, &header, &specs, None, None).unwrap();
11871 let model = Arc::new(CmfModel::open(&path).unwrap());
11872 let (a, b) = (
11873 QTensor::from_model(&model, "q").unwrap(),
11874 QTensor::from_model(&model, "kv").unwrap(),
11875 );
11876 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
11877 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
11878 let x: Vec<f32> = (0..cols)
11879 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
11880 .collect();
11881 let pool = Pool::new(3);
11882 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
11883 a.matvec(&x, &mut ea, Some(&pool));
11884 b.matvec(&x, &mut eb, Some(&pool));
11885 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
11886 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11887 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
11888 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
11889 let _ = std::fs::remove_dir_all(&dir);
11890 }
11891
11892 #[test]
11895 fn batched_matmat_equals_per_position_matvec() {
11896 let (rows, cols, b) = (8, 64, 5);
11897 let groups = rows * cols / GROUP_SIZE;
11899 let mut q4 = Vec::new();
11900 for i in 0..groups * 16 {
11901 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11902 }
11903 for g in 0..groups {
11904 q4.extend_from_slice(
11905 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11906 );
11907 }
11908 let ng = cols / GROUP_SIZE;
11910 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
11911 let mut vb = bits.clone();
11912 for g in 0..rows * ng {
11913 vb.extend_from_slice(
11914 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
11915 );
11916 }
11917 for r in 0..rows {
11918 let bw = bits[r] as usize;
11919 let (mut acc, mut nb) = (0u64, 0usize);
11920 let mut rowbytes = Vec::new();
11921 for i in 0..cols {
11922 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11923 acc = (acc << bw) | v;
11924 nb += bw;
11925 while nb >= 8 {
11926 nb -= 8;
11927 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11928 }
11929 }
11930 if nb > 0 {
11931 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11932 }
11933 vb.extend_from_slice(&rowbytes);
11934 }
11935 let offsets = vbit_row_offsets(&vb, rows, cols);
11936
11937 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11938
11939 let mut got = vec![0f32; b * rows];
11941 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
11942 for bi in 0..b {
11943 let mut expect = vec![0f32; rows];
11944 q4matvec(
11945 &q4,
11946 &xs[bi * cols..(bi + 1) * cols],
11947 rows,
11948 cols,
11949 &mut expect,
11950 None,
11951 );
11952 assert_eq!(
11953 &got[bi * rows..(bi + 1) * rows],
11954 &expect[..],
11955 "q4 batch pos {bi}"
11956 );
11957 }
11958
11959 let mut got = vec![0f32; b * rows];
11961 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
11962 for bi in 0..b {
11963 let mut expect = vec![0f32; rows];
11964 vbitmatvec(
11965 &vb,
11966 &offsets,
11967 &xs[bi * cols..(bi + 1) * cols],
11968 rows,
11969 cols,
11970 &mut expect,
11971 None,
11972 );
11973 assert_eq!(
11974 &got[bi * rows..(bi + 1) * rows],
11975 &expect[..],
11976 "vbit batch pos {bi}"
11977 );
11978 }
11979 }
11980
11981 #[test]
11985 fn q4_tiled_matches_q4_block_bitexact() {
11986 let (rows, cols, b) = (8usize, 128usize, 3usize);
11987 let groups = rows * cols / GROUP_SIZE;
11988 let mut split = Vec::with_capacity(groups * 18);
11989 for i in 0..groups * 16 {
11990 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11991 }
11992 for g in 0..groups {
11993 split.extend_from_slice(
11994 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11995 );
11996 }
11997 let (packed, scales) = split.split_at(groups * 16);
11999 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
12000 for g in 0..groups {
12001 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
12002 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
12003 }
12004
12005 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12006 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12008
12009 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
12010 q4matvec(&split, &x1, rows, cols, &mut a, None);
12011 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
12012 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
12013
12014 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12015 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
12016 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
12017 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
12018 assert_eq!(a1, t1);
12019 assert_eq!(a2, t2);
12020
12021 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12022 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12023 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
12024 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
12025 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
12026 }
12027
12028 #[test]
12035 fn q4matvec_sdot_outlier_exact() {
12036 let (rows, cols) = (4, 128);
12037 let groups = rows * cols / GROUP_SIZE;
12038 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12039 for i in 0..groups * 16 {
12040 bytes.push(((i * 11 + 5) % 256) as u8);
12041 }
12042 for g in 0..groups {
12043 let s = 0.02 + 0.002 * g as f32;
12044 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12045 }
12046 let mut x: Vec<f32> = (0..cols)
12047 .map(|i| match i % 3 {
12048 0 => 1.0,
12049 1 => -1.0,
12050 _ => 0.0,
12051 })
12052 .collect();
12053 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
12056 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12057 let mut expect = vec![0.0f32; rows];
12058 for r in 0..rows {
12059 expect[r] = reference[r * cols..(r + 1) * cols]
12060 .iter()
12061 .zip(&x)
12062 .map(|(w, xv)| w * xv)
12063 .sum();
12064 }
12065 let mut got = vec![0.0f32; rows];
12066 q4matvec(&bytes, &x, rows, cols, &mut got, None);
12067 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12068 for r in 0..rows {
12069 assert!(
12070 (got[r] - expect[r]).abs() < 2e-3 * scale,
12071 "row {r}: {} vs {} (outlier term must be exact)",
12072 got[r],
12073 expect[r]
12074 );
12075 }
12076 }
12077
12078 #[test]
12082 fn q1t_matvec_matches_reference() {
12083 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
12084 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
12086 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
12087 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
12089 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
12090 let mut bytes = Vec::new();
12091 for r in 0..rows {
12092 for g in 0..gpr {
12093 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
12094 let mut c = [0u8; 7];
12095 for k in 0..GROUP_SIZE {
12096 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
12098 0
12099 } else {
12100 ((k + r * 3 + g) % 3) as u8 };
12102 cortiq_core::quant::q1t_pack(&mut c, k, code);
12103 }
12104 bytes.extend_from_slice(&c);
12105 }
12106 }
12107 let mut row_ptr = vec![0u32; rows + 1];
12110 for &(idx, _) in &outliers {
12111 row_ptr[idx as usize / cols + 1] += 1;
12112 }
12113 for r in 0..rows {
12114 row_ptr[r + 1] += row_ptr[r];
12115 }
12116 for &p in &row_ptr {
12117 bytes.extend_from_slice(&p.to_le_bytes());
12118 }
12119 for &(idx, v) in &outliers {
12120 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
12121 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
12122 }
12123
12124 let mut refw = vec![0f32; rows * cols];
12125 dequant_q1t(&bytes, rows, cols, &mut refw);
12126 let x: Vec<f32> = (0..cols)
12129 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12130 .collect();
12131 let mut expect = vec![0f32; rows];
12132 for r in 0..rows {
12133 let mut a = 0.0f32;
12134 for j in 0..cols {
12135 a += refw[r * cols + j] * x[j];
12136 }
12137 expect[r] = a;
12138 }
12139 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
12140 let mut got = vec![0f32; rows];
12141 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
12142 for r in 0..rows {
12143 assert!(
12144 (got[r] - expect[r]).abs() < tol(expect[r]),
12145 "row {r}: {} vs {}",
12146 got[r],
12147 expect[r]
12148 );
12149 }
12150 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
12152 let mut gm = vec![0f32; 2 * rows];
12153 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
12154 for r in 0..rows {
12155 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
12156 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
12157 }
12158 let xb: Vec<f32> = (0..cols)
12162 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
12163 .collect();
12164 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12165 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
12166 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
12167 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12168 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
12169 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
12170 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
12171 }
12172
12173 #[test]
12176 fn q1t_matvec2_odd_gpr_matches_singles() {
12177 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12178 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
12180 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12181 for r in 0..rows {
12182 for g in 0..gpr {
12183 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
12184 let mut c = [0u8; 7];
12185 for k in 0..GROUP_SIZE {
12186 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
12187 }
12188 bytes.extend_from_slice(&c);
12189 }
12190 }
12191 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12192 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12193 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12194 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12195 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12196 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12197 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12198 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
12199 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
12200 }
12201
12202 #[test]
12206 #[ignore]
12207 fn q1t_matvec2_speed() {
12208 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12209 use std::time::Instant;
12210 let (rows, cols) = (8192usize, 4096usize);
12211 let gpr = cols / GROUP_SIZE;
12212 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12213 for r in 0..rows {
12214 for g in 0..gpr {
12215 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12216 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12217 let mut c = [0u8; 7];
12218 for k in 0..GROUP_SIZE {
12219 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12220 }
12221 bytes.extend_from_slice(&c);
12222 }
12223 }
12224 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12225 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12226 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12227 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12228 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12230 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12231 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
12232 for _ in 0..8 {
12233 let t0 = Instant::now();
12234 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12235 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
12236 let t1 = Instant::now();
12237 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12238 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12239 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
12240 }
12241 assert_eq!(p1, s1);
12242 assert_eq!(p2, s2);
12243 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
12244 }
12245
12246 #[test]
12250 #[ignore]
12251 fn q1t_matvec_speed() {
12252 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
12253 use std::time::Instant;
12254 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
12256 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
12257 for r in 0..rows {
12258 for g in 0..gpr {
12259 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12260 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12261 let mut c = [0u8; 7];
12262 for k in 0..GROUP_SIZE {
12263 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12264 }
12265 bytes.extend_from_slice(&c);
12266 }
12267 }
12268 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
12270 let mut idx = 0usize;
12271 while idx < n {
12272 row_ptr[idx / cols + 1] += 1;
12273 idx += stride;
12274 }
12275 for r in 0..rows {
12276 row_ptr[r + 1] += row_ptr[r];
12277 }
12278 for &p in &row_ptr {
12279 bytes.extend_from_slice(&p.to_le_bytes());
12280 }
12281 let mut idx = 0usize;
12282 while idx < n {
12283 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
12284 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
12285 idx += stride;
12286 }
12287 let x: Vec<f32> = (0..cols)
12290 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12291 .collect();
12292 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
12293
12294 let slow = |out: &mut [f32]| {
12296 let mut buf = vec![0f32; cols];
12297 for r in 0..rows {
12298 for g in 0..gpr {
12299 let off = (r * gpr + g) * Q1T_TILE;
12300 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
12301 let codes = &bytes[off + 2..off + Q1T_TILE];
12302 for k in 0..GROUP_SIZE {
12303 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
12304 1 => s,
12305 2 => -s,
12306 _ => 0.0,
12307 };
12308 }
12309 }
12310 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
12311 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
12312 }
12313 };
12314 let iters = 5;
12315 let mut a = vec![0f32; rows];
12316 slow(&mut a); let t = Instant::now();
12318 for _ in 0..iters {
12319 slow(&mut a);
12320 }
12321 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
12322
12323 let mut b = vec![0f32; rows];
12324 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
12326 for _ in 0..iters {
12327 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
12328 }
12329 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
12330
12331 for r in 0..rows {
12332 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
12333 }
12334 println!(
12335 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
12336 slow_ms / fast_ms
12337 );
12338 }
12339}
12340
12341#[cfg(test)]
12342mod gemm_bench {
12343 #[test]
12354 #[ignore]
12355 fn q4tp_matmat_throughput() {
12356 let b: usize = std::env::var("CMF_BENCH_B")
12360 .ok()
12361 .and_then(|v| v.parse().ok())
12362 .unwrap_or(296);
12363 let (rows, cols) = (9216usize, 2304usize);
12364 let (_, _, _) = (rows, cols, b);
12365 let total =
12366 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
12367 .unwrap();
12368 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
12373 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
12374 let lo = cortiq_core::quant::f32_to_f16(-4.0);
12375 let step = cortiq_core::quant::f32_to_f16(0.1);
12376 for r in 0..rows {
12377 let o = params_off + r * 4;
12378 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
12379 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
12380 }
12381 let _ = codes_off;
12382 let xs: Vec<f32> = (0..b * cols)
12383 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
12384 .collect();
12385 let mut out = vec![0f32; b * rows];
12386 let pool = crate::pool::Pool::from_env();
12387 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12393 let reps: usize = std::env::var("CMF_BENCH_REPS")
12394 .ok()
12395 .and_then(|v| v.parse().ok())
12396 .unwrap_or(10);
12397 let mut best = [f64::MAX; 2];
12398 let mut sums = [0f32; 2];
12399 for _ in 0..reps {
12400 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
12401 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
12402 let t = std::time::Instant::now();
12403 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12404 best[k] = best[k].min(t.elapsed().as_secs_f64());
12405 sums[k] = out.iter().take(64).sum::<f32>();
12406 }
12407 }
12408 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
12409 for (k, name) in ["previous", "tuned "].iter().enumerate() {
12410 println!(
12411 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
12412 best[k] * 1e3,
12413 flops / best[k] / 1e9,
12414 sums[k]
12415 );
12416 }
12417 assert!(
12418 (sums[0] - sums[1]).abs() < 1e-2,
12419 "the tuned kernel changed the result: {} vs {}",
12420 sums[0],
12421 sums[1]
12422 );
12423 }
12424
12425 #[test]
12431 fn q4tp_matmat_blocked_matches_scalar() {
12432 use std::sync::atomic::Ordering::Relaxed;
12433 for &(rows, cols, b) in &[
12441 (64usize, 128usize, 7usize),
12442 (33, 96, 4),
12443 (16, 256, 9),
12444 (192, 2304, 37),
12445 ] {
12446 let total = cortiq_core::quant::expected_nbytes(
12447 cortiq_core::TensorDtype::Q4TiledP,
12448 &[rows, cols],
12449 )
12450 .unwrap();
12451 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
12452 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
12453 let lo = cortiq_core::quant::f32_to_f16(-4.0);
12454 let step = cortiq_core::quant::f32_to_f16(0.1);
12455 for r in 0..rows {
12456 let o = params_off + r * 4;
12457 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
12458 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
12459 }
12460 let xs: Vec<f32> = (0..b * cols)
12461 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
12462 .collect();
12463 let mut got = vec![0f32; b * rows];
12464 let mut want = vec![0f32; b * rows];
12465 let gpr = cols / 32;
12466 let view = super::Q4tpView::new(&bytes, rows, cols);
12467 let pool = crate::pool::Pool::from_env();
12468 super::Q4TP_ALT.store(2, Relaxed);
12469 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
12470 super::Q4TP_ALT.store(1, Relaxed);
12471 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
12472 super::Q4TP_ALT.store(0, Relaxed);
12473 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
12480 let (mut worst, mut at) = (0f32, 0usize);
12481 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
12482 if (g - w).abs() > worst {
12483 worst = (g - w).abs();
12484 at = i;
12485 }
12486 }
12487 assert!(
12488 worst <= 1e-4 * scale,
12489 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
12490 (scale {scale:.3e}) at cell {at}: {} vs {}",
12491 got[at],
12492 want[at]
12493 );
12494
12495 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
12502 for bi in 0..b {
12503 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
12504 for r in 0..rows {
12505 let mut sc = vec![0f32; gpr];
12506 view.scales_into(r, gpr, &mut sc);
12507 let mut exact = 0f64;
12508 for j in 0..cols {
12509 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
12510 exact += w as f64 * sq as f64 * act.xq[j] as f64;
12511 }
12512 exact *= act.sx as f64;
12513 for &(j, xv) in &act.outliers {
12514 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
12515 exact += w as f64 * sq as f64 * xv as f64;
12516 }
12517 let i = bi * rows + r;
12518 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
12519 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
12520 }
12521 }
12522 println!(
12523 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
12524 per-column {e_scalar:.3e}"
12525 );
12526 assert!(
12531 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
12532 "{rows}x{cols} b={b}: error against f64 too large — blocked \
12533 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
12534 );
12535 }
12536 }
12537
12538 #[test]
12542 #[ignore]
12543 fn q4t_matmat_throughput() {
12544 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
12545 let total =
12546 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
12547 .unwrap();
12548 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
12551 let sc = cortiq_core::quant::f32_to_f16(0.02);
12552 for t in bytes.chunks_mut(super::Q4_TILE) {
12553 t[..2].copy_from_slice(&sc.to_le_bytes());
12554 }
12555 let xs: Vec<f32> = (0..b * cols)
12556 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
12557 .collect();
12558 let mut out = vec![0f32; b * rows];
12559 let pool = crate::pool::Pool::from_env();
12560 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12561 let reps: usize = std::env::var("CMF_BENCH_REPS")
12562 .ok()
12563 .and_then(|v| v.parse().ok())
12564 .unwrap_or(10);
12565 let mut best = f64::MAX;
12566 for _ in 0..reps {
12567 let t = std::time::Instant::now();
12568 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12569 best = best.min(t.elapsed().as_secs_f64());
12570 }
12571 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
12572 println!(
12573 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
12574 best * 1e3,
12575 flops / best / 1e9,
12576 out.iter().take(64).sum::<f32>()
12577 );
12578 }
12579}