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 if FULL_GPU_Q8.get() {
167 return 1.0;
168 }
169 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
170 *F.get_or_init(|| {
171 std::env::var("CMF_GPU_SPLIT")
172 .ok()
173 .and_then(|v| v.parse::<f32>().ok())
174 .unwrap_or(0.5)
175 .clamp(0.0, 1.0)
176 })
177}
178
179impl QTensor {
180 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
181 debug_assert_eq!(data.len(), rows * cols);
182 Self::F32 { data, rows, cols }
183 }
184
185 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
188 let idx = model
191 .tensor_index(name)
192 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
193 let entry = &model.tensors[idx];
194 if entry.shape.len() != 2 {
195 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
196 }
197 let (rows, cols) = (entry.shape[0], entry.shape[1]);
198 let bytes = model.entry_bytes(entry);
199
200 match entry.dtype {
201 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
202 let n = rows * cols;
203 let scales_off = n;
204 let row_scale: Vec<f32> = (0..rows)
205 .map(|o| {
206 f16_to_f32(u16::from_le_bytes([
207 bytes[scales_off + o * 2],
208 bytes[scales_off + o * 2 + 1],
209 ]))
210 })
211 .collect();
212 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
213 let col_off = n + rows * 2;
214 (0..cols)
215 .map(|i| {
216 f16_to_f32(u16::from_le_bytes([
217 bytes[col_off + i * 2],
218 bytes[col_off + i * 2 + 1],
219 ]))
220 })
221 .collect()
222 } else {
223 Vec::new()
224 };
225 Ok(Self::Mapped {
226 model: model.clone(),
227 idx,
228 dtype: entry.dtype,
229 rows,
230 cols,
231 row_scale,
232 col_field,
233 vbit_offsets: Vec::new(),
234 repack: q8_repack(bytes, rows, cols),
235 })
236 }
237 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
239 model: model.clone(),
240 idx,
241 dtype: entry.dtype,
242 rows,
243 cols,
244 row_scale: Vec::new(),
245 col_field: Vec::new(),
246 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
247 repack: Vec::new(),
248 }),
249 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
253 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
254 let offsets: Vec<usize> = (0..=rows)
255 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
256 .collect();
257 Ok(Self::Mapped {
258 model: model.clone(),
259 idx,
260 dtype: entry.dtype,
261 rows,
262 cols,
263 row_scale: Vec::new(),
264 col_field: Vec::new(),
265 vbit_offsets: offsets,
266 repack: Vec::new(),
267 })
268 }
269 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
275 model: model.clone(),
276 idx,
277 dtype: entry.dtype,
278 rows,
279 cols,
280 row_scale: Vec::new(),
281 col_field: Vec::new(),
282 vbit_offsets: Vec::new(),
283 repack: Vec::new(),
284 }),
285 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
288 model: model.clone(),
289 idx,
290 dtype: entry.dtype,
291 rows,
292 cols,
293 row_scale: Vec::new(),
294 col_field: Vec::new(),
295 vbit_offsets: Vec::new(),
296 repack: Vec::new(),
297 }),
298 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
300 model: model.clone(),
301 idx,
302 dtype: entry.dtype,
303 rows,
304 cols,
305 row_scale: Vec::new(),
306 col_field: Vec::new(),
307 vbit_offsets: Vec::new(),
308 repack: Vec::new(),
309 }),
310 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
311 model: model.clone(),
312 idx,
313 dtype: entry.dtype,
314 rows,
315 cols,
316 row_scale: Vec::new(),
317 col_field: Vec::new(),
318 vbit_offsets: Vec::new(),
319 repack: Vec::new(),
320 }),
321 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
323 model: model.clone(),
324 idx,
325 dtype: entry.dtype,
326 rows,
327 cols,
328 row_scale: Vec::new(),
329 col_field: Vec::new(),
330 vbit_offsets: Vec::new(),
331 repack: Vec::new(),
332 }),
333 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
337 model: model.clone(),
338 idx,
339 dtype: entry.dtype,
340 rows,
341 cols,
342 row_scale: Vec::new(),
343 col_field: Vec::new(),
344 vbit_offsets: Vec::new(),
345 repack: Vec::new(),
346 }),
347 _ => {
349 let mut data = vec![0.0f32; rows * cols];
350 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
351 Ok(Self::from_f32(data, rows, cols))
352 }
353 }
354 }
355
356 pub(crate) fn is_q1(&self) -> bool {
359 matches!(
360 self,
361 Self::Mapped {
362 dtype: TensorDtype::Q1,
363 ..
364 }
365 )
366 }
367
368 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
371 match self {
372 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
373 _ => None,
374 }
375 }
376
377 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
384 if self.has_prism_contract() {
385 return None;
386 }
387 match self {
388 #[cfg(target_os = "macos")]
389 Self::Mapped {
390 dtype: TensorDtype::Q1T,
391 ..
392 } if !crate::gpu::metal_q1t_enabled() => None,
393 Self::Mapped {
394 idx,
395 dtype:
396 TensorDtype::Q1
397 | TensorDtype::Q1T
398 | TensorDtype::Q4Block
399 | TensorDtype::Q4Tiled
400 | TensorDtype::Q4TiledP
404 | TensorDtype::Q8Row
405 | TensorDtype::Q8_2f,
406 rows,
407 cols,
408 ..
409 } => Some((*idx, *rows, *cols)),
410 _ => None,
411 }
412 }
413
414 #[cfg(target_os = "macos")]
421 pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
422 if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
423 let name = &model.tensors[idx].name;
424 let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
425 let affine = kind == 9 && crate::prism::is_affine_target(model, name);
426 if forward && affine {
427 let e = model.tensors.get(idx)?;
428 return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
429 }
430 }
431 self.q1_parts()
432 }
433
434 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
440 if self.has_prism_contract() {
441 return None;
442 }
443 match self {
444 Self::Mapped {
445 idx,
446 dtype: TensorDtype::Q4Tiled,
447 rows,
448 cols,
449 ..
450 } => Some((*idx, *rows, *cols)),
451 _ => None,
452 }
453 }
454
455 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
459 if self.has_prism_contract() {
460 return None;
461 }
462 match self {
463 Self::Mapped {
464 idx,
465 dtype: TensorDtype::Q4TiledP,
466 rows,
467 cols,
468 ..
469 } => Some((*idx, *rows, *cols)),
470 _ => None,
471 }
472 }
473
474 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
479 if self.has_prism_contract() {
480 return None;
481 }
482 match self {
483 Self::Mapped {
484 idx,
485 dtype: TensorDtype::Q8Row,
486 rows,
487 cols,
488 row_scale,
489 col_field,
490 ..
491 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
492 _ => None,
493 }
494 }
495
496 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
500 match self {
501 Self::Mapped { dtype, .. } => Some(*dtype),
502 _ => None,
503 }
504 }
505
506 pub fn model_idx(&self) -> Option<usize> {
511 match self {
512 Self::Mapped { idx, .. } => Some(*idx),
513 _ => None,
514 }
515 }
516
517 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
522 match self {
523 Self::Mapped { model, .. } => Some(model.clone()),
524 _ => None,
525 }
526 }
527
528 pub(crate) fn has_prism_contract(&self) -> bool {
533 matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
534 }
535
536 pub fn rows(&self) -> usize {
537 match self {
538 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
539 }
540 }
541
542 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
545 if self.has_prism_contract() {
546 return None;
547 }
548 match self {
549 Self::Mapped {
550 model,
551 idx,
552 dtype: TensorDtype::Q4Tiled,
553 ..
554 } => Some((model, *idx)),
555 _ => None,
556 }
557 }
558
559 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
562 if self.has_prism_contract() {
563 return None;
564 }
565 match self {
566 Self::Mapped {
567 model,
568 idx,
569 dtype: TensorDtype::Q4TiledP,
570 ..
571 } => Some((model, *idx)),
572 _ => None,
573 }
574 }
575
576 pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
584 if self.has_prism_contract() {
585 return None;
586 }
587 match self {
588 Self::Mapped {
589 model,
590 idx,
591 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
592 ..
593 } => Some((model, *idx)),
594 _ => None,
595 }
596 }
597
598 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
601 if self.has_prism_contract() {
602 return None;
603 }
604 match self {
605 Self::Mapped {
606 model,
607 idx,
608 dtype: TensorDtype::Q2TiledP,
609 ..
610 } => Some((model, *idx)),
611 _ => None,
612 }
613 }
614
615 pub fn cols(&self) -> usize {
616 match self {
617 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
618 }
619 }
620
621 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
624 if self.has_prism_contract() {
625 return None;
626 }
627 match self {
628 Self::Mapped {
629 model,
630 idx,
631 dtype: TensorDtype::Q1,
632 ..
633 } => Some((model, *idx)),
634 _ => None,
635 }
636 }
637
638 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
648 if self.has_prism_contract() {
649 return None;
650 }
651 self.graph_weight_descriptor()
652 }
653
654 pub(crate) fn graph_weight_descriptor(
659 &self,
660 ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
661 match self {
662 Self::Mapped {
663 model,
664 idx,
665 dtype: TensorDtype::Q8Row,
666 row_scale,
667 ..
668 } => Some((model, *idx, 0, row_scale.as_slice())),
669 Self::Mapped {
670 model,
671 idx,
672 dtype: TensorDtype::Q1,
673 ..
674 } => Some((model, *idx, 1, &[])),
675 Self::Mapped {
680 model,
681 idx,
682 dtype: TensorDtype::Q4Tiled,
683 ..
684 } => Some((model, *idx, 5, &[])),
685 Self::Mapped {
689 model,
690 idx,
691 dtype: TensorDtype::Q4TiledP,
692 ..
693 } => Some((model, *idx, 6, &[])),
694 Self::Mapped {
695 model,
696 idx,
697 dtype: TensorDtype::Q4Block,
698 ..
699 } => Some((model, *idx, 2, &[])),
700 Self::Mapped {
705 model,
706 idx,
707 dtype: TensorDtype::Q8_2f,
708 ..
709 } => Some((model, *idx, 7, &[])),
710 Self::Mapped {
711 model,
712 idx,
713 dtype: TensorDtype::Q1T,
714 ..
715 } => Some((model, *idx, 3, &[])),
716 Self::Mapped {
720 model,
721 idx,
722 dtype: TensorDtype::Q2TiledP,
723 ..
724 } => Some((model, *idx, 9, &[])),
725 _ => None,
726 }
727 }
728
729 pub fn as_f32(&self) -> Option<&[f32]> {
732 match self {
733 Self::F32 { data, .. } => Some(data),
734 Self::Mapped { .. } => None,
735 }
736 }
737
738 fn quant_bytes(&self) -> &[u8] {
739 match self {
740 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
741 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
742 }
743 }
744
745 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
747 let cols = self.cols();
748 debug_assert_eq!(dst.len(), cols);
749 match self {
750 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
751 Self::Mapped {
752 model,
753 idx,
754 dtype,
755 row_scale,
756 col_field,
757 vbit_offsets,
758 ..
759 } => {
760 if *dtype == TensorDtype::Q4Tiled {
761 let bytes = self.quant_bytes();
762 let gpr = cols / GROUP_SIZE;
763 for gi in 0..gpr {
764 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
765 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
766 for (k, &b) in tile[2..].iter().enumerate() {
767 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
768 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
769 }
770 }
771 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
772 crate::prism::inverse_embedding(model, dst);
773 }
774 return;
775 }
776 if *dtype == TensorDtype::Q4TiledP {
777 let bytes = self.quant_bytes();
778 let gpr = cols / GROUP_SIZE;
779 let v = Q4tpView::new(bytes, self.rows(), cols);
780 let mut sc = vec![0f32; gpr];
781 v.scales_into(r, gpr, &mut sc);
782 for gi in 0..gpr {
783 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
784 let s = sc[gi];
785 for (k, &b) in tile.iter().enumerate() {
786 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
787 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
788 }
789 }
790 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
791 crate::prism::inverse_embedding(model, dst);
792 }
793 return;
794 }
795 if *dtype == TensorDtype::Q2TiledP {
796 let bytes = self.quant_bytes();
797 let gpr = cols / GROUP_SIZE;
798 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
799 let mut sc = vec![0f32; gpr];
800 v.scales_into(r, gpr, &mut sc);
801 for gi in 0..gpr {
802 let ch =
803 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
804 let s = sc[gi];
805 for (k, &b) in ch.iter().enumerate() {
806 for j in 0..4 {
807 let center = if crate::prism::is_affine_target(
808 model,
809 &model.tensors[*idx].name,
810 ) {
811 1.0
812 } else {
813 1.5
814 };
815 dst[gi * GROUP_SIZE + k * 4 + j] =
816 (((b >> (2 * j)) & 3) as f32 - center) * s;
817 }
818 }
819 }
820 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
821 crate::prism::inverse_embedding(model, dst);
822 }
823 return;
824 }
825 if *dtype == TensorDtype::Q4Block {
826 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
827 let gpr = cols / GROUP_SIZE;
828 for gi in 0..gpr {
829 let g = r * gpr + gi;
830 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
831 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
832 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
833 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
834 }
835 }
836 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
837 crate::prism::inverse_embedding(model, dst);
838 }
839 return;
840 }
841 if *dtype == TensorDtype::Q1 {
842 let bytes = self.quant_bytes();
843 let gpr = cols / GROUP_SIZE;
844 for gi in 0..gpr {
845 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
846 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
847 for (j, &b) in tile[2..].iter().enumerate() {
848 for k in 0..8 {
849 dst[gi * GROUP_SIZE + j * 8 + k] =
850 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
851 }
852 }
853 }
854 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
855 crate::prism::inverse_embedding(model, dst);
856 }
857 return;
858 }
859 if *dtype == TensorDtype::Q1T {
860 let bytes = self.quant_bytes();
861 let gpr = cols / GROUP_SIZE;
862 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
863 for gi in 0..gpr {
864 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
865 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
866 bytes[off],
867 bytes[off + 1],
868 ]));
869 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
870 for k in 0..GROUP_SIZE {
871 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
872 {
873 1 => s,
874 2 => -s,
875 _ => 0.0,
876 };
877 }
878 }
879 let rows = self.rows();
881 let entries = base_len + (rows + 1) * 4;
882 if entries <= bytes.len() {
883 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
884 let r0 = u32::from_le_bytes([
885 ptrs[r * 4],
886 ptrs[r * 4 + 1],
887 ptrs[r * 4 + 2],
888 ptrs[r * 4 + 3],
889 ]) as usize;
890 let r1 = u32::from_le_bytes([
891 ptrs[(r + 1) * 4],
892 ptrs[(r + 1) * 4 + 1],
893 ptrs[(r + 1) * 4 + 2],
894 ptrs[(r + 1) * 4 + 3],
895 ]) as usize;
896 let off = entries + r0 * 4;
897 for i in 0..r1 - r0 {
898 let item = &bytes[off + i * 4..off + i * 4 + 4];
899 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
900 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
901 item[2], item[3],
902 ]));
903 if c < cols {
904 dst[c] = v;
905 }
906 }
907 }
908 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
909 crate::prism::inverse_embedding(model, dst);
910 }
911 return;
912 }
913 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
914 let bytes = self.quant_bytes();
915 let rows = self.rows();
916 let ng = cols / GROUP_SIZE;
917 let bits = &bytes[..rows];
918 let sc_off = rows;
919 let off = vbit_offsets[r];
922 let b = bits[r] as usize;
923 let l = ((1usize << (b - 1)) - 1) as f32;
924 let data = &bytes[off..];
925 let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
926 for (i, d) in dst.iter_mut().enumerate() {
927 while nbits < b {
928 acc = (acc << 8) | data[byte_idx] as u64;
929 byte_idx += 1;
930 nbits += 8;
931 }
932 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
933 nbits -= b;
934 let so = (r * ng + i / GROUP_SIZE) * 2;
935 let sv = f16_to_f32(u16::from_le_bytes([
936 bytes[sc_off + so],
937 bytes[sc_off + so + 1],
938 ]));
939 *d = (u - l) * sv;
940 }
941 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
942 crate::prism::inverse_embedding(model, dst);
943 }
944 return;
945 }
946 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
947 let s = row_scale[r];
948 match dtype {
949 TensorDtype::Q8Row => {
950 for (d, &b) in dst.iter_mut().zip(q) {
951 *d = (b as i8) as f32 * s;
952 }
953 }
954 TensorDtype::Q8_2f => {
955 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
956 *d = (b as i8) as f32 * s * col_field[i];
957 }
958 }
959 _ => unreachable!(),
960 }
961 if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
962 crate::prism::inverse_embedding(model, dst);
963 }
964 }
965 }
966 }
967
968 pub fn sparse_col_ok(&self) -> bool {
973 match self {
974 Self::F32 { .. } => true,
975 Self::Mapped { dtype, .. } => {
976 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
977 }
978 }
979 }
980
981 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
985 let inter = self.cols();
986 let hidden = self.rows();
987 debug_assert_eq!(out.len(), hidden);
988 match self {
989 Self::F32 { data, .. } => {
990 for (k, o) in out.iter_mut().enumerate() {
991 *o += w * data[k * inter + c];
992 }
993 }
994 Self::Mapped {
995 dtype,
996 row_scale,
997 col_field,
998 ..
999 } => {
1000 let q = self.quant_bytes();
1001 let colf = if *dtype == TensorDtype::Q8_2f {
1002 col_field[c]
1003 } else {
1004 1.0
1005 };
1006 let wc = w * colf;
1007 for (k, o) in out.iter_mut().enumerate() {
1008 let b = q[k * inter + c] as i8 as f32;
1009 *o += wc * b * row_scale[k];
1010 }
1011 }
1012 }
1013 }
1014
1015 #[inline]
1024 pub fn prefetch_row(&self, r: usize) {
1025 let Self::Mapped { dtype, .. } = self else {
1026 return;
1027 };
1028 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1029 return;
1030 }
1031 let cols = self.cols();
1032 let q = self.quant_bytes();
1033 let (a, b) = (r * cols, (r + 1) * cols);
1034 if b > q.len() {
1035 return;
1036 }
1037 let mut j = a;
1038 while j < b {
1039 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1040 j += 512;
1041 }
1042 }
1043
1044 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1053 let cols = self.cols();
1054 debug_assert_eq!(out.len(), cols);
1055 match self {
1056 Self::F32 { data, .. } => {
1057 let row = &data[r * cols..(r + 1) * cols];
1058 for (o, v) in out.iter_mut().zip(row) {
1059 *o += w * v;
1060 }
1061 }
1062 Self::Mapped {
1063 dtype,
1064 row_scale,
1065 col_field,
1066 ..
1067 } => match dtype {
1068 TensorDtype::Q8Row => {
1069 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1070 let ws = w * row_scale[r];
1071 let row: &[i8] =
1072 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1073 axpy_i8_f32(out, row, ws);
1074 }
1075 TensorDtype::Q8_2f => {
1076 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1077 let ws = w * row_scale[r];
1078 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1079 *o += ws * c * (*b as i8 as f32);
1080 }
1081 }
1082 _ => {
1083 self.row_f32(r, scratch);
1084 for (o, v) in out.iter_mut().zip(scratch.iter()) {
1085 *o += w * v;
1086 }
1087 }
1088 },
1089 }
1090 }
1091
1092 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1096 let cols = self.cols();
1097 match self {
1098 Self::F32 { data, .. } => {
1099 let row = &data[r * cols..(r + 1) * cols];
1100 row.iter().zip(x).map(|(w, v)| w * v).sum()
1101 }
1102 Self::Mapped {
1103 model,
1104 idx,
1105 dtype,
1106 row_scale,
1107 col_field,
1108 ..
1109 } => {
1110 let prism_forward =
1111 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1112 if prism_forward {
1113 let transformed = crate::prism::forward(model, &x[..cols]);
1114 let gpr = cols / GROUP_SIZE;
1115 match dtype {
1116 TensorDtype::Q2TiledP => {
1117 let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1118 let mut sc = vec![0f32; gpr];
1119 v.scales_into(r, gpr, &mut sc);
1120 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1121 return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1122 }
1123 return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1124 }
1125 _ => {
1126 self.row_f32(r, scratch);
1127 return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1128 }
1129 }
1130 }
1131 match dtype {
1132 TensorDtype::Q8Row => {
1133 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1134 dot_i8_f32(q, x) * row_scale[r]
1135 }
1136 TensorDtype::Q8_2f => {
1137 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1138 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1139 }
1140 _ => {
1141 self.row_f32(r, scratch);
1142 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1143 }
1144 }
1145 }
1146 }
1147 }
1148
1149 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1152 match self {
1153 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
1157 Self::Mapped {
1158 model,
1159 idx,
1160 dtype,
1161 rows,
1162 cols,
1163 row_scale,
1164 col_field,
1165 vbit_offsets,
1166 repack,
1167 } => {
1168 let _ = (model, idx);
1169 assert!(
1179 out.len() >= *rows && x.len() >= *cols,
1180 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1181 out.len(),
1182 x.len(),
1183 );
1184 let prism_forward =
1185 crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1186 if *dtype == TensorDtype::Q2TiledP
1187 && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1188 {
1189 use std::sync::atomic::{AtomicUsize, Ordering};
1190 static N: AtomicUsize = AtomicUsize::new(0);
1191 let n = N.fetch_add(1, Ordering::Relaxed);
1192 if n < 128 {
1193 eprintln!(
1194 "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1195 model.tensors[*idx].name,
1196 prism_forward,
1197 rows,
1198 cols,
1199 crate::gpu::enabled_here(),
1200 crate::gpu::q2tp_gpu_opt_in(),
1201 crate::gpu::cur_layer(),
1202 );
1203 }
1204 }
1205 if prism_forward {
1210 let transformed = crate::prism::forward(model, &x[..*cols]);
1211 match dtype {
1212 TensorDtype::Q4Block => {
1213 q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1214 }
1215 TensorDtype::Q4Tiled => {
1216 q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1217 }
1218 TensorDtype::Q4TiledP => {
1219 q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1220 }
1221 TensorDtype::Q2TiledP => {
1222 let affine =
1223 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1224 if *rows * *cols >= 8_388_608
1225 && crate::gpu::enabled_here()
1226 && crate::gpu::q2tp_gpu_opt_in()
1227 {
1228 let gpu_ok = if affine {
1229 crate::gpu::q2tp_affine_matvec(
1230 model,
1231 *idx,
1232 &transformed,
1233 *rows,
1234 *cols,
1235 out,
1236 )
1237 } else {
1238 crate::gpu::q2tp_matvec(
1239 model,
1240 *idx,
1241 &transformed,
1242 *rows,
1243 *cols,
1244 out,
1245 )
1246 };
1247 if gpu_ok {
1248 return;
1249 }
1250 }
1251 if affine {
1252 q2tp_affine_matvec(
1253 self.quant_bytes(),
1254 &transformed,
1255 *rows,
1256 *cols,
1257 out,
1258 pool,
1259 )
1260 } else {
1261 q2tp_matvec(
1262 self.quant_bytes(),
1263 &transformed,
1264 *rows,
1265 *cols,
1266 out,
1267 pool,
1268 )
1269 }
1270 }
1271 TensorDtype::Q1 => {
1272 q1_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1273 }
1274 TensorDtype::Q1T => {
1275 q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1276 }
1277 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1278 self.quant_bytes(),
1279 vbit_offsets,
1280 &transformed,
1281 *rows,
1282 *cols,
1283 out,
1284 pool,
1285 ),
1286 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1287 self.quant_bytes(),
1288 repack,
1289 row_scale,
1290 &transformed,
1291 col_field,
1292 *dtype,
1293 *rows,
1294 *cols,
1295 out,
1296 pool,
1297 ),
1298 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1299 }
1300 return;
1301 }
1302 if *dtype == TensorDtype::Q4Block {
1303 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1307 let t0 = std::time::Instant::now();
1308 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1309 crate::gpu::ProbeArm::Gpu => {
1310 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1311 crate::gpu::probe_record(
1312 crate::gpu::OpClass::Matvec,
1313 true,
1314 t0.elapsed(),
1315 );
1316 return;
1317 }
1318 }
1319 crate::gpu::ProbeArm::CpuTimed => {
1320 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1321 crate::gpu::probe_record(
1322 crate::gpu::OpClass::Matvec,
1323 false,
1324 t0.elapsed(),
1325 );
1326 return;
1327 }
1328 crate::gpu::ProbeArm::Cpu => {}
1329 }
1330 }
1331 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1332 return;
1333 }
1334 if *dtype == TensorDtype::Q4Tiled {
1335 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1340 let t0 = std::time::Instant::now();
1341 let cls = crate::gpu::matvec_class(*rows, *cols);
1342 match crate::gpu::probe_arm(cls) {
1343 crate::gpu::ProbeArm::Gpu => {
1344 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1345 crate::gpu::probe_record(cls, true, t0.elapsed());
1346 return;
1347 }
1348 }
1349 crate::gpu::ProbeArm::CpuTimed => {
1350 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1351 crate::gpu::probe_record(cls, false, t0.elapsed());
1352 return;
1353 }
1354 crate::gpu::ProbeArm::Cpu => {}
1355 }
1356 }
1357 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1358 return;
1359 }
1360 if *dtype == TensorDtype::Q4TiledP {
1361 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1367 let t0 = std::time::Instant::now();
1368 let cls = crate::gpu::matvec_class(*rows, *cols);
1369 match crate::gpu::probe_arm(cls) {
1370 crate::gpu::ProbeArm::Gpu => {
1371 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1372 crate::gpu::probe_record(cls, true, t0.elapsed());
1373 return;
1374 }
1375 }
1376 crate::gpu::ProbeArm::CpuTimed => {
1377 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1378 crate::gpu::probe_record(cls, false, t0.elapsed());
1379 return;
1380 }
1381 crate::gpu::ProbeArm::Cpu => {}
1382 }
1383 }
1384 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1385 return;
1386 }
1387 if *dtype == TensorDtype::Q2TiledP {
1388 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1389 return;
1390 }
1391 if *dtype == TensorDtype::Q1 {
1392 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1397 let t0 = std::time::Instant::now();
1398 let arm = if crate::gpu::q1_force() {
1399 crate::gpu::ProbeArm::Gpu
1400 } else {
1401 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1402 };
1403 match arm {
1404 crate::gpu::ProbeArm::Gpu => {
1405 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1406 crate::gpu::probe_record(
1407 crate::gpu::OpClass::Matvec,
1408 true,
1409 t0.elapsed(),
1410 );
1411 return;
1412 }
1413 }
1414 crate::gpu::ProbeArm::CpuTimed => {
1415 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1416 crate::gpu::probe_record(
1417 crate::gpu::OpClass::Matvec,
1418 false,
1419 t0.elapsed(),
1420 );
1421 return;
1422 }
1423 crate::gpu::ProbeArm::Cpu => {}
1424 }
1425 }
1426 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1427 return;
1428 }
1429 if *dtype == TensorDtype::Q1T {
1430 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1434 let t0 = std::time::Instant::now();
1435 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1436 crate::gpu::ProbeArm::Gpu => {
1437 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1438 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1439 crate::gpu::probe_record(
1440 crate::gpu::OpClass::Matvec,
1441 true,
1442 t0.elapsed(),
1443 );
1444 return;
1445 }
1446 }
1447 crate::gpu::ProbeArm::CpuTimed => {
1448 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1449 crate::gpu::probe_record(
1450 crate::gpu::OpClass::Matvec,
1451 false,
1452 t0.elapsed(),
1453 );
1454 return;
1455 }
1456 crate::gpu::ProbeArm::Cpu => {}
1457 }
1458 }
1459 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1460 return;
1461 }
1462 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1463 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1464 return;
1465 }
1466 let xs = prescale(x, col_field, *dtype);
1467 if *rows >= crate::gpu::min_rows()
1472 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1473 && gpu_lmhead_enabled()
1474 && crate::gpu::enabled_here()
1475 {
1476 let t0 = std::time::Instant::now();
1479 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1480 crate::gpu::ProbeArm::Gpu => {}
1481 crate::gpu::ProbeArm::CpuTimed => {
1482 qmatvec(
1483 self.quant_bytes(),
1484 repack,
1485 row_scale,
1486 x,
1487 col_field,
1488 *dtype,
1489 *rows,
1490 *cols,
1491 out,
1492 pool,
1493 );
1494 crate::gpu::probe_record(
1495 crate::gpu::OpClass::Matvec,
1496 false,
1497 t0.elapsed(),
1498 );
1499 return;
1500 }
1501 crate::gpu::ProbeArm::Cpu => {
1502 qmatvec(
1503 self.quant_bytes(),
1504 repack,
1505 row_scale,
1506 x,
1507 col_field,
1508 *dtype,
1509 *rows,
1510 *cols,
1511 out,
1512 pool,
1513 );
1514 return;
1515 }
1516 }
1517 let frac = gpu_split_frac();
1518 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1519 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1520 let bytes = self.quant_bytes();
1521 let ok = std::thread::scope(|sc| {
1522 let g = sc.spawn(|| {
1523 crate::gpu::q8_matvec_range(
1524 model,
1525 *idx,
1526 cpu_rows,
1527 &row_scale[cpu_rows..],
1528 &xs,
1529 *rows - cpu_rows,
1530 *cols,
1531 out_gpu,
1532 )
1533 });
1534 if cpu_rows > 0 {
1535 let rep_cpu = if repack.is_empty() {
1538 &[][..]
1539 } else {
1540 &repack[..(cpu_rows / 4) * 4 * *cols]
1541 };
1542 qmatvec(
1543 &bytes[..cpu_rows * *cols],
1544 rep_cpu,
1545 &row_scale[..cpu_rows],
1546 x,
1547 col_field,
1548 *dtype,
1549 cpu_rows,
1550 *cols,
1551 out_cpu,
1552 pool,
1553 );
1554 }
1555 g.join().unwrap_or(false)
1556 });
1557 if ok {
1558 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1559 return;
1560 }
1561 qmatvec(
1564 &bytes[cpu_rows * *cols..(*rows) * *cols],
1565 &[],
1566 &row_scale[cpu_rows..],
1567 x,
1568 col_field,
1569 *dtype,
1570 *rows - cpu_rows,
1571 *cols,
1572 out_gpu,
1573 pool,
1574 );
1575 return;
1576 }
1577 qmatvec(
1578 self.quant_bytes(),
1579 repack,
1580 row_scale,
1581 x,
1582 col_field,
1583 *dtype,
1584 *rows,
1585 *cols,
1586 out,
1587 pool,
1588 );
1589 }
1590 }
1591 }
1592
1593 pub fn matvec2(
1595 &self,
1596 x1: &[f32],
1597 x2: &[f32],
1598 o1: &mut [f32],
1599 o2: &mut [f32],
1600 pool: Option<&Pool>,
1601 ) {
1602 match self {
1603 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1604 Self::Mapped {
1605 model,
1606 idx,
1607 dtype,
1608 rows,
1609 cols,
1610 row_scale,
1611 col_field,
1612 vbit_offsets,
1613 ..
1614 } => {
1615 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1616 let tx1 = crate::prism::forward(model, &x1[..*cols]);
1617 let tx2 = crate::prism::forward(model, &x2[..*cols]);
1618 match dtype {
1619 TensorDtype::Q4Block => {
1620 q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1621 }
1622 TensorDtype::Q4Tiled => {
1623 q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1624 }
1625 TensorDtype::Q4TiledP => {
1626 q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1627 }
1628 TensorDtype::Q2TiledP => {
1629 if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1630 q2tp_affine_matvec2(
1631 self.quant_bytes(),
1632 &tx1,
1633 &tx2,
1634 *rows,
1635 *cols,
1636 o1,
1637 o2,
1638 pool,
1639 )
1640 } else {
1641 q2tp_matvec2(
1642 self.quant_bytes(),
1643 &tx1,
1644 &tx2,
1645 *rows,
1646 *cols,
1647 o1,
1648 o2,
1649 pool,
1650 )
1651 }
1652 }
1653 TensorDtype::Q1 => {
1654 q1_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1655 }
1656 TensorDtype::Q1T => {
1657 q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1658 }
1659 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1660 self.quant_bytes(),
1661 vbit_offsets,
1662 &tx1,
1663 &tx2,
1664 *rows,
1665 *cols,
1666 o1,
1667 o2,
1668 pool,
1669 ),
1670 TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1671 self.quant_bytes(),
1672 row_scale,
1673 &tx1,
1674 &tx2,
1675 col_field,
1676 *dtype,
1677 *rows,
1678 *cols,
1679 o1,
1680 o2,
1681 pool,
1682 ),
1683 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1684 }
1685 return;
1686 }
1687 if *dtype == TensorDtype::Q4Block {
1688 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1689 return;
1690 }
1691 if *dtype == TensorDtype::Q4Tiled {
1692 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1693 return;
1694 }
1695 if *dtype == TensorDtype::Q4TiledP {
1696 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1697 return;
1698 }
1699 if *dtype == TensorDtype::Q2TiledP {
1700 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1701 return;
1702 }
1703 if *dtype == TensorDtype::Q1 {
1704 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1705 return;
1706 }
1707 if *dtype == TensorDtype::Q1T {
1708 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1714 return;
1715 }
1716 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1717 vbitmatvec2(
1718 self.quant_bytes(),
1719 vbit_offsets,
1720 x1,
1721 x2,
1722 *rows,
1723 *cols,
1724 o1,
1725 o2,
1726 pool,
1727 );
1728 return;
1729 }
1730 qmatvec2(
1731 self.quant_bytes(),
1732 row_scale,
1733 x1,
1734 x2,
1735 col_field,
1736 *dtype,
1737 *rows,
1738 *cols,
1739 o1,
1740 o2,
1741 pool,
1742 );
1743 }
1744 }
1745 }
1746}
1747
1748impl QTensor {
1749 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1757 if self.has_prism_contract() {
1758 return None;
1759 }
1760 match self {
1761 Self::Mapped {
1762 model, idx, dtype, ..
1763 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1764 _ => None,
1765 }
1766 }
1767
1768 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1769 let cols = self.cols();
1770 let rows = self.rows();
1771 debug_assert_eq!(xs_all.len(), b * cols);
1772 debug_assert_eq!(out.len(), b * rows);
1773 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1774 if crate::gptq_capture::capturing() {
1778 if let Self::Mapped { model, idx, .. } = self {
1779 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1780 }
1781 }
1782 match self {
1783 Self::F32 { data, .. } => {
1784 let out_addr = SendMut(out.as_mut_ptr());
1785 let run = |start: usize, end: usize| {
1786 for o in start..end {
1787 let row = &data[o * cols..(o + 1) * cols];
1788 for bi in 0..b {
1789 let x = &xs_all[bi * cols..(bi + 1) * cols];
1790 let mut acc = 0f32;
1791 for j in 0..cols {
1792 acc += row[j] * x[j];
1793 }
1794 unsafe { *out_addr.at(bi * rows + o) = acc };
1795 }
1796 }
1797 };
1798 dispatch_rows(pool, rows, &run);
1799 }
1800 Self::Mapped {
1801 model,
1802 idx,
1803 dtype,
1804 row_scale,
1805 col_field,
1806 vbit_offsets,
1807 ..
1808 } => {
1809 if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1810 let mut transformed = Vec::with_capacity(xs_all.len());
1811 for bi in 0..b {
1812 transformed.extend_from_slice(&crate::prism::forward(
1813 model,
1814 &xs_all[bi * cols..(bi + 1) * cols],
1815 ));
1816 }
1817 match dtype {
1818 TensorDtype::Q4Block => {
1819 q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1820 }
1821 TensorDtype::Q4Tiled => {
1822 q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1823 }
1824 TensorDtype::Q4TiledP => {
1825 q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1826 }
1827 TensorDtype::Q2TiledP => {
1828 let affine =
1829 crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1830 let gpu_batch_ok = if affine {
1836 b >= 2
1837 } else {
1838 b >= 32 && b * rows * cols >= 128_000_000
1839 };
1840 if gpu_batch_ok
1841 && cols % 32 == 0
1842 && crate::gpu::enabled_here()
1843 && crate::gpu::q2tp_gpu_opt_in()
1844 {
1845 let gpu_ok = if affine {
1846 crate::gpu::q2tp_affine_matmat(
1847 model,
1848 *idx,
1849 &transformed,
1850 b,
1851 rows,
1852 cols,
1853 out,
1854 )
1855 } else {
1856 crate::gpu::q2tp_matmat(
1857 model,
1858 *idx,
1859 &transformed,
1860 b,
1861 rows,
1862 cols,
1863 out,
1864 )
1865 };
1866 if gpu_ok {
1867 return;
1868 }
1869 }
1870 if affine
1874 && b == 1
1875 && cols % 32 == 0
1876 && crate::gpu::enabled_here()
1877 && crate::gpu::q2tp_gpu_opt_in()
1878 && crate::gpu::q2tp_affine_matvec(
1879 model,
1880 *idx,
1881 &transformed[..cols],
1882 rows,
1883 cols,
1884 &mut out[..rows],
1885 )
1886 {
1887 return;
1888 }
1889 if affine {
1890 q2tp_affine_matmat(
1891 self.quant_bytes(),
1892 &transformed,
1893 b,
1894 rows,
1895 cols,
1896 out,
1897 pool,
1898 )
1899 } else {
1900 q2tp_matmat(
1901 self.quant_bytes(),
1902 &transformed,
1903 b,
1904 rows,
1905 cols,
1906 out,
1907 pool,
1908 )
1909 }
1910 }
1911 TensorDtype::Q1 => {
1912 q1_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1913 }
1914 TensorDtype::Q1T => {
1915 q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1916 }
1917 TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
1918 self.quant_bytes(),
1919 vbit_offsets,
1920 &transformed,
1921 b,
1922 rows,
1923 cols,
1924 out,
1925 pool,
1926 ),
1927 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1928 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1929 .map(|bi| {
1930 prescale(
1931 &transformed[bi * cols..(bi + 1) * cols],
1932 col_field,
1933 *dtype,
1934 )
1935 })
1936 .collect();
1937 qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
1938 }
1939 _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1940 }
1941 return;
1942 }
1943 if *dtype == TensorDtype::Q4Block {
1944 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1945 return;
1946 }
1947 if *dtype == TensorDtype::Q4TiledP {
1948 if b >= 32
1960 && b * rows * cols >= 128_000_000
1961 && cols % 32 == 0
1962 && !crate::gpu::mm_killed()
1963 && crate::gpu::enabled_here()
1964 {
1965 let class = if b >= 128 {
1966 crate::gpu::OpClass::MatmatWide
1967 } else {
1968 crate::gpu::OpClass::Matmat
1969 };
1970 if let Self::Mapped { model, idx, .. } = self {
1971 if crate::mm_ab::on() {
1983 let mut g = vec![0f32; b * rows];
1984 let t = std::time::Instant::now();
1985 let took = crate::gpu::q4tp_matmat(
1986 model, *idx, xs_all, b, rows, cols, &mut g,
1987 );
1988 let dg = t.elapsed();
1989 let t = std::time::Instant::now();
1990 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1991 let dc = t.elapsed();
1992 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1993 return;
1994 }
1995 let t0 = std::time::Instant::now();
1996 let resident = crate::gpu::weight_is_resident(model, *idx);
2000 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2001 crate::gpu::ProbeArm::Gpu => {
2002 if crate::gpu::q4tp_matmat(
2003 model, *idx, xs_all, b, rows, cols, out,
2004 ) {
2005 let el = t0.elapsed();
2006 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2016 let budget = std::time::Duration::from_secs_f64(
2017 flops / 1.5e12 * 8.0 + 0.020,
2018 );
2019 crate::gpu::mm_budget_check(
2020 "q4tp matmat",
2021 el,
2022 budget,
2023 crate::gpu::probe_was_cold() || !resident,
2024 );
2025 crate::gpu::probe_record(class, true, el);
2026 return;
2027 }
2028 }
2029 crate::gpu::ProbeArm::CpuTimed => {
2030 q4tp_matmat(
2031 self.quant_bytes(),
2032 xs_all,
2033 b,
2034 rows,
2035 cols,
2036 out,
2037 pool,
2038 );
2039 crate::gpu::probe_record(class, false, t0.elapsed());
2040 return;
2041 }
2042 crate::gpu::ProbeArm::Cpu => {}
2043 }
2044 }
2045 }
2046 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2047 return;
2048 }
2049 if *dtype == TensorDtype::Q2TiledP {
2050 if b >= 32
2056 && b * rows * cols >= 128_000_000
2057 && cols % 32 == 0
2058 && !crate::gpu::mm_killed()
2059 && crate::gpu::enabled_here()
2060 {
2061 let class = if b >= 128 {
2062 crate::gpu::OpClass::MatmatWide
2063 } else {
2064 crate::gpu::OpClass::Matmat
2065 };
2066 if let Self::Mapped { model, idx, .. } = self {
2067 let t0 = std::time::Instant::now();
2068 match crate::gpu::probe_arm(class) {
2069 crate::gpu::ProbeArm::Gpu => {
2070 if crate::gpu::q2tp_matmat(
2071 model, *idx, xs_all, b, rows, cols, out,
2072 ) {
2073 crate::gpu::probe_record(class, true, t0.elapsed());
2074 return;
2075 }
2076 }
2077 crate::gpu::ProbeArm::CpuTimed => {
2078 q2tp_matmat(
2079 self.quant_bytes(),
2080 xs_all,
2081 b,
2082 rows,
2083 cols,
2084 out,
2085 pool,
2086 );
2087 crate::gpu::probe_record(class, false, t0.elapsed());
2088 return;
2089 }
2090 crate::gpu::ProbeArm::Cpu => {}
2091 }
2092 }
2093 }
2094 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2099 return;
2100 }
2101 if *dtype == TensorDtype::Q4Tiled {
2102 if b >= 32
2114 && b * rows * cols >= 128_000_000
2115 && cols % 32 == 0
2116 && !crate::gpu::mm_killed()
2117 && crate::gpu::enabled_here()
2118 {
2119 let class = if b >= 128 {
2120 crate::gpu::OpClass::MatmatWide
2121 } else {
2122 crate::gpu::OpClass::Matmat
2123 };
2124 if let Self::Mapped { model, idx, .. } = self {
2125 let t0 = std::time::Instant::now();
2126 match crate::gpu::probe_arm(class) {
2127 crate::gpu::ProbeArm::Gpu => {
2128 if crate::gpu::q4t_matmat(
2129 model, *idx, xs_all, b, rows, cols, out,
2130 ) {
2131 let el = t0.elapsed();
2132 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2142 let budget = std::time::Duration::from_secs_f64(
2143 flops / 1.5e12 * 8.0 + 0.020,
2144 );
2145 crate::gpu::mm_budget_check(
2146 "q4t matmat",
2147 el,
2148 budget,
2149 crate::gpu::probe_was_cold(),
2150 );
2151 crate::gpu::probe_record(class, true, el);
2152 return;
2153 }
2154 }
2155 crate::gpu::ProbeArm::CpuTimed => {
2156 q4t_matmat(
2157 self.quant_bytes(),
2158 xs_all,
2159 b,
2160 rows,
2161 cols,
2162 out,
2163 pool,
2164 );
2165 crate::gpu::probe_record(class, false, t0.elapsed());
2166 return;
2167 }
2168 crate::gpu::ProbeArm::Cpu => {}
2169 }
2170 }
2171 }
2172 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2173 return;
2174 }
2175 if *dtype == TensorDtype::Q1 {
2176 if b >= 32
2179 && b * rows * cols >= 128_000_000
2180 && cols % 64 == 0
2181 && crate::gpu::enabled_here()
2182 {
2183 if let Self::Mapped { model, idx, .. } = self {
2184 let t0 = std::time::Instant::now();
2185 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2186 crate::gpu::ProbeArm::Gpu => {
2187 if crate::gpu::q1_matmat(
2188 model, *idx, xs_all, b, rows, cols, out,
2189 ) {
2190 crate::gpu::probe_record(
2191 crate::gpu::OpClass::Matmat,
2192 true,
2193 t0.elapsed(),
2194 );
2195 return;
2196 }
2197 }
2198 crate::gpu::ProbeArm::CpuTimed => {
2199 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2200 crate::gpu::probe_record(
2201 crate::gpu::OpClass::Matmat,
2202 false,
2203 t0.elapsed(),
2204 );
2205 return;
2206 }
2207 crate::gpu::ProbeArm::Cpu => {}
2208 }
2209 }
2210 }
2211 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2212 return;
2213 }
2214 if *dtype == TensorDtype::Q1T {
2215 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2218 if let Self::Mapped { model, idx, .. } = self {
2219 let t0 = std::time::Instant::now();
2220 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2221 crate::gpu::ProbeArm::Gpu => {
2222 if crate::gpu::q1t_matmat(
2223 model, *idx, xs_all, b, rows, cols, out,
2224 ) {
2225 crate::gpu::probe_record(
2226 crate::gpu::OpClass::Matmat,
2227 true,
2228 t0.elapsed(),
2229 );
2230 return;
2231 }
2232 }
2233 crate::gpu::ProbeArm::CpuTimed => {
2234 q1t_matmat(
2235 self.quant_bytes(),
2236 xs_all,
2237 b,
2238 rows,
2239 cols,
2240 out,
2241 pool,
2242 );
2243 crate::gpu::probe_record(
2244 crate::gpu::OpClass::Matmat,
2245 false,
2246 t0.elapsed(),
2247 );
2248 return;
2249 }
2250 crate::gpu::ProbeArm::Cpu => {}
2251 }
2252 }
2253 }
2254 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2255 return;
2256 }
2257 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2258 vbitmatmat(
2259 self.quant_bytes(),
2260 vbit_offsets,
2261 xs_all,
2262 b,
2263 rows,
2264 cols,
2265 out,
2266 pool,
2267 );
2268 return;
2269 }
2270 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2271 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2272 .collect();
2273 if row_exact()
2279 && (1..=4).contains(&b)
2280 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2281 && crate::gpu::enabled_here()
2282 && crate::gpu::wgpu_active()
2283 {
2284 let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2285 if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2286 return;
2287 }
2288 }
2289 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2295 if let Self::Mapped { model, idx, .. } = self {
2296 let t0 = std::time::Instant::now();
2297 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2298 crate::gpu::ProbeArm::Gpu
2299 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2300 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2301 {
2302 let q = self.quant_bytes();
2306 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2307 return;
2308 }
2309 crate::gpu::ProbeArm::Gpu => {
2310 let flat: Vec<f32> =
2311 pre.iter().flat_map(|v| v.iter().copied()).collect();
2312 if crate::gpu::q8_matmat(
2313 model, *idx, row_scale, &flat, b, rows, cols, out,
2314 ) {
2315 crate::gpu::probe_record(
2316 crate::gpu::OpClass::Matmat,
2317 true,
2318 t0.elapsed(),
2319 );
2320 return;
2321 }
2322 }
2323 crate::gpu::ProbeArm::CpuTimed => {
2324 let q = self.quant_bytes();
2325 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2326 crate::gpu::probe_record(
2327 crate::gpu::OpClass::Matmat,
2328 false,
2329 t0.elapsed(),
2330 );
2331 return;
2332 }
2333 crate::gpu::ProbeArm::Cpu => {}
2334 }
2335 }
2336 }
2337 let q = self.quant_bytes();
2338 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2339 }
2340 }
2341 }
2342}
2343
2344impl QTensor {
2345 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2355 let (rows, cols) = (self.rows(), self.cols());
2356 let Self::Mapped {
2357 model,
2358 idx,
2359 dtype,
2360 row_scale,
2361 col_field,
2362 ..
2363 } = self
2364 else {
2365 return false;
2366 };
2367 if crate::prism::has_contract(model) {
2368 return false;
2369 }
2370 match *dtype {
2371 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2372 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2376 if *dtype == TensorDtype::Q8_2f
2379 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2380 && crate::gpu::q8_matmat_2f(
2381 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2382 )
2383 {
2384 return true;
2385 }
2386 let flat: Vec<f32> = (0..b)
2387 .flat_map(|bi| {
2388 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2389 })
2390 .collect();
2391 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2392 }
2393 _ => false,
2394 }
2395 }
2396
2397 pub fn matvec_many<const N: usize>(
2404 ts: [&QTensor; N],
2405 x: &[f32],
2406 mut outs: [&mut [f32]; N],
2407 pool: Option<&Pool>,
2408 ) {
2409 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2410 if ts.iter().any(|t| t.has_prism_contract()) {
2411 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2415 t.matvec(x, o, pool);
2416 }
2417 return;
2418 }
2419 let uniform_q8 = ts.iter().all(|t| {
2420 matches!(
2421 t,
2422 Self::Mapped {
2423 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2424 ..
2425 }
2426 )
2427 });
2428 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2429 let uniform_q4 = ts.iter().all(|t| {
2430 matches!(
2431 t,
2432 Self::Mapped {
2433 dtype: TensorDtype::Q4Block,
2434 ..
2435 }
2436 )
2437 });
2438 let uniform_vbit = ts.iter().all(|t| {
2439 matches!(
2440 t,
2441 Self::Mapped {
2442 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2443 ..
2444 }
2445 )
2446 });
2447 let uniform_q1 = ts.iter().all(|t| {
2448 matches!(
2449 t,
2450 Self::Mapped {
2451 dtype: TensorDtype::Q1,
2452 ..
2453 }
2454 )
2455 });
2456 let uniform_q1t = ts.iter().all(|t| {
2457 matches!(
2458 t,
2459 Self::Mapped {
2460 dtype: TensorDtype::Q1T,
2461 ..
2462 }
2463 )
2464 });
2465 let uniform_q4tp = ts.iter().all(|t| {
2470 matches!(
2471 t,
2472 Self::Mapped {
2473 dtype: TensorDtype::Q4TiledP,
2474 ..
2475 }
2476 )
2477 }) && ts
2478 .iter()
2479 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2480 let Some(pool) = pool else {
2481 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2482 t.matvec(x, o, None);
2483 }
2484 return;
2485 };
2486 if total_rows < 256
2487 || !(uniform_q8
2488 || uniform_f32
2489 || uniform_q4
2490 || uniform_vbit
2491 || uniform_q1
2492 || uniform_q1t
2493 || uniform_q4tp)
2494 {
2495 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2496 t.matvec(x, o, Some(pool));
2497 }
2498 return;
2499 }
2500
2501 if uniform_q4tp {
2502 let cols = ts[0].cols();
2508 let gpr = cols / GROUP_SIZE;
2509 let views: Vec<Q4tpView> = ts
2510 .iter()
2511 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2512 .collect();
2513 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2514 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2515 let locate = |flat: usize| -> (usize, usize) {
2517 let mut acc = 0;
2518 for (i, &r) in rows_of.iter().enumerate() {
2519 if flat < acc + r {
2520 return (i, flat - acc);
2521 }
2522 acc += r;
2523 }
2524 (rows_of.len() - 1, 0)
2525 };
2526 let (views, outs_addr) = (&views, &outs_addr);
2527 if a8w8_enabled() {
2528 let act = split_act(x);
2529 let act = &act;
2530 let run = |start: usize, end: usize| {
2531 let mut sc = vec![0f32; gpr];
2532 for flat in start..end {
2533 let (t, r) = locate(flat);
2534 let v = &views[t];
2535 v.scales_into(r, gpr, &mut sc);
2536 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2537 for &(j, xv) in &act.outliers {
2538 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2539 acc += w * s * xv;
2540 }
2541 unsafe { *outs_addr[t].at(r) = acc };
2543 }
2544 };
2545 pool.run_rows(total_rows, &run);
2546 } else {
2547 let run = |start: usize, end: usize| {
2548 let mut sc = vec![0f32; gpr];
2549 for flat in start..end {
2550 let (t, r) = locate(flat);
2551 let v = &views[t];
2552 v.scales_into(r, gpr, &mut sc);
2553 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2555 }
2556 };
2557 pool.run_rows(total_rows, &run);
2558 }
2559 return;
2560 }
2561
2562 if uniform_q1 {
2563 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2566 if a8w8_enabled() {
2567 let act = split_act(x);
2568 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2569 let (act, gsum) = (&act, &gsum);
2570 let closures: [_; N] = std::array::from_fn(|i| {
2571 let (bytes, gpr, out) =
2572 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2573 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2574 });
2575 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2576 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2577 pool.run_many(&parts);
2578 } else {
2579 let closures: [_; N] = std::array::from_fn(|i| {
2580 let (bytes, gpr, out) =
2581 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2582 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2583 });
2584 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2585 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2586 pool.run_many(&parts);
2587 }
2588 return;
2589 }
2590
2591 if uniform_q1t {
2592 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2596 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2597 if a8w8_enabled() {
2598 let act = split_act(x);
2599 let act = &act;
2600 let x_ref = x;
2601 let closures: [_; N] = std::array::from_fn(|i| {
2602 let bytes = ts[i].quant_bytes();
2603 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2604 let gpr = cols / GROUP_SIZE;
2605 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2606 let out = outs_addr[i];
2607 move |s: usize, e: usize| {
2608 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2609 }
2610 });
2611 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2612 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2613 pool.run_many(&parts);
2614 } else {
2615 let x_ref = x;
2616 let closures: [_; N] = std::array::from_fn(|i| {
2617 let bytes = ts[i].quant_bytes();
2618 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2619 let gpr = cols / GROUP_SIZE;
2620 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2621 let out = outs_addr[i];
2622 move |s: usize, e: usize| {
2623 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2624 }
2625 });
2626 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2627 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2628 pool.run_many(&parts);
2629 }
2630 return;
2631 }
2632
2633 if uniform_q4 || uniform_vbit {
2634 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2635 if a8w8_enabled() {
2637 let act = split_act(x);
2638 let act = &act;
2639 if uniform_q4 {
2640 let closures: [_; N] = std::array::from_fn(|i| {
2641 let (packed, scales) =
2642 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2643 let (gpr, cols, out) =
2644 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2645 move |s: usize, e: usize| {
2646 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2647 }
2648 });
2649 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2650 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2651 pool.run_many(&parts);
2652 } else {
2653 let closures: [_; N] = std::array::from_fn(|i| {
2654 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2655 unreachable!()
2656 };
2657 let (bytes, rows, cols, out) = (
2658 ts[i].quant_bytes(),
2659 ts[i].rows(),
2660 ts[i].cols(),
2661 outs_addr[i],
2662 );
2663 move |s: usize, e: usize| {
2664 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2665 }
2666 });
2667 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2668 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2669 pool.run_many(&parts);
2670 }
2671 return;
2672 }
2673 if uniform_q4 {
2674 let closures: [_; N] = std::array::from_fn(|i| {
2675 let (packed, scales) =
2676 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2677 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2678 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2679 });
2680 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2681 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2682 pool.run_many(&parts);
2683 } else {
2684 let closures: [_; N] = std::array::from_fn(|i| {
2685 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2686 unreachable!()
2687 };
2688 let (bytes, rows, cols, out) = (
2689 ts[i].quant_bytes(),
2690 ts[i].rows(),
2691 ts[i].cols(),
2692 outs_addr[i],
2693 );
2694 move |s: usize, e: usize| {
2695 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2696 }
2697 });
2698 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2699 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2700 pool.run_many(&parts);
2701 }
2702 return;
2703 }
2704
2705 if uniform_f32 {
2706 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2707 let closures: [_; N] = std::array::from_fn(|i| {
2708 let Self::F32 { data, cols, .. } = ts[i] else {
2709 unreachable!()
2710 };
2711 let out = outs_addr[i];
2712 move |start: usize, end: usize| {
2713 for o in start..end {
2714 let row = &data[o * cols..(o + 1) * cols];
2715 let mut sum = 0.0f32;
2716 for j in 0..*cols {
2717 sum += row[j] * x[j];
2718 }
2719 unsafe { *out.at(o) = sum };
2721 }
2722 }
2723 });
2724 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2725 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2726 pool.run_many(&parts);
2727 return;
2728 }
2729
2730 struct Ctx<'a> {
2733 bytes: &'a [u8],
2734 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2735 rep: &'a [u8],
2736 row_scale: &'a [f32],
2737 cols: usize,
2738 xs: std::borrow::Cow<'a, [f32]>,
2739 }
2740 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2741 let Self::Mapped {
2742 dtype,
2743 cols,
2744 row_scale,
2745 col_field,
2746 repack,
2747 ..
2748 } = ts[i]
2749 else {
2750 unreachable!()
2751 };
2752 Ctx {
2753 bytes: ts[i].quant_bytes(),
2754 rep: repack,
2755 row_scale,
2756 cols: *cols,
2757 xs: prescale(x, col_field, *dtype),
2758 }
2759 });
2760 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2761 #[cfg(target_arch = "aarch64")]
2762 if sdot_enabled() {
2763 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2764 let closures: [_; N] = std::array::from_fn(|i| {
2765 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2766 move |start: usize, end: usize| {
2767 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2768 }
2769 });
2770 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2771 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2772 pool.run_many(&parts);
2773 return;
2774 }
2775 #[cfg(target_arch = "x86_64")]
2776 if avx2_a8w8_enabled() {
2777 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2778 let closures: [_; N] = std::array::from_fn(|i| {
2779 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2780 move |start: usize, end: usize| {
2781 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2782 }
2783 });
2784 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2785 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2786 pool.run_many(&parts);
2787 return;
2788 }
2789 let closures: [_; N] = std::array::from_fn(|i| {
2790 let (c, out) = (&ctxs[i], outs_addr[i]);
2791 move |start: usize, end: usize| {
2792 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2793 }
2794 });
2795 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2796 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2797 pool.run_many(&parts);
2798 }
2799}
2800
2801impl QTensor {
2802 #[allow(clippy::needless_range_loop)]
2807 pub fn matvec2_many<const N: usize>(
2808 ts: [&QTensor; N],
2809 x1: &[f32],
2810 x2: &[f32],
2811 mut o1s: [&mut [f32]; N],
2812 mut o2s: [&mut [f32]; N],
2813 pool: Option<&Pool>,
2814 ) {
2815 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2816 if ts.iter().any(|t| t.has_prism_contract()) {
2817 for i in 0..N {
2818 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2819 }
2820 return;
2821 }
2822 let uniform_q8 = ts.iter().all(|t| {
2823 matches!(
2824 t,
2825 Self::Mapped {
2826 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2827 ..
2828 }
2829 )
2830 });
2831 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2832 let uniform_q4 = ts.iter().all(|t| {
2833 matches!(
2834 t,
2835 Self::Mapped {
2836 dtype: TensorDtype::Q4Block,
2837 ..
2838 }
2839 )
2840 });
2841 let uniform_vbit = ts.iter().all(|t| {
2842 matches!(
2843 t,
2844 Self::Mapped {
2845 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2846 ..
2847 }
2848 )
2849 });
2850 let fusable = pool.is_some()
2851 && total_rows >= 256
2852 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2853 if !fusable {
2854 for i in 0..N {
2855 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2856 }
2857 return;
2858 }
2859 let pool = pool.unwrap();
2860
2861 if uniform_q4 || uniform_vbit {
2862 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2863 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2864 if a8w8_enabled() {
2866 let a1 = split_act(x1);
2867 let a2 = split_act(x2);
2868 let (a1, a2) = (&a1, &a2);
2869 if uniform_q4 {
2870 let closures: [_; N] = std::array::from_fn(|i| {
2871 let (packed, scales) =
2872 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2873 let (gpr, cols, o1, o2) =
2874 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2875 move |s: usize, e: usize| {
2876 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2877 }
2878 });
2879 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2880 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2881 pool.run_many(&parts);
2882 } else {
2883 let closures: [_; N] = std::array::from_fn(|i| {
2884 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2885 unreachable!()
2886 };
2887 let (bytes, rows, cols, o1, o2) = (
2888 ts[i].quant_bytes(),
2889 ts[i].rows(),
2890 ts[i].cols(),
2891 p1[i],
2892 p2[i],
2893 );
2894 move |s: usize, e: usize| {
2895 vbit_range2_a8w8(
2896 bytes,
2897 vbit_offsets,
2898 x1,
2899 x2,
2900 a1,
2901 a2,
2902 rows,
2903 cols,
2904 o1,
2905 o2,
2906 s,
2907 e,
2908 )
2909 }
2910 });
2911 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2912 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2913 pool.run_many(&parts);
2914 }
2915 return;
2916 }
2917 if uniform_q4 {
2918 let closures: [_; N] = std::array::from_fn(|i| {
2919 let (packed, scales) =
2920 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2921 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2922 move |s: usize, e: usize| {
2923 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2924 }
2925 });
2926 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2927 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2928 pool.run_many(&parts);
2929 } else {
2930 let closures: [_; N] = std::array::from_fn(|i| {
2931 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2932 unreachable!()
2933 };
2934 let (bytes, rows, cols, o1, o2) = (
2935 ts[i].quant_bytes(),
2936 ts[i].rows(),
2937 ts[i].cols(),
2938 p1[i],
2939 p2[i],
2940 );
2941 move |s: usize, e: usize| {
2942 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2943 }
2944 });
2945 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2946 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2947 pool.run_many(&parts);
2948 }
2949 return;
2950 }
2951
2952 if uniform_f32 {
2953 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2954 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2955 let closures: [_; N] = std::array::from_fn(|i| {
2956 let Self::F32 { data, cols, .. } = ts[i] else {
2957 unreachable!()
2958 };
2959 let (o1, o2) = (p1[i], p2[i]);
2960 move |start: usize, end: usize| {
2961 for o in start..end {
2962 let row = &data[o * cols..(o + 1) * cols];
2963 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2964 for j in 0..*cols {
2965 s1 += row[j] * x1[j];
2966 s2 += row[j] * x2[j];
2967 }
2968 unsafe {
2970 *o1.at(o) = s1;
2971 *o2.at(o) = s2;
2972 }
2973 }
2974 }
2975 });
2976 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2977 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2978 pool.run_many(&parts);
2979 return;
2980 }
2981
2982 struct Ctx<'a> {
2983 bytes: &'a [u8],
2984 row_scale: &'a [f32],
2985 cols: usize,
2986 xs1: std::borrow::Cow<'a, [f32]>,
2987 xs2: std::borrow::Cow<'a, [f32]>,
2988 }
2989 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2990 let Self::Mapped {
2991 dtype,
2992 cols,
2993 row_scale,
2994 col_field,
2995 ..
2996 } = ts[i]
2997 else {
2998 unreachable!()
2999 };
3000 Ctx {
3001 bytes: ts[i].quant_bytes(),
3002 row_scale,
3003 cols: *cols,
3004 xs1: prescale(x1, col_field, *dtype),
3005 xs2: prescale(x2, col_field, *dtype),
3006 }
3007 });
3008 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3009 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3010 #[cfg(target_arch = "aarch64")]
3011 if sdot_enabled() {
3012 let acts: [(SplitAct, SplitAct); N] =
3013 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3014 let closures: [_; N] = std::array::from_fn(|i| {
3015 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3016 move |start: usize, end: usize| {
3017 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3018 }
3019 });
3020 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3021 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3022 pool.run_many(&parts);
3023 return;
3024 }
3025 #[cfg(target_arch = "x86_64")]
3026 if avx2_a8w8_enabled() {
3027 let acts: [(SplitAct, SplitAct); N] =
3028 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3029 let closures: [_; N] = std::array::from_fn(|i| {
3030 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3031 move |start: usize, end: usize| {
3032 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3033 }
3034 });
3035 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3036 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3037 pool.run_many(&parts);
3038 return;
3039 }
3040 let closures: [_; N] = std::array::from_fn(|i| {
3041 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3042 move |start: usize, end: usize| {
3043 q8_range2_f32(
3044 c.bytes,
3045 c.row_scale,
3046 &c.xs1,
3047 &c.xs2,
3048 c.cols,
3049 o1,
3050 o2,
3051 start,
3052 end,
3053 )
3054 }
3055 });
3056 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3057 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3058 pool.run_many(&parts);
3059 }
3060
3061 pub fn matvec_silu_mul(
3066 gate: &QTensor,
3067 up: &QTensor,
3068 x: &[f32],
3069 out: &mut [f32],
3070 pool: Option<&Pool>,
3071 ) -> bool {
3072 Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3073 }
3074
3075 pub fn matvec_silu_mul_limited(
3081 gate: &QTensor,
3082 up: &QTensor,
3083 x: &[f32],
3084 out: &mut [f32],
3085 limit: f32,
3086 pool: Option<&Pool>,
3087 ) -> bool {
3088 if gate.has_prism_contract() || up.has_prism_contract() {
3089 return false;
3093 }
3094 let inter = gate.rows();
3095 debug_assert_eq!(up.rows(), inter);
3096 debug_assert_eq!(out.len(), inter);
3097 debug_assert_eq!(gate.cols(), up.cols());
3098 if !a8w8_enabled() {
3099 return false;
3100 }
3101 let act = split_act(x);
3102 let act = &act;
3103 let x_ref = x;
3104 let out_addr = SendMut(out.as_mut_ptr());
3105
3106 match (gate, up) {
3107 (
3109 Self::Mapped {
3110 dtype: TensorDtype::Q4Block,
3111 ..
3112 },
3113 Self::Mapped {
3114 dtype: TensorDtype::Q4Block,
3115 ..
3116 },
3117 ) => {
3118 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3119 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3120 let gpr = gate.cols() / GROUP_SIZE;
3121 let cols = gate.cols();
3122 let run = move |start: usize, end: usize| {
3123 for r in start..end {
3124 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3125 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3126 for &(j, xv) in &act.outliers {
3127 let flat = r * cols + j;
3128 let gb = gp[flat / 2];
3129 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3130 let gsc = f16_to_f32(u16::from_le_bytes([
3131 gs[(flat / GROUP_SIZE) * 2],
3132 gs[(flat / GROUP_SIZE) * 2 + 1],
3133 ]));
3134 gv += ((gn as i32 - 8) as f32) * gsc * xv;
3135 let ub = up_p[flat / 2];
3136 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3137 let usc = f16_to_f32(u16::from_le_bytes([
3138 up_s[(flat / GROUP_SIZE) * 2],
3139 up_s[(flat / GROUP_SIZE) * 2 + 1],
3140 ]));
3141 uv += ((un as i32 - 8) as f32) * usc * xv;
3142 }
3143 let silu_g = gv / (1.0 + (-gv).exp());
3144 unsafe { *out_addr.at(r) = silu_g * uv };
3146 }
3147 };
3148 dispatch_rows(pool, inter, &run);
3149 true
3150 }
3151 (
3155 Self::Mapped {
3156 dtype: TensorDtype::Q4Tiled,
3157 ..
3158 },
3159 Self::Mapped {
3160 dtype: TensorDtype::Q4Tiled,
3161 ..
3162 },
3163 ) => {
3164 let g_bytes = gate.quant_bytes();
3165 let u_bytes = up.quant_bytes();
3166 let gpr = gate.cols() / GROUP_SIZE;
3167 let run = move |start: usize, end: usize| {
3168 for r in start..end {
3169 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3170 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3171 for &(j, xv) in &act.outliers {
3172 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3173 gv += w * s * xv;
3174 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3175 uv += w * s * xv;
3176 }
3177 let silu_g = gv / (1.0 + (-gv).exp());
3178 unsafe { *out_addr.at(r) = silu_g * uv };
3180 }
3181 };
3182 dispatch_rows(pool, inter, &run);
3183 true
3184 }
3185 (
3188 Self::Mapped {
3189 dtype: TensorDtype::Q4TiledP,
3190 ..
3191 },
3192 Self::Mapped {
3193 dtype: TensorDtype::Q4TiledP,
3194 ..
3195 },
3196 ) => {
3197 let cols = gate.cols();
3198 let gpr = cols / GROUP_SIZE;
3199 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3200 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3201 let run = |start: usize, end: usize| {
3202 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3203 for r in start..end {
3204 gv_view.scales_into(r, gpr, &mut gsc);
3205 uv_view.scales_into(r, gpr, &mut usc);
3206 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3207 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3208 for &(j, xv) in &act.outliers {
3209 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3210 gv += w * s * xv;
3211 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3212 uv += w * s * xv;
3213 }
3214 let silu_g = gv / (1.0 + (-gv).exp());
3215 unsafe { *out_addr.at(r) = silu_g * uv };
3217 }
3218 };
3219 dispatch_rows(pool, inter, &run);
3220 true
3221 }
3222 (
3228 Self::Mapped {
3229 dtype: TensorDtype::Q1,
3230 ..
3231 },
3232 Self::Mapped {
3233 dtype: TensorDtype::Q1,
3234 ..
3235 },
3236 ) => {
3237 let g_bytes = gate.quant_bytes();
3238 let u_bytes = up.quant_bytes();
3239 let gpr = gate.cols() / GROUP_SIZE;
3240 let gsum = q1_group_sums(&act.xq, gpr);
3241 let gsum = &gsum;
3242 let run = move |start: usize, end: usize| {
3243 for r in start..end {
3244 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
3245 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
3246 for &(j, xv) in &act.outliers {
3247 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
3248 gv += w * s * xv;
3249 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
3250 uv += w * s * xv;
3251 }
3252 let silu_g = gv / (1.0 + (-gv).exp());
3253 unsafe { *out_addr.at(r) = silu_g * uv };
3255 }
3256 };
3257 dispatch_rows(pool, inter, &run);
3258 true
3259 }
3260 (
3264 Self::Mapped {
3265 dtype: TensorDtype::Q2TiledP,
3266 ..
3267 },
3268 Self::Mapped {
3269 dtype: TensorDtype::Q2TiledP,
3270 ..
3271 },
3272 ) => {
3273 let cols = gate.cols();
3274 let gpr = cols / GROUP_SIZE;
3275 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3276 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3277 let gsum = q1_group_sums(&act.xq, gpr);
3278 let gsum = &gsum;
3279 let run = move |start: usize, end: usize| {
3280 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3281 for r in start..end {
3282 gv_view.scales_into(r, gpr, &mut gsc);
3283 uv_view.scales_into(r, gpr, &mut usc);
3284 let mut gv =
3285 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
3286 let mut uv =
3287 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
3288 for &(j, xv) in &act.outliers {
3289 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3290 gv += w * s * xv;
3291 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
3292 uv += w * s * xv;
3293 }
3294 let silu_g = gv / (1.0 + (-gv).exp());
3295 unsafe { *out_addr.at(r) = silu_g * uv };
3297 }
3298 };
3299 dispatch_rows(pool, inter, &run);
3300 true
3301 }
3302 (
3307 Self::Mapped {
3308 dtype: TensorDtype::Q8Row,
3309 row_scale: g_rs,
3310 ..
3311 },
3312 Self::Mapped {
3313 dtype: TensorDtype::Q8Row,
3314 row_scale: u_rs,
3315 ..
3316 },
3317 ) => {
3318 let g_bytes = gate.quant_bytes();
3319 let u_bytes = up.quant_bytes();
3320 let cols = gate.cols();
3321 let run = move |start: usize, end: usize| {
3322 for r in start..end {
3323 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3324 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3325 let silu_g = gv / (1.0 + (-gv).exp());
3326 unsafe { *out_addr.at(r) = silu_g * uv };
3328 }
3329 };
3330 dispatch_rows(pool, inter, &run);
3331 true
3332 }
3333 (
3335 Self::Mapped {
3336 dtype: TensorDtype::Q1T,
3337 ..
3338 },
3339 Self::Mapped {
3340 dtype: TensorDtype::Q1T,
3341 ..
3342 },
3343 ) => {
3344 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3345 let g_bytes = gate.quant_bytes();
3346 let u_bytes = up.quant_bytes();
3347 let gpr = gate.cols() / GROUP_SIZE;
3348 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3349 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3350 let run = move |start: usize, end: usize| {
3351 for r in start..end {
3352 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3353 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3354 for &(j, xv) in &act.outliers {
3355 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3356 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3357 }
3358 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3359 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3360 let silu_g = gv / (1.0 + (-gv).exp());
3361 unsafe { *out_addr.at(r) = silu_g * uv };
3363 }
3364 };
3365 dispatch_rows(pool, inter, &run);
3366 true
3367 }
3368 _ => false,
3369 }
3370 }
3371
3372 pub fn moe_gate_up_many(
3387 pairs: &[(&QTensor, &QTensor)],
3388 x: &[f32],
3389 outs: &mut [Vec<f32>],
3390 pool: Option<&Pool>,
3391 ) -> bool {
3392 if pairs.is_empty() || pairs.len() != outs.len() {
3393 return false;
3394 }
3395 if !a8w8_enabled() {
3396 let groups = vec![vec![0]; pairs.len()];
3397 return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3398 }
3399 let inter = pairs[0].0.rows();
3400 let cols = pairs[0].0.cols();
3401 if cols % GROUP_SIZE != 0 {
3402 return false;
3403 }
3404 let gpr = cols / GROUP_SIZE;
3405 let q2 = matches!(
3408 pairs[0].0,
3409 Self::Mapped {
3410 dtype: TensorDtype::Q2TiledP,
3411 ..
3412 }
3413 );
3414 let want = if q2 {
3415 TensorDtype::Q2TiledP
3416 } else {
3417 TensorDtype::Q4TiledP
3418 };
3419 let mut views = Vec::with_capacity(pairs.len() * 2);
3420 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3421 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3422 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3423 if !both
3424 || g.rows() != inter
3425 || u.rows() != inter
3426 || g.cols() != cols
3427 || u.cols() != cols
3428 || o.len() != inter
3429 {
3430 return false;
3431 }
3432 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3433 views.push(mk(g.quant_bytes(), inter, cols));
3434 views.push(mk(u.quant_bytes(), inter, cols));
3435 }
3436 let act = split_act(x);
3437 let gsum = if q2 {
3438 q1_group_sums(&act.xq, gpr)
3439 } else {
3440 Vec::new()
3441 };
3442 let (act, gsum) = (&act, &gsum);
3443 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3444 let (views, ptrs) = (&views, &ptrs);
3445 let run = |start: usize, end: usize| {
3446 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3447 for flat in start..end {
3448 let (e, r) = (flat / inter, flat % inter);
3449 let gv_view = &views[e * 2];
3450 let uv_view = &views[e * 2 + 1];
3451 gv_view.scales_into(r, gpr, &mut gsc);
3452 uv_view.scales_into(r, gpr, &mut usc);
3453 let (mut gv, mut uv) = if q2 {
3454 (
3455 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
3456 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
3457 )
3458 } else {
3459 (
3460 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
3461 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
3462 )
3463 };
3464 for &(j, xv) in &act.outliers {
3465 let (og, ou) = if q2 {
3466 (
3467 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3468 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
3469 )
3470 } else {
3471 (
3472 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3473 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
3474 )
3475 };
3476 gv += og.0 * og.1 * xv;
3477 uv += ou.0 * ou.1 * xv;
3478 }
3479 let silu_g = gv / (1.0 + (-gv).exp());
3480 unsafe { *ptrs[e].at(r) = silu_g * uv };
3482 }
3483 };
3484 dispatch_rows(pool, pairs.len() * inter, &run);
3485 true
3486 }
3487
3488 pub fn moe_down_many(
3497 downs: &[&QTensor],
3498 gs: &[Vec<f32>],
3499 weights: &[f32],
3500 out: &mut [f32],
3501 pool: Option<&Pool>,
3502 ) -> bool {
3503 if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3504 return false;
3505 }
3506 if !a8w8_enabled() {
3507 let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3508 if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3509 return false;
3510 }
3511 out.fill(0.0);
3512 for (row, &w) in terms.iter().zip(weights) {
3513 for (o, &v) in out.iter_mut().zip(row) {
3514 *o += w * v;
3515 }
3516 }
3517 return true;
3518 }
3519 let rows = out.len();
3520 let cols = downs[0].cols();
3521 if cols % GROUP_SIZE != 0 {
3522 return false;
3523 }
3524 let gpr = cols / GROUP_SIZE;
3525 let mut views = Vec::with_capacity(downs.len());
3526 for (d, g) in downs.iter().zip(gs.iter()) {
3527 if !matches!(
3528 d,
3529 Self::Mapped {
3530 dtype: TensorDtype::Q4TiledP,
3531 ..
3532 }
3533 ) || d.rows() != rows
3534 || d.cols() != cols
3535 || g.len() != cols
3536 {
3537 return false;
3538 }
3539 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3540 }
3541 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3543 let out_addr = SendMut(out.as_mut_ptr());
3551 let (views, acts, weights) = (&views, &acts, &weights);
3552 let run = |start: usize, end: usize| {
3553 let mut sc = vec![0f32; gpr];
3554 for r in start..end {
3555 let mut acc = 0f32;
3556 for (e, v) in views.iter().enumerate() {
3557 v.scales_into(r, gpr, &mut sc);
3558 let a = &acts[e];
3559 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3560 for &(j, xv) in &a.outliers {
3561 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3562 d += w * s * xv;
3563 }
3564 acc += weights[e] * d;
3565 }
3566 unsafe { *out_addr.at(r) = acc };
3568 }
3569 };
3570 dispatch_rows(pool, rows, &run);
3571 true
3572 }
3573
3574 pub fn moe_gate_up_rows(
3586 pairs: &[(&QTensor, &QTensor)],
3587 groups: &[Vec<usize>],
3588 xs: &[f32],
3589 outs: &mut [Vec<f32>],
3590 pool: Option<&Pool>,
3591 ) -> bool {
3592 if pairs.is_empty() || pairs.len() != groups.len() {
3593 return false;
3594 }
3595 let inter = pairs[0].0.rows();
3596 let cols = pairs[0].0.cols();
3597 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
3598 if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
3599 return false;
3600 }
3601 let b = xs.len() / cols;
3602 let gpr = cols / GROUP_SIZE;
3603 let mut views = Vec::with_capacity(pairs.len() * 2);
3604 for (g, u) in pairs {
3605 let q4tp = |t: &QTensor| {
3606 matches!(
3607 t,
3608 Self::Mapped {
3609 dtype: TensorDtype::Q4TiledP,
3610 ..
3611 }
3612 )
3613 };
3614 if g.has_prism_contract()
3615 || u.has_prism_contract()
3616 || !q4tp(g)
3617 || !q4tp(u)
3618 || g.rows() != inter
3619 || u.rows() != inter
3620 || g.cols() != cols
3621 || u.cols() != cols
3622 {
3623 return false;
3624 }
3625 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
3626 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
3627 }
3628 if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
3629 return false;
3630 }
3631 let quantized = a8w8_enabled();
3632 let acts: Vec<SplitAct> = if quantized {
3633 (0..b)
3634 .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
3635 .collect()
3636 } else {
3637 Vec::new()
3638 };
3639 let mut offs = Vec::with_capacity(groups.len());
3640 let mut o = 0usize;
3641 for g in groups {
3642 offs.push(o);
3643 o += g.len();
3644 }
3645 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3646 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3647 let run = |start: usize, end: usize| {
3648 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3649 for flat in start..end {
3650 let (e, r) = (flat / inter, flat % inter);
3651 let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
3652 gv_view.scales_into(r, gpr, &mut gsc);
3653 uv_view.scales_into(r, gpr, &mut usc);
3654 for (k, &t) in groups[e].iter().enumerate() {
3655 if !quantized {
3656 let x = &xs[t * cols..(t + 1) * cols];
3657 let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
3658 let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
3659 unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
3660 continue;
3661 }
3662 let act = &acts[t];
3663 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3664 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3665 for &(j, xv) in &act.outliers {
3666 let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3667 let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3668 gv += og.0 * og.1 * xv;
3669 uv += ou.0 * ou.1 * xv;
3670 }
3671 let silu_g = gv / (1.0 + (-gv).exp());
3672 unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
3675 }
3676 }
3677 };
3678 dispatch_rows(pool, pairs.len() * inter, &run);
3679 true
3680 }
3681
3682 pub fn moe_down_rows(
3689 downs: &[&QTensor],
3690 group_lens: &[usize],
3691 gs: &[Vec<f32>],
3692 outs: &mut [Vec<f32>],
3693 pool: Option<&Pool>,
3694 ) -> bool {
3695 if downs.is_empty() || downs.len() != group_lens.len() {
3696 return false;
3697 }
3698 let rows = downs[0].rows();
3699 let cols = downs[0].cols();
3700 let n_pairs: usize = group_lens.iter().sum();
3701 if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
3702 return false;
3703 }
3704 let gpr = cols / GROUP_SIZE;
3705 let mut views = Vec::with_capacity(downs.len());
3706 for d in downs {
3707 if d.has_prism_contract()
3708 || !matches!(
3709 d,
3710 Self::Mapped {
3711 dtype: TensorDtype::Q4TiledP,
3712 ..
3713 }
3714 )
3715 || d.rows() != rows
3716 || d.cols() != cols
3717 {
3718 return false;
3719 }
3720 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3721 }
3722 if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
3723 return false;
3724 }
3725 let quantized = a8w8_enabled();
3726 let acts: Vec<SplitAct> = if quantized {
3727 gs.iter().map(|g| split_act(g)).collect()
3728 } else {
3729 Vec::new()
3730 };
3731 let mut offs = Vec::with_capacity(group_lens.len());
3732 let mut o = 0usize;
3733 for &l in group_lens {
3734 offs.push(o);
3735 o += l;
3736 }
3737 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3738 let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3739 let run = |start: usize, end: usize| {
3740 let mut sc = vec![0f32; gpr];
3741 for flat in start..end {
3742 let (e, r) = (flat / rows, flat % rows);
3743 let v = &views[e];
3744 v.scales_into(r, gpr, &mut sc);
3745 for k in 0..group_lens[e] {
3746 if !quantized {
3747 let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
3748 unsafe { *ptrs[offs[e] + k].at(r) = d };
3749 continue;
3750 }
3751 let a = &acts[offs[e] + k];
3752 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3753 for &(j, xv) in &a.outliers {
3754 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3755 d += w * s * xv;
3756 }
3757 unsafe { *ptrs[offs[e] + k].at(r) = d };
3759 }
3760 }
3761 };
3762 dispatch_rows(pool, downs.len() * rows, &run);
3763 true
3764 }
3765}
3766
3767#[cfg(target_os = "macos")]
3772mod accel_blas {
3773 #[link(name = "Accelerate", kind = "framework")]
3774 unsafe extern "C" {
3775 pub fn cblas_sgemm(
3776 order: i32,
3777 trans_a: i32,
3778 trans_b: i32,
3779 m: i32,
3780 n: i32,
3781 k: i32,
3782 alpha: f32,
3783 a: *const f32,
3784 lda: i32,
3785 b: *const f32,
3786 ldb: i32,
3787 beta: f32,
3788 c: *mut f32,
3789 ldc: i32,
3790 );
3791 }
3792}
3793
3794#[cfg(target_os = "macos")]
3795pub(crate) fn accel_gemm_enabled() -> bool {
3796 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3797 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3798}
3799
3800#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3803pub(crate) fn accel_gemm_enabled() -> bool {
3804 true
3805}
3806
3807#[cfg(target_arch = "aarch64")]
3814#[allow(clippy::too_many_arguments)]
3815pub(crate) fn neon_gemm_rm(
3816 m: usize,
3817 n: usize,
3818 k: usize,
3819 alpha: f32,
3820 a: &[f32],
3821 lda: usize,
3822 b_mat: &[f32],
3823 ldb: usize,
3824 b_rows_are_n: bool,
3825 c: &mut [f32],
3826 ldc: usize,
3827) {
3828 debug_assert!(a.len() >= (m - 1) * lda + k);
3829 debug_assert!(c.len() >= (m - 1) * ldc + n);
3830 unsafe {
3832 use core::arch::aarch64::*;
3833 let mut i = 0usize;
3834 while i < m {
3835 let mi = (m - i).min(4);
3836 let mut j = 0usize;
3837 while j < n {
3838 let nj = (n - j).min(8);
3839 if mi == 4 && nj == 8 {
3840 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3841 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3842 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3843 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3844 for p in 0..k {
3845 let (b0, b1) = if b_rows_are_n {
3846 let base = b_mat.as_ptr().add(j * ldb + p);
3849 let g = |o: usize| *base.add(o * ldb);
3850 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3851 } else {
3852 let base = b_mat.as_ptr().add(p * ldb + j);
3853 (
3854 [*base, *base.add(1), *base.add(2), *base.add(3)],
3855 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3856 )
3857 };
3858 let bv0 = vld1q_f32(b0.as_ptr());
3859 let bv1 = vld1q_f32(b1.as_ptr());
3860 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3861 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3862 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3863 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3864 c0a = vfmaq_f32(c0a, a0, bv0);
3865 c0b = vfmaq_f32(c0b, a0, bv1);
3866 c1a = vfmaq_f32(c1a, a1, bv0);
3867 c1b = vfmaq_f32(c1b, a1, bv1);
3868 c2a = vfmaq_f32(c2a, a2, bv0);
3869 c2b = vfmaq_f32(c2b, a2, bv1);
3870 c3a = vfmaq_f32(c3a, a3, bv0);
3871 c3b = vfmaq_f32(c3b, a3, bv1);
3872 }
3873 let al = vdupq_n_f32(alpha);
3874 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3875 .iter()
3876 .enumerate()
3877 {
3878 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3879 vst1q_f32(dst, vmulq_f32(*ca, al));
3880 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3881 }
3882 } else {
3883 for r in 0..mi {
3884 for q in 0..nj {
3885 let mut acc = 0f32;
3886 for p in 0..k {
3887 let bv = if b_rows_are_n {
3888 b_mat[(j + q) * ldb + p]
3889 } else {
3890 b_mat[p * ldb + j + q]
3891 };
3892 acc += a[(i + r) * lda + p] * bv;
3893 }
3894 c[(i + r) * ldc + j + q] = acc * alpha;
3895 }
3896 }
3897 }
3898 j += nj;
3899 }
3900 i += mi;
3901 }
3902 }
3903}
3904
3905#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3907#[allow(clippy::too_many_arguments)]
3908pub(crate) fn sgemm_rm(
3909 m: usize,
3910 n: usize,
3911 k: usize,
3912 alpha: f32,
3913 a: &[f32],
3914 lda: usize,
3915 b_mat: &[f32],
3916 ldb: usize,
3917 b_rows_are_n: bool,
3918 c: &mut [f32],
3919 ldc: usize,
3920) {
3921 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3922}
3923
3924#[allow(clippy::too_many_arguments)]
3928pub fn sgemm_public(
3929 m: usize,
3930 n: usize,
3931 k: usize,
3932 alpha: f32,
3933 a: &[f32],
3934 lda: usize,
3935 b_mat: &[f32],
3936 ldb: usize,
3937 b_rows_are_n: bool,
3938 c: &mut [f32],
3939 ldc: usize,
3940) {
3941 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3942 {
3943 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3944 }
3945 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3950 {
3951 for i in 0..m {
3952 for j in 0..n {
3953 let mut acc = 0f32;
3954 for p in 0..k {
3955 let bv = if b_rows_are_n {
3956 b_mat[j * ldb + p]
3957 } else {
3958 b_mat[p * ldb + j]
3959 };
3960 acc += a[i * lda + p] * bv;
3961 }
3962 c[i * ldc + j] = alpha * acc;
3963 }
3964 }
3965 }
3966}
3967
3968#[cfg(target_os = "macos")]
3971#[allow(clippy::too_many_arguments)]
3972pub(crate) fn sgemm_rm(
3973 m: usize,
3974 n: usize,
3975 k: usize,
3976 alpha: f32,
3977 a: &[f32],
3978 lda: usize,
3979 b_mat: &[f32],
3980 ldb: usize,
3981 b_rows_are_n: bool,
3982 c: &mut [f32],
3983 ldc: usize,
3984) {
3985 debug_assert!(a.len() >= (m - 1) * lda + k);
3986 debug_assert!(c.len() >= (m - 1) * ldc + n);
3987 #[cfg(target_arch = "aarch64")]
3992 if std::env::var("CMF_FORCE_NEON_GEMM")
3993 .map(|v| v == "1")
3994 .unwrap_or(false)
3995 {
3996 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3997 }
3998 unsafe {
3999 accel_blas::cblas_sgemm(
4000 101, 111, if b_rows_are_n { 112 } else { 111 },
4003 m as i32,
4004 n as i32,
4005 k as i32,
4006 alpha,
4007 a.as_ptr(),
4008 lda as i32,
4009 b_mat.as_ptr(),
4010 ldb as i32,
4011 0.0,
4012 c.as_mut_ptr(),
4013 ldc as i32,
4014 );
4015 }
4016}
4017
4018#[cfg(target_os = "macos")]
4025fn qmatmat_accel(
4026 q: &[u8],
4027 row_scale: &[f32],
4028 pre: &[std::borrow::Cow<'_, [f32]>],
4029 rows: usize,
4030 cols: usize,
4031 out: &mut [f32],
4032 pool: Option<&Pool>,
4033) {
4034 const TR: usize = 2048;
4039 let b = pre.len();
4040 thread_local! {
4041 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4042 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4043 }
4044 XPANEL.with(|xp| {
4045 WTILE.with(|wt| {
4046 let mut xpanel = xp.borrow_mut();
4047 xpanel.clear();
4048 for x in pre {
4049 xpanel.extend_from_slice(x);
4050 }
4051 let mut wtile = wt.borrow_mut();
4052 wtile.resize(TR * cols, 0.0);
4053 let mut r0 = 0usize;
4054 while r0 < rows {
4055 let tr = TR.min(rows - r0);
4056 let wt_addr = SendMut(wtile.as_mut_ptr());
4058 let run = |start: usize, end: usize| {
4059 for r in start..end {
4060 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4061 let s = row_scale[r0 + r];
4062 let dst =
4064 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4065 for (d, &v) in dst.iter_mut().zip(row) {
4066 *d = (v as i8) as f32 * s;
4067 }
4068 }
4069 };
4070 dispatch_rows(pool, tr, &run);
4071 unsafe {
4073 accel_blas::cblas_sgemm(
4074 101, 111, 112, b as i32,
4078 tr as i32,
4079 cols as i32,
4080 1.0,
4081 xpanel.as_ptr(),
4082 cols as i32,
4083 wtile.as_ptr(),
4084 cols as i32,
4085 0.0,
4086 out.as_mut_ptr().add(r0),
4087 rows as i32,
4088 );
4089 }
4090 r0 += tr;
4091 }
4092 })
4093 });
4094}
4095
4096fn qmatmat(
4097 q: &[u8],
4098 row_scale: &[f32],
4099 pre: &[std::borrow::Cow<'_, [f32]>],
4100 rows: usize,
4101 cols: usize,
4102 out: &mut [f32],
4103 pool: Option<&Pool>,
4104) {
4105 let b = pre.len();
4106 debug_assert_eq!(out.len(), b * rows);
4107 #[cfg(target_os = "macos")]
4112 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4113 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4114 return;
4115 }
4116 #[cfg(target_arch = "aarch64")]
4117 if sdot_enabled() {
4118 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4119 let out_addr = SendMut(out.as_mut_ptr());
4120 let blocked_ok = blocked_enabled();
4123 let use_i8mm = i8mm_enabled();
4124 if blocked_ok {
4125 let run = |start: usize, end: usize| {
4126 let mut o = start;
4127 while o < end {
4128 if o + 2 <= end {
4129 let r0 = &q[o * cols..(o + 1) * cols];
4130 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4131 let mut bi = 0usize;
4132 while bi + 4 <= acts.len() {
4133 let xs = [
4134 acts[bi].xq.as_slice(),
4135 acts[bi + 1].xq.as_slice(),
4136 acts[bi + 2].xq.as_slice(),
4137 acts[bi + 3].xq.as_slice(),
4138 ];
4139 let d = if use_i8mm {
4140 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4141 } else {
4142 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4143 };
4144 for (r, row) in [r0, r1].into_iter().enumerate() {
4145 for k in 0..4 {
4146 let act = &acts[bi + k];
4147 let mut v = d[r][k] as f32 * act.sx;
4148 for &(j, xv) in &act.outliers {
4149 v += (row[j] as i8) as f32 * xv;
4150 }
4151 unsafe {
4152 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4153 };
4154 }
4155 }
4156 bi += 4;
4157 }
4158 while bi < acts.len() {
4159 for (r, row) in [r0, r1].into_iter().enumerate() {
4160 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4161 unsafe { *out_addr.at(bi * rows + o + r) = v };
4162 }
4163 bi += 1;
4164 }
4165 o += 2;
4166 } else {
4167 let row = &q[o * cols..(o + 1) * cols];
4168 for (bi, act) in acts.iter().enumerate() {
4169 let v = row_dot_sdot(row, act) * row_scale[o];
4170 unsafe { *out_addr.at(bi * rows + o) = v };
4171 }
4172 o += 1;
4173 }
4174 }
4175 };
4176 dispatch_rows(pool, rows, &run);
4177 return;
4178 }
4179 let run = |start: usize, end: usize| {
4180 for o in start..end {
4181 let row = &q[o * cols..(o + 1) * cols];
4182 for (bi, act) in acts.iter().enumerate() {
4183 let v = row_dot_sdot(row, act) * row_scale[o];
4184 unsafe { *out_addr.at(bi * rows + o) = v };
4185 }
4186 }
4187 };
4188 dispatch_rows(pool, rows, &run);
4189 return;
4190 }
4191 #[cfg(target_arch = "x86_64")]
4196 if avx2_a8w8_enabled() {
4197 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4198 let out_addr = SendMut(out.as_mut_ptr());
4199 let blocked_ok = blocked_enabled();
4202 if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4203 let run = |start: usize, end: usize| {
4204 let mut o = start;
4205 while o < end {
4206 if o + 2 <= end {
4207 let r0 = &q[o * cols..(o + 1) * cols];
4208 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4209 let mut bi = 0usize;
4210 while bi + 4 <= acts.len() {
4211 let xs = [
4212 acts[bi].xq.as_slice(),
4213 acts[bi + 1].xq.as_slice(),
4214 acts[bi + 2].xq.as_slice(),
4215 acts[bi + 3].xq.as_slice(),
4216 ];
4217 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4218 for (r, row) in [r0, r1].into_iter().enumerate() {
4219 for k in 0..4 {
4220 let act = &acts[bi + k];
4221 let mut v = d[r][k] as f32 * act.sx;
4222 for &(j, xv) in &act.outliers {
4223 v += (row[j] as i8) as f32 * xv;
4224 }
4225 unsafe {
4226 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4227 };
4228 }
4229 }
4230 bi += 4;
4231 }
4232 while bi < acts.len() {
4233 for (r, row) in [r0, r1].into_iter().enumerate() {
4234 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4235 unsafe { *out_addr.at(bi * rows + o + r) = v };
4236 }
4237 bi += 1;
4238 }
4239 o += 2;
4240 } else {
4241 let row = &q[o * cols..(o + 1) * cols];
4242 for (bi, act) in acts.iter().enumerate() {
4243 let v = row_dot_avx2(row, act) * row_scale[o];
4244 unsafe { *out_addr.at(bi * rows + o) = v };
4245 }
4246 o += 1;
4247 }
4248 }
4249 };
4250 dispatch_rows(pool, rows, &run);
4251 return;
4252 }
4253 let run = |start: usize, end: usize| {
4254 for o in start..end {
4255 let row = &q[o * cols..(o + 1) * cols];
4256 for (bi, act) in acts.iter().enumerate() {
4257 let v = row_dot_avx2(row, act) * row_scale[o];
4258 unsafe { *out_addr.at(bi * rows + o) = v };
4259 }
4260 }
4261 };
4262 dispatch_rows(pool, rows, &run);
4263 return;
4264 }
4265 let out_addr = SendMut(out.as_mut_ptr());
4266 let run = |start: usize, end: usize| {
4267 for o in start..end {
4268 let row = &q[o * cols..(o + 1) * cols];
4269 for (bi, x) in pre.iter().enumerate() {
4270 let mut acc = 0f32;
4271 for j in 0..cols {
4272 acc += (row[j] as i8) as f32 * x[j];
4273 }
4274 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4275 }
4276 }
4277 };
4278 dispatch_rows(pool, rows, &run);
4279}
4280
4281fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4284 match pool {
4285 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4286 _ => run(0, rows),
4287 }
4288}
4289
4290fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4292 let groups = rows * cols / GROUP_SIZE;
4293 bytes.split_at(groups * 16)
4294}
4295
4296#[inline]
4301fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4302 #[cfg(target_arch = "aarch64")]
4303 unsafe {
4304 return vbit_fill4_neon(data, buf);
4305 }
4306 #[cfg(target_arch = "x86_64")]
4307 if avx2_enabled() {
4308 return unsafe { vbit_fill4_avx2(data, buf) };
4309 }
4310 #[allow(unreachable_code)]
4311 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4312 let u = unpack8::<4>(&data[blk * 4..]);
4313 for k in 0..8 {
4314 chunk[k] = (u[k] - 7) as i8 as u8;
4315 }
4316 }
4317}
4318
4319#[cfg(target_arch = "aarch64")]
4320#[target_feature(enable = "neon")]
4321unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4322 unsafe {
4325 use core::arch::aarch64::*;
4326 let n = buf.len();
4327 let mask = vdupq_n_u8(0x0F);
4328 let seven = vdupq_n_s8(7);
4329 let mut g = 0usize;
4330 while g * 32 + 32 <= n {
4331 let b = vld1q_u8(data.as_ptr().add(g * 16));
4332 let hi = vshrq_n_u8::<4>(b);
4333 let lo = vandq_u8(b, mask);
4334 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4335 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4336 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4337 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4338 g += 1;
4339 }
4340 }
4341}
4342
4343#[cfg(target_arch = "x86_64")]
4344#[target_feature(enable = "avx2")]
4345unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4346 unsafe {
4348 use core::arch::x86_64::*;
4349 let n = buf.len();
4350 let mask = _mm_set1_epi8(0x0F);
4351 let seven = _mm256_set1_epi8(7);
4352 let mut g = 0usize;
4353 while g * 32 + 32 <= n {
4354 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4355 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4356 let lo = _mm_and_si128(b, mask);
4357 let z = _mm256_sub_epi8(
4358 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4359 seven,
4360 );
4361 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4362 g += 1;
4363 }
4364 }
4365}
4366
4367#[inline(always)]
4372fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4373 let mut acc = 0u64;
4374 for i in 0..B {
4375 acc = (acc << 8) | data[i] as u64;
4376 }
4377 let mask = (1u64 << B) - 1;
4378 let mut out = [0i32; 8];
4379 for (k, o) in out.iter_mut().enumerate() {
4380 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4381 }
4382 out
4383}
4384
4385#[allow(clippy::too_many_arguments)]
4391fn vbitmatvec(
4392 bytes: &[u8],
4393 offsets: &[usize],
4394 x: &[f32],
4395 rows: usize,
4396 cols: usize,
4397 out: &mut [f32],
4398 pool: Option<&Pool>,
4399) {
4400 debug_assert_eq!(out.len(), rows);
4401 debug_assert_eq!(offsets.len(), rows + 1);
4402
4403 if a8w8_enabled() {
4407 let act = split_act(x);
4408 let out_addr = SendMut(out.as_mut_ptr());
4409 let run = move |start: usize, end: usize| {
4410 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4411 };
4412 dispatch_rows(pool, rows, &run);
4413 return;
4414 }
4415
4416 let out_addr = SendMut(out.as_mut_ptr());
4417 let run = move |start: usize, end: usize| {
4418 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4419 };
4420 dispatch_rows(pool, rows, &run);
4421}
4422
4423#[allow(clippy::too_many_arguments)]
4427fn vbit_range_a8w8(
4428 bytes: &[u8],
4429 offsets: &[usize],
4430 x: &[f32],
4431 act: &SplitAct,
4432 rows: usize,
4433 cols: usize,
4434 out: SendMut,
4435 start: usize,
4436 end: usize,
4437) {
4438 let ng = cols / GROUP_SIZE;
4439 let bits = &bytes[..rows];
4440 let sc_off = rows;
4441 let row_dot = |r: usize| -> f32 {
4442 let b = bits[r] as usize;
4443 let l = (1i32 << (b - 1)) - 1;
4444 let mask = (1u64 << b) - 1;
4445 let data = &bytes[offsets[r]..offsets[r + 1]];
4446 if b == 8 {
4447 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4449 let mut dot = 0f32;
4450 for g in 0..ng {
4451 let so = (r * ng + g) * 2;
4452 let sgf = f16_to_f32(u16::from_le_bytes([
4453 bytes[sc_off + so],
4454 bytes[sc_off + so + 1],
4455 ]));
4456 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4457 let mut gd = 0f32;
4458 for &xv in xg.iter() {
4459 if nbits < 8 {
4460 acc = (acc << 8) | data[idx] as u64;
4461 idx += 1;
4462 nbits += 8;
4463 }
4464 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4465 nbits -= 8;
4466 gd += (u - l) as f32 * xv;
4467 }
4468 dot += gd * sgf;
4469 }
4470 return dot;
4471 }
4472 thread_local! {
4476 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4477 const { std::cell::RefCell::new(Vec::new()) };
4478 }
4479 #[inline(always)]
4480 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4481 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4482 let u = unpack8::<B>(&data[blk * B..]);
4483 for k in 0..8 {
4484 chunk[k] = (u[k] - l) as i8 as u8;
4485 }
4486 }
4487 }
4488 let _ = mask;
4489 VBIT_SCRATCH.with(|scratch| {
4490 let mut buf = scratch.borrow_mut();
4491 buf.resize(cols, 0);
4492 match b {
4493 3 => fill::<3>(data, l, &mut buf),
4494 4 => vbit_fill4(data, &mut buf),
4495 5 => fill::<5>(data, l, &mut buf),
4496 6 => fill::<6>(data, l, &mut buf),
4497 _ => unreachable!(),
4498 }
4499 let mut dot = 0f32;
4500 for g in 0..ng {
4501 let so = (r * ng + g) * 2;
4502 let s = f16_to_f32(u16::from_le_bytes([
4503 bytes[sc_off + so],
4504 bytes[sc_off + so + 1],
4505 ]));
4506 let d = dot_i8_i8(
4507 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4508 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4509 ) as f32
4510 * act.sx;
4511 dot += d * s;
4512 }
4513 for &(j, xv) in &act.outliers {
4514 let so = (r * ng + j / GROUP_SIZE) * 2;
4515 let s = f16_to_f32(u16::from_le_bytes([
4516 bytes[sc_off + so],
4517 bytes[sc_off + so + 1],
4518 ]));
4519 dot += (buf[j] as i8) as f32 * s * xv;
4521 }
4522 dot
4523 })
4524 };
4525 for r in start..end {
4526 unsafe { *out.at(r) = row_dot(r) };
4528 }
4529}
4530
4531#[allow(clippy::too_many_arguments)]
4533fn vbit_range_f32(
4534 bytes: &[u8],
4535 offsets: &[usize],
4536 x: &[f32],
4537 rows: usize,
4538 cols: usize,
4539 out: SendMut,
4540 start: usize,
4541 end: usize,
4542) {
4543 let ng = cols / GROUP_SIZE;
4544 let bits = &bytes[..rows];
4545 let sc_off = rows;
4546 #[inline(always)]
4550 fn dot_row<const B: usize>(
4551 data: &[u8],
4552 bytes: &[u8],
4553 sc_off: usize,
4554 r: usize,
4555 ng: usize,
4556 x: &[f32],
4557 ) -> f32 {
4558 let l = ((1i32 << (B - 1)) - 1) as f32;
4559 let gbytes = GROUP_SIZE * B / 8;
4560 let mut dot = 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 xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4568 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4569 let mut gd = 0f32;
4570 for blk in 0..GROUP_SIZE / 8 {
4571 let u = unpack8::<B>(&gd0[blk * B..]);
4572 let xb = &xg[blk * 8..blk * 8 + 8];
4573 for k in 0..8 {
4574 gd += (u[k] as f32 - l) * xb[k];
4575 }
4576 }
4577 dot += gd * s;
4578 }
4579 dot
4580 }
4581 for r in start..end {
4582 let data = &bytes[offsets[r]..offsets[r + 1]];
4583 let v = match bits[r] {
4584 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4585 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4586 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4587 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4588 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4589 b => unreachable!("vbit bit-width {b} (validated at load)"),
4590 };
4591 unsafe { *out.at(r) = v };
4593 }
4594}
4595
4596#[allow(clippy::too_many_arguments)]
4601fn vbitmatvec2(
4602 bytes: &[u8],
4603 offsets: &[usize],
4604 x1: &[f32],
4605 x2: &[f32],
4606 rows: usize,
4607 cols: usize,
4608 o1: &mut [f32],
4609 o2: &mut [f32],
4610 pool: Option<&Pool>,
4611) {
4612 debug_assert_eq!(o1.len(), rows);
4613 debug_assert_eq!(o2.len(), rows);
4614
4615 if a8w8_enabled() {
4616 let a1 = split_act(x1);
4617 let a2 = split_act(x2);
4618 let p1 = SendMut(o1.as_mut_ptr());
4619 let p2 = SendMut(o2.as_mut_ptr());
4620 let run = move |start: usize, end: usize| {
4621 vbit_range2_a8w8(
4622 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4623 )
4624 };
4625 dispatch_rows(pool, rows, &run);
4626 return;
4627 }
4628
4629 let p1 = SendMut(o1.as_mut_ptr());
4630 let p2 = SendMut(o2.as_mut_ptr());
4631 let run = move |start: usize, end: usize| {
4632 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4633 };
4634 dispatch_rows(pool, rows, &run);
4635}
4636
4637#[allow(clippy::too_many_arguments)]
4641fn vbit_range2_a8w8(
4642 bytes: &[u8],
4643 offsets: &[usize],
4644 x1: &[f32],
4645 x2: &[f32],
4646 a1: &SplitAct,
4647 a2: &SplitAct,
4648 rows: usize,
4649 cols: usize,
4650 p1: SendMut,
4651 p2: SendMut,
4652 start: usize,
4653 end: usize,
4654) {
4655 let ng = cols / GROUP_SIZE;
4656 let bits = &bytes[..rows];
4657 let sc_off = rows;
4658 let row_dots = |r: usize| -> (f32, f32) {
4659 let b = bits[r] as usize;
4660 let l = (1i32 << (b - 1)) - 1;
4661 let data = &bytes[offsets[r]..offsets[r + 1]];
4662 if b == 8 {
4663 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4666 let (mut d1, mut d2) = (0f32, 0f32);
4667 for g in 0..ng {
4668 let so = (r * ng + g) * 2;
4669 let sgf = f16_to_f32(u16::from_le_bytes([
4670 bytes[sc_off + so],
4671 bytes[sc_off + so + 1],
4672 ]));
4673 let (mut g1, mut g2) = (0f32, 0f32);
4674 for k in 0..GROUP_SIZE {
4675 if nbits < 8 {
4676 acc = (acc << 8) | data[idx] as u64;
4677 idx += 1;
4678 nbits += 8;
4679 }
4680 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4681 nbits -= 8;
4682 let w = (u - l) as f32;
4683 g1 += w * x1[g * GROUP_SIZE + k];
4684 g2 += w * x2[g * GROUP_SIZE + k];
4685 }
4686 d1 += g1 * sgf;
4687 d2 += g2 * sgf;
4688 }
4689 return (d1, d2);
4690 }
4691 thread_local! {
4692 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4693 const { std::cell::RefCell::new(Vec::new()) };
4694 }
4695 #[inline(always)]
4696 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4697 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4698 let u = unpack8::<B>(&data[blk * B..]);
4699 for k in 0..8 {
4700 chunk[k] = (u[k] - l) as i8 as u8;
4701 }
4702 }
4703 }
4704 VBIT_SCRATCH2.with(|scratch| {
4705 let mut buf = scratch.borrow_mut();
4706 buf.resize(cols, 0);
4707 match b {
4708 3 => fill::<3>(data, l, &mut buf),
4709 4 => vbit_fill4(data, &mut buf),
4710 5 => fill::<5>(data, l, &mut buf),
4711 6 => fill::<6>(data, l, &mut buf),
4712 _ => unreachable!(),
4713 }
4714 let (mut d1, mut d2) = (0f32, 0f32);
4715 for g in 0..ng {
4716 let so = (r * ng + g) * 2;
4717 let s = f16_to_f32(u16::from_le_bytes([
4718 bytes[sc_off + so],
4719 bytes[sc_off + so + 1],
4720 ]));
4721 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4722 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
4723 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
4724 d1 += v1 * s;
4725 d2 += v2 * s;
4726 }
4727 for &(j, xv) in &a1.outliers {
4728 let so = (r * ng + j / GROUP_SIZE) * 2;
4729 let s = f16_to_f32(u16::from_le_bytes([
4730 bytes[sc_off + so],
4731 bytes[sc_off + so + 1],
4732 ]));
4733 d1 += (buf[j] as i8) as f32 * s * xv;
4734 }
4735 for &(j, xv) in &a2.outliers {
4736 let so = (r * ng + j / GROUP_SIZE) * 2;
4737 let s = f16_to_f32(u16::from_le_bytes([
4738 bytes[sc_off + so],
4739 bytes[sc_off + so + 1],
4740 ]));
4741 d2 += (buf[j] as i8) as f32 * s * xv;
4742 }
4743 (d1, d2)
4744 })
4745 };
4746 for r in start..end {
4747 let (v1, v2) = row_dots(r);
4748 unsafe {
4750 *p1.at(r) = v1;
4751 *p2.at(r) = v2;
4752 }
4753 }
4754}
4755
4756#[allow(clippy::too_many_arguments)]
4760fn vbit_range2_f32(
4761 bytes: &[u8],
4762 offsets: &[usize],
4763 x1: &[f32],
4764 x2: &[f32],
4765 rows: usize,
4766 cols: usize,
4767 p1: SendMut,
4768 p2: SendMut,
4769 start: usize,
4770 end: usize,
4771) {
4772 let ng = cols / GROUP_SIZE;
4773 let bits = &bytes[..rows];
4774 let sc_off = rows;
4775 #[inline(always)]
4776 #[allow(clippy::too_many_arguments)]
4777 fn dot_row2<const B: usize>(
4778 data: &[u8],
4779 bytes: &[u8],
4780 sc_off: usize,
4781 r: usize,
4782 ng: usize,
4783 x1: &[f32],
4784 x2: &[f32],
4785 ) -> (f32, f32) {
4786 let l = ((1i32 << (B - 1)) - 1) as f32;
4787 let gbytes = GROUP_SIZE * B / 8;
4788 let (mut d1, mut d2) = (0f32, 0f32);
4789 for g in 0..ng {
4790 let so = (r * ng + g) * 2;
4791 let s = f16_to_f32(u16::from_le_bytes([
4792 bytes[sc_off + so],
4793 bytes[sc_off + so + 1],
4794 ]));
4795 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4796 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4797 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4798 let (mut g1, mut g2) = (0f32, 0f32);
4799 for blk in 0..GROUP_SIZE / 8 {
4800 let u = unpack8::<B>(&gd0[blk * B..]);
4801 for k in 0..8 {
4802 let w = u[k] as f32 - l;
4803 g1 += w * x1g[blk * 8 + k];
4804 g2 += w * x2g[blk * 8 + k];
4805 }
4806 }
4807 d1 += g1 * s;
4808 d2 += g2 * s;
4809 }
4810 (d1, d2)
4811 }
4812 for r in start..end {
4813 let data = &bytes[offsets[r]..offsets[r + 1]];
4814 let (v1, v2) = match bits[r] {
4815 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4816 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4817 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4818 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4819 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4820 b => unreachable!("vbit bit-width {b} (validated at load)"),
4821 };
4822 unsafe {
4824 *p1.at(r) = v1;
4825 *p2.at(r) = v2;
4826 }
4827 }
4828}
4829
4830#[inline]
4837#[allow(unreachable_code)]
4838fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4839 #[cfg(target_arch = "aarch64")]
4840 unsafe {
4841 return dot_q4t_row_sdot(bytes, r, gpr, xq);
4842 }
4843 #[cfg(target_arch = "x86_64")]
4844 unsafe {
4845 if vnni_tiles_enabled() {
4846 return dot_q4t_row_vnni(bytes, r, gpr, xq);
4847 }
4848 return dot_q4t_row_avx2(bytes, r, gpr, xq);
4849 }
4850 let mut acc = 0f32;
4851 for gi in 0..gpr {
4852 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4853 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4854 let mut d = 0i32;
4855 for (k, &b) in tile[2..].iter().enumerate() {
4856 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4857 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4858 }
4859 acc += d as f32 * s;
4860 }
4861 acc
4862}
4863
4864#[cfg(target_arch = "aarch64")]
4865#[target_feature(enable = "neon,dotprod")]
4866unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4867 unsafe {
4870 use core::arch::aarch64::*;
4871 use core::arch::asm;
4872 let lomask = vdupq_n_u8(0x0F);
4873 let eight = vdupq_n_s8(8);
4874 let mut acc = 0f32;
4875 for gi in 0..gpr {
4876 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4877 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4878 let b = vld1q_u8(t.add(2));
4879 let lo = vandq_u8(b, lomask);
4880 let hi = vshrq_n_u8::<4>(b);
4881 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4882 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4883 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4884 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4885 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4886 asm!(
4887 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4888 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4889 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4890 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4891 options(pure, nomem, nostack),
4892 );
4893 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4894 }
4895 acc
4896 }
4897}
4898
4899#[cfg(target_arch = "x86_64")]
4900#[target_feature(enable = "avx2")]
4901unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4902 unsafe {
4904 use core::arch::x86_64::*;
4905 let lomask = _mm_set1_epi8(0x0F);
4906 let eight = _mm256_set1_epi8(8);
4907 let ones = _mm256_set1_epi16(1);
4908 let mut acc = 0f32;
4909 for gi in 0..gpr {
4910 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4911 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4912 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4913 let lo = _mm_and_si128(b, lomask);
4914 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4915 let w = _mm256_sub_epi8(
4916 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4917 eight,
4918 );
4919 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4920 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4921 let d = _mm256_madd_epi16(p16, ones);
4922 let hi128 = _mm256_extracti128_si256::<1>(d);
4923 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4924 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4925 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4926 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4927 }
4928 acc
4929 }
4930}
4931
4932#[cfg(target_arch = "x86_64")]
4936#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4937unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4938 unsafe {
4940 use core::arch::x86_64::*;
4941 let lomask = _mm_set1_epi8(0x0F);
4942 let eight = _mm256_set1_epi8(8);
4943 let mut acc = 0f32;
4944 for gi in 0..gpr {
4945 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4946 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4947 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4948 let lo = _mm_and_si128(b, lomask);
4949 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4950 let w = _mm256_sub_epi8(
4951 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4952 eight,
4953 );
4954 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4955 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4956 acc += d as f32 * s;
4957 }
4958 acc
4959 }
4960}
4961
4962#[cfg(target_arch = "x86_64")]
4967#[target_feature(enable = "avx2,fma")]
4972unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4973 unsafe {
4975 use core::arch::x86_64::*;
4976 let lomask = _mm_set1_epi8(0x0F);
4977 let eight = _mm256_set1_epi8(8);
4978 let ones = _mm256_set1_epi16(1);
4979 let mut f0 = _mm256_setzero_ps();
4992 let mut f1 = _mm256_setzero_ps();
4993 let mut f2 = _mm256_setzero_ps();
4994 let mut f3 = _mm256_setzero_ps();
4995 for gi in 0..gpr {
4996 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4997 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4998 let sv = _mm256_set1_ps(s);
4999 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5000 let lo = _mm_and_si128(bb, lomask);
5001 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5002 let w = _mm256_sub_epi8(
5003 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5004 eight,
5005 );
5006 let aw = _mm256_abs_epi8(w);
5007 let off = gi * GROUP_SIZE;
5008 let dot = |xq: &[i8]| {
5009 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5010 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5011 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5012 };
5013 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5014 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5015 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5016 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5017 }
5018 [
5019 hsum256_ps(f0),
5020 hsum256_ps(f1),
5021 hsum256_ps(f2),
5022 hsum256_ps(f3),
5023 ]
5024 }
5025}
5026
5027#[cfg(target_arch = "x86_64")]
5030#[target_feature(enable = "avx2")]
5031#[inline]
5032unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5033 unsafe {
5035 use core::arch::x86_64::*;
5036 let hi = _mm256_extractf128_ps::<1>(v);
5037 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5038 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5039 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5040 _mm_cvtss_f32(s)
5041 }
5042}
5043
5044#[cfg(target_arch = "x86_64")]
5046#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5047unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5048 unsafe {
5050 use core::arch::x86_64::*;
5051 let lomask = _mm_set1_epi8(0x0F);
5052 let eight = _mm256_set1_epi8(8);
5053 let mut f0 = _mm256_setzero_ps();
5056 let mut f1 = _mm256_setzero_ps();
5057 let mut f2 = _mm256_setzero_ps();
5058 let mut f3 = _mm256_setzero_ps();
5059 for gi in 0..gpr {
5060 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5061 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5062 let sv = _mm256_set1_ps(s);
5063 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5064 let lo = _mm_and_si128(bb, lomask);
5065 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5066 let w = _mm256_sub_epi8(
5067 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5068 eight,
5069 );
5070 let aw = _mm256_abs_epi8(w);
5071 let off = gi * GROUP_SIZE;
5072 let dot = |xq: &[i8]| {
5073 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5074 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5075 _mm256_setzero_si256(),
5076 aw,
5077 _mm256_sign_epi8(x, w),
5078 ))
5079 };
5080 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5081 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5082 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5083 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5084 }
5085 let acc = [
5086 hsum256_ps(f0),
5087 hsum256_ps(f1),
5088 hsum256_ps(f2),
5089 hsum256_ps(f3),
5090 ];
5091 acc
5092 }
5093}
5094
5095#[cfg(target_arch = "aarch64")]
5100#[target_feature(enable = "neon,dotprod")]
5101unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5102 unsafe {
5104 use core::arch::aarch64::*;
5105 use core::arch::asm;
5106 let lomask = vdupq_n_u8(0x0F);
5107 let eight = vdupq_n_s8(8);
5108 let mut acc = [0f32; 4];
5109 for gi in 0..gpr {
5110 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5111 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5112 let b = vld1q_u8(t.add(2));
5113 let lo = vandq_u8(b, lomask);
5114 let hi = vshrq_n_u8::<4>(b);
5115 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5116 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5117 for (k, xq) in xs.iter().enumerate() {
5118 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5119 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5120 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5121 asm!(
5122 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5123 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5124 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5125 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5126 options(pure, nomem, nostack),
5127 );
5128 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5129 }
5130 }
5131 acc
5132 }
5133}
5134
5135#[inline]
5137fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5138 let gi = j / GROUP_SIZE;
5139 let k = j % GROUP_SIZE;
5140 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5141 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5142 let byte = tile[2 + k / 2];
5143 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5144 ((nib as i32 - 8) as f32, s)
5145}
5146
5147#[inline]
5150fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5151 let mut acc = 0f32;
5152 for gi in 0..gpr {
5153 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5154 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5155 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5156 let mut ga = 0f32;
5157 for (k, &b) in tile[2..].iter().enumerate() {
5158 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5159 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5160 }
5161 acc += ga * s;
5162 }
5163 acc
5164}
5165
5166struct Q4tpView<'a> {
5170 nib: &'a [u8],
5171 params: &'a [u8],
5172 codes: &'a [u8],
5173 stride: usize,
5174 zero_rung: bool,
5176}
5177
5178impl<'a> Q4tpView<'a> {
5179 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5180 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5181 Self {
5182 nib: &bytes[..params_off],
5183 params: &bytes[params_off..codes_off],
5184 codes: &bytes[codes_off..],
5185 stride,
5186 zero_rung: false,
5187 }
5188 }
5189
5190 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5192 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5193 Self {
5194 nib: &bytes[..params_off],
5195 params: &bytes[params_off..codes_off],
5196 codes: &bytes[codes_off..],
5197 stride,
5198 zero_rung: true,
5199 }
5200 }
5201
5202 #[inline]
5218 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5219 let tab = if self.zero_rung {
5220 q2tp_ladder(self.params, r)
5221 } else {
5222 q4tp_ladder(self.params, r)
5223 };
5224 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5225 let out = &mut out[..gpr];
5226 let mut chunks = out.chunks_exact_mut(8);
5227 let mut ci = 0usize;
5228 for c in &mut chunks {
5229 let w = u64::from(codes[ci])
5230 | u64::from(codes[ci + 1]) << 8
5231 | u64::from(codes[ci + 2]) << 16
5232 | u64::from(codes[ci + 3]) << 24
5233 | u64::from(codes[ci + 4]) << 32;
5234 for (k, o) in c.iter_mut().enumerate() {
5235 *o = tab[((w >> (5 * k)) & 31) as usize];
5236 }
5237 ci += 5;
5238 }
5239 let tail = &codes[ci..];
5242 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5243 *o = tab[q4tp_code(tail, k)];
5244 }
5245 }
5246}
5247
5248#[inline]
5249fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5250 #[cfg(target_arch = "aarch64")]
5251 unsafe {
5252 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5253 }
5254 #[cfg(target_arch = "x86_64")]
5255 unsafe {
5256 if vnni_tiles_enabled() {
5257 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5258 }
5259 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5260 }
5261 #[allow(unreachable_code)]
5262 {
5263 let mut acc = 0f32;
5264 for gi in 0..gpr {
5265 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5266 let s = scales[gi];
5267 let mut d = 0i32;
5268 for (k, &b) in tile.iter().enumerate() {
5269 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5270 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5271 }
5272 acc += d as f32 * s;
5273 }
5274 acc
5275 }
5276}
5277
5278#[cfg(target_arch = "aarch64")]
5281#[target_feature(enable = "neon,dotprod")]
5282unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5283 unsafe {
5286 use core::arch::aarch64::*;
5287 use core::arch::asm;
5288 let lomask = vdupq_n_u8(0x0F);
5289 let eight = vdupq_n_s8(8);
5290 let mut acc = 0f32;
5291 for gi in 0..gpr {
5292 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5293 let s = *scales.get_unchecked(gi);
5294 let b = vld1q_u8(t);
5295 let lo = vandq_u8(b, lomask);
5296 let hi = vshrq_n_u8::<4>(b);
5297 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5298 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5299 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5300 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5301 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5302 asm!(
5303 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5304 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5305 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5306 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5307 options(pure, nomem, nostack),
5308 );
5309 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5310 }
5311 acc
5312 }
5313}
5314
5315#[cfg(target_arch = "x86_64")]
5316#[target_feature(enable = "avx2")]
5317unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5318 unsafe {
5320 use core::arch::x86_64::*;
5321 let lomask = _mm_set1_epi8(0x0F);
5322 let eight = _mm256_set1_epi8(8);
5323 let ones = _mm256_set1_epi16(1);
5324 let mut acc = 0f32;
5325 for gi in 0..gpr {
5326 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5327 let s = *scales.get_unchecked(gi);
5328 let b = _mm_loadu_si128(t as *const __m128i);
5329 let lo = _mm_and_si128(b, lomask);
5330 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5331 let w = _mm256_sub_epi8(
5332 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5333 eight,
5334 );
5335 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5336 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5337 let d = _mm256_madd_epi16(p16, ones);
5338 let hi128 = _mm256_extracti128_si256::<1>(d);
5339 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5340 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5341 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5342 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5343 }
5344 acc
5345 }
5346}
5347
5348#[cfg(target_arch = "x86_64")]
5351#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5352unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5353 unsafe {
5355 use core::arch::x86_64::*;
5356 let lomask = _mm_set1_epi8(0x0F);
5357 let eight = _mm256_set1_epi8(8);
5358 let mut acc = 0f32;
5359 for gi in 0..gpr {
5360 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5361 let s = *scales.get_unchecked(gi);
5362 let b = _mm_loadu_si128(t as *const __m128i);
5363 let lo = _mm_and_si128(b, lomask);
5364 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5365 let w = _mm256_sub_epi8(
5366 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5367 eight,
5368 );
5369 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5370 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5371 }
5372 acc
5373 }
5374}
5375
5376#[inline]
5379fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5380 #[cfg(target_arch = "x86_64")]
5381 if avx2_enabled() {
5382 return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5384 }
5385 q4tp_row_float_scalar(nib, r, gpr, x, scales)
5386}
5387
5388#[inline]
5389fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5390 let mut acc = 0f32;
5391 for gi in 0..gpr {
5392 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5393 let s = scales[gi];
5394 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5395 let mut ga = 0f32;
5396 for (k, &b) in tile.iter().enumerate() {
5397 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5398 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5399 }
5400 acc += ga * s;
5401 }
5402 acc
5403}
5404
5405#[cfg(target_arch = "x86_64")]
5409#[target_feature(enable = "avx2")]
5410unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5411 unsafe {
5414 use core::arch::x86_64::*;
5415 let mask = _mm_set1_epi8(15);
5416 let eight = _mm_set1_epi8(8);
5417 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5418 let mut acc = 0.0f32;
5419 for gi in 0..gpr {
5420 let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5421 let lo = _mm_and_si128(packed, mask);
5422 let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5423 let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5424 let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5425 let xp = x.as_ptr().add(gi * GROUP_SIZE);
5426 let a = _mm256_mul_ps(
5427 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5428 _mm256_loadu_ps(xp),
5429 );
5430 let b = _mm256_mul_ps(
5431 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5432 _mm256_loadu_ps(xp.add(8)),
5433 );
5434 let c = _mm256_mul_ps(
5435 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5436 _mm256_loadu_ps(xp.add(16)),
5437 );
5438 let d = _mm256_mul_ps(
5439 _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5440 _mm256_loadu_ps(xp.add(24)),
5441 );
5442 let mut pairs = [0.0f32; 16];
5443 _mm256_storeu_ps(
5444 pairs.as_mut_ptr(),
5445 _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5446 );
5447 _mm256_storeu_ps(
5448 pairs.as_mut_ptr().add(8),
5449 _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5450 );
5451 let mut ga = 0.0f32;
5452 for v in pairs {
5453 ga += v;
5454 }
5455 acc += ga * scales[gi];
5456 }
5457 acc
5458 }
5459}
5460
5461#[inline]
5464fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5465 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5466 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5467 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5468 ((n as i32 - 8) as f32, scales[gi])
5469}
5470
5471fn q4tp_matvec(
5473 bytes: &[u8],
5474 x: &[f32],
5475 rows: usize,
5476 cols: usize,
5477 out: &mut [f32],
5478 pool: Option<&Pool>,
5479) {
5480 debug_assert_eq!(out.len(), rows);
5481 let gpr = cols / GROUP_SIZE;
5482 let v = Q4tpView::new(bytes, rows, cols);
5483 let out_addr = SendMut(out.as_mut_ptr());
5484 if a8w8_enabled() {
5485 let act = split_act(x);
5486 let run = |start: usize, end: usize| {
5487 with_krow(gpr, |sc| {
5489 for r in start..end {
5490 v.scales_into(r, gpr, sc);
5491 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5492 for &(j, xv) in &act.outliers {
5493 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5494 acc += w * s * xv;
5495 }
5496 unsafe { *out_addr.at(r) = acc };
5498 }
5499 })
5500 };
5501 dispatch_rows(pool, rows, &run);
5502 return;
5503 }
5504 let run = |start: usize, end: usize| {
5505 with_krow(gpr, |sc| {
5506 for r in start..end {
5507 v.scales_into(r, gpr, sc);
5508 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5510 }
5511 })
5512 };
5513 dispatch_rows(pool, rows, &run);
5514}
5515
5516#[allow(clippy::too_many_arguments)]
5519fn q4tp_matvec2(
5520 bytes: &[u8],
5521 x1: &[f32],
5522 x2: &[f32],
5523 rows: usize,
5524 cols: usize,
5525 o1: &mut [f32],
5526 o2: &mut [f32],
5527 pool: Option<&Pool>,
5528) {
5529 let gpr = cols / GROUP_SIZE;
5530 let v = Q4tpView::new(bytes, rows, cols);
5531 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5532 let run = |start: usize, end: usize| {
5533 let mut sc = vec![0f32; gpr];
5534 for r in start..end {
5535 v.scales_into(r, gpr, &mut sc);
5536 unsafe {
5538 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5539 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5540 }
5541 }
5542 };
5543 dispatch_rows(pool, rows, &run);
5544}
5545
5546#[inline]
5549fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5550 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5551 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5552 let c = (byte >> (2 * (k % 4))) & 3;
5553 (c as f32 - 1.5, scales[gi])
5554}
5555
5556#[cfg(target_arch = "x86_64")]
5557const Q2TP_DECODE_U32: [u32; 256] = {
5558 let mut tab = [0u32; 256];
5559 let mut b = 0usize;
5560 while b < 256 {
5561 tab[b] = ((b as u32) & 3)
5562 | ((((b as u32) >> 2) & 3) << 8)
5563 | ((((b as u32) >> 4) & 3) << 16)
5564 | ((((b as u32) >> 6) & 3) << 24);
5565 b += 1;
5566 }
5567 tab
5568};
5569
5570#[cfg(target_arch = "x86_64")]
5574#[target_feature(enable = "avx2")]
5575unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5576 use core::arch::x86_64::*;
5577 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5578 let codes = _mm256_setr_epi32(
5579 Q2TP_DECODE_U32[ch[0] as usize] as i32,
5580 Q2TP_DECODE_U32[ch[1] as usize] as i32,
5581 Q2TP_DECODE_U32[ch[2] as usize] as i32,
5582 Q2TP_DECODE_U32[ch[3] as usize] as i32,
5583 Q2TP_DECODE_U32[ch[4] as usize] as i32,
5584 Q2TP_DECODE_U32[ch[5] as usize] as i32,
5585 Q2TP_DECODE_U32[ch[6] as usize] as i32,
5586 Q2TP_DECODE_U32[ch[7] as usize] as i32,
5587 );
5588 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5589 let pair = _mm256_maddubs_epi16(codes, xv);
5590 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5591 let sum128 = _mm_add_epi32(
5592 _mm256_castsi256_si128(quad),
5593 _mm256_extracti128_si256(quad, 1),
5594 );
5595 let sum64 = _mm_hadd_epi32(sum128, sum128);
5596 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5597}
5598
5599#[inline]
5606fn dot_q2tp_row_i8(
5607 chunks: &[u8],
5608 r: usize,
5609 gpr: usize,
5610 xq: &[i8],
5611 gsum: &[i32],
5612 scales: &[f32],
5613) -> f32 {
5614 let mut acc = 0f32;
5615 let base = r * gpr * Q2TP_CHUNK;
5616 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5617 let mut codes = [0i8; GROUP_SIZE];
5618 #[cfg(target_arch = "x86_64")]
5619 let avx2 = std::arch::is_x86_feature_detected!("avx2");
5620 for gi in 0..gpr {
5621 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5622 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5623 #[cfg(target_arch = "aarch64")]
5624 let dot = unsafe {
5630 use core::arch::aarch64::*;
5631 let b = vld1_u8(ch.as_ptr());
5632 let three = vdup_n_u8(3);
5633 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5634 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5635 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5636 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5637 let x4 = vld4_s8(xg.as_ptr());
5638 let mut acc4 = vdupq_n_s32(0);
5639 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5640 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5641 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5642 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5643 vaddvq_s32(acc4)
5644 };
5645 #[cfg(target_arch = "x86_64")]
5646 let dot: i32 = if avx2 {
5647 unsafe { q2tp_code_dot_avx2(ch, xg) }
5650 } else {
5651 ch.iter()
5652 .enumerate()
5653 .map(|(k, &b)| {
5654 ((b & 3) as i32) * xg[k * 4] as i32
5655 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5656 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5657 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5658 })
5659 .sum()
5660 };
5661 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5662 let dot: i32 = {
5663 for (k, &b) in ch.iter().enumerate() {
5664 codes[k * 4] = (b & 3) as i8;
5665 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5666 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5667 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5668 }
5669 codes
5670 .iter()
5671 .zip(xg)
5672 .map(|(&c, &x)| c as i32 * x as i32)
5673 .sum()
5674 };
5675 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5676 }
5677 acc
5678}
5679
5680fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5684 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5685}
5686
5687#[inline]
5691fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5692 q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
5693}
5694
5695#[inline]
5696fn q2tp_row_exact_center(
5697 chunks: &[u8],
5698 r: usize,
5699 gpr: usize,
5700 x: &[f32],
5701 scales: &[f32],
5702 center: f32,
5703) -> f32 {
5704 let mut acc = 0f32;
5705 for gi in 0..gpr {
5706 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
5707 let s = scales[gi];
5708 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5709 let mut g = 0f32;
5710 for (k, &b) in ch.iter().enumerate() {
5711 g += ((b & 3) as f32 - center) * xb[k * 4]
5712 + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
5713 + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
5714 + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
5715 }
5716 acc += s * g;
5717 }
5718 acc
5719}
5720
5721fn q2tp_matvec(
5722 bytes: &[u8],
5723 x: &[f32],
5724 rows: usize,
5725 cols: usize,
5726 out: &mut [f32],
5727 pool: Option<&Pool>,
5728) {
5729 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
5730}
5731
5732fn q2tp_affine_matvec(
5733 bytes: &[u8],
5734 x: &[f32],
5735 rows: usize,
5736 cols: usize,
5737 out: &mut [f32],
5738 pool: Option<&Pool>,
5739) {
5740 q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
5741}
5742
5743fn q2tp_matvec_mode(
5744 bytes: &[u8],
5745 x: &[f32],
5746 rows: usize,
5747 cols: usize,
5748 out: &mut [f32],
5749 pool: Option<&Pool>,
5750 affine: bool,
5751) {
5752 debug_assert_eq!(out.len(), rows);
5753 let gpr = cols / GROUP_SIZE;
5754 let v = Q4tpView::new_q2(bytes, rows, cols);
5755 let out_addr = SendMut(out.as_mut_ptr());
5756 if !affine && a8w8_enabled() {
5761 let act = split_act(x);
5762 let gsum = q1_group_sums(&act.xq, gpr);
5763 let (act, gsum) = (&act, &gsum);
5764 let run = move |start: usize, end: usize| {
5765 with_krow(gpr, |sc| {
5766 for r in start..end {
5767 v.scales_into(r, gpr, sc);
5768 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
5769 for &(j, xv) in &act.outliers {
5770 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
5771 acc += w * s * xv;
5772 }
5773 unsafe { *out_addr.at(r) = acc };
5775 }
5776 })
5777 };
5778 dispatch_rows(pool, rows, &run);
5779 return;
5780 }
5781 let run = |start: usize, end: usize| {
5782 with_krow(gpr, |sc| {
5783 for r in start..end {
5784 v.scales_into(r, gpr, sc);
5785 unsafe {
5787 *out_addr.at(r) = if affine {
5788 q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
5789 } else {
5790 q2tp_row_exact(v.nib, r, gpr, x, sc)
5791 }
5792 };
5793 }
5794 })
5795 };
5796 dispatch_rows(pool, rows, &run);
5797}
5798
5799#[allow(clippy::too_many_arguments)]
5801fn q2tp_matvec2(
5802 bytes: &[u8],
5803 x1: &[f32],
5804 x2: &[f32],
5805 rows: usize,
5806 cols: usize,
5807 o1: &mut [f32],
5808 o2: &mut [f32],
5809 pool: Option<&Pool>,
5810) {
5811 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
5812}
5813
5814#[allow(clippy::too_many_arguments)]
5815fn q2tp_affine_matvec2(
5816 bytes: &[u8],
5817 x1: &[f32],
5818 x2: &[f32],
5819 rows: usize,
5820 cols: usize,
5821 o1: &mut [f32],
5822 o2: &mut [f32],
5823 pool: Option<&Pool>,
5824) {
5825 q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
5826}
5827
5828#[allow(clippy::too_many_arguments)]
5829fn q2tp_matvec2_mode(
5830 bytes: &[u8],
5831 x1: &[f32],
5832 x2: &[f32],
5833 rows: usize,
5834 cols: usize,
5835 o1: &mut [f32],
5836 o2: &mut [f32],
5837 pool: Option<&Pool>,
5838 affine: bool,
5839) {
5840 let gpr = cols / GROUP_SIZE;
5841 let v = Q4tpView::new_q2(bytes, rows, cols);
5842 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5843 let run = |start: usize, end: usize| {
5844 let mut sc = vec![0f32; gpr];
5845 for r in start..end {
5846 v.scales_into(r, gpr, &mut sc);
5847 unsafe {
5849 *p1.at(r) = if affine {
5850 q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
5851 } else {
5852 q2tp_row_exact(v.nib, r, gpr, x1, &sc)
5853 };
5854 *p2.at(r) = if affine {
5855 q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
5856 } else {
5857 q2tp_row_exact(v.nib, r, gpr, x2, &sc)
5858 };
5859 }
5860 }
5861 };
5862 dispatch_rows(pool, rows, &run);
5863}
5864
5865pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
5872 let gpr = cols / GROUP_SIZE;
5877 let v = Q4tpView::new_q2(bytes, rows, cols);
5878 with_krow(gpr, |sc| {
5879 for r in 0..rows {
5880 v.scales_into(r, gpr, sc);
5881 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
5882 }
5883 });
5884}
5885
5886pub fn q2tp_affine_matvec_for_test(
5889 bytes: &[u8],
5890 x: &[f32],
5891 rows: usize,
5892 cols: usize,
5893 out: &mut [f32],
5894) {
5895 q2tp_affine_matvec(bytes, x, rows, cols, out, None);
5896}
5897
5898pub fn q2tp_matmat_for_test(
5899 bytes: &[u8],
5900 xs_all: &[f32],
5901 b: usize,
5902 rows: usize,
5903 cols: usize,
5904 out: &mut [f32],
5905) {
5906 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
5907}
5908
5909fn q2tp_matmat(
5910 bytes: &[u8],
5911 xs_all: &[f32],
5912 b: usize,
5913 rows: usize,
5914 cols: usize,
5915 out: &mut [f32],
5916 pool: Option<&Pool>,
5917) {
5918 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
5919}
5920
5921fn q2tp_affine_matmat(
5922 bytes: &[u8],
5923 xs_all: &[f32],
5924 b: usize,
5925 rows: usize,
5926 cols: usize,
5927 out: &mut [f32],
5928 pool: Option<&Pool>,
5929) {
5930 q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
5931}
5932
5933fn q2tp_matmat_mode(
5934 bytes: &[u8],
5935 xs_all: &[f32],
5936 b: usize,
5937 rows: usize,
5938 cols: usize,
5939 out: &mut [f32],
5940 pool: Option<&Pool>,
5941 affine: bool,
5942) {
5943 debug_assert_eq!(out.len(), b * rows);
5944 let gpr = cols / GROUP_SIZE;
5945 let v = Q4tpView::new_q2(bytes, rows, cols);
5946 let out_addr = SendMut(out.as_mut_ptr());
5947 let run = |start: usize, end: usize| {
5948 let mut sc = vec![0f32; gpr];
5949 for r in start..end {
5950 v.scales_into(r, gpr, &mut sc);
5951 for bi in 0..b {
5952 let x = &xs_all[bi * cols..(bi + 1) * cols];
5953 unsafe {
5955 *out_addr.at(bi * rows + r) = if affine {
5956 q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
5957 } else {
5958 q2tp_row_exact(v.nib, r, gpr, x, &sc)
5959 }
5960 };
5961 }
5962 }
5963 };
5964 dispatch_rows(pool, rows, &run);
5965}
5966
5967#[cfg(target_arch = "aarch64")]
5971#[target_feature(enable = "neon,dotprod")]
5972unsafe fn dot_q4tp_row_1x4_sdot_v1(
5973 nib: &[u8],
5974 r: usize,
5975 gpr: usize,
5976 xs: [&[i8]; 4],
5977 scales: &[f32],
5978) -> [f32; 4] {
5979 unsafe {
5980 use core::arch::aarch64::*;
5981 use core::arch::asm;
5982 let lomask = vdupq_n_u8(0x0F);
5983 let eight = vdupq_n_s8(8);
5984 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
5985 for gi in 0..gpr {
5986 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5987 let s = *scales.get_unchecked(gi);
5988 let bb = vld1q_u8(t);
5989 let lo = vandq_u8(bb, lomask);
5990 let hi = vshrq_n_u8::<4>(bb);
5991 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5992 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5993 let mut d = [0f32; 4];
5994 for (k, dk) in d.iter_mut().enumerate() {
5995 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5996 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5997 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5998 asm!(
5999 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6000 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6001 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6002 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6003 options(pure, nomem, nostack),
6004 );
6005 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6006 }
6007 f0 += d[0];
6008 f1 += d[1];
6009 f2 += d[2];
6010 f3 += d[3];
6011 }
6012 [f0, f1, f2, f3]
6013 }
6014}
6015
6016#[allow(dead_code)]
6023static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6024
6025#[cfg(target_arch = "x86_64")]
6031fn q4tp_blocked_x86() -> bool {
6032 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6033 1 => false,
6034 2 => avx512vnni_enabled(),
6039 _ => blocked_enabled() && avx512vnni_enabled(),
6043 }
6044}
6045
6046#[cfg(target_arch = "aarch64")]
6048#[allow(dead_code)]
6049fn q4tp_v1() -> bool {
6050 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6051 1 => true,
6052 2 => false,
6053 _ => {
6054 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6055 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6056 }
6057 }
6058}
6059
6060#[cfg(target_arch = "x86_64")]
6072#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6073unsafe fn dot_q4tp_2x8_avx512(
6074 nib: &[u8],
6075 r0: usize,
6076 gpr: usize,
6077 xs: [&[i8]; 8],
6078 sc0: &[f32],
6079 sc1: &[f32],
6080) -> [[f32; 8]; 2] {
6081 unsafe {
6084 use core::arch::x86_64::*;
6085 let lomask = _mm256_set1_epi8(0x0F);
6086 let eight = _mm256_set1_epi8(8);
6087 let zero = _mm512_setzero_si512();
6088 let mut v0 = [_mm512_setzero_ps(); 8];
6089 let mut v1 = [_mm512_setzero_ps(); 8];
6090 let pairs = gpr / 2;
6091 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6092 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6093 let bb = _mm256_loadu_si256(t as *const __m256i);
6094 let lo = _mm256_and_si256(bb, lomask);
6095 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6096 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6097 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6098 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6099 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6100 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6101 };
6102 for gp in 0..pairs {
6103 let gi = gp * 2;
6104 let (wa0, neg0) = unpack(r0, gi);
6105 let (wa1, neg1) = unpack(r0 + 1, gi);
6106 let off = gi * GROUP_SIZE;
6107 let sv = |sc: &[f32]| {
6108 _mm512_insertf32x8::<1>(
6109 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6110 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6111 )
6112 };
6113 let s0 = sv(sc0);
6114 let s1 = sv(sc1);
6115 for k in 0..8 {
6116 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6117 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6118 zero,
6119 wa0,
6120 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6121 ));
6122 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6123 zero,
6124 wa1,
6125 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6126 ));
6127 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6128 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6129 }
6130 }
6131 let mut acc = [[0f32; 8]; 2];
6132 for k in 0..8 {
6133 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6134 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6135 }
6136 if gpr % 2 == 1 {
6137 let off = (gpr - 1) * GROUP_SIZE;
6138 for j in off..off + GROUP_SIZE {
6139 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6140 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6141 for k in 0..8 {
6142 let x = *xs[k].get_unchecked(j) as f32;
6143 acc[0][k] += w0 * sa * x;
6144 acc[1][k] += w1 * sb * x;
6145 }
6146 }
6147 }
6148 acc
6149 }
6150}
6151
6152#[cfg(target_arch = "x86_64")]
6157#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6158unsafe fn dot_q4tp_row_1x8_avx512(
6159 nib: &[u8],
6160 r: usize,
6161 gpr: usize,
6162 xs: [&[i8]; 8],
6163 scales: &[f32],
6164) -> [f32; 8] {
6165 unsafe {
6167 use core::arch::x86_64::*;
6168 let lomask = _mm256_set1_epi8(0x0F);
6169 let eight = _mm256_set1_epi8(8);
6170 let zero = _mm512_setzero_si512();
6171 let (mut v0, mut v1, mut v2, mut v3) = (
6172 _mm512_setzero_ps(),
6173 _mm512_setzero_ps(),
6174 _mm512_setzero_ps(),
6175 _mm512_setzero_ps(),
6176 );
6177 let (mut v4, mut v5, mut v6, mut v7) = (
6178 _mm512_setzero_ps(),
6179 _mm512_setzero_ps(),
6180 _mm512_setzero_ps(),
6181 _mm512_setzero_ps(),
6182 );
6183 let pairs = gpr / 2;
6184 for gp in 0..pairs {
6185 let gi = gp * 2;
6186 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6187 let bb = _mm256_loadu_si256(t as *const __m256i);
6188 let lo = _mm256_and_si256(bb, lomask);
6189 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6190 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6195 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6196 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6197 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6198 let wabs = _mm512_abs_epi8(w);
6199 let neg = _mm512_movepi8_mask(w);
6200 let off = gi * GROUP_SIZE;
6201 let sv = _mm512_insertf32x8::<1>(
6202 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6203 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6204 );
6205 let dot = |x: &[i8]| -> __m512 {
6206 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6207 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6208 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6209 };
6210 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6211 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6212 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6213 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6214 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6215 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6216 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6217 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6218 }
6219 let mut acc = [
6220 _mm512_reduce_add_ps(v0),
6221 _mm512_reduce_add_ps(v1),
6222 _mm512_reduce_add_ps(v2),
6223 _mm512_reduce_add_ps(v3),
6224 _mm512_reduce_add_ps(v4),
6225 _mm512_reduce_add_ps(v5),
6226 _mm512_reduce_add_ps(v6),
6227 _mm512_reduce_add_ps(v7),
6228 ];
6229 if gpr % 2 == 1 {
6232 let off = (gpr - 1) * GROUP_SIZE;
6233 for j in off..off + GROUP_SIZE {
6234 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6235 let ws = w * s;
6236 for k in 0..8 {
6237 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6238 }
6239 }
6240 }
6241 acc
6242 }
6243}
6244
6245#[cfg(target_arch = "x86_64")]
6258#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6259unsafe fn dot_q4tp_row_1x4_avx512(
6260 nib: &[u8],
6261 r: usize,
6262 gpr: usize,
6263 xs: [&[i8]; 4],
6264 scales: &[f32],
6265) -> [f32; 4] {
6266 unsafe {
6268 use core::arch::x86_64::*;
6269 let lomask = _mm256_set1_epi8(0x0F);
6270 let eight = _mm256_set1_epi8(8);
6271 let zero = _mm512_setzero_si512();
6272 let (mut v0, mut v1, mut v2, mut v3) = (
6273 _mm512_setzero_ps(),
6274 _mm512_setzero_ps(),
6275 _mm512_setzero_ps(),
6276 _mm512_setzero_ps(),
6277 );
6278 let pairs = gpr / 2;
6279 for gp in 0..pairs {
6280 let gi = gp * 2;
6281 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6282 let bb = _mm256_loadu_si256(t as *const __m256i);
6283 let lo = _mm256_and_si256(bb, lomask);
6284 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6285 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6290 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6291 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6292 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6293 let wabs = _mm512_abs_epi8(w);
6294 let neg = _mm512_movepi8_mask(w);
6295 let off = gi * GROUP_SIZE;
6296 let sv = _mm512_insertf32x8::<1>(
6297 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6298 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6299 );
6300 let dot = |x: &[i8]| -> __m512 {
6301 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6302 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6303 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6304 };
6305 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6306 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6307 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6308 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6309 }
6310 let mut acc = [
6311 _mm512_reduce_add_ps(v0),
6312 _mm512_reduce_add_ps(v1),
6313 _mm512_reduce_add_ps(v2),
6314 _mm512_reduce_add_ps(v3),
6315 ];
6316 if gpr % 2 == 1 {
6319 let off = (gpr - 1) * GROUP_SIZE;
6320 for j in off..off + GROUP_SIZE {
6321 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6322 let ws = w * s;
6323 for k in 0..4 {
6324 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6325 }
6326 }
6327 }
6328 acc
6329 }
6330}
6331
6332#[cfg(target_arch = "aarch64")]
6336#[target_feature(enable = "neon,dotprod")]
6337unsafe fn dot_q4tp_row_1x4_sdot(
6338 nib: &[u8],
6339 r: usize,
6340 gpr: usize,
6341 xs: [&[i8]; 4],
6342 scales: &[f32],
6343) -> [f32; 4] {
6344 unsafe {
6346 use core::arch::aarch64::*;
6347 use core::arch::asm;
6348 let lomask = vdupq_n_u8(0x0F);
6349 let eight = vdupq_n_s8(8);
6350 let (mut v0, mut v1, mut v2, mut v3) = (
6366 vdupq_n_f32(0.0),
6367 vdupq_n_f32(0.0),
6368 vdupq_n_f32(0.0),
6369 vdupq_n_f32(0.0),
6370 );
6371 for gi in 0..gpr {
6372 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6373 let s = *scales.get_unchecked(gi);
6374 let bb = vld1q_u8(t);
6375 let lo = vandq_u8(bb, lomask);
6376 let hi = vshrq_n_u8::<4>(bb);
6377 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6378 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6379 let off = gi * GROUP_SIZE;
6380 let dot4 = |x: &[i8]| -> int32x4_t {
6381 let x0 = vld1q_s8(x.as_ptr().add(off));
6382 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6383 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6384 asm!(
6385 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6386 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6387 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6388 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6389 options(pure, nomem, nostack),
6390 );
6391 vaddq_s32(a0, a1)
6392 };
6393 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6394 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6395 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6396 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6397 }
6398 [
6399 vaddvq_f32(v0),
6400 vaddvq_f32(v1),
6401 vaddvq_f32(v2),
6402 vaddvq_f32(v3),
6403 ]
6404 }
6405}
6406
6407fn q4tp_matmat(
6411 bytes: &[u8],
6412 xs_all: &[f32],
6413 b: usize,
6414 rows: usize,
6415 cols: usize,
6416 out: &mut [f32],
6417 pool: Option<&Pool>,
6418) {
6419 debug_assert_eq!(out.len(), b * rows);
6420 let gpr = cols / GROUP_SIZE;
6421 let v = Q4tpView::new(bytes, rows, cols);
6422
6423 #[cfg(target_os = "macos")]
6425 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6426 dequant_matmat_accel(
6427 &|r, dst| {
6428 let mut sc = [0f32; 32];
6429 let mut scv;
6430 let s: &[f32] = if gpr <= 32 {
6431 v.scales_into(r, gpr, &mut sc);
6432 &sc[..gpr]
6433 } else {
6434 scv = vec![0f32; gpr];
6435 v.scales_into(r, gpr, &mut scv);
6436 &scv
6437 };
6438 for gi in 0..gpr {
6439 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6440 for (k, &bb) in tile.iter().enumerate() {
6441 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6442 dst[gi * GROUP_SIZE + k * 2 + 1] =
6443 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6444 }
6445 }
6446 },
6447 xs_all,
6448 b,
6449 rows,
6450 cols,
6451 out,
6452 pool,
6453 );
6454 return;
6455 }
6456
6457 let out_addr = SendMut(out.as_mut_ptr());
6458 if a8w8_enabled() {
6459 let acts: Vec<SplitAct> = (0..b)
6460 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6461 .collect();
6462 let acts = &acts;
6463 #[cfg(target_arch = "aarch64")]
6464 let blocked_ok = sdot_enabled() && blocked_enabled();
6465 #[cfg(target_arch = "x86_64")]
6472 let blocked_ok = q4tp_blocked_x86() && !row_exact();
6473 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6474 let blocked_ok = false;
6475 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6485 .ok()
6486 .and_then(|v| v.parse().ok())
6487 .filter(|v| *v > 0)
6488 .unwrap_or(256);
6489 let run = |start: usize, end: usize| {
6490 for abase in (0..acts.len()).step_by(panel_cols) {
6491 let alen = (acts.len() - abase).min(panel_cols);
6492 let mut sc = vec![0f32; gpr];
6493 #[cfg(target_arch = "x86_64")]
6494 let mut r_lo = start;
6495 #[cfg(target_arch = "x86_64")]
6496 if blocked_ok && alen >= 8 {
6497 let mut sc1 = vec![0f32; gpr];
6498 while r_lo + 2 <= end {
6499 v.scales_into(r_lo, gpr, &mut sc);
6500 v.scales_into(r_lo + 1, gpr, &mut sc1);
6501 let mut bi = 0usize;
6502 while bi + 8 <= alen {
6503 let xs = [
6504 acts[abase + bi].xq.as_slice(),
6505 acts[abase + bi + 1].xq.as_slice(),
6506 acts[abase + bi + 2].xq.as_slice(),
6507 acts[abase + bi + 3].xq.as_slice(),
6508 acts[abase + bi + 4].xq.as_slice(),
6509 acts[abase + bi + 5].xq.as_slice(),
6510 acts[abase + bi + 6].xq.as_slice(),
6511 acts[abase + bi + 7].xq.as_slice(),
6512 ];
6513 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6514 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6515 for k in 0..8 {
6516 let act = &acts[abase + bi + k];
6517 let mut acc = dr[k] * act.sx;
6518 for &(j, xv) in &act.outliers {
6519 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6520 acc += w * s * xv;
6521 }
6522 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6524 }
6525 }
6526 bi += 8;
6527 }
6528 for row in [r_lo, r_lo + 1] {
6531 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6532 for b2 in bi..alen {
6533 let act = &acts[abase + b2];
6534 let xs4 = [
6535 act.xq.as_slice(),
6536 act.xq.as_slice(),
6537 act.xq.as_slice(),
6538 act.xq.as_slice(),
6539 ];
6540 let d =
6541 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6542 let mut acc = d[0] * act.sx;
6543 for &(j, xv) in &act.outliers {
6544 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6545 acc += w * s * xv;
6546 }
6547 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6549 }
6550 }
6551 r_lo += 2;
6552 }
6553 }
6554 #[cfg(target_arch = "x86_64")]
6555 let row_start = r_lo;
6556 #[cfg(not(target_arch = "x86_64"))]
6557 let row_start = start;
6558 for r in row_start..end {
6559 v.scales_into(r, gpr, &mut sc);
6560 let mut bi = 0usize;
6561 #[cfg(target_arch = "x86_64")]
6562 if blocked_ok {
6563 while bi + 8 <= alen {
6564 let xs = [
6565 acts[abase + bi].xq.as_slice(),
6566 acts[abase + bi + 1].xq.as_slice(),
6567 acts[abase + bi + 2].xq.as_slice(),
6568 acts[abase + bi + 3].xq.as_slice(),
6569 acts[abase + bi + 4].xq.as_slice(),
6570 acts[abase + bi + 5].xq.as_slice(),
6571 acts[abase + bi + 6].xq.as_slice(),
6572 acts[abase + bi + 7].xq.as_slice(),
6573 ];
6574 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6575 for k in 0..8 {
6576 let act = &acts[abase + bi + k];
6577 let mut acc = d[k] * act.sx;
6578 for &(j, xv) in &act.outliers {
6579 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6580 acc += w * s * xv;
6581 }
6582 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6584 }
6585 bi += 8;
6586 }
6587 while bi + 4 <= alen {
6588 let xs = [
6589 acts[abase + bi].xq.as_slice(),
6590 acts[abase + bi + 1].xq.as_slice(),
6591 acts[abase + bi + 2].xq.as_slice(),
6592 acts[abase + bi + 3].xq.as_slice(),
6593 ];
6594 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6595 for k in 0..4 {
6596 let act = &acts[abase + bi + k];
6597 let mut acc = d[k] * act.sx;
6598 for &(j, xv) in &act.outliers {
6599 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6600 acc += w * s * xv;
6601 }
6602 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6604 }
6605 bi += 4;
6606 }
6607 }
6608 #[cfg(target_arch = "aarch64")]
6609 if blocked_ok {
6610 while bi + 4 <= alen {
6611 let xs = [
6612 acts[abase + bi].xq.as_slice(),
6613 acts[abase + bi + 1].xq.as_slice(),
6614 acts[abase + bi + 2].xq.as_slice(),
6615 acts[abase + bi + 3].xq.as_slice(),
6616 ];
6617 let d = unsafe {
6618 if q4tp_v1() {
6619 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6620 } else {
6621 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6622 }
6623 };
6624 for k in 0..4 {
6625 let act = &acts[abase + bi + k];
6626 let mut acc = d[k] * act.sx;
6627 for &(j, xv) in &act.outliers {
6628 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6629 acc += w * s * xv;
6630 }
6631 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6633 }
6634 bi += 4;
6635 }
6636 }
6637 let _ = blocked_ok;
6638 while bi < alen {
6639 let act = &acts[abase + bi];
6640 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6641 for &(j, xv) in &act.outliers {
6642 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6643 acc += w * s * xv;
6644 }
6645 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
6647 bi += 1;
6648 }
6649 }
6650 }
6651 };
6652 dispatch_rows(pool, rows, &run);
6653 return;
6654 }
6655
6656 let run = |start: usize, end: usize| {
6657 let mut sc = vec![0f32; gpr];
6658 for r in start..end {
6659 v.scales_into(r, gpr, &mut sc);
6660 for bi in 0..b {
6661 let x = &xs_all[bi * cols..(bi + 1) * cols];
6662 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
6664 }
6665 }
6666 };
6667 dispatch_rows(pool, rows, &run);
6668}
6669
6670fn q4t_matvec(
6672 bytes: &[u8],
6673 x: &[f32],
6674 rows: usize,
6675 cols: usize,
6676 out: &mut [f32],
6677 pool: Option<&Pool>,
6678) {
6679 debug_assert_eq!(out.len(), rows);
6680 let gpr = cols / GROUP_SIZE;
6681 let out_addr = SendMut(out.as_mut_ptr());
6682 if a8w8_enabled() {
6683 let act = split_act(x);
6684 let run = move |start: usize, end: usize| {
6685 for r in start..end {
6686 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6687 for &(j, xv) in &act.outliers {
6688 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6689 acc += w * s * xv;
6690 }
6691 unsafe { *out_addr.at(r) = acc };
6693 }
6694 };
6695 dispatch_rows(pool, rows, &run);
6696 return;
6697 }
6698 let run = move |start: usize, end: usize| {
6699 for r in start..end {
6700 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
6702 }
6703 };
6704 dispatch_rows(pool, rows, &run);
6705}
6706
6707#[allow(clippy::too_many_arguments)]
6709fn q4t_matvec2(
6710 bytes: &[u8],
6711 x1: &[f32],
6712 x2: &[f32],
6713 rows: usize,
6714 cols: usize,
6715 o1: &mut [f32],
6716 o2: &mut [f32],
6717 pool: Option<&Pool>,
6718) {
6719 let gpr = cols / GROUP_SIZE;
6720 let p1 = SendMut(o1.as_mut_ptr());
6721 let p2 = SendMut(o2.as_mut_ptr());
6722 if a8w8_enabled() {
6723 let a1 = split_act(x1);
6724 let a2 = split_act(x2);
6725 let run = move |start: usize, end: usize| {
6726 for r in start..end {
6727 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
6728 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
6729 for &(j, xv) in &a1.outliers {
6730 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6731 v1 += w * s * xv;
6732 }
6733 for &(j, xv) in &a2.outliers {
6734 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6735 v2 += w * s * xv;
6736 }
6737 unsafe {
6739 *p1.at(r) = v1;
6740 *p2.at(r) = v2;
6741 }
6742 }
6743 };
6744 dispatch_rows(pool, rows, &run);
6745 return;
6746 }
6747 let run = move |start: usize, end: usize| {
6748 for r in start..end {
6749 unsafe {
6751 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
6752 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
6753 }
6754 }
6755 };
6756 dispatch_rows(pool, rows, &run);
6757}
6758
6759#[allow(clippy::too_many_arguments)]
6761#[cfg(target_os = "macos")]
6767fn dequant_matmat_accel(
6768 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
6769 xs_all: &[f32],
6770 b: usize,
6771 rows: usize,
6772 cols: usize,
6773 out: &mut [f32],
6774 pool: Option<&Pool>,
6775) {
6776 const TR: usize = 2048;
6777 thread_local! {
6778 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6779 }
6780 WTILE.with(|wt| {
6781 let mut wtile = wt.borrow_mut();
6782 wtile.resize(TR * cols, 0.0);
6783 let mut r0 = 0usize;
6784 while r0 < rows {
6785 let tr = TR.min(rows - r0);
6786 let wt_addr = SendMut(wtile.as_mut_ptr());
6787 let run = |start: usize, end: usize| {
6788 for r in start..end {
6789 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
6791 dequant_row(r0 + r, dst);
6792 }
6793 };
6794 dispatch_rows(pool, tr, &run);
6795 unsafe {
6796 accel_blas::cblas_sgemm(
6797 101, 111, 112, b as i32,
6801 tr as i32,
6802 cols as i32,
6803 1.0,
6804 xs_all.as_ptr(),
6805 cols as i32,
6806 wtile.as_ptr(),
6807 cols as i32,
6808 0.0,
6809 out.as_mut_ptr().add(r0),
6810 rows as i32,
6811 );
6812 }
6813 r0 += tr;
6814 }
6815 });
6816}
6817
6818fn q4t_matmat(
6819 bytes: &[u8],
6820 xs_all: &[f32],
6821 b: usize,
6822 rows: usize,
6823 cols: usize,
6824 out: &mut [f32],
6825 pool: Option<&Pool>,
6826) {
6827 debug_assert_eq!(out.len(), b * rows);
6828 let gpr = cols / GROUP_SIZE;
6829 #[cfg(target_os = "macos")]
6833 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6834 dequant_matmat_accel(
6835 &|r, dst| {
6836 for gi in 0..gpr {
6837 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6838 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6839 for (k, &bb) in tile[2..].iter().enumerate() {
6840 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
6841 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
6842 }
6843 }
6844 },
6845 xs_all,
6846 b,
6847 rows,
6848 cols,
6849 out,
6850 pool,
6851 );
6852 return;
6853 }
6854 let out_addr = SendMut(out.as_mut_ptr());
6855 if a8w8_enabled() {
6856 let acts: Vec<SplitAct> = (0..b)
6857 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6858 .collect();
6859 let acts = &acts;
6860 #[cfg(target_arch = "x86_64")]
6861 let blocked_ok = avx2_enabled() && blocked_enabled();
6862 #[cfg(target_arch = "aarch64")]
6863 let blocked_ok = sdot_enabled() && blocked_enabled();
6864 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
6865 let blocked_ok = false;
6866 let run = move |start: usize, end: usize| {
6867 for r in start..end {
6868 let mut bi = 0usize;
6869 #[cfg(target_arch = "aarch64")]
6870 if blocked_ok {
6871 while bi + 4 <= acts.len() {
6872 let xs = [
6873 acts[bi].xq.as_slice(),
6874 acts[bi + 1].xq.as_slice(),
6875 acts[bi + 2].xq.as_slice(),
6876 acts[bi + 3].xq.as_slice(),
6877 ];
6878 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
6879 for k in 0..4 {
6880 let act = &acts[bi + k];
6881 let mut acc = d[k] * act.sx;
6882 for &(j, xv) in &act.outliers {
6883 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6884 acc += w * sc * xv;
6885 }
6886 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6888 }
6889 bi += 4;
6890 }
6891 }
6892 #[cfg(target_arch = "x86_64")]
6893 if blocked_ok {
6894 while bi + 4 <= acts.len() {
6895 let xs = [
6896 acts[bi].xq.as_slice(),
6897 acts[bi + 1].xq.as_slice(),
6898 acts[bi + 2].xq.as_slice(),
6899 acts[bi + 3].xq.as_slice(),
6900 ];
6901 let d = unsafe {
6902 if vnni_tiles_enabled() {
6903 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
6904 } else {
6905 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
6906 }
6907 };
6908 for k in 0..4 {
6909 let act = &acts[bi + k];
6910 let mut acc = d[k] * act.sx;
6911 for &(j, xv) in &act.outliers {
6912 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6913 acc += w * sc * xv;
6914 }
6915 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6917 }
6918 bi += 4;
6919 }
6920 }
6921 let _ = blocked_ok;
6922 while bi < acts.len() {
6923 let act = &acts[bi];
6924 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6925 for &(j, xv) in &act.outliers {
6926 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6927 acc += w * s * xv;
6928 }
6929 unsafe { *out_addr.at(bi * rows + r) = acc };
6931 bi += 1;
6932 }
6933 }
6934 };
6935 dispatch_rows(pool, rows, &run);
6936 return;
6937 }
6938 let run = move |start: usize, end: usize| {
6939 for r in start..end {
6940 for bi in 0..b {
6941 let x = &xs_all[bi * cols..(bi + 1) * cols];
6942 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
6944 }
6945 }
6946 };
6947 dispatch_rows(pool, rows, &run);
6948}
6949
6950fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
6959 (0..gpr)
6960 .map(|gi| {
6961 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
6962 .iter()
6963 .map(|&v| v as i32)
6964 .sum()
6965 })
6966 .collect()
6967}
6968
6969#[inline]
6973#[allow(unreachable_code)]
6974#[cfg(target_arch = "x86_64")]
6979#[target_feature(enable = "avx2")]
6980unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6981 unsafe {
6983 use core::arch::x86_64::*;
6984 let expand = _mm256_setr_epi8(
6986 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,
6987 3, 3, 3,
6988 );
6989 let bitsel = _mm256_setr_epi8(
6990 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6991 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6992 );
6993 let ones8 = _mm256_set1_epi8(1);
6994 let ones16 = _mm256_set1_epi16(1);
6995 let mut acc = 0f32;
6996 for gi in 0..gpr {
6997 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6998 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6999 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7000 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7001 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7002 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7003 let sel = _mm256_and_si256(x, mask);
7004 let p16 = _mm256_maddubs_epi16(ones8, sel);
7006 let d32 = _mm256_madd_epi16(p16, ones16);
7007 let hi128 = _mm256_extracti128_si256::<1>(d32);
7008 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7009 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7010 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7011 let msum = _mm_cvtsi128_si32(s32);
7012 let d = 2 * msum - gsum[gi];
7015 acc += d as f32 * s;
7016 }
7017 acc
7018 }
7019}
7020
7021#[cfg(target_arch = "x86_64")]
7024#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7025unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7026 unsafe {
7028 use core::arch::x86_64::*;
7029 let expand = _mm256_setr_epi8(
7030 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,
7031 3, 3, 3,
7032 );
7033 let bitsel = _mm256_setr_epi8(
7034 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7035 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7036 );
7037 let ones8 = _mm256_set1_epi8(1);
7038 let mut acc = 0f32;
7039 for gi in 0..gpr {
7040 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7041 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7042 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7043 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7044 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7045 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7046 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7047 let d = 2 * msum - gsum[gi];
7048 acc += d as f32 * s;
7049 }
7050 acc
7051 }
7052}
7053
7054#[cfg(target_arch = "x86_64")]
7056#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7057unsafe fn dot_q1_row_1x4_vnni(
7058 bytes: &[u8],
7059 r: usize,
7060 gpr: usize,
7061 xs: [&[i8]; 4],
7062 gsums: [&[i32]; 4],
7063) -> [f32; 4] {
7064 unsafe {
7066 use core::arch::x86_64::*;
7067 let expand = _mm256_setr_epi8(
7068 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,
7069 3, 3, 3,
7070 );
7071 let bitsel = _mm256_setr_epi8(
7072 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7073 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7074 );
7075 let ones8 = _mm256_set1_epi8(1);
7076 let mut acc = [0f32; 4];
7077 for gi in 0..gpr {
7078 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7079 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7080 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7081 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7082 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7083 for (k, xq) in xs.iter().enumerate() {
7084 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7085 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7086 let d = 2 * msum - gsums[k][gi];
7087 acc[k] += d as f32 * s;
7088 }
7089 }
7090 acc
7091 }
7092}
7093
7094#[cfg(target_arch = "x86_64")]
7097#[target_feature(enable = "avx2")]
7098unsafe fn dot_q1_row_1x4_avx2(
7099 bytes: &[u8],
7100 r: usize,
7101 gpr: usize,
7102 xs: [&[i8]; 4],
7103 gsums: [&[i32]; 4],
7104) -> [f32; 4] {
7105 unsafe {
7107 use core::arch::x86_64::*;
7108 let expand = _mm256_setr_epi8(
7109 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,
7110 3, 3, 3,
7111 );
7112 let bitsel = _mm256_setr_epi8(
7113 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7114 -128, 1, 2, 4, 8, 16, 32, 64, -128,
7115 );
7116 let ones8 = _mm256_set1_epi8(1);
7117 let ones16 = _mm256_set1_epi16(1);
7118 let mut acc = [0f32; 4];
7119 for gi in 0..gpr {
7120 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7121 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7122 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7123 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7124 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7125 for (k, xq) in xs.iter().enumerate() {
7126 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7127 let sel = _mm256_and_si256(x, mask);
7128 let p16 = _mm256_maddubs_epi16(ones8, sel);
7129 let d32 = _mm256_madd_epi16(p16, ones16);
7130 let hi128 = _mm256_extracti128_si256::<1>(d32);
7131 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7132 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7133 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7134 let msum = _mm_cvtsi128_si32(s32);
7135 let d = 2 * msum - gsums[k][gi];
7136 acc[k] += d as f32 * s;
7137 }
7138 }
7139 acc
7140 }
7141}
7142
7143#[allow(unreachable_code)]
7144fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7145 #[cfg(target_arch = "aarch64")]
7146 unsafe {
7147 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7148 }
7149 #[cfg(target_arch = "x86_64")]
7150 if avx2_enabled() {
7151 unsafe {
7152 if vnni_tiles_enabled() {
7153 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7154 }
7155 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7156 }
7157 }
7158 let _ = gsum;
7159 let mut acc = 0f32;
7160 for gi in 0..gpr {
7161 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7162 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7163 let mut d = 0i32;
7164 for (j, &b) in tile[2..].iter().enumerate() {
7165 for k in 0..8 {
7166 let w = ((b >> k) & 1) as i32 * 2 - 1;
7167 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7168 }
7169 }
7170 acc += d as f32 * s;
7171 }
7172 acc
7173}
7174
7175#[cfg(target_arch = "aarch64")]
7184#[target_feature(enable = "neon,dotprod")]
7185unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7186 unsafe {
7189 use core::arch::aarch64::*;
7190 use core::arch::asm;
7191 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7192 let m = vld1q_u8(MASKS.as_ptr());
7193 macro_rules! tile_dot {
7195 ($t:expr, $x:expr) => {{
7196 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7197 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7198 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7199 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7200 let x0 = vld1q_s8($x);
7201 let x1 = vld1q_s8($x.add(16));
7202 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7203 asm!(
7204 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7205 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7206 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7207 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7208 options(pure, nomem, nostack),
7209 );
7210 vaddq_s32(a0, a1)
7211 }};
7212 }
7213 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7222 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7223 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7224 const IW11: [u8; 16] = [
7225 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7226 ];
7227 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7228 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7229 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7230 let isc = vld1_u8(ISC.as_ptr());
7231 macro_rules! tile_dot_tbl {
7233 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7234 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7235 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7236 let x0 = vld1q_s8($x);
7237 let x1 = vld1q_s8($x.add(16));
7238 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7239 asm!(
7240 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7241 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7242 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7243 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7244 options(pure, nomem, nostack),
7245 );
7246 vaddq_s32(a0, a1)
7247 }};
7248 }
7249 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7250 let row_base = r * gpr * Q1_TILE;
7251 let abs_end = bytes.len();
7252 let xp = xq.as_ptr();
7253 let gp = gsum.as_ptr();
7254 let mut accv = vdupq_n_f32(0.0);
7255 let mut gi = 0;
7256 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7259 let t0 = base.add(gi * Q1_TILE);
7260 let ld_a = vld1q_u8(t0);
7261 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7262 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7263 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7264 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7265 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7266 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7268 let g = vld1q_s32(gp.add(gi));
7269 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7270 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7271 let scf: float32x4_t;
7272 asm!(
7273 "fcvtl {o:v}.4s, {i:v}.4h",
7274 o = out(vreg) scf, i = in(vreg) sc16,
7275 options(pure, nomem, nostack),
7276 );
7277 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7278 gi += 4;
7279 }
7280 let mut acc = vaddvq_f32(accv);
7281 while gi < gpr {
7282 let t = base.add(gi * Q1_TILE);
7283 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7284 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7285 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7286 gi += 1;
7287 }
7288 acc
7289 }
7290}
7291
7292#[cfg(target_arch = "aarch64")]
7297#[target_feature(enable = "neon,dotprod")]
7298unsafe fn dot_q1_row_1x4_sdot(
7299 bytes: &[u8],
7300 r: usize,
7301 gpr: usize,
7302 xs: [&[i8]; 4],
7303 gs: [&[i32]; 4],
7304) -> [f32; 4] {
7305 unsafe {
7307 use core::arch::aarch64::*;
7308 use core::arch::asm;
7309 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7310 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7311 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7312 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7313 const IW11: [u8; 16] = [
7314 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7315 ];
7316 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7317 let m = vld1q_u8(MASKS.as_ptr());
7318 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7319 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7320 let isc = vld1_u8(ISC.as_ptr());
7321 macro_rules! sdot2 {
7322 ($w0:expr, $w1:expr, $x:expr) => {{
7323 let x0 = vld1q_s8($x);
7324 let x1 = vld1q_s8($x.add(16));
7325 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7326 asm!(
7327 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7328 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7329 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7330 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7331 options(pure, nomem, nostack),
7332 );
7333 vaddq_s32(a0, a1)
7334 }};
7335 }
7336 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7337 let row_base = r * gpr * Q1_TILE;
7338 let abs_end = bytes.len();
7339 let mut accv = [vdupq_n_f32(0.0); 4];
7340 let mut gi = 0;
7341 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7342 let t0 = base.add(gi * Q1_TILE);
7343 let ld_a = vld1q_u8(t0);
7344 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7345 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7347 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7348 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7349 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7350 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7351 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7352 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7353 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7354 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7355 let scf: float32x4_t;
7356 asm!(
7357 "fcvtl {o:v}.4s, {i:v}.4h",
7358 o = out(vreg) scf, i = in(vreg) sc16,
7359 options(pure, nomem, nostack),
7360 );
7361 for k in 0..4 {
7362 let xp = xs[k].as_ptr();
7363 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7364 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7365 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7366 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7367 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7368 let g = vld1q_s32(gs[k].as_ptr().add(gi));
7369 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7370 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7371 }
7372 gi += 4;
7373 }
7374 let mut acc = [
7375 vaddvq_f32(accv[0]),
7376 vaddvq_f32(accv[1]),
7377 vaddvq_f32(accv[2]),
7378 vaddvq_f32(accv[3]),
7379 ];
7380 while gi < gpr {
7381 let t = base.add(gi * Q1_TILE);
7382 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7383 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7384 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7385 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7386 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7387 for k in 0..4 {
7388 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7389 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7390 }
7391 gi += 1;
7392 }
7393 acc
7394 }
7395}
7396
7397#[inline]
7399fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7400 let gi = j / GROUP_SIZE;
7401 let k = j % GROUP_SIZE;
7402 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7403 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7404 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7405 ((bit as i32 * 2 - 1) as f32, s)
7406}
7407
7408#[inline]
7410fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7411 let mut acc = 0f32;
7412 for gi in 0..gpr {
7413 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7414 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7415 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7416 let mut ga = 0f32;
7417 for (j, &b) in tile[2..].iter().enumerate() {
7418 for k in 0..8 {
7419 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7420 }
7421 }
7422 acc += ga * s;
7423 }
7424 acc
7425}
7426
7427#[allow(clippy::too_many_arguments)]
7430fn q1_range_a8w8(
7431 bytes: &[u8],
7432 gpr: usize,
7433 act: &SplitAct,
7434 gsum: &[i32],
7435 out: SendMut,
7436 start: usize,
7437 end: usize,
7438) {
7439 for r in start..end {
7440 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7441 for &(j, xv) in &act.outliers {
7442 let (w, s) = q1_outlier(bytes, r, gpr, j);
7443 acc += w * s * xv;
7444 }
7445 unsafe { *out.at(r) = acc };
7447 }
7448}
7449
7450fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7452 for r in start..end {
7453 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
7455 }
7456}
7457
7458fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
7463 let entries = base_len + (rows + 1) * 4;
7464 (base_len, entries, entries <= bytes.len())
7465}
7466
7467#[inline]
7469fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
7470 let o = rp_off + r * 4;
7471 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
7472}
7473
7474const SIGN5: [[f32; 5]; 256] = {
7478 let mut lut = [[0.0f32; 5]; 256];
7479 let pow3 = [1u16, 3, 9, 27, 81];
7480 let mut byte = 0usize;
7481 while byte < 256 {
7482 let mut i = 0usize;
7483 while i < 5 {
7484 let code = (byte as u16 / pow3[i]) % 3;
7485 lut[byte][i] = if code == 1 {
7486 1.0
7487 } else if code == 2 {
7488 -1.0
7489 } else {
7490 0.0
7491 };
7492 i += 1;
7493 }
7494 byte += 1;
7495 }
7496 lut
7497};
7498
7499const SIGN5_I8: [[i8; 5]; 256] = {
7501 let mut lut = [[0i8; 5]; 256];
7502 let pow3 = [1u16, 3, 9, 27, 81];
7503 let mut byte = 0usize;
7504 while byte < 256 {
7505 let mut i = 0usize;
7506 while i < 5 {
7507 let code = (byte as u16 / pow3[i]) % 3;
7508 lut[byte][i] = if code == 1 {
7509 1
7510 } else if code == 2 {
7511 -1
7512 } else {
7513 0
7514 };
7515 i += 1;
7516 }
7517 byte += 1;
7518 }
7519 lut
7520};
7521
7522const SIGN5_U64: [u64; 256] = {
7528 let mut lut = [0u64; 256];
7529 let pow3 = [1u16, 3, 9, 27, 81];
7530 let mut byte = 0usize;
7531 while byte < 256 {
7532 let mut v = 0u64;
7533 let mut i = 0usize;
7534 while i < 5 {
7535 let code = (byte as u16 / pow3[i]) % 3;
7536 let s: u8 = if code == 1 {
7537 1
7538 } else if code == 2 {
7539 0xFF
7540 } else {
7541 0
7542 };
7543 v |= (s as u64) << (i * 8);
7544 i += 1;
7545 }
7546 lut[byte] = v;
7547 byte += 1;
7548 }
7549 lut
7550};
7551
7552#[inline]
7557fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
7558 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7559 let off = (r * gpr + j / GROUP_SIZE) * TILE;
7560 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7561 let within = j % GROUP_SIZE;
7562 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
7563}
7564
7565#[cfg(target_arch = "aarch64")]
7568#[target_feature(enable = "neon,dotprod")]
7569#[inline]
7570unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
7571 unsafe {
7573 use core::arch::aarch64::*;
7574 use core::arch::asm;
7575 let w0 = vld1q_s8(w);
7576 let w1 = vld1q_s8(w.add(16));
7577 let x0 = vld1q_s8(x);
7578 let x1 = vld1q_s8(x.add(16));
7579 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7580 asm!(
7581 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7582 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7583 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7584 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7585 options(pure, nomem, nostack),
7586 );
7587 vaddvq_s32(vaddq_s32(a0, a1))
7588 }
7589}
7590
7591#[cfg(target_arch = "x86_64")]
7594#[target_feature(enable = "avx2")]
7595#[inline]
7596unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
7597 unsafe {
7599 use core::arch::x86_64::*;
7600 let wv = _mm256_loadu_si256(w as *const __m256i);
7601 let xv = _mm256_loadu_si256(x as *const __m256i);
7602 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7603 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
7604 let hi128 = _mm256_extracti128_si256::<1>(d);
7605 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7606 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7607 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7608 _mm_cvtsi128_si32(s32)
7609 }
7610}
7611
7612#[inline]
7617fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
7618 debug_assert!(dst.len() >= 40);
7619 unsafe {
7622 let p = dst.as_mut_ptr();
7623 for bi in 0..7 {
7624 core::ptr::write_unaligned(
7625 p.add(bi * 5) as *mut u64,
7626 SIGN5_U64[*codes.add(bi) as usize],
7627 );
7628 }
7629 }
7630}
7631
7632#[inline]
7637fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
7638 #[cfg(target_arch = "aarch64")]
7639 unsafe {
7640 return sdot32_i8(w, x);
7641 }
7642 #[cfg(target_arch = "x86_64")]
7643 unsafe {
7644 return i8dot32_avx2(w, x);
7645 }
7646 #[allow(unreachable_code)]
7647 unsafe {
7648 let mut s = 0i32;
7649 for k in 0..GROUP_SIZE {
7650 s += *w.add(k) as i32 * *x.add(k) as i32;
7651 }
7652 s
7653 }
7654}
7655
7656#[inline]
7657unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
7658 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
7659 (
7660 SIGN5_U64[*codes as usize],
7661 SIGN5_U64[*codes.add(1) as usize],
7662 SIGN5_U64[*codes.add(2) as usize],
7663 SIGN5_U64[*codes.add(3) as usize],
7664 SIGN5_U64[*codes.add(4) as usize],
7665 SIGN5_U64[*codes.add(5) as usize],
7666 SIGN5_U64[*codes.add(6) as usize],
7667 )
7668 };
7669
7670 let u0 = s0 | (s1 << 40);
7671 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
7672 let u2 = (s3 >> 8) | (s4 << 32);
7673 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
7674
7675 (u0, u1, u2, u3)
7676}
7677
7678#[cfg(target_arch = "aarch64")]
7682#[target_feature(enable = "neon,dotprod")]
7683unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7684 use core::arch::aarch64::*;
7685 use core::arch::asm;
7686 unsafe {
7687 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7688 let mut acc = 0f32;
7689 let bytes_ptr = bytes.as_ptr();
7690 let xq_ptr = xq.as_ptr();
7691 let row_off = r * gpr * TILE;
7692
7693 let gpr2 = gpr & !1;
7694 let mut gi = 0;
7695 while gi < gpr2 {
7696 let off0 = row_off + gi * TILE;
7697 let off1 = off0 + TILE;
7698 let s0 = f16_to_f32(u16::from_le_bytes([
7699 *bytes_ptr.add(off0),
7700 *bytes_ptr.add(off0 + 1),
7701 ]));
7702 let s1 = f16_to_f32(u16::from_le_bytes([
7703 *bytes_ptr.add(off1),
7704 *bytes_ptr.add(off1 + 1),
7705 ]));
7706
7707 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7708 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7709
7710 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7711 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7712 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7713 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7714
7715 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7716 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7717 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
7718 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
7719
7720 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
7721 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7722 asm!(
7723 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
7724 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
7725 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
7726 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
7727 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
7728 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
7729 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
7730 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
7731 options(pure, nomem, nostack),
7732 );
7733 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
7734 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
7735 acc += d0 as f32 * s0 + d1 as f32 * s1;
7736 gi += 2;
7737 }
7738
7739 if gi < gpr {
7740 let off = row_off + gi * TILE;
7741 let s = f16_to_f32(u16::from_le_bytes([
7742 *bytes_ptr.add(off),
7743 *bytes_ptr.add(off + 1),
7744 ]));
7745 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7746 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7747 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7748 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7749 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7750 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7751 asm!(
7752 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7753 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7754 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7755 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7756 options(pure, nomem, nostack),
7757 );
7758 let d = vaddvq_s32(vaddq_s32(a0, a1));
7759 acc += d as f32 * s;
7760 }
7761 acc
7762 }
7763}
7764
7765#[cfg(target_arch = "x86_64")]
7767#[target_feature(enable = "avx2")]
7768unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7769 use core::arch::x86_64::*;
7770 unsafe {
7771 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7772 let mut acc = 0f32;
7773 let bytes_ptr = bytes.as_ptr();
7774 let xq_ptr = xq.as_ptr();
7775 let row_off = r * gpr * TILE;
7776
7777 let ones = _mm256_set1_epi16(1);
7778 for gi in 0..gpr {
7779 let off = row_off + gi * TILE;
7780 let s = f16_to_f32(u16::from_le_bytes([
7781 *bytes_ptr.add(off),
7782 *bytes_ptr.add(off + 1),
7783 ]));
7784 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7785 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7786 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7787 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7788 let d256 = _mm256_madd_epi16(p16, ones);
7789 let d128 = _mm_add_epi32(
7790 _mm256_castsi256_si128(d256),
7791 _mm256_extracti128_si256(d256, 1),
7792 );
7793 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
7794 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
7795 acc += d32 as f32 * s;
7796 }
7797 acc
7798 }
7799}
7800
7801#[cfg(target_arch = "x86_64")]
7803#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7804unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7805 use core::arch::x86_64::*;
7806 unsafe {
7808 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7809 let mut acc = 0f32;
7810 let bytes_ptr = bytes.as_ptr();
7811 let xq_ptr = xq.as_ptr();
7812 let row_off = r * gpr * TILE;
7813 for gi in 0..gpr {
7814 let off = row_off + gi * TILE;
7815 let s = f16_to_f32(u16::from_le_bytes([
7816 *bytes_ptr.add(off),
7817 *bytes_ptr.add(off + 1),
7818 ]));
7819 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7820 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7821 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7822 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7823 acc += d as f32 * s;
7824 }
7825 acc
7826 }
7827}
7828
7829#[inline]
7833fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7834 #[cfg(target_arch = "aarch64")]
7835 unsafe {
7836 return q1t_dot_row_sdot(bytes, r, gpr, xq);
7837 }
7838 #[cfg(target_arch = "x86_64")]
7839 unsafe {
7840 if vnni_tiles_enabled() {
7841 return q1t_dot_row_vnni(bytes, r, gpr, xq);
7842 }
7843 return q1t_dot_row_avx2(bytes, r, gpr, xq);
7844 }
7845 #[allow(unreachable_code)]
7846 {
7847 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7848 let mut acc = 0f32;
7849 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
7851 let off = (r * gpr + gi) * TILE;
7852 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7853 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
7854 let mut d = 0i32;
7855 for k in 0..GROUP_SIZE {
7856 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
7857 }
7858 acc += d as f32 * s;
7859 }
7860 acc
7861 }
7862}
7863
7864fn q1t_row_outlier_correction(
7871 bytes: &[u8],
7872 r: usize,
7873 rp_off: usize,
7874 entries_off: usize,
7875 has_ov: bool,
7876 x: &[f32],
7877) -> f32 {
7878 if !has_ov {
7879 return 0.0;
7880 }
7881 let (c0, c1) = (
7882 q1t_rowptr(bytes, rp_off, r),
7883 q1t_rowptr(bytes, rp_off, r + 1),
7884 );
7885 let mut corr = 0f32;
7886 for p in c0..c1 {
7887 let e = entries_off + p * 4;
7888 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7889 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7890 corr += val * x[col];
7891 }
7892 corr
7893}
7894
7895fn q1t_dequant_row(
7899 bytes: &[u8],
7900 r: usize,
7901 gpr: usize,
7902 rp_off: usize,
7903 entries_off: usize,
7904 has_ov: bool,
7905 buf: &mut [f32],
7906) {
7907 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7908 for g in 0..gpr {
7909 let off = (r * gpr + g) * TILE;
7910 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7911 let codes = &bytes[off + 2..off + TILE];
7912 let bc = g * GROUP_SIZE;
7913 for bi in 0..6 {
7915 let lut = &SIGN5[codes[bi] as usize];
7916 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
7917 for i in 0..5 {
7918 d[i] = lut[i] * s;
7919 }
7920 }
7921 let lut = &SIGN5[codes[6] as usize];
7922 buf[bc + 30] = lut[0] * s;
7923 buf[bc + 31] = lut[1] * s;
7924 }
7925 if !has_ov {
7926 return;
7927 }
7928 let (c0, c1) = (
7929 q1t_rowptr(bytes, rp_off, r),
7930 q1t_rowptr(bytes, rp_off, r + 1),
7931 );
7932 for p in c0..c1 {
7933 let e = entries_off + p * 4;
7934 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7935 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7936 }
7937}
7938
7939fn q1t_add_overlay(
7943 bytes: &[u8],
7944 x: &[f32],
7945 rows: usize,
7946 cols: usize,
7947 out: &mut [f32],
7948 pool: Option<&Pool>,
7949) {
7950 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7951 let gpr = cols / GROUP_SIZE;
7952 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7953 if !has_ov {
7954 return;
7955 }
7956 let out_addr = SendMut(out.as_mut_ptr());
7957 let run = move |start: usize, end: usize| {
7958 for r in start..end {
7959 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7960 unsafe { *out_addr.at(r) += corr };
7962 }
7963 };
7964 dispatch_rows(pool, rows, &run);
7965}
7966
7967#[allow(clippy::too_many_arguments)]
7970fn q1t_range_a8w8(
7971 bytes: &[u8],
7972 gpr: usize,
7973 rp_off: usize,
7974 ent_off: usize,
7975 has_ov: bool,
7976 act: &SplitAct,
7977 x: &[f32],
7978 out: SendMut,
7979 start: usize,
7980 end: usize,
7981) {
7982 for r in start..end {
7983 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7984 for &(j, xv) in &act.outliers {
7985 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7986 }
7987 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7988 unsafe { *out.at(r) = acc };
7990 }
7991}
7992
7993#[allow(clippy::too_many_arguments)]
7996fn q1t_range_f32_batch(
7997 bytes: &[u8],
7998 gpr: usize,
7999 rp_off: usize,
8000 ent_off: usize,
8001 has_ov: bool,
8002 x: &[f32],
8003 out: SendMut,
8004 start: usize,
8005 end: usize,
8006) {
8007 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8008 let mut sg = [0f32; GROUP_SIZE];
8009 for r in start..end {
8010 let mut acc = 0f32;
8011 for g in 0..gpr {
8012 let off = (r * gpr + g) * TILE;
8013 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8014 let codes = &bytes[off + 2..off + TILE];
8015 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8016 for bi in 0..6 {
8017 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8018 }
8019 let lut = &SIGN5[codes[6] as usize];
8020 sg[30] = lut[0];
8021 sg[31] = lut[1];
8022 let mut gsum = 0f32;
8023 for k in 0..GROUP_SIZE {
8024 gsum += sg[k] * xg[k];
8025 }
8026 acc += s * gsum;
8027 }
8028 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8029 unsafe { *out.at(r) = acc };
8031 }
8032}
8033
8034fn q1t_matvec(
8038 bytes: &[u8],
8039 x: &[f32],
8040 rows: usize,
8041 cols: usize,
8042 out: &mut [f32],
8043 pool: Option<&Pool>,
8044) {
8045 debug_assert_eq!(out.len(), rows);
8046 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8047 let gpr = cols / GROUP_SIZE;
8048 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8049 let out_addr = SendMut(out.as_mut_ptr());
8050 if a8w8_enabled() {
8054 let act = split_act(x);
8055 let act = &act;
8056 let run = move |start: usize, end: usize| {
8057 for r in start..end {
8058 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8059 for &(j, xv) in &act.outliers {
8060 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8061 }
8062 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8063 unsafe { *out_addr.at(r) = acc };
8065 }
8066 };
8067 dispatch_rows(pool, rows, &run);
8068 return;
8069 }
8070 let run = move |start: usize, end: usize| {
8071 let mut sg = [0f32; GROUP_SIZE];
8075 for r in start..end {
8076 let mut acc = 0f32;
8077 for g in 0..gpr {
8078 let off = (r * gpr + g) * TILE;
8079 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8080 let codes = &bytes[off + 2..off + TILE];
8081 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8082 for bi in 0..6 {
8083 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8084 }
8085 let lut = &SIGN5[codes[6] as usize];
8086 sg[30] = lut[0];
8087 sg[31] = lut[1];
8088 let mut gsum = 0f32;
8089 for k in 0..GROUP_SIZE {
8090 gsum += sg[k] * xg[k];
8091 }
8092 acc += s * gsum;
8093 }
8094 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8095 unsafe { *out_addr.at(r) = acc };
8096 }
8097 };
8098 dispatch_rows(pool, rows, &run);
8099}
8100
8101#[cfg(target_arch = "aarch64")]
8107#[target_feature(enable = "neon,dotprod")]
8108unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
8109 use core::arch::aarch64::*;
8110 use core::arch::asm;
8111 unsafe {
8113 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8114 let bytes_ptr = bytes.as_ptr();
8115 let row_off = r * gpr * TILE;
8116 let xp = [xa.as_ptr(), xb.as_ptr()];
8117 let mut acc = [0f32; 2];
8118 macro_rules! sdot2 {
8119 ($w0:expr, $w1:expr, $x:expr) => {{
8120 let x0 = vld1q_s8($x);
8121 let x1 = vld1q_s8($x.add(16));
8122 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8123 asm!(
8124 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8125 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8126 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8127 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8128 options(pure, nomem, nostack),
8129 );
8130 vaddvq_s32(vaddq_s32(a0, a1))
8131 }};
8132 }
8133 let gpr2 = gpr & !1;
8134 let mut gi = 0;
8135 while gi < gpr2 {
8136 let off0 = row_off + gi * TILE;
8137 let off1 = off0 + TILE;
8138 let s0 = f16_to_f32(u16::from_le_bytes([
8139 *bytes_ptr.add(off0),
8140 *bytes_ptr.add(off0 + 1),
8141 ]));
8142 let s1 = f16_to_f32(u16::from_le_bytes([
8143 *bytes_ptr.add(off1),
8144 *bytes_ptr.add(off1 + 1),
8145 ]));
8146 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8147 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8148 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8149 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8150 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8151 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8152 for k in 0..2 {
8153 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
8154 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
8155 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
8156 }
8157 gi += 2;
8158 }
8159 if gi < gpr {
8160 let off = row_off + gi * TILE;
8161 let s = f16_to_f32(u16::from_le_bytes([
8162 *bytes_ptr.add(off),
8163 *bytes_ptr.add(off + 1),
8164 ]));
8165 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8166 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8167 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8168 for k in 0..2 {
8169 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
8170 acc[k] += d as f32 * s;
8171 }
8172 }
8173 acc
8174 }
8175}
8176
8177fn q1t_matvec2(
8183 bytes: &[u8],
8184 x1: &[f32],
8185 x2: &[f32],
8186 rows: usize,
8187 cols: usize,
8188 o1: &mut [f32],
8189 o2: &mut [f32],
8190 pool: Option<&Pool>,
8191) {
8192 debug_assert_eq!(o1.len(), rows);
8193 debug_assert_eq!(o2.len(), rows);
8194 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8195 let gpr = cols / GROUP_SIZE;
8196 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8197 let out1 = SendMut(o1.as_mut_ptr());
8198 let out2 = SendMut(o2.as_mut_ptr());
8199 if a8w8_enabled() {
8200 let a1 = split_act(x1);
8201 let a2 = split_act(x2);
8202 let (a1, a2) = (&a1, &a2);
8203 let run = move |start: usize, end: usize| {
8204 for r in start..end {
8205 #[cfg(target_arch = "aarch64")]
8206 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
8209 #[cfg(not(target_arch = "aarch64"))]
8210 let ds = [
8211 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
8212 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
8213 ];
8214 let mut acc1 = ds[0] * a1.sx;
8215 for &(j, xv) in &a1.outliers {
8216 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
8217 }
8218 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
8219 let mut acc2 = ds[1] * a2.sx;
8220 for &(j, xv) in &a2.outliers {
8221 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
8222 }
8223 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
8224 unsafe {
8226 *out1.at(r) = acc1;
8227 *out2.at(r) = acc2;
8228 }
8229 }
8230 };
8231 dispatch_rows(pool, rows, &run);
8232 return;
8233 }
8234 let run = move |start: usize, end: usize| {
8235 let mut sg = [0f32; GROUP_SIZE];
8238 for r in start..end {
8239 let mut acc1 = 0f32;
8240 let mut acc2 = 0f32;
8241 for g in 0..gpr {
8242 let off = (r * gpr + g) * TILE;
8243 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8244 let codes = &bytes[off + 2..off + TILE];
8245 for bi in 0..6 {
8246 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8247 }
8248 let lut = &SIGN5[codes[6] as usize];
8249 sg[30] = lut[0];
8250 sg[31] = lut[1];
8251 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8252 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8253 let mut gsum1 = 0f32;
8254 for k in 0..GROUP_SIZE {
8255 gsum1 += sg[k] * xg1[k];
8256 }
8257 acc1 += s * gsum1;
8258 let mut gsum2 = 0f32;
8259 for k in 0..GROUP_SIZE {
8260 gsum2 += sg[k] * xg2[k];
8261 }
8262 acc2 += s * gsum2;
8263 }
8264 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
8265 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
8266 unsafe {
8268 *out1.at(r) = acc1;
8269 *out2.at(r) = acc2;
8270 }
8271 }
8272 };
8273 dispatch_rows(pool, rows, &run);
8274}
8275
8276fn q1t_matmat(
8279 bytes: &[u8],
8280 xs: &[f32],
8281 b: usize,
8282 rows: usize,
8283 cols: usize,
8284 out: &mut [f32],
8285 pool: Option<&Pool>,
8286) {
8287 debug_assert_eq!(out.len(), b * rows);
8288 const TILE: usize = cortiq_core::quant::Q1T_TILE;
8289 let gpr = cols / GROUP_SIZE;
8290 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8291 let out_addr = SendMut(out.as_mut_ptr());
8292 if a8w8_enabled() {
8296 let acts: Vec<SplitAct> = (0..b)
8297 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
8298 .collect();
8299 let acts = &acts;
8300 let run = move |start: usize, end: usize| {
8301 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
8305 for g in 0..gpr {
8306 let off = (r * gpr + g) * TILE;
8307 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8308 q1t_unpack_group_i8(
8309 bytes.as_ptr().wrapping_add(off + 2),
8310 &mut sg[g * GROUP_SIZE..],
8311 );
8312 }
8313 for bi in 0..b {
8314 let act = &acts[bi];
8315 let mut isum = 0f32;
8316 for g in 0..gpr {
8317 let d = q1t_i8dot32(
8318 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
8319 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
8320 );
8321 isum += d as f32 * sc[g];
8322 }
8323 let mut acc = isum * act.sx;
8324 for &(j, xv) in &act.outliers {
8325 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8326 }
8327 accs[bi] = acc;
8328 }
8329 if has_ov {
8333 let (c0, c1) = (
8334 q1t_rowptr(bytes, rp_off, r),
8335 q1t_rowptr(bytes, rp_off, r + 1),
8336 );
8337 for p in c0..c1 {
8338 let e = ent_off + p * 4;
8339 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8340 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8341 for bi in 0..b {
8342 accs[bi] += val * xs[bi * cols + col];
8343 }
8344 }
8345 }
8346 for bi in 0..b {
8347 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
8348 }
8349 }
8350 };
8351 dispatch_rows(pool, rows, &run);
8352 return;
8353 }
8354 let run = move |start: usize, end: usize| {
8355 let mut buf = vec![0f32; cols];
8356 for r in start..end {
8357 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
8358 for bi in 0..b {
8359 let xr = &xs[bi * cols..(bi + 1) * cols];
8360 let mut acc = 0f32;
8361 for j in 0..cols {
8362 acc += buf[j] * xr[j];
8363 }
8364 unsafe { *out_addr.at(bi * rows + r) = acc };
8365 }
8366 }
8367 };
8368 dispatch_rows(pool, rows, &run);
8369}
8370
8371fn q1_matvec(
8372 bytes: &[u8],
8373 x: &[f32],
8374 rows: usize,
8375 cols: usize,
8376 out: &mut [f32],
8377 pool: Option<&Pool>,
8378) {
8379 debug_assert_eq!(out.len(), rows);
8380 let gpr = cols / GROUP_SIZE;
8381 let out_addr = SendMut(out.as_mut_ptr());
8382 if a8w8_enabled() {
8383 let act = split_act(x);
8384 let gsum = q1_group_sums(&act.xq, gpr);
8385 let (act, gsum) = (&act, &gsum);
8386 let run = move |start: usize, end: usize| {
8387 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
8388 };
8389 dispatch_rows(pool, rows, &run);
8390 return;
8391 }
8392 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
8393 dispatch_rows(pool, rows, &run);
8394}
8395
8396#[allow(clippy::too_many_arguments)]
8398fn q1_matvec2(
8399 bytes: &[u8],
8400 x1: &[f32],
8401 x2: &[f32],
8402 rows: usize,
8403 cols: usize,
8404 o1: &mut [f32],
8405 o2: &mut [f32],
8406 pool: Option<&Pool>,
8407) {
8408 let gpr = cols / GROUP_SIZE;
8409 let p1 = SendMut(o1.as_mut_ptr());
8410 let p2 = SendMut(o2.as_mut_ptr());
8411 if a8w8_enabled() {
8412 let a1 = split_act(x1);
8413 let a2 = split_act(x2);
8414 let g1 = q1_group_sums(&a1.xq, gpr);
8415 let g2 = q1_group_sums(&a2.xq, gpr);
8416 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
8417 let run = move |start: usize, end: usize| {
8418 for r in start..end {
8419 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
8420 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
8421 for &(j, xv) in &a1.outliers {
8422 let (w, s) = q1_outlier(bytes, r, gpr, j);
8423 v1 += w * s * xv;
8424 }
8425 for &(j, xv) in &a2.outliers {
8426 let (w, s) = q1_outlier(bytes, r, gpr, j);
8427 v2 += w * s * xv;
8428 }
8429 unsafe {
8431 *p1.at(r) = v1;
8432 *p2.at(r) = v2;
8433 }
8434 }
8435 };
8436 dispatch_rows(pool, rows, &run);
8437 return;
8438 }
8439 let run = move |start: usize, end: usize| {
8440 for r in start..end {
8441 unsafe {
8443 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
8444 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
8445 }
8446 }
8447 };
8448 dispatch_rows(pool, rows, &run);
8449}
8450
8451#[allow(clippy::too_many_arguments)]
8453fn q1_matmat(
8454 bytes: &[u8],
8455 xs_all: &[f32],
8456 b: usize,
8457 rows: usize,
8458 cols: usize,
8459 out: &mut [f32],
8460 pool: Option<&Pool>,
8461) {
8462 debug_assert_eq!(out.len(), b * rows);
8463 let gpr = cols / GROUP_SIZE;
8464 let out_addr = SendMut(out.as_mut_ptr());
8465 if a8w8_enabled() {
8466 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
8467 .map(|bi| {
8468 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
8469 let gsum = q1_group_sums(&act.xq, gpr);
8470 (act, gsum)
8471 })
8472 .collect();
8473 let acts = &acts;
8474 #[cfg(target_arch = "x86_64")]
8475 let blocked_ok = avx2_enabled() && blocked_enabled();
8476 #[cfg(target_arch = "aarch64")]
8477 let blocked_ok = sdot_enabled() && blocked_enabled();
8478 let run = move |start: usize, end: usize| {
8479 for r in start..end {
8480 let mut bi = 0usize;
8481 #[cfg(target_arch = "aarch64")]
8484 if blocked_ok {
8485 while bi + 4 <= acts.len() {
8486 let xs = [
8487 acts[bi].0.xq.as_slice(),
8488 acts[bi + 1].0.xq.as_slice(),
8489 acts[bi + 2].0.xq.as_slice(),
8490 acts[bi + 3].0.xq.as_slice(),
8491 ];
8492 let gs = [
8493 acts[bi].1.as_slice(),
8494 acts[bi + 1].1.as_slice(),
8495 acts[bi + 2].1.as_slice(),
8496 acts[bi + 3].1.as_slice(),
8497 ];
8498 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
8499 for k in 0..4 {
8500 let (act, _) = &acts[bi + k];
8501 let mut acc = d[k] * act.sx;
8502 for &(j, xv) in &act.outliers {
8503 let (w, sc) = q1_outlier(bytes, r, gpr, j);
8504 acc += w * sc * xv;
8505 }
8506 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8508 }
8509 bi += 4;
8510 }
8511 }
8512 #[cfg(target_arch = "x86_64")]
8513 if blocked_ok {
8514 while bi + 4 <= acts.len() {
8515 let xs = [
8516 acts[bi].0.xq.as_slice(),
8517 acts[bi + 1].0.xq.as_slice(),
8518 acts[bi + 2].0.xq.as_slice(),
8519 acts[bi + 3].0.xq.as_slice(),
8520 ];
8521 let gs = [
8522 acts[bi].1.as_slice(),
8523 acts[bi + 1].1.as_slice(),
8524 acts[bi + 2].1.as_slice(),
8525 acts[bi + 3].1.as_slice(),
8526 ];
8527 let d = unsafe {
8528 if vnni_tiles_enabled() {
8529 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
8530 } else {
8531 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
8532 }
8533 };
8534 for k in 0..4 {
8535 let (act, _) = &acts[bi + k];
8536 let mut acc = d[k] * act.sx;
8537 for &(j, xv) in &act.outliers {
8538 let (w, sc) = q1_outlier(bytes, r, gpr, j);
8539 acc += w * sc * xv;
8540 }
8541 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8543 }
8544 bi += 4;
8545 }
8546 }
8547 while bi < acts.len() {
8548 let (act, gsum) = &acts[bi];
8549 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
8550 for &(j, xv) in &act.outliers {
8551 let (w, s) = q1_outlier(bytes, r, gpr, j);
8552 acc += w * s * xv;
8553 }
8554 unsafe { *out_addr.at(bi * rows + r) = acc };
8556 bi += 1;
8557 }
8558 }
8559 };
8560 dispatch_rows(pool, rows, &run);
8561 return;
8562 }
8563 let run = move |start: usize, end: usize| {
8564 for r in start..end {
8565 for bi in 0..b {
8566 let x = &xs_all[bi * cols..(bi + 1) * cols];
8567 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
8569 }
8570 }
8571 };
8572 dispatch_rows(pool, rows, &run);
8573}
8574
8575fn q4matvec(
8581 bytes: &[u8],
8582 x: &[f32],
8583 rows: usize,
8584 cols: usize,
8585 out: &mut [f32],
8586 pool: Option<&Pool>,
8587) {
8588 debug_assert_eq!(out.len(), rows);
8589 let (packed, scales) = q4_split(bytes, rows, cols);
8590 let gpr = cols / GROUP_SIZE;
8591 let out_addr = SendMut(out.as_mut_ptr());
8592
8593 if a8w8_enabled() {
8594 let act = split_act(x);
8595 let run = move |start: usize, end: usize| {
8596 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
8597 };
8598 dispatch_rows(pool, rows, &run);
8599 return;
8600 }
8601
8602 let run =
8603 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
8604 dispatch_rows(pool, rows, &run);
8605}
8606
8607#[inline]
8610#[allow(unreachable_code)]
8611#[cfg(target_arch = "x86_64")]
8616#[target_feature(enable = "avx2")]
8617unsafe fn dot_q4b_row_1x4_avx2(
8618 buf: &[u8],
8619 scales: &[u8],
8620 g0: usize,
8621 gpr: usize,
8622 xs: [&[i8]; 4],
8623) -> [f32; 4] {
8624 unsafe {
8626 use core::arch::x86_64::*;
8627 let ones = _mm256_set1_epi16(1);
8628 let mut acc = [0f32; 4];
8629 for gi in 0..gpr {
8630 let s = f16_to_f32(u16::from_le_bytes([
8631 scales[(g0 + gi) * 2],
8632 scales[(g0 + gi) * 2 + 1],
8633 ]));
8634 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8635 let aw = _mm256_abs_epi8(w);
8636 for (k, xq) in xs.iter().enumerate() {
8637 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8638 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8639 let d = _mm256_madd_epi16(p16, ones);
8640 let hi128 = _mm256_extracti128_si256::<1>(d);
8641 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8642 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8643 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8644 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
8645 }
8646 }
8647 acc
8648 }
8649}
8650
8651#[cfg(target_arch = "x86_64")]
8653#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8654unsafe fn dot_q4b_row_1x4_vnni(
8655 buf: &[u8],
8656 scales: &[u8],
8657 g0: usize,
8658 gpr: usize,
8659 xs: [&[i8]; 4],
8660) -> [f32; 4] {
8661 unsafe {
8663 use core::arch::x86_64::*;
8664 let mut acc = [0f32; 4];
8665 for gi in 0..gpr {
8666 let s = f16_to_f32(u16::from_le_bytes([
8667 scales[(g0 + gi) * 2],
8668 scales[(g0 + gi) * 2 + 1],
8669 ]));
8670 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8671 let aw = _mm256_abs_epi8(w);
8672 for (k, xq) in xs.iter().enumerate() {
8673 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8674 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8675 acc[k] += d as f32 * s;
8676 }
8677 }
8678 acc
8679 }
8680}
8681
8682#[cfg(target_arch = "x86_64")]
8688#[target_feature(enable = "avx2")]
8689unsafe fn dot_q4b_row_1x4_sx_avx2(
8690 buf: &[u8],
8691 scales: &[u8],
8692 g0: usize,
8693 gpr: usize,
8694 xs: [&[i8]; 4],
8695 sxs: [f32; 4],
8696) -> [f32; 4] {
8697 unsafe {
8699 use core::arch::x86_64::*;
8700 let ones = _mm256_set1_epi16(1);
8701 let mut acc = [0f32; 4];
8702 for gi in 0..gpr {
8703 let s = f16_to_f32(u16::from_le_bytes([
8704 scales[(g0 + gi) * 2],
8705 scales[(g0 + gi) * 2 + 1],
8706 ]));
8707 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8708 let aw = _mm256_abs_epi8(w);
8709 for (k, xq) in xs.iter().enumerate() {
8710 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8711 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8712 let d = _mm256_madd_epi16(p16, ones);
8713 let hi128 = _mm256_extracti128_si256::<1>(d);
8714 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8715 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8716 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8717 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
8718 }
8719 }
8720 acc
8721 }
8722}
8723
8724#[cfg(target_arch = "x86_64")]
8727#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8728unsafe fn dot_q4b_row_1x4_sx_vnni(
8729 buf: &[u8],
8730 scales: &[u8],
8731 g0: usize,
8732 gpr: usize,
8733 xs: [&[i8]; 4],
8734 sxs: [f32; 4],
8735) -> [f32; 4] {
8736 unsafe {
8738 use core::arch::x86_64::*;
8739 let mut acc = [0f32; 4];
8740 for gi in 0..gpr {
8741 let s = f16_to_f32(u16::from_le_bytes([
8742 scales[(g0 + gi) * 2],
8743 scales[(g0 + gi) * 2 + 1],
8744 ]));
8745 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8746 let aw = _mm256_abs_epi8(w);
8747 for (k, xq) in xs.iter().enumerate() {
8748 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8749 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8750 acc[k] += (d as f32 * sxs[k]) * s;
8751 }
8752 }
8753 acc
8754 }
8755}
8756
8757#[allow(unreachable_code)]
8758fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8759 #[cfg(target_arch = "aarch64")]
8760 unsafe {
8761 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
8762 }
8763 #[cfg(target_arch = "x86_64")]
8764 unsafe {
8765 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
8766 }
8767 let mut acc = 0f32;
8768 for gi in 0..gpr {
8769 let g = g0 + gi;
8770 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8771 let mut d = 0i32;
8772 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8773 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
8774 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
8775 }
8776 acc += d as f32 * s;
8777 }
8778 acc
8779}
8780
8781#[inline]
8783#[allow(unreachable_code)]
8784fn dot_q4_row_i8_2(
8785 packed: &[u8],
8786 scales: &[u8],
8787 g0: usize,
8788 gpr: usize,
8789 xq1: &[i8],
8790 xq2: &[i8],
8791) -> (f32, f32) {
8792 #[cfg(target_arch = "aarch64")]
8793 unsafe {
8794 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
8795 }
8796 #[cfg(target_arch = "x86_64")]
8797 unsafe {
8798 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
8799 }
8800 (
8801 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
8802 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
8803 )
8804}
8805
8806#[allow(clippy::too_many_arguments)]
8809fn q4_range_a8w8(
8810 packed: &[u8],
8811 scales: &[u8],
8812 gpr: usize,
8813 cols: usize,
8814 act: &SplitAct,
8815 out: SendMut,
8816 start: usize,
8817 end: usize,
8818) {
8819 for r in start..end {
8820 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
8821 for &(j, xv) in &act.outliers {
8823 let flat = r * cols + j;
8824 let byte = packed[flat / 2];
8825 let nib = if flat & 1 == 0 {
8826 byte & 0x0F
8827 } else {
8828 byte >> 4
8829 };
8830 let s = f16_to_f32(u16::from_le_bytes([
8831 scales[(flat / GROUP_SIZE) * 2],
8832 scales[(flat / GROUP_SIZE) * 2 + 1],
8833 ]));
8834 acc += ((nib as i32 - 8) as f32) * s * xv;
8835 }
8836 unsafe { *out.at(r) = acc };
8838 }
8839}
8840
8841#[allow(clippy::too_many_arguments)]
8844fn q4_range2_a8w8(
8845 packed: &[u8],
8846 scales: &[u8],
8847 gpr: usize,
8848 cols: usize,
8849 a1: &SplitAct,
8850 a2: &SplitAct,
8851 p1: SendMut,
8852 p2: SendMut,
8853 start: usize,
8854 end: usize,
8855) {
8856 for r in start..end {
8857 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
8858 let mut acc1 = s1 * a1.sx;
8859 let mut acc2 = s2 * a2.sx;
8860 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
8862 for &(j, xv) in outliers {
8863 let flat = r * cols + j;
8864 let byte = packed[flat / 2];
8865 let nib = if flat & 1 == 0 {
8866 byte & 0x0F
8867 } else {
8868 byte >> 4
8869 };
8870 let s = f16_to_f32(u16::from_le_bytes([
8871 scales[(flat / GROUP_SIZE) * 2],
8872 scales[(flat / GROUP_SIZE) * 2 + 1],
8873 ]));
8874 *acc += ((nib as i32 - 8) as f32) * s * xv;
8875 }
8876 };
8877 fix(&a1.outliers, &mut acc1);
8878 fix(&a2.outliers, &mut acc2);
8879 unsafe {
8881 *p1.at(r) = acc1;
8882 *p2.at(r) = acc2;
8883 }
8884 }
8885}
8886
8887fn q4_range_f32(
8889 packed: &[u8],
8890 scales: &[u8],
8891 gpr: usize,
8892 x: &[f32],
8893 out: SendMut,
8894 start: usize,
8895 end: usize,
8896) {
8897 for r in start..end {
8898 let mut acc = 0f32;
8899 for gi in 0..gpr {
8900 let g = r * gpr + gi;
8901 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8902 let pk = &packed[g * 16..(g + 1) * 16];
8903 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8904 let mut ga = 0f32;
8905 for (k, &b) in pk.iter().enumerate() {
8906 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
8907 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
8908 }
8909 acc += ga * s;
8910 }
8911 unsafe { *out.at(r) = acc };
8913 }
8914}
8915
8916#[allow(clippy::too_many_arguments)]
8920fn q4matvec2(
8921 bytes: &[u8],
8922 x1: &[f32],
8923 x2: &[f32],
8924 rows: usize,
8925 cols: usize,
8926 o1: &mut [f32],
8927 o2: &mut [f32],
8928 pool: Option<&Pool>,
8929) {
8930 debug_assert_eq!(o1.len(), rows);
8931 debug_assert_eq!(o2.len(), rows);
8932 let (packed, scales) = q4_split(bytes, rows, cols);
8933 let gpr = cols / GROUP_SIZE;
8934
8935 if a8w8_enabled() {
8936 let a1 = split_act(x1);
8937 let a2 = split_act(x2);
8938 let p1 = SendMut(o1.as_mut_ptr());
8939 let p2 = SendMut(o2.as_mut_ptr());
8940 let run = move |start: usize, end: usize| {
8941 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
8942 };
8943 dispatch_rows(pool, rows, &run);
8944 return;
8945 }
8946
8947 let p1 = SendMut(o1.as_mut_ptr());
8948 let p2 = SendMut(o2.as_mut_ptr());
8949 let run = move |start: usize, end: usize| {
8950 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
8951 };
8952 dispatch_rows(pool, rows, &run);
8953}
8954
8955#[allow(clippy::too_many_arguments)]
8957fn q4_range2_f32(
8958 packed: &[u8],
8959 scales: &[u8],
8960 gpr: usize,
8961 x1: &[f32],
8962 x2: &[f32],
8963 p1: SendMut,
8964 p2: SendMut,
8965 start: usize,
8966 end: usize,
8967) {
8968 for r in start..end {
8969 let (mut acc1, mut acc2) = (0f32, 0f32);
8970 for gi in 0..gpr {
8971 let g = r * gpr + gi;
8972 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8973 let pk = &packed[g * 16..(g + 1) * 16];
8974 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8975 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8976 let (mut g1, mut g2) = (0f32, 0f32);
8977 for (k, &b) in pk.iter().enumerate() {
8978 let wl = (b & 0x0F) as f32 - 8.0;
8979 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
8980 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
8981 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
8982 }
8983 acc1 += g1 * s;
8984 acc2 += g2 * s;
8985 }
8986 unsafe {
8988 *p1.at(r) = acc1;
8989 *p2.at(r) = acc2;
8990 }
8991 }
8992}
8993
8994thread_local! {
8995 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8998 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8999}
9000
9001#[allow(clippy::too_many_arguments)]
9007fn q4matmat(
9008 bytes: &[u8],
9009 xs_all: &[f32],
9010 b: usize,
9011 rows: usize,
9012 cols: usize,
9013 out: &mut [f32],
9014 pool: Option<&Pool>,
9015) {
9016 debug_assert_eq!(xs_all.len(), b * cols);
9017 debug_assert_eq!(out.len(), b * rows);
9018 let (packed, scales) = q4_split(bytes, rows, cols);
9019 let gpr = cols / GROUP_SIZE;
9020 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9021
9022 if a8w8_enabled() {
9023 let acts: Vec<SplitAct> = (0..b)
9024 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
9025 .collect();
9026 let acts = &acts;
9027 let out_addr = SendMut(out.as_mut_ptr());
9028 let run = move |start: usize, end: usize| {
9029 ROW_I8.with(|rb| {
9030 let mut buf = rb.borrow_mut();
9031 buf.resize(cols, 0);
9032 for r in start..end {
9033 for gi in 0..gpr {
9037 let g = r * gpr + gi;
9038 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9039 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
9040 buf[gi * GROUP_SIZE + k * 2 + 1] =
9041 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
9042 }
9043 }
9044 let mut bi = 0usize;
9045 #[cfg(target_arch = "x86_64")]
9046 if avx2_enabled() && blocked_enabled() {
9047 while bi + 4 <= acts.len() {
9048 let xs = [
9049 acts[bi].xq.as_slice(),
9050 acts[bi + 1].xq.as_slice(),
9051 acts[bi + 2].xq.as_slice(),
9052 acts[bi + 3].xq.as_slice(),
9053 ];
9054 let d = unsafe {
9055 if vnni_tiles_enabled() {
9056 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
9057 } else {
9058 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
9059 }
9060 };
9061 for k in 0..4 {
9062 let act = &acts[bi + k];
9063 let mut acc = d[k] * act.sx;
9064 for &(j, xv) in &act.outliers {
9065 acc += (buf[j] as i8) as f32
9066 * gscale((r * cols + j) / GROUP_SIZE)
9067 * xv;
9068 }
9069 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9071 }
9072 bi += 4;
9073 }
9074 }
9075 while bi < acts.len() {
9076 let act = &acts[bi];
9077 let mut acc = 0f32;
9078 for gi in 0..gpr {
9079 let d = dot_i8_i8(
9080 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
9081 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
9082 );
9083 acc += d as f32 * gscale(r * gpr + gi);
9084 }
9085 acc *= act.sx;
9086 for &(j, xv) in &act.outliers {
9088 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
9089 }
9090 unsafe { *out_addr.at(bi * rows + r) = acc };
9092 bi += 1;
9093 }
9094 }
9095 })
9096 };
9097 dispatch_rows(pool, rows, &run);
9098 return;
9099 }
9100
9101 let out_addr = SendMut(out.as_mut_ptr());
9102 let run = move |start: usize, end: usize| {
9103 ROW_F32.with(|rb| {
9104 let mut buf = rb.borrow_mut();
9105 buf.resize(cols, 0.0);
9106 for r in start..end {
9107 for gi in 0..gpr {
9110 let g = r * gpr + gi;
9111 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9112 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
9113 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
9114 }
9115 }
9116 for bi in 0..b {
9117 let x = &xs_all[bi * cols..(bi + 1) * cols];
9118 let mut acc = 0f32;
9119 for gi in 0..gpr {
9120 let mut ga = 0f32;
9121 for k in 0..GROUP_SIZE / 2 {
9126 let e = gi * GROUP_SIZE + k * 2;
9127 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
9128 }
9129 acc += ga * gscale(r * gpr + gi);
9130 }
9131 unsafe { *out_addr.at(bi * rows + r) = acc };
9133 }
9134 }
9135 })
9136 };
9137 dispatch_rows(pool, rows, &run);
9138}
9139
9140#[allow(clippy::too_many_arguments)]
9145fn vbitmatmat(
9146 bytes: &[u8],
9147 offsets: &[usize],
9148 xs_all: &[f32],
9149 b: usize,
9150 rows: usize,
9151 cols: usize,
9152 out: &mut [f32],
9153 pool: Option<&Pool>,
9154) {
9155 debug_assert_eq!(xs_all.len(), b * cols);
9156 debug_assert_eq!(out.len(), b * rows);
9157 debug_assert_eq!(offsets.len(), rows + 1);
9158 let ng = cols / GROUP_SIZE;
9159 let bits = &bytes[..rows];
9160 let sc_off = rows;
9161 let gscale = |r: usize, g: usize| {
9162 let so = (r * ng + g) * 2;
9163 f16_to_f32(u16::from_le_bytes([
9164 bytes[sc_off + so],
9165 bytes[sc_off + so + 1],
9166 ]))
9167 };
9168
9169 let decode_f32 = |r: usize, dst: &mut [f32]| {
9171 let bw = bits[r] as usize;
9172 let l = ((1i32 << (bw - 1)) - 1) as f32;
9173 let data = &bytes[offsets[r]..offsets[r + 1]];
9174 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
9175 for d in dst.iter_mut() {
9176 while nbits < bw {
9177 acc = (acc << 8) | data[idx] as u64;
9178 idx += 1;
9179 nbits += 8;
9180 }
9181 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
9182 nbits -= bw;
9183 *d = u - l;
9184 }
9185 };
9186
9187 if a8w8_enabled() {
9188 let acts: Vec<SplitAct> = (0..b)
9189 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
9190 .collect();
9191 let acts = &acts;
9192 let out_addr = SendMut(out.as_mut_ptr());
9193 let run = move |start: usize, end: usize| {
9194 for r in start..end {
9195 let bw = bits[r] as usize;
9196 if bw == 8 {
9197 ROW_F32.with(|rb| {
9200 let mut buf = rb.borrow_mut();
9201 buf.resize(cols, 0.0);
9202 decode_f32(r, &mut buf);
9203 for bi in 0..b {
9204 let x = &xs_all[bi * cols..(bi + 1) * cols];
9205 let mut dot = 0f32;
9206 for g in 0..ng {
9207 let mut gd = 0f32;
9208 for k in 0..GROUP_SIZE {
9209 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9210 }
9211 dot += gd * gscale(r, g);
9212 }
9213 unsafe { *out_addr.at(bi * rows + r) = dot };
9215 }
9216 });
9217 continue;
9218 }
9219 let l = (1i32 << (bw - 1)) - 1;
9220 let data = &bytes[offsets[r]..offsets[r + 1]];
9221 ROW_I8.with(|rb| {
9222 let mut buf = rb.borrow_mut();
9223 buf.resize(cols, 0);
9224 #[inline(always)]
9225 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
9226 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
9227 let u = unpack8::<B>(&data[blk * B..]);
9228 for k in 0..8 {
9229 chunk[k] = (u[k] - l) as i8 as u8;
9230 }
9231 }
9232 }
9233 match bw {
9234 3 => fill::<3>(data, l, &mut buf),
9235 4 => vbit_fill4(data, &mut buf),
9236 5 => fill::<5>(data, l, &mut buf),
9237 6 => fill::<6>(data, l, &mut buf),
9238 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
9239 }
9240 let mut bi = 0usize;
9241 #[cfg(target_arch = "x86_64")]
9245 if avx2_enabled() && blocked_enabled() {
9246 while bi + 4 <= acts.len() {
9247 let xs = [
9248 acts[bi].xq.as_slice(),
9249 acts[bi + 1].xq.as_slice(),
9250 acts[bi + 2].xq.as_slice(),
9251 acts[bi + 3].xq.as_slice(),
9252 ];
9253 let sxs = [
9254 acts[bi].sx,
9255 acts[bi + 1].sx,
9256 acts[bi + 2].sx,
9257 acts[bi + 3].sx,
9258 ];
9259 let d = unsafe {
9260 if vnni_tiles_enabled() {
9261 dot_q4b_row_1x4_sx_vnni(
9262 &buf,
9263 &bytes[sc_off..],
9264 r * ng,
9265 ng,
9266 xs,
9267 sxs,
9268 )
9269 } else {
9270 dot_q4b_row_1x4_sx_avx2(
9271 &buf,
9272 &bytes[sc_off..],
9273 r * ng,
9274 ng,
9275 xs,
9276 sxs,
9277 )
9278 }
9279 };
9280 for k in 0..4 {
9281 let act = &acts[bi + k];
9282 let mut dot = d[k];
9283 for &(j, xv) in &act.outliers {
9284 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9285 }
9286 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
9288 }
9289 bi += 4;
9290 }
9291 }
9292 while bi < acts.len() {
9293 let act = &acts[bi];
9294 let mut dot = 0f32;
9295 for g in 0..ng {
9296 let d = dot_i8_i8(
9297 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9298 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9299 ) as f32
9300 * act.sx;
9301 dot += d * gscale(r, g);
9302 }
9303 for &(j, xv) in &act.outliers {
9304 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9305 }
9306 unsafe { *out_addr.at(bi * rows + r) = dot };
9308 bi += 1;
9309 }
9310 });
9311 }
9312 };
9313 dispatch_rows(pool, rows, &run);
9314 return;
9315 }
9316
9317 let out_addr = SendMut(out.as_mut_ptr());
9318 let run = move |start: usize, end: usize| {
9319 ROW_F32.with(|rb| {
9320 let mut buf = rb.borrow_mut();
9321 buf.resize(cols, 0.0);
9322 for r in start..end {
9323 decode_f32(r, &mut buf);
9324 for bi in 0..b {
9325 let x = &xs_all[bi * cols..(bi + 1) * cols];
9326 let mut dot = 0f32;
9327 for g in 0..ng {
9328 let mut gd = 0f32;
9329 for k in 0..GROUP_SIZE {
9330 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9331 }
9332 dot += gd * gscale(r, g);
9333 }
9334 unsafe { *out_addr.at(bi * rows + r) = dot };
9336 }
9337 }
9338 })
9339 };
9340 dispatch_rows(pool, rows, &run);
9341}
9342
9343pub(crate) fn gpu_batch_job<'a>(
9347 t: &'a QTensor,
9348 x: &[f32],
9349) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
9350 match t {
9351 QTensor::Mapped {
9352 model,
9353 idx,
9354 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
9355 rows,
9356 cols,
9357 row_scale,
9358 col_field,
9359 ..
9360 } => Some((
9361 model.clone(),
9362 crate::gpu::BatchJob {
9363 idx: *idx,
9364 rows: *rows,
9365 cols: *cols,
9366 row_scale,
9367 xs: prescale(x, col_field, *dt).into_owned(),
9368 layout: crate::gpu::BatchLayout::Q8,
9369 },
9370 )),
9371 QTensor::Mapped {
9373 model,
9374 idx,
9375 dtype: TensorDtype::Q1,
9376 rows,
9377 cols,
9378 ..
9379 } => Some((
9380 model.clone(),
9381 crate::gpu::BatchJob {
9382 idx: *idx,
9383 rows: *rows,
9384 cols: *cols,
9385 row_scale: &[],
9386 xs: x.to_vec(),
9387 layout: crate::gpu::BatchLayout::Q1,
9388 },
9389 )),
9390 QTensor::Mapped {
9395 model,
9396 idx,
9397 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
9398 rows,
9399 cols,
9400 ..
9401 } => Some((
9402 model.clone(),
9403 crate::gpu::BatchJob {
9404 idx: *idx,
9405 rows: *rows,
9406 cols: *cols,
9407 row_scale: &[],
9408 xs: x.to_vec(),
9409 layout: if *dt == TensorDtype::Q4Tiled {
9410 crate::gpu::BatchLayout::Q4t
9411 } else {
9412 crate::gpu::BatchLayout::Q4tp
9413 },
9414 },
9415 )),
9416 _ => None,
9417 }
9418}
9419
9420thread_local! {
9421 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9422 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9423}
9424
9425pub(crate) fn prescale<'a>(
9426 x: &'a [f32],
9427 col_field: &[f32],
9428 dtype: TensorDtype,
9429) -> std::borrow::Cow<'a, [f32]> {
9430 if dtype == TensorDtype::Q8_2f {
9431 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
9432 } else {
9433 std::borrow::Cow::Borrowed(x)
9434 }
9435}
9436
9437pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
9440 x: &[f32],
9441 col_field: &[f32],
9442 dtype: TensorDtype,
9443 buf_id: u8,
9444 f: F,
9445) -> R {
9446 if dtype == TensorDtype::Q8_2f {
9447 if buf_id == 1 {
9448 PRESCALE_BUF1.with(|b| {
9449 let mut buf = b.borrow_mut();
9450 buf.clear();
9451 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9452 f(&buf)
9453 })
9454 } else {
9455 PRESCALE_BUF2.with(|b| {
9456 let mut buf = b.borrow_mut();
9457 buf.clear();
9458 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9459 f(&buf)
9460 })
9461 }
9462 } else {
9463 f(x)
9464 }
9465}
9466
9467#[cfg(target_arch = "x86_64")]
9472pub(crate) fn avx2_enabled() -> bool {
9473 use std::sync::OnceLock;
9474 static ON: OnceLock<bool> = OnceLock::new();
9475 *ON.get_or_init(|| {
9476 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
9477 && std::arch::is_x86_feature_detected!("avx2")
9478 && std::arch::is_x86_feature_detected!("fma")
9479 })
9480}
9481
9482#[cfg(target_arch = "x86_64")]
9487fn avx2_a8w8_enabled() -> bool {
9488 if FLOAT_ACTIVATIONS.get() {
9489 return false;
9490 }
9491 use std::sync::OnceLock;
9492 static ON: OnceLock<bool> = OnceLock::new();
9493 *ON.get_or_init(|| {
9494 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
9495 })
9496}
9497
9498thread_local! {
9499 static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
9500}
9501
9502pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
9505 struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
9506 impl Drop for Restore {
9507 fn drop(&mut self) {
9508 FULL_GPU_Q8.set(self.0);
9509 }
9510 }
9511 Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
9512}
9513
9514thread_local! {
9518 static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
9519}
9520
9521pub(crate) fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
9522 struct Restore(bool);
9523 impl Drop for Restore {
9524 fn drop(&mut self) {
9525 FLOAT_ACTIVATIONS.set(self.0);
9526 }
9527 }
9528 let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
9529 f()
9530}
9531
9532static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
9542
9543pub(crate) fn row_exact() -> bool {
9544 ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
9545}
9546
9547fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
9548 struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
9549 impl Drop for Restore<'_> {
9550 fn drop(&mut self) {
9551 self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
9552 }
9553 }
9554 active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
9555 let _restore = Restore(active);
9556 f()
9557}
9558
9559pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
9561 counted_row_exact_scope(&ROW_EXACT, f)
9562}
9563
9564#[inline]
9568pub(crate) fn a8w8_enabled() -> bool {
9569 #[cfg(target_arch = "aarch64")]
9570 {
9571 sdot_enabled()
9572 }
9573 #[cfg(target_arch = "x86_64")]
9574 {
9575 avx2_a8w8_enabled()
9576 }
9577 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
9578 {
9579 false
9580 }
9581}
9582
9583#[inline]
9586#[allow(unreachable_code)]
9587fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
9588 #[cfg(target_arch = "aarch64")]
9589 unsafe {
9590 return dot_i8_sdot(w, xq);
9591 }
9592 #[cfg(target_arch = "x86_64")]
9593 unsafe {
9594 if avx512vnni_enabled() {
9595 return dot_i8_i8_vnni(w, xq);
9596 }
9597 return dot_i8_i8_avx2(w, xq);
9598 }
9599 w.iter()
9600 .zip(xq)
9601 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
9602 .sum()
9603}
9604
9605#[cfg(target_arch = "x86_64")]
9609fn avx512vnni_enabled() -> bool {
9610 use std::sync::OnceLock;
9611 static ON: OnceLock<bool> = OnceLock::new();
9612 *ON.get_or_init(|| {
9613 std::env::var("CMF_AVX512")
9614 .map(|v| v != "0")
9615 .unwrap_or(true)
9616 && std::arch::is_x86_feature_detected!("avx512f")
9617 && std::arch::is_x86_feature_detected!("avx512bw")
9618 && std::arch::is_x86_feature_detected!("avx512vl")
9619 && std::arch::is_x86_feature_detected!("avx512vnni")
9620 })
9621}
9622
9623#[cfg(target_arch = "x86_64")]
9631fn vnni_tiles_enabled() -> bool {
9632 use std::sync::OnceLock;
9633 static ON: OnceLock<bool> = OnceLock::new();
9634 *ON.get_or_init(|| {
9635 std::env::var("CMF_VNNI_TILES")
9636 .map(|v| v != "0")
9637 .unwrap_or(true)
9638 && avx512vnni_enabled()
9639 })
9640}
9641
9642#[cfg(target_arch = "x86_64")]
9647#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9648#[inline]
9649unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
9650 unsafe {
9652 use core::arch::x86_64::*;
9653 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
9654 let hi128 = _mm256_extracti128_si256::<1>(d);
9655 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9656 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9657 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9658 _mm_cvtsi128_si32(s32)
9659 }
9660}
9661
9662#[cfg(target_arch = "x86_64")]
9667#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9668unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9669 unsafe {
9671 use core::arch::x86_64::*;
9672 let n = w.len();
9673 let mut j = 0usize;
9674 let mut total: i32;
9675 {
9680 #[inline(always)]
9681 unsafe fn step(
9682 w: *const u8,
9683 x: *const i8,
9684 acc: core::arch::x86_64::__m512i,
9685 ) -> core::arch::x86_64::__m512i {
9686 unsafe {
9687 use core::arch::x86_64::*;
9688 let wv = _mm512_loadu_si512(w as *const _);
9689 let xv = _mm512_loadu_si512(x as *const _);
9690 let aw = _mm512_abs_epi8(wv);
9691 let neg = _mm512_movepi8_mask(wv);
9692 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
9693 _mm512_dpbusd_epi32(acc, aw, sx)
9694 }
9695 }
9696 let (mut a0, mut a1, mut a2, mut a3) = (
9697 _mm512_setzero_si512(),
9698 _mm512_setzero_si512(),
9699 _mm512_setzero_si512(),
9700 _mm512_setzero_si512(),
9701 );
9702 while j + 256 <= n {
9703 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9704 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
9705 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
9706 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
9707 j += 256;
9708 }
9709 while j + 64 <= n {
9710 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9711 j += 64;
9712 }
9713 let s01 = _mm512_add_epi32(a0, a1);
9714 let s23 = _mm512_add_epi32(a2, a3);
9715 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
9716 }
9717 if j + 32 <= n {
9719 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9720 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9721 let d = _mm256_dpbusd_epi32(
9722 _mm256_setzero_si256(),
9723 _mm256_abs_epi8(wv),
9724 _mm256_sign_epi8(xv, wv),
9725 );
9726 let hi128 = _mm256_extracti128_si256::<1>(d);
9727 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9728 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9729 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9730 total += _mm_cvtsi128_si32(s32);
9731 j += 32;
9732 }
9733 while j < n {
9734 total += (w[j] as i8) as i32 * xq[j] as i32;
9735 j += 1;
9736 }
9737 total
9738 }
9739}
9740
9741#[cfg(target_arch = "x86_64")]
9743#[target_feature(enable = "avx2,fma")]
9744unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
9745 unsafe {
9747 use core::arch::x86_64::*;
9748 let n = x.len();
9749 let wp = w.as_ptr();
9750 let xp = x.as_ptr();
9751 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
9752 let mut j = 0usize;
9753 while j + 16 <= n {
9754 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
9755 let lo = _mm256_cvtepi8_epi32(wb);
9756 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
9757 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
9758 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
9759 j += 16;
9760 }
9761 let acc = _mm256_add_ps(a0, a1);
9762 let hi128 = _mm256_extractf128_ps::<1>(acc);
9763 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
9764 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
9765 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
9766 let mut sum = _mm_cvtss_f32(s32);
9767 while j < n {
9768 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
9769 j += 1;
9770 }
9771 sum
9772 }
9773}
9774
9775#[cfg(target_arch = "x86_64")]
9780#[target_feature(enable = "avx2")]
9781unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
9782 unsafe {
9784 use core::arch::x86_64::*;
9785 let n = w.len();
9786 let ones = _mm256_set1_epi16(1);
9787 let mut acc = _mm256_setzero_si256();
9788 let mut j = 0usize;
9789 while j + 32 <= n {
9790 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9791 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9792 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9793 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
9794 j += 32;
9795 }
9796 let hi128 = _mm256_extracti128_si256::<1>(acc);
9797 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
9798 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9799 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9800 let mut s = _mm_cvtsi128_si32(s32);
9801 while j < n {
9802 s += (w[j] as i8) as i32 * xq[j] as i32;
9803 j += 1;
9804 }
9805 s
9806 }
9807}
9808
9809#[cfg(target_arch = "aarch64")]
9813#[target_feature(enable = "neon,i8mm")]
9814unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9815 unsafe {
9817 use core::arch::aarch64::*;
9818 use core::arch::asm;
9819 let n = w0.len();
9820 let w0p = w0.as_ptr() as *const i8;
9821 let w1p = w1.as_ptr() as *const i8;
9822 let mut acc01 = vdupq_n_s32(0);
9825 let mut acc23 = vdupq_n_s32(0);
9826 let mut i = 0usize;
9827 while i + 8 <= n {
9828 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
9829 let xb01 = vcombine_s8(
9830 vld1_s8(xs[0].as_ptr().add(i)),
9831 vld1_s8(xs[1].as_ptr().add(i)),
9832 );
9833 let xb23 = vcombine_s8(
9834 vld1_s8(xs[2].as_ptr().add(i)),
9835 vld1_s8(xs[3].as_ptr().add(i)),
9836 );
9837 asm!(
9838 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
9839 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
9840 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
9841 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
9842 options(pure, nomem, nostack),
9843 );
9844 i += 8;
9845 }
9846 let mut out = [[0i32; 4]; 2];
9847 let a01: [i32; 4] = core::mem::transmute(acc01);
9848 let a23: [i32; 4] = core::mem::transmute(acc23);
9849 out[0][0] = a01[0];
9850 out[0][1] = a01[1];
9851 out[1][0] = a01[2];
9852 out[1][1] = a01[3];
9853 out[0][2] = a23[0];
9854 out[0][3] = a23[1];
9855 out[1][2] = a23[2];
9856 out[1][3] = a23[3];
9857 if i < n {
9858 for (k, x) in xs.iter().enumerate() {
9859 for j in i..n {
9860 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9861 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9862 }
9863 }
9864 }
9865 out
9866 }
9867}
9868
9869#[cfg(target_arch = "aarch64")]
9873#[target_feature(enable = "neon,dotprod")]
9874unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9875 unsafe {
9877 use core::arch::aarch64::*;
9878 use core::arch::asm;
9879 let n = w0.len();
9880 let w0p = w0.as_ptr() as *const i8;
9881 let w1p = w1.as_ptr() as *const i8;
9882 let mut acc = [[vdupq_n_s32(0); 4]; 2];
9883 let mut i = 0usize;
9884 while i + 16 <= n {
9885 let wv0 = vld1q_s8(w0p.add(i));
9886 let wv1 = vld1q_s8(w1p.add(i));
9887 for (k, x) in xs.iter().enumerate() {
9888 let xv = vld1q_s8(x.as_ptr().add(i));
9889 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
9890 asm!(
9891 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
9892 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
9893 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9894 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
9895 options(pure, nomem, nostack),
9896 );
9897 acc[0][k] = a0;
9898 acc[1][k] = a1;
9899 }
9900 i += 16;
9901 }
9902 let mut out = [[0i32; 4]; 2];
9903 for r in 0..2 {
9904 for k in 0..4 {
9905 out[r][k] = vaddvq_s32(acc[r][k]);
9906 }
9907 }
9908 if i < n {
9909 for (k, x) in xs.iter().enumerate() {
9910 for j in i..n {
9911 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9912 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9913 }
9914 }
9915 }
9916 out
9917 }
9918}
9919
9920#[cfg(target_arch = "x86_64")]
9926#[target_feature(enable = "avx2")]
9927unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9928 unsafe {
9930 use core::arch::x86_64::*;
9931 let n = w0.len();
9932 let ones = _mm256_set1_epi16(1);
9933 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
9934 let mut j = 0usize;
9935 while j + 32 <= n {
9936 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
9937 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
9938 let aw0 = _mm256_abs_epi8(wv0);
9939 let aw1 = _mm256_abs_epi8(wv1);
9940 for (k, x) in xs.iter().enumerate() {
9941 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
9942 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
9943 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
9944 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
9945 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
9946 }
9947 j += 32;
9948 }
9949 let mut out = [[0i32; 4]; 2];
9950 for r in 0..2 {
9951 for k in 0..4 {
9952 let a = acc[r][k];
9953 let hi128 = _mm256_extracti128_si256::<1>(a);
9954 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
9955 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9956 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9957 out[r][k] = _mm_cvtsi128_si32(s32);
9958 }
9959 }
9960 if j < n {
9961 for (k, x) in xs.iter().enumerate() {
9962 for i in j..n {
9963 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
9964 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
9965 }
9966 }
9967 }
9968 out
9969 }
9970}
9971
9972#[cfg(target_arch = "x86_64")]
9977#[inline]
9978fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
9979 let dot = if avx512vnni_enabled() && row.len() >= 64 {
9980 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
9981 } else {
9982 unsafe { dot_i8_i8_avx2(row, &act.xq) }
9983 };
9984 let mut acc = dot as f32 * act.sx;
9985 for &(j, xv) in &act.outliers {
9986 acc += (row[j] as i8) as f32 * xv;
9987 }
9988 acc
9989}
9990
9991#[cfg(target_arch = "x86_64")]
9995#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9996unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9997 unsafe {
9999 use core::arch::x86_64::*;
10000 let n = w.len();
10001 let flip = _mm512_set1_epi8(-128); #[inline(always)]
10003 unsafe fn step(
10004 w: *const u8,
10005 x: *const i8,
10006 flip: core::arch::x86_64::__m512i,
10007 acc: core::arch::x86_64::__m512i,
10008 ) -> core::arch::x86_64::__m512i {
10009 unsafe {
10010 use core::arch::x86_64::*;
10011 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
10012 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
10013 }
10014 }
10015 let (mut a0, mut a1, mut a2, mut a3) = (
10016 _mm512_setzero_si512(),
10017 _mm512_setzero_si512(),
10018 _mm512_setzero_si512(),
10019 _mm512_setzero_si512(),
10020 );
10021 let mut j = 0usize;
10022 while j + 256 <= n {
10023 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
10024 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
10025 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
10026 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
10027 j += 256;
10028 }
10029 while j + 64 <= n {
10030 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
10031 j += 64;
10032 }
10033 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
10034 _mm512_add_epi32(a0, a1),
10035 _mm512_add_epi32(a2, a3),
10036 ));
10037 while j < n {
10039 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
10040 j += 1;
10041 }
10042 total
10043 }
10044}
10045
10046#[cfg(target_arch = "x86_64")]
10052#[target_feature(enable = "avx2")]
10053unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10054 unsafe {
10057 use core::arch::x86_64::*;
10058 let lomask = _mm_set1_epi8(0x0F);
10059 let eight = _mm256_set1_epi8(8);
10060 let ones = _mm256_set1_epi16(1);
10061 let mut acc = 0f32;
10062 for gi in 0..gpr {
10063 let g = g0 + gi;
10064 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10065 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
10066 let lo = _mm_and_si128(b, lomask);
10067 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
10068 let w = _mm256_sub_epi8(
10069 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
10070 eight,
10071 );
10072 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10073 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
10074 let d = _mm256_madd_epi16(p16, ones);
10075 let hi128 = _mm256_extracti128_si256::<1>(d);
10076 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10077 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10078 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10079 acc += _mm_cvtsi128_si32(s32) as f32 * s;
10080 }
10081 acc
10082 }
10083}
10084
10085#[cfg(target_arch = "x86_64")]
10088#[target_feature(enable = "avx2")]
10089unsafe fn dot_q4_row_avx2_2(
10090 packed: &[u8],
10091 scales: &[u8],
10092 g0: usize,
10093 gpr: usize,
10094 xq1: &[i8],
10095 xq2: &[i8],
10096) -> (f32, f32) {
10097 unsafe {
10099 use core::arch::x86_64::*;
10100 let lomask = _mm_set1_epi8(0x0F);
10101 let eight = _mm256_set1_epi8(8);
10102 let ones = _mm256_set1_epi16(1);
10103 let (mut acc1, mut acc2) = (0f32, 0f32);
10104 #[inline(always)]
10105 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
10106 unsafe {
10107 use core::arch::x86_64::*;
10108 let hi128 = _mm256_extracti128_si256::<1>(d);
10109 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10110 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10111 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10112 _mm_cvtsi128_si32(s32)
10113 }
10114 }
10115 for gi in 0..gpr {
10116 let g = g0 + gi;
10117 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10118 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
10119 let lo = _mm_and_si128(b, lomask);
10120 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
10121 let w = _mm256_sub_epi8(
10122 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
10123 eight,
10124 );
10125 let aw = _mm256_abs_epi8(w);
10126 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10127 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10128 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
10129 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
10130 acc1 += hsum(d1) as f32 * s;
10131 acc2 += hsum(d2) as f32 * s;
10132 }
10133 (acc1, acc2)
10134 }
10135}
10136
10137#[cfg(target_arch = "x86_64")]
10139fn q8_range_avx2(
10140 q: &[u8],
10141 row_scale: &[f32],
10142 act: &SplitAct,
10143 cols: usize,
10144 out_addr: SendMut,
10145 start: usize,
10146 end: usize,
10147) {
10148 for o in start..end {
10149 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10150 unsafe { *out_addr.at(o) = v };
10152 }
10153}
10154
10155#[cfg(target_arch = "x86_64")]
10157#[allow(clippy::too_many_arguments)]
10158fn q8_range2_avx2(
10159 q: &[u8],
10160 row_scale: &[f32],
10161 a1: &SplitAct,
10162 a2: &SplitAct,
10163 cols: usize,
10164 p1: SendMut,
10165 p2: SendMut,
10166 start: usize,
10167 end: usize,
10168) {
10169 for o in start..end {
10170 let row = &q[o * cols..(o + 1) * cols];
10171 unsafe {
10173 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
10174 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
10175 }
10176 }
10177}
10178
10179#[cfg(target_arch = "aarch64")]
10190fn i8mm_enabled() -> bool {
10191 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10192 *ON.get_or_init(|| {
10193 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
10194 && std::arch::is_aarch64_feature_detected!("i8mm")
10195 })
10196}
10197
10198#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
10202fn sdot_enabled() -> bool {
10203 if FLOAT_ACTIVATIONS.get() {
10204 return false;
10205 }
10206 use std::sync::OnceLock;
10207 static ON: OnceLock<bool> = OnceLock::new();
10208 *ON.get_or_init(|| {
10209 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
10210 if !want {
10211 return false;
10212 }
10213
10214 #[cfg(target_arch = "aarch64")]
10215 {
10216 if std::arch::is_aarch64_feature_detected!("dotprod") {
10217 return true;
10218 }
10219 #[cfg(target_os = "android")]
10220 {
10221 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
10222 if cpuinfo.lines().any(|l| {
10223 (l.starts_with("Features") || l.starts_with("features"))
10224 && l.contains("asimddp")
10225 }) {
10226 return true;
10227 }
10228 }
10229 }
10230 false
10231 }
10232 #[cfg(not(target_arch = "aarch64"))]
10233 {
10234 false
10235 }
10236 })
10237}
10238
10239struct SplitAct {
10244 xq: Vec<i8>,
10245 sx: f32,
10246 outliers: Vec<(usize, f32)>,
10247 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
10250 xsum: i32,
10251}
10252
10253thread_local! {
10254 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
10257 const { std::cell::RefCell::new(Vec::new()) };
10258}
10259
10260impl Drop for SplitAct {
10261 fn drop(&mut self) {
10262 let buf = std::mem::take(&mut self.xq);
10263 if buf.capacity() > 0 {
10264 XQ_FREE.with(|f| {
10265 let mut f = f.borrow_mut();
10266 if f.len() < 16 {
10267 f.push(buf);
10268 }
10269 });
10270 }
10271 }
10272}
10273
10274thread_local! {
10275 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10282}
10283
10284#[inline]
10287fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
10288 KROW.with(|s| {
10289 let mut b = s.borrow_mut();
10290 if b.len() < n {
10291 b.resize(n, 0.0);
10292 }
10293 f(&mut b[..n])
10294 })
10295}
10296
10297#[inline(always)]
10308fn q8_round(t: f32) -> i8 {
10309 let t = t.clamp(-127.0, 127.0);
10310 let i = t as i32;
10311 let f = t - i as f32;
10312 let r = if f >= 0.5 {
10313 i + 1
10314 } else if f <= -0.5 {
10315 i - 1
10316 } else {
10317 i
10318 };
10319 r as i8
10320}
10321
10322fn split_act(x: &[f32]) -> SplitAct {
10323 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10324 let n = x.len();
10325 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
10326 let thr = 8.0 * rms;
10327 let mut outliers: Vec<(usize, f32)> = Vec::new();
10331 let mut amax = 0f32;
10332 for (j, &v) in x.iter().enumerate() {
10333 let a = v.abs();
10334 if a > thr {
10335 outliers.push((j, v));
10336 } else if a > amax {
10337 amax = a;
10338 }
10339 }
10340 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10341 let inv = 1.0 / sx;
10342 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10343 xq.clear();
10344 xq.reserve(n);
10345 if outliers.is_empty() {
10346 xq.extend(
10347 x.iter()
10348 .map(|&v| q8_round(v * inv)),
10349 );
10350 } else {
10351 xq.extend(x.iter().map(|&v| {
10353 if v.abs() > thr {
10354 0
10355 } else {
10356 q8_round(v * inv)
10357 }
10358 }));
10359 }
10360 let xsum = xq.iter().map(|&v| v as i32).sum();
10361 SplitAct {
10362 xq,
10363 sx,
10364 outliers,
10365 xsum,
10366 }
10367}
10368
10369fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
10370 let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10371 let n = x.len();
10372 let rms = (x
10373 .iter()
10374 .zip(col)
10375 .map(|(&a, &c)| {
10376 let v = a * c;
10377 (v * v) as f64
10378 })
10379 .sum::<f64>()
10380 / n.max(1) as f64)
10381 .sqrt() as f32;
10382 let thr = 8.0 * rms;
10383
10384 let mut outliers = Vec::new();
10385 let mut amax = 0f32;
10386 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
10387 let v = a * c;
10388 let s = v.abs();
10389 if s > thr {
10390 outliers.push((j, v));
10391 } else if s > amax {
10392 amax = s;
10393 }
10394 }
10395
10396 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10397 let inv = 1.0 / sx;
10398 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10399 xq.clear();
10400 xq.reserve(n);
10401 if outliers.is_empty() {
10402 xq.extend(
10403 x.iter()
10404 .zip(col)
10405 .map(|(&a, &c)| q8_round((a * c) * inv)),
10406 );
10407 } else {
10408 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
10409 let v = a * c;
10410 if v.abs() > thr {
10411 0
10412 } else {
10413 q8_round(v * inv)
10414 }
10415 }));
10416 }
10417 let xsum = xq.iter().map(|&v| v as i32).sum();
10418 SplitAct {
10419 xq,
10420 sx,
10421 outliers,
10422 xsum,
10423 }
10424}
10425
10426#[cfg(target_arch = "aarch64")]
10429#[target_feature(enable = "neon,dotprod")]
10430unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
10431 unsafe {
10433 use core::arch::aarch64::*;
10434 use core::arch::asm;
10435 let wp = w.as_ptr() as *const i8;
10436 let n = w.len();
10437 let (mut a0, mut a1, mut a2, mut a3) = (
10438 vdupq_n_s32(0),
10439 vdupq_n_s32(0),
10440 vdupq_n_s32(0),
10441 vdupq_n_s32(0),
10442 );
10443 let mut i = 0;
10444 while i + 64 <= n {
10445 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10446 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
10447 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
10448 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
10449 asm!(
10450 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10451 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10452 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
10453 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
10454 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10455 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10456 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
10457 options(pure, nomem, nostack),
10458 );
10459 i += 64;
10460 }
10461 while i + 16 <= n {
10462 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10463 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
10464 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
10465 i += 16;
10466 }
10467 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
10468 while i < n {
10469 s += (*wp.add(i)) as i32 * xq[i] as i32;
10470 i += 1;
10471 }
10472 s
10473 }
10474}
10475
10476#[cfg(target_arch = "aarch64")]
10480#[target_feature(enable = "neon,dotprod")]
10481unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
10482 unsafe {
10484 use core::arch::aarch64::*;
10485 use core::arch::asm;
10486 let n = xq.len();
10487 let px = xq.as_ptr();
10488 let (p0, p1, p2, p3) = (
10489 w0.as_ptr() as *const i8,
10490 w1.as_ptr() as *const i8,
10491 w2.as_ptr() as *const i8,
10492 w3.as_ptr() as *const i8,
10493 );
10494 let (mut a0, mut a1, mut a2, mut a3) = (
10495 vdupq_n_s32(0),
10496 vdupq_n_s32(0),
10497 vdupq_n_s32(0),
10498 vdupq_n_s32(0),
10499 );
10500 let mut i = 0;
10501 while i + 16 <= n {
10502 let x = vld1q_s8(px.add(i));
10503 let v0 = vld1q_s8(p0.add(i));
10504 let v1 = vld1q_s8(p1.add(i));
10505 let v2 = vld1q_s8(p2.add(i));
10506 let v3 = vld1q_s8(p3.add(i));
10507 asm!(
10508 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10509 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10510 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10511 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10512 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10513 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10514 options(pure, nomem, nostack),
10515 );
10516 i += 16;
10517 }
10518 let mut r = [
10519 vaddvq_s32(a0),
10520 vaddvq_s32(a1),
10521 vaddvq_s32(a2),
10522 vaddvq_s32(a3),
10523 ];
10524 while i < n {
10525 let xi = *px.add(i) as i32;
10526 r[0] += (*p0.add(i)) as i32 * xi;
10527 r[1] += (*p1.add(i)) as i32 * xi;
10528 r[2] += (*p2.add(i)) as i32 * xi;
10529 r[3] += (*p3.add(i)) as i32 * xi;
10530 i += 1;
10531 }
10532 r
10533 }
10534}
10535
10536#[cfg(target_arch = "aarch64")]
10543#[target_feature(enable = "neon,dotprod")]
10544unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
10545 unsafe {
10548 use core::arch::aarch64::*;
10549 use core::arch::asm;
10550 let n = xq.len();
10551 let px = xq.as_ptr();
10552 let pg = g.as_ptr() as *const i8;
10553 let (mut a0, mut a1, mut a2, mut a3) = (
10554 vdupq_n_s32(0),
10555 vdupq_n_s32(0),
10556 vdupq_n_s32(0),
10557 vdupq_n_s32(0),
10558 );
10559 let mut i = 0;
10560 while i + 16 <= n {
10561 let x = vld1q_s8(px.add(i));
10562 let base = pg.add(4 * i);
10563 let v0 = vld1q_s8(base);
10564 let v1 = vld1q_s8(base.add(16));
10565 let v2 = vld1q_s8(base.add(32));
10566 let v3 = vld1q_s8(base.add(48));
10567 asm!(
10568 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10569 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10570 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10571 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10572 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10573 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10574 options(pure, nomem, nostack),
10575 );
10576 i += 16;
10577 }
10578 [
10579 vaddvq_s32(a0),
10580 vaddvq_s32(a1),
10581 vaddvq_s32(a2),
10582 vaddvq_s32(a3),
10583 ]
10584 }
10585}
10586
10587#[cfg(target_arch = "aarch64")]
10593fn q8_range_sdot(
10594 q: &[u8],
10595 rep: &[u8],
10596 row_scale: &[f32],
10597 act: &SplitAct,
10598 cols: usize,
10599 out_addr: SendMut,
10600 start: usize,
10601 end: usize,
10602) {
10603 let mut o = start;
10604 if !rep.is_empty() {
10607 while o < end && o % 4 != 0 {
10608 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10609 unsafe { *out_addr.at(o) = v };
10610 o += 1;
10611 }
10612 }
10613 while o + 4 <= end {
10614 let r = if rep.is_empty() {
10615 unsafe {
10616 dot_i8_sdot_4rows(
10617 &q[o * cols..(o + 1) * cols],
10618 &q[(o + 1) * cols..(o + 2) * cols],
10619 &q[(o + 2) * cols..(o + 3) * cols],
10620 &q[(o + 3) * cols..(o + 4) * cols],
10621 &act.xq,
10622 )
10623 }
10624 } else {
10625 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
10626 };
10627 for k in 0..4 {
10628 let mut acc = r[k] as f32 * act.sx;
10629 for &(j, xv) in &act.outliers {
10630 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
10631 }
10632 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
10634 }
10635 o += 4;
10636 }
10637 while o < end {
10638 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10639 unsafe { *out_addr.at(o) = v };
10640 o += 1;
10641 }
10642}
10643
10644#[cfg(target_arch = "aarch64")]
10647#[allow(clippy::too_many_arguments)]
10648fn q8_range2_sdot(
10649 q: &[u8],
10650 row_scale: &[f32],
10651 a1: &SplitAct,
10652 a2: &SplitAct,
10653 cols: usize,
10654 p1: SendMut,
10655 p2: SendMut,
10656 start: usize,
10657 end: usize,
10658) {
10659 for o in start..end {
10660 let row = &q[o * cols..(o + 1) * cols];
10661 unsafe {
10663 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
10664 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
10665 }
10666 }
10667}
10668
10669#[allow(clippy::too_many_arguments)]
10671fn q8_range2_f32(
10672 q: &[u8],
10673 row_scale: &[f32],
10674 x1: &[f32],
10675 x2: &[f32],
10676 cols: usize,
10677 p1: SendMut,
10678 p2: SendMut,
10679 start: usize,
10680 end: usize,
10681) {
10682 for o in start..end {
10683 let row = &q[o * cols..(o + 1) * cols];
10684 unsafe {
10686 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
10687 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
10688 }
10689 }
10690}
10691
10692fn q8_range_f32(
10694 q: &[u8],
10695 row_scale: &[f32],
10696 xs: &[f32],
10697 cols: usize,
10698 out_addr: SendMut,
10699 start: usize,
10700 end: usize,
10701) {
10702 for o in start..end {
10703 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10704 unsafe { *out_addr.at(o) = v };
10706 }
10707}
10708
10709#[inline]
10713fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
10714 #[cfg(target_arch = "aarch64")]
10715 return row_dot_sdot(row, act);
10716 #[cfg(target_arch = "x86_64")]
10717 return row_dot_avx2(row, act);
10718 #[allow(unreachable_code)]
10719 q8_row_dot_scalar(row, act)
10720}
10721
10722#[allow(dead_code)]
10723fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
10724 let mut acc = 0i32;
10725 for (k, &b) in row.iter().enumerate() {
10726 acc += (b as i8) as i32 * act.xq[k] as i32;
10727 }
10728 let mut acc = acc as f32 * act.sx;
10729 for &(j, xv) in &act.outliers {
10730 acc += (row[j] as i8) as f32 * xv;
10731 }
10732 acc
10733}
10734
10735#[cfg(target_arch = "aarch64")]
10738#[inline]
10739fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
10740 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
10741 for &(j, xv) in &act.outliers {
10742 acc += (row[j] as i8) as f32 * xv;
10743 }
10744 acc
10745}
10746
10747#[cfg(target_arch = "aarch64")]
10755#[target_feature(enable = "neon,dotprod")]
10756unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10757 unsafe {
10760 use core::arch::aarch64::*;
10761 use core::arch::asm;
10762 let lomask = vdupq_n_u8(0x0F);
10763 let eight = vdupq_n_s8(8);
10764 let mut acc = 0f32;
10765 for gi in 0..gpr {
10766 let g = g0 + gi;
10767 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10768 let b = vld1q_u8(packed.as_ptr().add(g * 16));
10769 let lo = vandq_u8(b, lomask);
10770 let hi = vshrq_n_u8::<4>(b);
10771 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10772 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10773 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
10774 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
10775 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10776 asm!(
10777 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
10778 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
10779 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10780 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
10781 options(pure, nomem, nostack),
10782 );
10783 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10784 }
10785 acc
10786 }
10787}
10788
10789#[cfg(target_arch = "aarch64")]
10794#[target_feature(enable = "neon,dotprod")]
10795unsafe fn dot_q4_row_sdot2(
10796 packed: &[u8],
10797 scales: &[u8],
10798 g0: usize,
10799 gpr: usize,
10800 xq1: &[i8],
10801 xq2: &[i8],
10802) -> (f32, f32) {
10803 unsafe {
10806 use core::arch::aarch64::*;
10807 use core::arch::asm;
10808 let lomask = vdupq_n_u8(0x0F);
10809 let eight = vdupq_n_s8(8);
10810 let (mut acc1, mut acc2) = (0f32, 0f32);
10811 for gi in 0..gpr {
10812 let g = g0 + gi;
10813 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10814 let b = vld1q_u8(packed.as_ptr().add(g * 16));
10815 let lo = vandq_u8(b, lomask);
10816 let hi = vshrq_n_u8::<4>(b);
10817 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10818 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10819 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
10820 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
10821 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
10822 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
10823 let (mut a0, mut a1, mut b0, mut b1) = (
10824 vdupq_n_s32(0),
10825 vdupq_n_s32(0),
10826 vdupq_n_s32(0),
10827 vdupq_n_s32(0),
10828 );
10829 asm!(
10830 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
10831 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
10832 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
10833 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
10834 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10835 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
10836 e0 = in(vreg) e0, e1 = in(vreg) e1,
10837 x10 = in(vreg) x10, x11 = in(vreg) x11,
10838 x20 = in(vreg) x20, x21 = in(vreg) x21,
10839 options(pure, nomem, nostack),
10840 );
10841 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10842 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
10843 }
10844 (acc1, acc2)
10845 }
10846}
10847
10848#[inline]
10853pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
10854 #[cfg(target_arch = "aarch64")]
10855 unsafe {
10856 return axpy_i8_f32_neon(acc, row, w);
10857 }
10858 #[cfg(target_arch = "x86_64")]
10859 if avx2_enabled() {
10860 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
10861 }
10862 #[allow(unreachable_code)]
10863 {
10864 for (a, &b) in acc.iter_mut().zip(row) {
10865 *a += w * b as f32;
10866 }
10867 }
10868}
10869
10870#[cfg(target_arch = "x86_64")]
10872#[target_feature(enable = "avx2,fma")]
10873unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
10874 unsafe {
10876 use core::arch::x86_64::*;
10877 let n = acc.len().min(row.len());
10878 let ap = acc.as_mut_ptr();
10879 let rp = row.as_ptr();
10880 let wv = _mm256_set1_ps(w);
10881 let mut j = 0usize;
10882 while j + 16 <= n {
10883 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
10884 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
10885 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
10886 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
10887 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
10888 _mm256_storeu_ps(ap.add(j), v0);
10889 _mm256_storeu_ps(ap.add(j + 8), v1);
10890 j += 16;
10891 }
10892 while j < n {
10893 *ap.add(j) += w * (*rp.add(j)) as f32;
10894 j += 1;
10895 }
10896 }
10897}
10898
10899#[cfg(target_arch = "aarch64")]
10900#[target_feature(enable = "neon")]
10901unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
10902 unsafe {
10904 use core::arch::aarch64::*;
10905 let n = acc.len().min(row.len());
10906 let ap = acc.as_mut_ptr();
10907 let rp = row.as_ptr();
10908 let wv = vdupq_n_f32(w);
10909 let mut j = 0usize;
10910 while j + 16 <= n {
10911 let rb = vld1q_s8(rp.add(j));
10912 let lo = vmovl_s8(vget_low_s8(rb));
10913 let hi = vmovl_s8(vget_high_s8(rb));
10914 for (off, half) in [(0, lo), (8, hi)] {
10915 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
10916 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
10917 let o = j + off;
10918 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
10919 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
10920 }
10921 j += 16;
10922 }
10923 while j < n {
10924 *ap.add(j) += w * (*rp.add(j)) as f32;
10925 j += 1;
10926 }
10927 }
10928}
10929
10930#[inline]
10933pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
10934 #[cfg(target_arch = "aarch64")]
10935 unsafe {
10936 return dot_i8_f32_neon(w, x);
10937 }
10938 #[cfg(target_arch = "x86_64")]
10939 if avx2_enabled() {
10940 return unsafe { dot_i8_f32_avx2(w, x) };
10941 }
10942 #[allow(unreachable_code)]
10943 {
10944 let mut sum = 0.0f32;
10945 for (j, &b) in w.iter().enumerate() {
10946 sum += (b as i8) as f32 * x[j];
10947 }
10948 sum
10949 }
10950}
10951
10952#[inline]
10956fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10957 #[cfg(target_arch = "aarch64")]
10958 unsafe {
10959 return dot_i8_col_f32_neon(w, x, col);
10960 }
10961 #[allow(unreachable_code)]
10962 {
10963 let mut sum = 0.0f32;
10964 for (j, &b) in w.iter().enumerate() {
10965 sum += (b as i8) as f32 * x[j] * col[j];
10966 }
10967 sum
10968 }
10969}
10970
10971#[cfg(target_arch = "aarch64")]
10972#[target_feature(enable = "neon")]
10973unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10974 unsafe {
10976 use core::arch::aarch64::*;
10977 let n = x.len();
10978 let wp = w.as_ptr() as *const i8;
10979 let xp = x.as_ptr();
10980 let cp = col.as_ptr();
10981 let (mut a0, mut a1, mut a2, mut a3) = (
10982 vdupq_n_f32(0.0),
10983 vdupq_n_f32(0.0),
10984 vdupq_n_f32(0.0),
10985 vdupq_n_f32(0.0),
10986 );
10987 let mut j = 0usize;
10988 while j + 16 <= n {
10989 let wb = vld1q_s8(wp.add(j));
10990 let lo = vmovl_s8(vget_low_s8(wb));
10991 let hi = vmovl_s8(vget_high_s8(wb));
10992 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10993 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10994 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10995 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10996 a0 = vfmaq_f32(
10997 a0,
10998 w0,
10999 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
11000 );
11001 a1 = vfmaq_f32(
11002 a1,
11003 w1,
11004 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
11005 );
11006 a2 = vfmaq_f32(
11007 a2,
11008 w2,
11009 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
11010 );
11011 a3 = vfmaq_f32(
11012 a3,
11013 w3,
11014 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
11015 );
11016 j += 16;
11017 }
11018 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
11019 while j < n {
11020 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
11021 j += 1;
11022 }
11023 sum
11024 }
11025}
11026
11027#[cfg(target_arch = "aarch64")]
11028#[target_feature(enable = "neon")]
11029unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
11030 unsafe {
11032 use core::arch::aarch64::*;
11033 let n = x.len();
11034 let wp = w.as_ptr() as *const i8;
11035 let xp = x.as_ptr();
11036 let (mut a0, mut a1, mut a2, mut a3) = (
11037 vdupq_n_f32(0.0),
11038 vdupq_n_f32(0.0),
11039 vdupq_n_f32(0.0),
11040 vdupq_n_f32(0.0),
11041 );
11042 let mut j = 0usize;
11043 while j + 16 <= n {
11044 let wb = vld1q_s8(wp.add(j));
11045 let lo = vmovl_s8(vget_low_s8(wb));
11046 let hi = vmovl_s8(vget_high_s8(wb));
11047 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
11048 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
11049 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
11050 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
11051 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
11052 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
11053 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
11054 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
11055 j += 16;
11056 }
11057 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
11058 while j < n {
11059 sum += (*wp.add(j)) as f32 * *xp.add(j);
11060 j += 1;
11061 }
11062 sum
11063 }
11064}
11065
11066#[allow(clippy::too_many_arguments)]
11067fn qmatvec(
11068 q: &[u8],
11069 rep: &[u8],
11070 row_scale: &[f32],
11071 x: &[f32],
11072 col_field: &[f32],
11073 dtype: TensorDtype,
11074 rows: usize,
11075 cols: usize,
11076 out: &mut [f32],
11077 pool: Option<&Pool>,
11078) {
11079 debug_assert_eq!(out.len(), rows);
11080 #[cfg(not(target_arch = "aarch64"))]
11081 let _ = rep;
11082
11083 #[cfg(target_arch = "aarch64")]
11084 if sdot_enabled() {
11085 let act = if dtype == TensorDtype::Q8_2f {
11086 split_act_q8_2f(x, col_field)
11087 } else {
11088 split_act(x)
11089 };
11090 let out_addr = SendMut(out.as_mut_ptr());
11091 let run_range = |start: usize, end: usize| {
11092 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
11093 };
11094 match pool {
11095 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11096 _ => run_range(0, rows),
11097 }
11098 return;
11099 }
11100 #[cfg(target_arch = "x86_64")]
11103 if avx2_a8w8_enabled() {
11104 let act = if dtype == TensorDtype::Q8_2f {
11105 split_act_q8_2f(x, col_field)
11106 } else {
11107 split_act(x)
11108 };
11109 let out_addr = SendMut(out.as_mut_ptr());
11110 let run_range = |start: usize, end: usize| {
11111 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
11112 };
11113 match pool {
11114 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11115 _ => run_range(0, rows),
11116 }
11117 return;
11118 }
11119
11120 prescale_with(x, col_field, dtype, 1, |xs| {
11121 let out_addr = SendMut(out.as_mut_ptr());
11122 let run_range = move |start: usize, end: usize| {
11123 for o in start..end {
11124 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
11125 unsafe { *out_addr.at(o) = v };
11127 }
11128 };
11129 match pool {
11130 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11131 _ => run_range(0, rows),
11132 }
11133 });
11134}
11135
11136#[allow(clippy::too_many_arguments)]
11137fn qmatvec2(
11138 q: &[u8],
11139 row_scale: &[f32],
11140 x1: &[f32],
11141 x2: &[f32],
11142 col_field: &[f32],
11143 dtype: TensorDtype,
11144 rows: usize,
11145 cols: usize,
11146 o1: &mut [f32],
11147 o2: &mut [f32],
11148 pool: Option<&Pool>,
11149) {
11150 #[cfg(target_arch = "aarch64")]
11151 if sdot_enabled() {
11152 let a1s = if dtype == TensorDtype::Q8_2f {
11153 split_act_q8_2f(x1, col_field)
11154 } else {
11155 split_act(x1)
11156 };
11157 let a2s = if dtype == TensorDtype::Q8_2f {
11158 split_act_q8_2f(x2, col_field)
11159 } else {
11160 split_act(x2)
11161 };
11162 let p1 = SendMut(o1.as_mut_ptr());
11163 let p2 = SendMut(o2.as_mut_ptr());
11164 let run_range = |start: usize, end: usize| {
11165 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
11166 };
11167 match pool {
11168 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11169 _ => run_range(0, rows),
11170 }
11171 return;
11172 }
11173 #[cfg(target_arch = "x86_64")]
11174 if avx2_a8w8_enabled() {
11175 let a1s = if dtype == TensorDtype::Q8_2f {
11176 split_act_q8_2f(x1, col_field)
11177 } else {
11178 split_act(x1)
11179 };
11180 let a2s = if dtype == TensorDtype::Q8_2f {
11181 split_act_q8_2f(x2, col_field)
11182 } else {
11183 split_act(x2)
11184 };
11185 let p1 = SendMut(o1.as_mut_ptr());
11186 let p2 = SendMut(o2.as_mut_ptr());
11187 let run_range = |start: usize, end: usize| {
11188 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
11189 };
11190 match pool {
11191 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11192 _ => run_range(0, rows),
11193 }
11194 return;
11195 }
11196
11197 prescale_with(x1, col_field, dtype, 1, |x1s| {
11198 prescale_with(x2, col_field, dtype, 2, |x2s| {
11199 let p1 = SendMut(o1.as_mut_ptr());
11200 let p2 = SendMut(o2.as_mut_ptr());
11201 let run_range = move |start: usize, end: usize| {
11202 for o in start..end {
11203 let row = &q[o * cols..(o + 1) * cols];
11204 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
11205 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
11206 unsafe {
11208 *p1.at(o) = s1;
11209 *p2.at(o) = s2;
11210 }
11211 }
11212 };
11213 match pool {
11214 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11215 _ => run_range(0, rows),
11216 }
11217 });
11218 });
11219}
11220
11221#[derive(Clone, Copy)]
11222struct SendMut(*mut f32);
11223unsafe impl Send for SendMut {}
11224unsafe impl Sync for SendMut {}
11225
11226impl SendMut {
11227 #[inline]
11228 fn at(self, i: usize) -> *mut f32 {
11229 unsafe { self.0.add(i) }
11230 }
11231}
11232
11233#[cfg(test)]
11234mod tests {
11235 #[test]
11239 fn q8_round_is_round_clamp() {
11240 let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
11241 let mut probe = vec![
11242 0.0f32,
11243 -0.0,
11244 f32::NAN,
11245 f32::INFINITY,
11246 f32::NEG_INFINITY,
11247 f32::MAX,
11248 f32::MIN,
11249 1e30,
11250 -1e30,
11251 f32::MIN_POSITIVE,
11252 -f32::MIN_POSITIVE,
11253 ];
11254 for k in -300i32..=300 {
11255 let h = k as f32 * 0.5;
11256 let up = f32::from_bits(h.to_bits() + 1);
11257 let down = f32::from_bits(h.to_bits().wrapping_sub(1));
11258 for t in [h, up, down] {
11259 probe.push(t);
11260 probe.push(-t);
11261 }
11262 }
11263 let mut t = -140.0f32;
11264 while t < 140.0 {
11265 probe.push(t);
11266 t += 0.000_731;
11267 }
11268 for t in probe {
11269 assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
11270 }
11271 }
11272
11273 use super::*;
11274
11275 #[test]
11276 fn q2tp_i8_dot_matches_exact_on_grid() {
11277 let (rows, cols) = (5, 64);
11281 let gpr = cols / GROUP_SIZE;
11282 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
11286 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
11287 .collect();
11288 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
11289 let x: Vec<f32> = (0..cols)
11290 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
11291 .collect();
11292 let act = split_act(&x);
11293 assert!(
11294 act.outliers.is_empty(),
11295 "on-grid input must have no outliers"
11296 );
11297 let gsum = q1_group_sums(&act.xq, gpr);
11298 for r in 0..rows {
11299 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
11300 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
11301 assert!(
11302 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
11303 "row {r}: exact {exact} vs i8 {fast}"
11304 );
11305 }
11306 }
11307
11308 #[test]
11309 fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
11310 let (rows, cols) = (1usize, GROUP_SIZE);
11311 let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
11312 bytes[..Q2TP_CHUNK].fill(0x24); bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
11316 bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
11317 bytes[Q2TP_CHUNK + 4] = 1; let x = vec![1.0f32; cols];
11319 let mut raw = vec![0.0f32; rows];
11320 let mut affine = vec![0.0f32; rows];
11321 q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
11322 q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
11323 assert_eq!(raw, vec![-24.0]);
11324 assert_eq!(affine, vec![-8.0]);
11325 assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
11326 }
11327
11328 #[cfg(target_arch = "x86_64")]
11329 #[test]
11330 fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
11331 if !std::arch::is_x86_feature_detected!("avx2") {
11336 return;
11337 }
11338 let mut seed = 0x9e3779b9u32;
11339 let mut next = || {
11340 seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
11341 seed
11342 };
11343 for _ in 0..20_000 {
11344 let mut ch = [0u8; Q2TP_CHUNK];
11345 let mut x = [0i8; GROUP_SIZE];
11346 for b in &mut ch {
11347 *b = next() as u8;
11348 }
11349 for v in &mut x {
11350 *v = (next() >> 24) as i8;
11351 }
11352 let mut reference = 0i32;
11353 for (k, &b) in ch.iter().enumerate() {
11354 reference += (b & 3) as i32 * x[k * 4] as i32;
11355 reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
11356 reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
11357 reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
11358 }
11359 let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
11362 assert_eq!(got, reference, "packed q2 lane mismatch");
11363 }
11364 }
11365
11366 #[test]
11367 fn q8_row_dot_fast_matches_scalar() {
11368 let cols = 96;
11371 let row: Vec<u8> = (0..cols)
11372 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
11373 .collect();
11374 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
11375 let act = split_act(&x);
11376 let fast = q8_row_dot(&row, &act);
11377 let scalar = q8_row_dot_scalar(&row, &act);
11378 assert!(
11379 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
11380 "fast {fast} vs scalar {scalar}"
11381 );
11382 }
11383
11384 #[test]
11385 fn f32_matvec_matches_matvec_rows_bitexact() {
11386 let (rows, cols) = (300, 40);
11387 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
11388 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
11389 let qt = QTensor::from_f32(w.clone(), rows, cols);
11390
11391 let mut a = vec![0.0f32; rows];
11392 matvec_rows(None, &w, &x, &mut a);
11393 let mut b = vec![0.0f32; rows];
11394 qt.matvec(&x, &mut b, None);
11395 assert_eq!(a, b);
11396 }
11397
11398 #[test]
11399 fn sdot_kernel_exact_on_grid() {
11400 eprintln!("sdot_enabled = {}", sdot_enabled());
11405 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
11407 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11408 .collect();
11409 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
11410 let x: Vec<f32> = (0..cols)
11411 .map(|i| match i % 3 {
11412 0 => 1.0,
11413 1 => -1.0,
11414 _ => 0.0,
11415 })
11416 .collect();
11417 let mut a = vec![0.0f32; rows];
11418 qmatvec(
11419 &w,
11420 &[],
11421 &scales,
11422 &x,
11423 &[],
11424 TensorDtype::Q8Row,
11425 rows,
11426 cols,
11427 &mut a,
11428 None,
11429 );
11430 for o in 0..rows {
11431 let mut acc = 0.0f32;
11432 for j in 0..cols {
11433 acc += (w[o * cols + j] as i8) as f32 * x[j];
11434 }
11435 let expect = acc * scales[o];
11436 assert!(
11437 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11438 "row {o}: {} vs {expect}",
11439 a[o]
11440 );
11441 }
11442 }
11443
11444 #[test]
11445 fn q1_tbl_fast_path_matches_reference() {
11446 let (rows, cols) = (5, 256);
11451 let gpr = cols / GROUP_SIZE;
11452 let mut bytes = Vec::new();
11453 for t in 0..rows * gpr {
11454 let s = 0.007 + (t % 11) as f32 * 0.004;
11455 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11456 for j in 0..4 {
11457 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
11458 }
11459 }
11460 let x: Vec<f32> = (0..cols)
11461 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
11462 .collect();
11463 let mut w = vec![0.0f32; rows * cols];
11464 cortiq_core::quant::dequant_q1(&bytes, &mut w);
11465 let mut got = vec![0.0f32; rows];
11466 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11467 for o in 0..rows {
11468 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11469 assert!(
11470 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11471 "row {o}: {} vs {expect}",
11472 got[o]
11473 );
11474 }
11475 let b = 5usize;
11478 let mut xs_all = Vec::new();
11479 for bi in 0..b {
11480 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
11481 }
11482 let mut mm = vec![0.0f32; b * rows];
11483 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
11484 for bi in 0..b {
11485 let mut single = vec![0.0f32; rows];
11486 q1_matvec(
11487 &bytes,
11488 &xs_all[bi * cols..(bi + 1) * cols],
11489 rows,
11490 cols,
11491 &mut single,
11492 None,
11493 );
11494 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
11495 }
11496 }
11497
11498 #[test]
11499 fn q1_kernels_match_exact_reference() {
11500 let (rows, cols) = (7, 96);
11502 let gpr = cols / GROUP_SIZE;
11503 let mut bytes = Vec::new();
11504 for t in 0..rows * gpr {
11505 let s = 0.01 + (t % 13) as f32 * 0.003;
11506 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11507 for j in 0..4 {
11508 bytes.push(((t * 31 + j * 97) % 251) as u8);
11509 }
11510 }
11511 let x: Vec<f32> = (0..cols)
11513 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
11514 .collect();
11515 let mut w = vec![0.0f32; rows * cols];
11517 cortiq_core::quant::dequant_q1(&bytes, &mut w);
11518 let mut expect = vec![0.0f32; rows];
11519 for o in 0..rows {
11520 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11521 }
11522 let mut got = vec![0.0f32; rows];
11523 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11524 for o in 0..rows {
11525 assert!(
11526 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
11527 "row {o}: {} vs {}",
11528 got[o],
11529 expect[o]
11530 );
11531 }
11532 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
11534 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
11535 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
11536 assert_eq!(a1, got);
11537 let mut xs = x.clone();
11538 xs.extend_from_slice(&x2);
11539 let mut mm = vec![0.0f32; 2 * rows];
11540 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
11541 assert_eq!(&mm[..rows], got.as_slice());
11542 assert_eq!(&mm[rows..], a2.as_slice());
11543 }
11544
11545 #[test]
11546 fn repack_is_bit_identical() {
11547 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
11553 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
11554 .collect();
11555 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
11556 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
11557 let rep = q8_repack_layout(&w, rows, cols);
11558 for g in 0..rows / 4 {
11560 for c in 0..cols / 16 {
11561 for lane in 0..4 {
11562 assert_eq!(
11563 &rep[g * 4 * cols + c * 64 + lane * 16
11564 ..g * 4 * cols + c * 64 + lane * 16 + 16],
11565 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
11566 );
11567 }
11568 }
11569 }
11570 let mut a = vec![0.0f32; rows];
11571 qmatvec(
11572 &w,
11573 &[],
11574 &scales,
11575 &x,
11576 &[],
11577 TensorDtype::Q8Row,
11578 rows,
11579 cols,
11580 &mut a,
11581 None,
11582 );
11583 let mut b = vec![0.0f32; rows];
11584 qmatvec(
11585 &w,
11586 &rep,
11587 &scales,
11588 &x,
11589 &[],
11590 TensorDtype::Q8Row,
11591 rows,
11592 cols,
11593 &mut b,
11594 None,
11595 );
11596 assert_eq!(a, b, "full-range repack output diverged");
11597
11598 #[cfg(target_arch = "aarch64")]
11599 if sdot_enabled() {
11600 let act = split_act(&x);
11602 let mut c1 = vec![0.0f32; rows];
11603 let mut c2 = vec![0.0f32; rows];
11604 q8_range_sdot(
11605 &w,
11606 &[],
11607 &scales,
11608 &act,
11609 cols,
11610 SendMut(c1.as_mut_ptr()),
11611 3,
11612 rows - 2,
11613 );
11614 q8_range_sdot(
11615 &w,
11616 &rep,
11617 &scales,
11618 &act,
11619 cols,
11620 SendMut(c2.as_mut_ptr()),
11621 3,
11622 rows - 2,
11623 );
11624 assert_eq!(c1, c2, "unaligned-range repack output diverged");
11625 }
11626 }
11627
11628 #[test]
11629 fn sdot_a8w8_noise_is_bounded() {
11630 let (rows, cols) = (16, 512);
11634 let w: Vec<u8> = (0..rows * cols)
11635 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11636 .collect();
11637 let scales = vec![0.01f32; rows];
11638 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
11639 let mut a = vec![0.0f32; rows];
11640 qmatvec(
11641 &w,
11642 &[],
11643 &scales,
11644 &x,
11645 &[],
11646 TensorDtype::Q8Row,
11647 rows,
11648 cols,
11649 &mut a,
11650 None,
11651 );
11652 let (mut num, mut den) = (0f64, 0f64);
11653 for o in 0..rows {
11654 let mut acc = 0.0f32;
11655 for j in 0..cols {
11656 acc += (w[o * cols + j] as i8) as f32 * x[j];
11657 }
11658 let expect = acc * scales[o];
11659 num += ((a[o] - expect) as f64).powi(2);
11660 den += (expect as f64).powi(2);
11661 }
11662 let rel = (num / den.max(1e-12)).sqrt();
11663 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
11664 }
11665
11666 #[test]
11667 fn i8_dot_neon_matches_scalar() {
11668 let n = 100;
11669 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
11670 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
11671 let mut scalar = 0.0f32;
11672 for j in 0..n {
11673 scalar += (w[j] as i8) as f32 * x[j];
11674 }
11675 let fast = dot_i8_f32(&w, &x);
11676 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
11677 }
11678
11679 #[test]
11681 fn vbitmatvec_matches_full_dequant() {
11682 let (rows, cols) = (6, 64);
11683 let ng = cols / GROUP_SIZE;
11684 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11686 let mut bytes = bits.clone();
11687 for g in 0..rows * ng {
11688 let s = 0.02 + 0.001 * g as f32;
11689 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11690 }
11691 for r in 0..rows {
11692 let b = bits[r] as usize;
11693 let (mut acc, mut nb) = (0u64, 0usize);
11694 let mut rowbytes = Vec::new();
11695 for i in 0..cols {
11696 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11697 acc = (acc << b) | v;
11698 nb += b;
11699 while nb >= 8 {
11700 nb -= 8;
11701 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11702 }
11703 }
11704 if nb > 0 {
11705 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11706 }
11707 bytes.extend_from_slice(&rowbytes);
11708 }
11709 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11710
11711 let mut reference = vec![0f32; rows * cols];
11712 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
11713 let mut expect = vec![0f32; rows];
11714 for r in 0..rows {
11715 expect[r] = reference[r * cols..(r + 1) * cols]
11716 .iter()
11717 .zip(&x)
11718 .map(|(w, xv)| w * xv)
11719 .sum();
11720 }
11721 let mut got = vec![0f32; rows];
11722 let offsets = vbit_row_offsets(&bytes, rows, cols);
11723 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
11724 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11728 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11729 for r in 0..rows {
11730 assert!(
11731 (got[r] - expect[r]).abs() < tol * scale,
11732 "row {r}: {} vs {}",
11733 got[r],
11734 expect[r]
11735 );
11736 }
11737 }
11738
11739 #[test]
11744 #[cfg(target_arch = "x86_64")]
11745 fn vbit_matmat_blocked_matches_per_row() {
11746 let (rows, cols, b) = (64usize, 128usize, 9usize);
11747 let ng = cols / GROUP_SIZE;
11748 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
11749 let mut bytes = bits.clone();
11750 for g in 0..rows * ng {
11751 let sc = 0.02 + 0.0005 * g as f32;
11752 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11753 }
11754 for r in 0..rows {
11755 let bw = bits[r] as usize;
11756 let (mut acc, mut nb) = (0u64, 0usize);
11757 let mut rowbytes = Vec::new();
11758 for i in 0..cols {
11759 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11760 acc = (acc << bw) | v;
11761 nb += bw;
11762 while nb >= 8 {
11763 nb -= 8;
11764 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11765 }
11766 }
11767 if nb > 0 {
11768 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11769 }
11770 bytes.extend_from_slice(&rowbytes);
11771 }
11772 let x: Vec<f32> = (0..b * cols)
11773 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11774 .collect();
11775 let offsets = vbit_row_offsets(&bytes, rows, cols);
11776 let mut y_a = vec![0f32; b * rows];
11777 let mut y_b = vec![0f32; b * rows];
11778 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11779 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
11780 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11781 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
11782 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11783 let max_d = y_a
11784 .iter()
11785 .zip(&y_b)
11786 .map(|(p, q)| (p - q).abs())
11787 .fold(0.0f32, f32::max);
11788 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
11789 }
11790
11791 #[test]
11799 fn q4t_matmat_blocked_matches_per_row() {
11800 let (rows, cols, b) = (16usize, 64usize, 9usize);
11801 let gpr = cols / GROUP_SIZE;
11802 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11803 for r in 0..rows {
11804 for g in 0..gpr {
11805 let t = (r * gpr + g) * Q4_TILE;
11806 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
11807 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11808 for k in 0..16 {
11809 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11810 }
11811 }
11812 }
11813 let x: Vec<f32> = (0..b * cols)
11814 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11815 .collect();
11816 let mut y_blk = vec![0f32; b * rows];
11817 let mut y_row = vec![0f32; b * rows];
11818 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11819 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
11820 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11821 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
11822 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11823 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
11824 }
11825
11826 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
11833 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
11834 let gpr = cols / GROUP_SIZE;
11835 let stride = q4tp_code_stride(gpr);
11836 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
11837 let mut b = vec![0u8; codes_off + rows * stride];
11838 for r in 0..rows {
11839 for g in 0..gpr {
11840 let t = (r * gpr + g) * Q4TP_NIB;
11841 for k in 0..16 {
11842 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11843 }
11844 }
11845 let lo = -6.0 - 0.03 * (r % 17) as f32;
11846 let step = 0.01 + 0.004 * (r % 11) as f32;
11847 let p = params_off + r * 4;
11848 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
11849 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
11850 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
11851 for g in 0..gpr {
11852 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
11853 }
11854 }
11855 b
11856 }
11857
11858 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
11862 let gpr = cols / GROUP_SIZE;
11863 let v = Q4tpView::new(bytes, rows, cols);
11864 let mut out = vec![0u8; rows * gpr * Q4_TILE];
11865 let mut sc = vec![0f32; gpr];
11866 for r in 0..rows {
11867 v.scales_into(r, gpr, &mut sc);
11868 for g in 0..gpr {
11869 let t = (r * gpr + g) * Q4_TILE;
11870 let s = sc[g];
11871 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11872 let src = (r * gpr + g) * Q4TP_NIB;
11873 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
11874 }
11875 }
11876 out
11877 }
11878
11879 #[test]
11885 fn q4tp_exact_path_matches_dequant_reference() {
11886 let (rows, cols) = (256usize, 512usize);
11887 let gpr = cols / GROUP_SIZE;
11888 let bytes = synth_q4tp(rows, cols);
11889 let mut w = vec![0f32; rows * cols];
11890 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11891
11892 let x: Vec<f32> = (0..cols)
11893 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11894 .collect();
11895 let v = Q4tpView::new(&bytes, rows, cols);
11896 let mut sc = vec![0f32; gpr];
11897 for r in 0..rows {
11898 v.scales_into(r, gpr, &mut sc);
11899 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
11900 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
11901 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11905 assert!(
11906 (got - want).abs() <= 1e-5 * mag,
11907 "row {r}: kernel {got} vs dequant {want}"
11908 );
11909 }
11910 }
11911
11912 #[test]
11918 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
11919 let (rows, cols) = (256usize, 512usize);
11920 let bytes = synth_q4tp(rows, cols);
11921 let twin = q4tp_as_q4t(&bytes, rows, cols);
11922 let x: Vec<f32> = (0..cols)
11923 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11924 .collect();
11925
11926 let mut got = vec![0f32; rows];
11927 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
11928 let mut want = vec![0f32; rows];
11929 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
11930
11931 let mut w = vec![0f32; rows * cols];
11934 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11935 for r in 0..rows {
11936 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11937 assert!(
11938 (got[r] - want[r]).abs() <= 1e-3 * mag,
11939 "row {r}: q4tp {} vs q4t {}",
11940 got[r],
11941 want[r]
11942 );
11943 }
11944 }
11945
11946 #[test]
11951 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
11952 let (rows, cols, b) = (256usize, 512usize, 5usize);
11953 let bytes = synth_q4tp(rows, cols);
11954 let twin = q4tp_as_q4t(&bytes, rows, cols);
11955 let xs: Vec<f32> = (0..b * cols)
11956 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11957 .collect();
11958
11959 let mut got = vec![0f32; b * rows];
11960 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
11961 let mut want = vec![0f32; b * rows];
11962 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
11963
11964 let mut w = vec![0f32; rows * cols];
11965 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11966 for t in 0..b {
11967 for r in 0..rows {
11968 let mag: f32 = (0..cols)
11969 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
11970 .sum();
11971 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
11972 assert!(
11973 (g - wa).abs() <= 1e-3 * mag,
11974 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
11975 );
11976 }
11977 }
11978 }
11979
11980 #[test]
11981 fn q4tp_matvec2_matches_the_single_stream_kernel() {
11982 let (rows, cols) = (128usize, 256usize);
11983 let gpr = cols / GROUP_SIZE;
11984 let bytes = synth_q4tp(rows, cols);
11985 let xs: Vec<f32> = (0..2 * cols)
11986 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11987 .collect();
11988
11989 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
11990 q4tp_matvec2(
11991 &bytes,
11992 &xs[..cols],
11993 &xs[cols..],
11994 rows,
11995 cols,
11996 &mut o1,
11997 &mut o2,
11998 None,
11999 );
12000
12001 let v = Q4tpView::new(&bytes, rows, cols);
12004 let mut sc = vec![0f32; gpr];
12005 for r in 0..rows {
12006 v.scales_into(r, gpr, &mut sc);
12007 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
12008 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
12009 }
12010 }
12011
12012 #[test]
12019 fn q4tp_matvec_keeps_pace_with_q4t() {
12020 let (rows, cols) = (4096usize, 3072usize);
12021 let bytes = synth_q4tp(rows, cols);
12022 let twin = q4tp_as_q4t(&bytes, rows, cols);
12023 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
12024 let mut o = vec![0f32; rows];
12025 let n = 12;
12026 let mut best = (f64::MAX, f64::MAX);
12027 for _ in 0..3 {
12030 let t0 = std::time::Instant::now();
12031 for _ in 0..n {
12032 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
12033 }
12034 best.0 = best.0.min(t0.elapsed().as_secs_f64());
12035 let t0 = std::time::Instant::now();
12036 for _ in 0..n {
12037 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
12038 }
12039 best.1 = best.1.min(t0.elapsed().as_secs_f64());
12040 }
12041 let ratio = best.1 / best.0;
12042 println!(
12043 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
12044 best.0 * 1e3 / n as f64,
12045 best.1 * 1e3 / n as f64
12046 );
12047 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
12048 }
12049
12050 #[cfg(target_os = "macos")]
12051 #[test]
12052 fn q4t_matmat_accel_matches_dequant_reference() {
12053 if !accel_gemm_enabled() {
12054 return; }
12056 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
12058 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
12059 for r in 0..rows {
12060 for g in 0..gpr {
12061 let t = (r * gpr + g) * Q4_TILE;
12062 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
12063 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
12064 for k in 0..16 {
12065 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
12066 }
12067 }
12068 }
12069 let x: Vec<f32> = (0..b * cols)
12070 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12071 .collect();
12072 let mut got = vec![0f32; b * rows];
12073 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
12074 let mut w = vec![0f32; rows * cols];
12076 for r in 0..rows {
12077 for g in 0..gpr {
12078 let t = (r * gpr + g) * Q4_TILE;
12079 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
12080 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
12081 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
12082 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
12083 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
12084 }
12085 }
12086 }
12087 for bi in 0..b {
12088 for r in 0..rows {
12089 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
12090 let d = (got[bi * rows + r] - want).abs();
12091 assert!(
12092 d <= want.abs().max(1.0) * 1e-4,
12093 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
12094 got[bi * rows + r]
12095 );
12096 }
12097 }
12098 }
12099
12100 #[test]
12101 fn q4matvec_matches_full_dequant() {
12102 let (rows, cols) = (8, 64);
12103 let groups = rows * cols / GROUP_SIZE;
12104 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12106 for i in 0..groups * 16 {
12107 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12108 }
12109 for g in 0..groups {
12110 let s = 0.01 + 0.003 * g as f32;
12111 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12112 }
12113 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12114
12115 let mut reference = vec![0.0f32; rows * cols];
12116 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12117 let mut expect = vec![0.0f32; rows];
12118 for r in 0..rows {
12119 expect[r] = reference[r * cols..(r + 1) * cols]
12120 .iter()
12121 .zip(&x)
12122 .map(|(w, xv)| w * xv)
12123 .sum();
12124 }
12125
12126 let mut got = vec![0.0f32; rows];
12127 q4matvec(&bytes, &x, rows, cols, &mut got, None);
12128 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
12132 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12133 for r in 0..rows {
12134 assert!(
12135 (got[r] - expect[r]).abs() < tol * scale,
12136 "row {r}: {} vs {}",
12137 got[r],
12138 expect[r]
12139 );
12140 }
12141 }
12142
12143 #[test]
12146 fn vbitmatvec2_equals_two_singles() {
12147 let (rows, cols) = (6, 64);
12148 let ng = cols / GROUP_SIZE;
12149 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
12150 let mut bytes = bits.clone();
12151 for g in 0..rows * ng {
12152 let s = 0.02 + 0.001 * g as f32;
12153 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12154 }
12155 for r in 0..rows {
12156 let b = bits[r] as usize;
12157 let (mut acc, mut nb) = (0u64, 0usize);
12158 let mut rowbytes = Vec::new();
12159 for i in 0..cols {
12160 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
12161 acc = (acc << b) | v;
12162 nb += b;
12163 while nb >= 8 {
12164 nb -= 8;
12165 rowbytes.push(((acc >> nb) & 0xFF) as u8);
12166 }
12167 }
12168 if nb > 0 {
12169 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
12170 }
12171 bytes.extend_from_slice(&rowbytes);
12172 }
12173 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
12174 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
12175 let offsets = vbit_row_offsets(&bytes, rows, cols);
12176
12177 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12178 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
12179 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
12180 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
12181 vbitmatvec2(
12182 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
12183 );
12184 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
12185 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
12186 }
12187
12188 #[test]
12190 fn q4matvec2_equals_two_singles() {
12191 let (rows, cols) = (8, 128);
12192 let groups = rows * cols / GROUP_SIZE;
12193 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12194 for i in 0..groups * 16 {
12195 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12196 }
12197 for g in 0..groups {
12198 let s = 0.01 + 0.003 * g as f32;
12199 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12200 }
12201 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12204 x1[9] = 250.0;
12205 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12206
12207 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12208 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
12209 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
12210 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
12211 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
12212 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
12213 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
12214 }
12215
12216 #[test]
12219 fn matvec_many_equals_separate_matvecs() {
12220 use crate::pool::Pool;
12221 let (r1, r2, cols) = (300, 200, 64);
12222 let mk = |salt: usize, rows: usize| {
12223 QTensor::from_f32(
12224 (0..rows * cols)
12225 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
12226 .collect(),
12227 rows,
12228 cols,
12229 )
12230 };
12231 let (a, b) = (mk(1, r1), mk(5, r2));
12232 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
12233 let pool = Pool::new(3);
12234
12235 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
12236 a.matvec(&x, &mut ea, Some(&pool));
12237 b.matvec(&x, &mut eb, Some(&pool));
12238 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
12239 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
12240 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
12241 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
12242 }
12243
12244 #[test]
12249 fn q4tp_matvec_many_equals_separate_matvecs() {
12250 use crate::pool::Pool;
12251 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
12252
12253 let (r1, r2, cols) = (300usize, 200usize, 64usize);
12254 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
12255 "arch_name": "tiny-q4tp",
12256 "hidden_size": cols,
12257 "intermediate_size": cols * 2,
12258 "num_layers": 1,
12259 "num_attention_heads": 2,
12260 "num_kv_heads": 1,
12261 "head_dim": 32,
12262 "vocab_size": r1,
12263 "layer_types": ["FullAttention"],
12264 "rms_norm_eps": 1e-6,
12265 "max_position_embeddings": 8,
12266 "linear_conv_kernel_dim": 0,
12267 "linear_num_key_heads": 0,
12268 "linear_num_value_heads": 0
12269 }))
12270 .unwrap();
12271 let header = CmfHeader {
12272 format: "cmf".into(),
12273 version: CMF_VERSION,
12274 arch,
12275 quant_type: QuantType::Q4Block,
12276 provenance: None,
12277 tokenizer_config: None,
12278 section_hashes: None,
12279 skills: Vec::new(),
12280 shard: None,
12281 calibration: None,
12282 routing: None,
12283 };
12284 let specs = [
12285 TensorSpec {
12286 name: "q".into(),
12287 dtype: TensorDtype::Q4TiledP,
12288 shape: vec![r1, cols],
12289 data: synth_q4tp(r1, cols),
12290 },
12291 TensorSpec {
12292 name: "kv".into(),
12293 dtype: TensorDtype::Q4TiledP,
12294 shape: vec![r2, cols],
12295 data: synth_q4tp(r2, cols),
12296 },
12297 ];
12298 let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
12299 std::fs::create_dir_all(&dir).unwrap();
12300 let path = dir.join("m.cmf");
12301 CmfModel::write(&path, &header, &specs, None, None).unwrap();
12302 let model = Arc::new(CmfModel::open(&path).unwrap());
12303 let (a, b) = (
12304 QTensor::from_model(&model, "q").unwrap(),
12305 QTensor::from_model(&model, "kv").unwrap(),
12306 );
12307 assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
12308 assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
12309 let x: Vec<f32> = (0..cols)
12310 .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
12311 .collect();
12312 let pool = Pool::new(3);
12313 let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
12314 a.matvec(&x, &mut ea, Some(&pool));
12315 b.matvec(&x, &mut eb, Some(&pool));
12316 let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
12317 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
12318 assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
12319 assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
12320 let _ = std::fs::remove_dir_all(&dir);
12321 }
12322
12323 #[test]
12329 fn multi_token_moe_rows_equal_single_token_decode() {
12330 use crate::pool::Pool;
12331 use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
12332
12333 let (h, inter, ne) = (64usize, 128usize, 3usize);
12334 let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
12335 "arch_name": "tiny-q4tp-moe",
12336 "hidden_size": h,
12337 "intermediate_size": inter,
12338 "num_layers": 1,
12339 "num_attention_heads": 2,
12340 "num_kv_heads": 1,
12341 "head_dim": 32,
12342 "vocab_size": 8,
12343 "layer_types": ["FullAttention"],
12344 "rms_norm_eps": 1e-6,
12345 "max_position_embeddings": 8,
12346 "linear_conv_kernel_dim": 0,
12347 "linear_num_key_heads": 0,
12348 "linear_num_value_heads": 0
12349 }))
12350 .unwrap();
12351 let header = CmfHeader {
12352 format: "cmf".into(),
12353 version: CMF_VERSION,
12354 arch,
12355 quant_type: QuantType::Q4Block,
12356 provenance: None,
12357 tokenizer_config: None,
12358 section_hashes: None,
12359 skills: Vec::new(),
12360 shard: None,
12361 calibration: None,
12362 routing: None,
12363 };
12364 let mut specs = Vec::new();
12365 for e in 0..ne {
12366 for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
12367 .into_iter()
12368 .enumerate()
12369 {
12370 let mut data = synth_q4tp(r, c);
12373 for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
12374 .iter_mut()
12375 .enumerate()
12376 {
12377 *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
12378 }
12379 specs.push(TensorSpec {
12380 name: format!("{n}{e}"),
12381 dtype: TensorDtype::Q4TiledP,
12382 shape: vec![r, c],
12383 data,
12384 });
12385 }
12386 }
12387 let dir = std::env::temp_dir().join(format!(
12388 "cmf-moe-rows-{}-{}",
12389 std::process::id(),
12390 FLOAT_ACTIVATIONS.get()
12391 ));
12392 std::fs::create_dir_all(&dir).unwrap();
12393 let path = dir.join("m.cmf");
12394 CmfModel::write(&path, &header, &specs, None, None).unwrap();
12395 let model = Arc::new(CmfModel::open(&path).unwrap());
12396 let t = |n: String| QTensor::from_model(&model, &n).unwrap();
12397 let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
12398 let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
12399 let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
12400 let b = 4usize;
12401 let mut xs: Vec<f32> = (0..b * h)
12402 .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
12403 .collect();
12404 xs[5] = 9.0; let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
12407 (vec![2, 0], vec![0.6, 0.4]),
12408 (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
12409 (vec![1], vec![1.0]),
12410 (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
12411 ];
12412 let pool = Pool::new(3);
12413 let mut want = vec![0f32; b * h];
12415 for (tk, (idx, w)) in routes.iter().enumerate() {
12416 let x = &xs[tk * h..(tk + 1) * h];
12417 let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
12418 let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
12419 assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
12420 if FLOAT_ACTIVATIONS.get() {
12421 for (slot, &e) in idx.iter().enumerate() {
12422 let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
12423 g[e].matvec(x, &mut gate, Some(&pool));
12424 u[e].matvec(x, &mut up, Some(&pool));
12425 for (v, u) in gate.iter_mut().zip(up) {
12426 *v = (*v / (1.0 + (-*v).exp())) * u;
12427 }
12428 assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
12429 }
12430 }
12431 let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
12432 assert!(QTensor::moe_down_many(
12433 &downs,
12434 &gs,
12435 w,
12436 &mut want[tk * h..(tk + 1) * h],
12437 Some(&pool)
12438 ));
12439 }
12440 if FLOAT_ACTIVATIONS.get() {
12441 for (tk, (idx, w)) in routes.iter().enumerate() {
12442 let mut scalar = vec![0.0; h];
12443 for (&e, &weight) in idx.iter().zip(w) {
12444 let (mut gate, mut up, mut down) =
12445 (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
12446 g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
12447 u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
12448 for (v, u) in gate.iter_mut().zip(up) {
12449 *v = (*v / (1.0 + (-*v).exp())) * u;
12450 }
12451 d[e].matvec(&gate, &mut down, Some(&pool));
12452 for (v, d) in scalar.iter_mut().zip(down) {
12453 *v += weight * d;
12454 }
12455 }
12456 assert_eq!(
12457 &want[tk * h..(tk + 1) * h],
12458 scalar,
12459 "float many equals scalar experts"
12460 );
12461 }
12462 }
12463 let mut experts: Vec<usize> = Vec::new();
12465 let mut groups: Vec<Vec<usize>> = Vec::new();
12466 for (tk, (idx, _)) in routes.iter().enumerate() {
12467 for &e in idx {
12468 match experts.iter().position(|&x| x == e) {
12469 Some(k) => groups[k].push(tk),
12470 None => {
12471 experts.push(e);
12472 groups.push(vec![tk]);
12473 }
12474 }
12475 }
12476 }
12477 let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
12478 let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
12479 let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
12480 assert!(QTensor::moe_gate_up_rows(
12481 &pairs,
12482 &groups,
12483 &xs,
12484 &mut gs,
12485 Some(&pool)
12486 ));
12487 let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
12488 let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
12489 let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
12490 assert!(QTensor::moe_down_rows(
12491 &downs,
12492 &lens,
12493 &gs,
12494 &mut ds,
12495 Some(&pool)
12496 ));
12497 let slot = |tk: usize, e: usize| {
12498 let k = experts.iter().position(|&x| x == e).unwrap();
12499 groups[..k].iter().map(|g| g.len()).sum::<usize>()
12500 + groups[k].iter().position(|&x| x == tk).unwrap()
12501 };
12502 let mut got = vec![0f32; b * h];
12503 for (tk, (idx, w)) in routes.iter().enumerate() {
12504 for i in 0..h {
12505 let mut acc = 0f32;
12506 for (&e, &we) in idx.iter().zip(w) {
12507 acc += we * ds[slot(tk, e)][i];
12508 }
12509 got[tk * h + i] = acc;
12510 }
12511 }
12512 assert!(want.iter().any(|v| *v != 0.0));
12513 assert_eq!(
12514 want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12515 got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12516 "multi-token MoE must equal decode bit for bit"
12517 );
12518
12519 let b5 = 5usize;
12522 let x5: Vec<f32> = (0..b5 * h)
12523 .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
12524 .collect();
12525 let mut mm = vec![0f32; b5 * inter];
12526 row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
12527 for tk in 0..b5 {
12528 let mut mv = vec![0f32; inter];
12529 g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
12530 assert_eq!(
12531 mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12532 mm[tk * inter..(tk + 1) * inter]
12533 .iter()
12534 .map(|v| v.to_bits())
12535 .collect::<Vec<_>>(),
12536 "row-exact matmat token {tk}"
12537 );
12538 }
12539 let _ = std::fs::remove_dir_all(&dir);
12542 }
12543
12544 #[test]
12545 fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
12546 use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
12547 let active = AtomicUsize::new(0);
12550 counted_row_exact_scope(&active, || {
12551 assert_eq!(active.load(Ordering::Acquire), 1);
12552 counted_row_exact_scope(&active, || {
12553 assert_eq!(active.load(Ordering::Acquire), 2);
12554 });
12555 assert_eq!(active.load(Ordering::Acquire), 1);
12556 });
12557 assert_eq!(active.load(Ordering::Acquire), 0);
12558
12559 let both_entered = Barrier::new(2);
12560 let release_last = Barrier::new(2);
12561 std::thread::scope(|s| {
12562 let first = s.spawn(|| counted_row_exact_scope(&active, || {
12563 both_entered.wait();
12564 }));
12565 let last = s.spawn(|| counted_row_exact_scope(&active, || {
12566 both_entered.wait();
12567 release_last.wait();
12568 }));
12569 first.join().unwrap();
12570 let after_first = active.load(Ordering::Acquire);
12571 release_last.wait();
12572 last.join().unwrap();
12573 assert_eq!(after_first, 1, "second request must remain exact");
12574 });
12575 assert_eq!(active.load(Ordering::Acquire), 0);
12576 let panic = std::panic::catch_unwind(|| {
12577 counted_row_exact_scope(&active, || panic!("scope unwind"));
12578 });
12579 assert!(panic.is_err());
12580 assert_eq!(active.load(Ordering::Acquire), 0);
12581 }
12582
12583 #[test]
12584 #[cfg(target_arch = "x86_64")]
12585 fn q4tp_float_avx2_is_bitwise_scalar() {
12586 if !avx2_enabled() {
12587 return;
12588 }
12589 for cols in [32, 64, 96, 2048, 4096] {
12590 let rows = 9;
12591 let bytes = synth_q4tp(rows, cols);
12592 let v = Q4tpView::new(&bytes, rows, cols);
12593 let gpr = cols / GROUP_SIZE;
12594 let mut sc = vec![0.0; gpr];
12595 for seed in 1..=5 {
12596 let xs: Vec<f32> = (0..cols)
12597 .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
12598 .collect();
12599 for r in 0..rows {
12600 v.scales_into(r, gpr, &mut sc);
12601 let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
12602 let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
12603 assert_eq!(
12604 scalar.to_bits(),
12605 vector.to_bits(),
12606 "cols={cols} row={r} seed={seed}"
12607 );
12608 }
12609 }
12610 }
12611 }
12612
12613 #[test]
12614 fn multi_token_moe_rows_float_equal_single_token_decode() {
12615 float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
12616 }
12617
12618 #[test]
12619 fn full_gpu_q8_scope_is_nested_and_thread_local() {
12620 assert!(!FULL_GPU_Q8.get());
12621 let before = gpu_split_frac();
12622 {
12623 let _guard = enter_full_gpu_q8_scope();
12624 assert_eq!(gpu_split_frac(), 1.0);
12625 {
12626 let _nested = enter_full_gpu_q8_scope();
12627 }
12628 assert_eq!(gpu_split_frac(), 1.0);
12629 std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
12630 }
12631 assert!(!FULL_GPU_Q8.get());
12632 assert_eq!(gpu_split_frac(), before);
12633 }
12634
12635 #[test]
12636 fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
12637 assert!(!FLOAT_ACTIVATIONS.get());
12638 let before = a8w8_enabled();
12639 float_activations_scope(|| {
12640 assert!(!a8w8_enabled());
12641 float_activations_scope(|| assert!(!a8w8_enabled()));
12642 assert!(FLOAT_ACTIVATIONS.get());
12643 std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
12644 .join()
12645 .unwrap();
12646 });
12647 assert!(!FLOAT_ACTIVATIONS.get());
12648 assert_eq!(a8w8_enabled(), before);
12649 let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
12650 assert!(!FLOAT_ACTIVATIONS.get());
12651 }
12652
12653 #[test]
12656 fn batched_matmat_equals_per_position_matvec() {
12657 let (rows, cols, b) = (8, 64, 5);
12658 let groups = rows * cols / GROUP_SIZE;
12660 let mut q4 = Vec::new();
12661 for i in 0..groups * 16 {
12662 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12663 }
12664 for g in 0..groups {
12665 q4.extend_from_slice(
12666 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
12667 );
12668 }
12669 let ng = cols / GROUP_SIZE;
12671 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
12672 let mut vb = bits.clone();
12673 for g in 0..rows * ng {
12674 vb.extend_from_slice(
12675 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
12676 );
12677 }
12678 for r in 0..rows {
12679 let bw = bits[r] as usize;
12680 let (mut acc, mut nb) = (0u64, 0usize);
12681 let mut rowbytes = Vec::new();
12682 for i in 0..cols {
12683 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
12684 acc = (acc << bw) | v;
12685 nb += bw;
12686 while nb >= 8 {
12687 nb -= 8;
12688 rowbytes.push(((acc >> nb) & 0xFF) as u8);
12689 }
12690 }
12691 if nb > 0 {
12692 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
12693 }
12694 vb.extend_from_slice(&rowbytes);
12695 }
12696 let offsets = vbit_row_offsets(&vb, rows, cols);
12697
12698 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12699
12700 let mut got = vec![0f32; b * rows];
12702 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
12703 for bi in 0..b {
12704 let mut expect = vec![0f32; rows];
12705 q4matvec(
12706 &q4,
12707 &xs[bi * cols..(bi + 1) * cols],
12708 rows,
12709 cols,
12710 &mut expect,
12711 None,
12712 );
12713 assert_eq!(
12714 &got[bi * rows..(bi + 1) * rows],
12715 &expect[..],
12716 "q4 batch pos {bi}"
12717 );
12718 }
12719
12720 let mut got = vec![0f32; b * rows];
12722 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
12723 for bi in 0..b {
12724 let mut expect = vec![0f32; rows];
12725 vbitmatvec(
12726 &vb,
12727 &offsets,
12728 &xs[bi * cols..(bi + 1) * cols],
12729 rows,
12730 cols,
12731 &mut expect,
12732 None,
12733 );
12734 assert_eq!(
12735 &got[bi * rows..(bi + 1) * rows],
12736 &expect[..],
12737 "vbit batch pos {bi}"
12738 );
12739 }
12740 }
12741
12742 #[test]
12746 fn q4_tiled_matches_q4_block_bitexact() {
12747 let (rows, cols, b) = (8usize, 128usize, 3usize);
12748 let groups = rows * cols / GROUP_SIZE;
12749 let mut split = Vec::with_capacity(groups * 18);
12750 for i in 0..groups * 16 {
12751 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12752 }
12753 for g in 0..groups {
12754 split.extend_from_slice(
12755 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
12756 );
12757 }
12758 let (packed, scales) = split.split_at(groups * 16);
12760 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
12761 for g in 0..groups {
12762 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
12763 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
12764 }
12765
12766 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12767 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12769
12770 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
12771 q4matvec(&split, &x1, rows, cols, &mut a, None);
12772 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
12773 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
12774
12775 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12776 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
12777 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
12778 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
12779 assert_eq!(a1, t1);
12780 assert_eq!(a2, t2);
12781
12782 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12783 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12784 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
12785 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
12786 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
12787 }
12788
12789 #[test]
12796 fn q4matvec_sdot_outlier_exact() {
12797 let (rows, cols) = (4, 128);
12798 let groups = rows * cols / GROUP_SIZE;
12799 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12800 for i in 0..groups * 16 {
12801 bytes.push(((i * 11 + 5) % 256) as u8);
12802 }
12803 for g in 0..groups {
12804 let s = 0.02 + 0.002 * g as f32;
12805 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12806 }
12807 let mut x: Vec<f32> = (0..cols)
12808 .map(|i| match i % 3 {
12809 0 => 1.0,
12810 1 => -1.0,
12811 _ => 0.0,
12812 })
12813 .collect();
12814 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
12817 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12818 let mut expect = vec![0.0f32; rows];
12819 for r in 0..rows {
12820 expect[r] = reference[r * cols..(r + 1) * cols]
12821 .iter()
12822 .zip(&x)
12823 .map(|(w, xv)| w * xv)
12824 .sum();
12825 }
12826 let mut got = vec![0.0f32; rows];
12827 q4matvec(&bytes, &x, rows, cols, &mut got, None);
12828 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12829 for r in 0..rows {
12830 assert!(
12831 (got[r] - expect[r]).abs() < 2e-3 * scale,
12832 "row {r}: {} vs {} (outlier term must be exact)",
12833 got[r],
12834 expect[r]
12835 );
12836 }
12837 }
12838
12839 #[test]
12843 fn q1t_matvec_matches_reference() {
12844 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
12845 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
12847 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
12848 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
12850 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
12851 let mut bytes = Vec::new();
12852 for r in 0..rows {
12853 for g in 0..gpr {
12854 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
12855 let mut c = [0u8; 7];
12856 for k in 0..GROUP_SIZE {
12857 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
12859 0
12860 } else {
12861 ((k + r * 3 + g) % 3) as u8 };
12863 cortiq_core::quant::q1t_pack(&mut c, k, code);
12864 }
12865 bytes.extend_from_slice(&c);
12866 }
12867 }
12868 let mut row_ptr = vec![0u32; rows + 1];
12871 for &(idx, _) in &outliers {
12872 row_ptr[idx as usize / cols + 1] += 1;
12873 }
12874 for r in 0..rows {
12875 row_ptr[r + 1] += row_ptr[r];
12876 }
12877 for &p in &row_ptr {
12878 bytes.extend_from_slice(&p.to_le_bytes());
12879 }
12880 for &(idx, v) in &outliers {
12881 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
12882 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
12883 }
12884
12885 let mut refw = vec![0f32; rows * cols];
12886 dequant_q1t(&bytes, rows, cols, &mut refw);
12887 let x: Vec<f32> = (0..cols)
12890 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12891 .collect();
12892 let mut expect = vec![0f32; rows];
12893 for r in 0..rows {
12894 let mut a = 0.0f32;
12895 for j in 0..cols {
12896 a += refw[r * cols + j] * x[j];
12897 }
12898 expect[r] = a;
12899 }
12900 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
12901 let mut got = vec![0f32; rows];
12902 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
12903 for r in 0..rows {
12904 assert!(
12905 (got[r] - expect[r]).abs() < tol(expect[r]),
12906 "row {r}: {} vs {}",
12907 got[r],
12908 expect[r]
12909 );
12910 }
12911 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
12913 let mut gm = vec![0f32; 2 * rows];
12914 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
12915 for r in 0..rows {
12916 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
12917 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
12918 }
12919 let xb: Vec<f32> = (0..cols)
12923 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
12924 .collect();
12925 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12926 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
12927 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
12928 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12929 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
12930 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
12931 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
12932 }
12933
12934 #[test]
12937 fn q1t_matvec2_odd_gpr_matches_singles() {
12938 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12939 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
12941 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12942 for r in 0..rows {
12943 for g in 0..gpr {
12944 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
12945 let mut c = [0u8; 7];
12946 for k in 0..GROUP_SIZE {
12947 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
12948 }
12949 bytes.extend_from_slice(&c);
12950 }
12951 }
12952 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12953 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12954 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12955 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12956 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12957 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12958 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12959 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
12960 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
12961 }
12962
12963 #[test]
12967 #[ignore]
12968 fn q1t_matvec2_speed() {
12969 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12970 use std::time::Instant;
12971 let (rows, cols) = (8192usize, 4096usize);
12972 let gpr = cols / GROUP_SIZE;
12973 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12974 for r in 0..rows {
12975 for g in 0..gpr {
12976 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12977 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12978 let mut c = [0u8; 7];
12979 for k in 0..GROUP_SIZE {
12980 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12981 }
12982 bytes.extend_from_slice(&c);
12983 }
12984 }
12985 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12986 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12987 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12988 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12989 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12991 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12992 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
12993 for _ in 0..8 {
12994 let t0 = Instant::now();
12995 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12996 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
12997 let t1 = Instant::now();
12998 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12999 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
13000 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
13001 }
13002 assert_eq!(p1, s1);
13003 assert_eq!(p2, s2);
13004 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
13005 }
13006
13007 #[test]
13011 #[ignore]
13012 fn q1t_matvec_speed() {
13013 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
13014 use std::time::Instant;
13015 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
13017 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
13018 for r in 0..rows {
13019 for g in 0..gpr {
13020 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
13021 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
13022 let mut c = [0u8; 7];
13023 for k in 0..GROUP_SIZE {
13024 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
13025 }
13026 bytes.extend_from_slice(&c);
13027 }
13028 }
13029 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
13031 let mut idx = 0usize;
13032 while idx < n {
13033 row_ptr[idx / cols + 1] += 1;
13034 idx += stride;
13035 }
13036 for r in 0..rows {
13037 row_ptr[r + 1] += row_ptr[r];
13038 }
13039 for &p in &row_ptr {
13040 bytes.extend_from_slice(&p.to_le_bytes());
13041 }
13042 let mut idx = 0usize;
13043 while idx < n {
13044 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
13045 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
13046 idx += stride;
13047 }
13048 let x: Vec<f32> = (0..cols)
13051 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
13052 .collect();
13053 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
13054
13055 let slow = |out: &mut [f32]| {
13057 let mut buf = vec![0f32; cols];
13058 for r in 0..rows {
13059 for g in 0..gpr {
13060 let off = (r * gpr + g) * Q1T_TILE;
13061 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
13062 let codes = &bytes[off + 2..off + Q1T_TILE];
13063 for k in 0..GROUP_SIZE {
13064 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
13065 1 => s,
13066 2 => -s,
13067 _ => 0.0,
13068 };
13069 }
13070 }
13071 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
13072 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
13073 }
13074 };
13075 let iters = 5;
13076 let mut a = vec![0f32; rows];
13077 slow(&mut a); let t = Instant::now();
13079 for _ in 0..iters {
13080 slow(&mut a);
13081 }
13082 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
13083
13084 let mut b = vec![0f32; rows];
13085 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
13087 for _ in 0..iters {
13088 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
13089 }
13090 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
13091
13092 for r in 0..rows {
13093 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
13094 }
13095 println!(
13096 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
13097 slow_ms / fast_ms
13098 );
13099 }
13100}
13101
13102#[cfg(test)]
13103mod gemm_bench {
13104 #[test]
13115 #[ignore]
13116 fn q4tp_matmat_throughput() {
13117 let b: usize = std::env::var("CMF_BENCH_B")
13121 .ok()
13122 .and_then(|v| v.parse().ok())
13123 .unwrap_or(296);
13124 let (rows, cols) = (9216usize, 2304usize);
13125 let (_, _, _) = (rows, cols, b);
13126 let total =
13127 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
13128 .unwrap();
13129 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13134 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13135 let lo = cortiq_core::quant::f32_to_f16(-4.0);
13136 let step = cortiq_core::quant::f32_to_f16(0.1);
13137 for r in 0..rows {
13138 let o = params_off + r * 4;
13139 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13140 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13141 }
13142 let _ = codes_off;
13143 let xs: Vec<f32> = (0..b * cols)
13144 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13145 .collect();
13146 let mut out = vec![0f32; b * rows];
13147 let pool = crate::pool::Pool::from_env();
13148 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13154 let reps: usize = std::env::var("CMF_BENCH_REPS")
13155 .ok()
13156 .and_then(|v| v.parse().ok())
13157 .unwrap_or(10);
13158 let mut best = [f64::MAX; 2];
13159 let mut sums = [0f32; 2];
13160 for _ in 0..reps {
13161 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
13162 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
13163 let t = std::time::Instant::now();
13164 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13165 best[k] = best[k].min(t.elapsed().as_secs_f64());
13166 sums[k] = out.iter().take(64).sum::<f32>();
13167 }
13168 }
13169 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
13170 for (k, name) in ["previous", "tuned "].iter().enumerate() {
13171 println!(
13172 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
13173 best[k] * 1e3,
13174 flops / best[k] / 1e9,
13175 sums[k]
13176 );
13177 }
13178 assert!(
13179 (sums[0] - sums[1]).abs() < 1e-2,
13180 "the tuned kernel changed the result: {} vs {}",
13181 sums[0],
13182 sums[1]
13183 );
13184 }
13185
13186 #[test]
13192 fn q4tp_matmat_blocked_matches_scalar() {
13193 use std::sync::atomic::Ordering::Relaxed;
13194 for &(rows, cols, b) in &[
13202 (64usize, 128usize, 7usize),
13203 (33, 96, 4),
13204 (16, 256, 9),
13205 (192, 2304, 37),
13206 ] {
13207 let total = cortiq_core::quant::expected_nbytes(
13208 cortiq_core::TensorDtype::Q4TiledP,
13209 &[rows, cols],
13210 )
13211 .unwrap();
13212 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13213 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
13214 let lo = cortiq_core::quant::f32_to_f16(-4.0);
13215 let step = cortiq_core::quant::f32_to_f16(0.1);
13216 for r in 0..rows {
13217 let o = params_off + r * 4;
13218 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13219 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13220 }
13221 let xs: Vec<f32> = (0..b * cols)
13222 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
13223 .collect();
13224 let mut got = vec![0f32; b * rows];
13225 let mut want = vec![0f32; b * rows];
13226 let gpr = cols / 32;
13227 let view = super::Q4tpView::new(&bytes, rows, cols);
13228 let pool = crate::pool::Pool::from_env();
13229 super::Q4TP_ALT.store(2, Relaxed);
13230 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
13231 super::Q4TP_ALT.store(1, Relaxed);
13232 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
13233 super::Q4TP_ALT.store(0, Relaxed);
13234 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13241 let (mut worst, mut at) = (0f32, 0usize);
13242 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
13243 if (g - w).abs() > worst {
13244 worst = (g - w).abs();
13245 at = i;
13246 }
13247 }
13248 assert!(
13249 worst <= 1e-4 * scale,
13250 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
13251 (scale {scale:.3e}) at cell {at}: {} vs {}",
13252 got[at],
13253 want[at]
13254 );
13255
13256 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
13263 for bi in 0..b {
13264 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
13265 for r in 0..rows {
13266 let mut sc = vec![0f32; gpr];
13267 view.scales_into(r, gpr, &mut sc);
13268 let mut exact = 0f64;
13269 for j in 0..cols {
13270 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
13271 exact += w as f64 * sq as f64 * act.xq[j] as f64;
13272 }
13273 exact *= act.sx as f64;
13274 for &(j, xv) in &act.outliers {
13275 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
13276 exact += w as f64 * sq as f64 * xv as f64;
13277 }
13278 let i = bi * rows + r;
13279 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
13280 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
13281 }
13282 }
13283 println!(
13284 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
13285 per-column {e_scalar:.3e}"
13286 );
13287 assert!(
13292 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
13293 "{rows}x{cols} b={b}: error against f64 too large — blocked \
13294 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
13295 );
13296 }
13297 }
13298
13299 #[test]
13303 #[ignore]
13304 fn q4t_matmat_throughput() {
13305 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
13306 let total =
13307 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
13308 .unwrap();
13309 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13312 let sc = cortiq_core::quant::f32_to_f16(0.02);
13313 for t in bytes.chunks_mut(super::Q4_TILE) {
13314 t[..2].copy_from_slice(&sc.to_le_bytes());
13315 }
13316 let xs: Vec<f32> = (0..b * cols)
13317 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13318 .collect();
13319 let mut out = vec![0f32; b * rows];
13320 let pool = crate::pool::Pool::from_env();
13321 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13322 let reps: usize = std::env::var("CMF_BENCH_REPS")
13323 .ok()
13324 .and_then(|v| v.parse().ok())
13325 .unwrap_or(10);
13326 let mut best = f64::MAX;
13327 for _ in 0..reps {
13328 let t = std::time::Instant::now();
13329 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13330 best = best.min(t.elapsed().as_secs_f64());
13331 }
13332 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
13333 println!(
13334 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
13335 best * 1e3,
13336 flops / best / 1e9,
13337 out.iter().take(64).sum::<f32>()
13338 );
13339 }
13340}