1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16 GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17 q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23 F32 {
24 data: Vec<f32>,
25 rows: usize,
26 cols: usize,
27 },
28 Mapped {
29 model: Arc<CmfModel>,
30 idx: usize,
32 dtype: TensorDtype,
33 rows: usize,
34 cols: usize,
35 row_scale: Vec<f32>,
37 col_field: Vec<f32>,
39 vbit_offsets: Vec<usize>,
43 repack: Vec<u8>,
51 },
52}
53
54fn repack_enabled() -> bool {
61 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62 *ON.get_or_init(|| {
63 std::env::var("CMF_REPACK")
64 .map(|v| v == "1")
65 .unwrap_or(cfg!(target_os = "android"))
66 })
67}
68
69fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73 #[cfg(target_arch = "aarch64")]
74 let arch_ok = sdot_enabled();
75 #[cfg(not(target_arch = "aarch64"))]
76 let arch_ok = false;
77 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78 return Vec::new();
79 }
80 q8_repack_layout(bytes, rows, cols)
81}
82
83fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86 let groups = rows / 4;
87 let mut rep = vec![0u8; groups * 4 * cols];
88 for g in 0..groups {
89 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90 for c in 0..cols / 16 {
91 for lane in 0..4 {
92 let src = (g * 4 + lane) * cols + c * 16;
93 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94 .copy_from_slice(&bytes[src..src + 16]);
95 }
96 }
97 }
98 rep
99}
100
101fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104 let ng = cols / GROUP_SIZE;
105 let bits = &bytes[..rows];
106 let mut offsets = Vec::with_capacity(rows + 1);
107 let mut off = rows + rows * ng * 2;
108 for r in 0..rows {
109 offsets.push(off);
110 off += (cols * bits[r] as usize).div_ceil(8);
111 }
112 offsets.push(off);
113 offsets
114}
115
116fn blocked_enabled() -> bool {
122 use std::sync::atomic::Ordering::Relaxed;
123 match BLOCKED_OVERRIDE.load(Relaxed) {
124 1 => false,
125 2 => true,
126 _ => {
127 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128 *ON.get_or_init(|| {
129 std::env::var("CMF_X86_BLOCKED")
130 .map(|v| v != "0")
131 .unwrap_or(true)
132 })
133 }
134 }
135}
136
137static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
138
139pub fn set_blocked_override(on: Option<bool>) {
146 let v = match on {
147 None => 0,
148 Some(false) => 1,
149 Some(true) => 2,
150 };
151 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
152}
153
154fn gpu_lmhead_enabled() -> bool {
155 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
156 *ON.get_or_init(|| {
157 std::env::var("CMF_GPU_LMHEAD")
158 .map(|v| v != "0")
159 .unwrap_or(true)
160 })
161}
162
163fn gpu_split_frac() -> f32 {
164 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
165 *F.get_or_init(|| {
166 std::env::var("CMF_GPU_SPLIT")
167 .ok()
168 .and_then(|v| v.parse::<f32>().ok())
169 .unwrap_or(0.5)
170 .clamp(0.0, 1.0)
171 })
172}
173
174impl QTensor {
175 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
176 debug_assert_eq!(data.len(), rows * cols);
177 Self::F32 { data, rows, cols }
178 }
179
180 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
183 let idx = model
186 .tensor_index(name)
187 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
188 let entry = &model.tensors[idx];
189 if entry.shape.len() != 2 {
190 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
191 }
192 let (rows, cols) = (entry.shape[0], entry.shape[1]);
193 let bytes = model.entry_bytes(entry);
194
195 match entry.dtype {
196 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
197 let n = rows * cols;
198 let scales_off = n;
199 let row_scale: Vec<f32> = (0..rows)
200 .map(|o| {
201 f16_to_f32(u16::from_le_bytes([
202 bytes[scales_off + o * 2],
203 bytes[scales_off + o * 2 + 1],
204 ]))
205 })
206 .collect();
207 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
208 let col_off = n + rows * 2;
209 (0..cols)
210 .map(|i| {
211 f16_to_f32(u16::from_le_bytes([
212 bytes[col_off + i * 2],
213 bytes[col_off + i * 2 + 1],
214 ]))
215 })
216 .collect()
217 } else {
218 Vec::new()
219 };
220 Ok(Self::Mapped {
221 model: model.clone(),
222 idx,
223 dtype: entry.dtype,
224 rows,
225 cols,
226 row_scale,
227 col_field,
228 vbit_offsets: Vec::new(),
229 repack: q8_repack(bytes, rows, cols),
230 })
231 }
232 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
234 model: model.clone(),
235 idx,
236 dtype: entry.dtype,
237 rows,
238 cols,
239 row_scale: Vec::new(),
240 col_field: Vec::new(),
241 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
242 repack: Vec::new(),
243 }),
244 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
248 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
249 let offsets: Vec<usize> = (0..=rows)
250 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
251 .collect();
252 Ok(Self::Mapped {
253 model: model.clone(),
254 idx,
255 dtype: entry.dtype,
256 rows,
257 cols,
258 row_scale: Vec::new(),
259 col_field: Vec::new(),
260 vbit_offsets: offsets,
261 repack: Vec::new(),
262 })
263 }
264 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
270 model: model.clone(),
271 idx,
272 dtype: entry.dtype,
273 rows,
274 cols,
275 row_scale: Vec::new(),
276 col_field: Vec::new(),
277 vbit_offsets: Vec::new(),
278 repack: Vec::new(),
279 }),
280 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
283 model: model.clone(),
284 idx,
285 dtype: entry.dtype,
286 rows,
287 cols,
288 row_scale: Vec::new(),
289 col_field: Vec::new(),
290 vbit_offsets: Vec::new(),
291 repack: Vec::new(),
292 }),
293 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
295 model: model.clone(),
296 idx,
297 dtype: entry.dtype,
298 rows,
299 cols,
300 row_scale: Vec::new(),
301 col_field: Vec::new(),
302 vbit_offsets: Vec::new(),
303 repack: Vec::new(),
304 }),
305 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
306 model: model.clone(),
307 idx,
308 dtype: entry.dtype,
309 rows,
310 cols,
311 row_scale: Vec::new(),
312 col_field: Vec::new(),
313 vbit_offsets: Vec::new(),
314 repack: Vec::new(),
315 }),
316 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
318 model: model.clone(),
319 idx,
320 dtype: entry.dtype,
321 rows,
322 cols,
323 row_scale: Vec::new(),
324 col_field: Vec::new(),
325 vbit_offsets: Vec::new(),
326 repack: Vec::new(),
327 }),
328 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
332 model: model.clone(),
333 idx,
334 dtype: entry.dtype,
335 rows,
336 cols,
337 row_scale: Vec::new(),
338 col_field: Vec::new(),
339 vbit_offsets: Vec::new(),
340 repack: Vec::new(),
341 }),
342 _ => {
344 let mut data = vec![0.0f32; rows * cols];
345 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
346 Ok(Self::from_f32(data, rows, cols))
347 }
348 }
349 }
350
351 pub(crate) fn is_q1(&self) -> bool {
354 matches!(
355 self,
356 Self::Mapped {
357 dtype: TensorDtype::Q1,
358 ..
359 }
360 )
361 }
362
363 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
366 match self {
367 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
368 _ => None,
369 }
370 }
371
372 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
379 match self {
380 #[cfg(target_os = "macos")]
381 Self::Mapped {
382 dtype: TensorDtype::Q1T,
383 ..
384 } if !crate::gpu::metal_q1t_enabled() => None,
385 Self::Mapped {
386 idx,
387 dtype:
388 TensorDtype::Q1
389 | TensorDtype::Q1T
390 | TensorDtype::Q4Block
391 | TensorDtype::Q4Tiled
392 | TensorDtype::Q4TiledP
396 | TensorDtype::Q8Row
397 | TensorDtype::Q8_2f,
398 rows,
399 cols,
400 ..
401 } => Some((*idx, *rows, *cols)),
402 _ => None,
403 }
404 }
405
406 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
412 match self {
413 Self::Mapped {
414 idx,
415 dtype: TensorDtype::Q4Tiled,
416 rows,
417 cols,
418 ..
419 } => Some((*idx, *rows, *cols)),
420 _ => None,
421 }
422 }
423
424 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
428 match self {
429 Self::Mapped {
430 idx,
431 dtype: TensorDtype::Q4TiledP,
432 rows,
433 cols,
434 ..
435 } => Some((*idx, *rows, *cols)),
436 _ => None,
437 }
438 }
439
440 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
445 match self {
446 Self::Mapped {
447 idx,
448 dtype: TensorDtype::Q8Row,
449 rows,
450 cols,
451 row_scale,
452 col_field,
453 ..
454 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
455 _ => None,
456 }
457 }
458
459 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
463 match self {
464 Self::Mapped { dtype, .. } => Some(*dtype),
465 _ => None,
466 }
467 }
468
469 pub fn model_idx(&self) -> Option<usize> {
474 match self {
475 Self::Mapped { idx, .. } => Some(*idx),
476 _ => None,
477 }
478 }
479
480 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
485 match self {
486 Self::Mapped { model, .. } => Some(model.clone()),
487 _ => None,
488 }
489 }
490
491 pub fn rows(&self) -> usize {
492 match self {
493 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
494 }
495 }
496
497 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
500 match self {
501 Self::Mapped {
502 model,
503 idx,
504 dtype: TensorDtype::Q4Tiled,
505 ..
506 } => Some((model, *idx)),
507 _ => None,
508 }
509 }
510
511 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
514 match self {
515 Self::Mapped {
516 model,
517 idx,
518 dtype: TensorDtype::Q4TiledP,
519 ..
520 } => Some((model, *idx)),
521 _ => None,
522 }
523 }
524
525 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
528 match self {
529 Self::Mapped {
530 model,
531 idx,
532 dtype: TensorDtype::Q2TiledP,
533 ..
534 } => Some((model, *idx)),
535 _ => None,
536 }
537 }
538
539 pub fn cols(&self) -> usize {
540 match self {
541 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
542 }
543 }
544
545 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
548 match self {
549 Self::Mapped {
550 model,
551 idx,
552 dtype: TensorDtype::Q1,
553 ..
554 } => Some((model, *idx)),
555 _ => None,
556 }
557 }
558
559 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
569 match self {
570 Self::Mapped {
571 model,
572 idx,
573 dtype: TensorDtype::Q8Row,
574 row_scale,
575 ..
576 } => Some((model, *idx, 0, row_scale.as_slice())),
577 Self::Mapped {
578 model,
579 idx,
580 dtype: TensorDtype::Q1,
581 ..
582 } => Some((model, *idx, 1, &[])),
583 Self::Mapped {
588 model,
589 idx,
590 dtype: TensorDtype::Q4Tiled,
591 ..
592 } => Some((model, *idx, 5, &[])),
593 Self::Mapped {
597 model,
598 idx,
599 dtype: TensorDtype::Q4TiledP,
600 ..
601 } => Some((model, *idx, 6, &[])),
602 Self::Mapped {
603 model,
604 idx,
605 dtype: TensorDtype::Q4Block,
606 ..
607 } => Some((model, *idx, 2, &[])),
608 Self::Mapped {
613 model,
614 idx,
615 dtype: TensorDtype::Q8_2f,
616 ..
617 } => Some((model, *idx, 7, &[])),
618 Self::Mapped {
619 model,
620 idx,
621 dtype: TensorDtype::Q1T,
622 ..
623 } => Some((model, *idx, 3, &[])),
624 Self::Mapped {
628 model,
629 idx,
630 dtype: TensorDtype::Q2TiledP,
631 ..
632 } => Some((model, *idx, 9, &[])),
633 _ => None,
634 }
635 }
636
637 pub fn as_f32(&self) -> Option<&[f32]> {
640 match self {
641 Self::F32 { data, .. } => Some(data),
642 Self::Mapped { .. } => None,
643 }
644 }
645
646 fn quant_bytes(&self) -> &[u8] {
647 match self {
648 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
649 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
650 }
651 }
652
653 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
655 let cols = self.cols();
656 debug_assert_eq!(dst.len(), cols);
657 match self {
658 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
659 Self::Mapped {
660 dtype,
661 row_scale,
662 col_field,
663 vbit_offsets,
664 ..
665 } => {
666 if *dtype == TensorDtype::Q4Tiled {
667 let bytes = self.quant_bytes();
668 let gpr = cols / GROUP_SIZE;
669 for gi in 0..gpr {
670 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
671 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
672 for (k, &b) in tile[2..].iter().enumerate() {
673 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
674 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
675 }
676 }
677 return;
678 }
679 if *dtype == TensorDtype::Q4TiledP {
680 let bytes = self.quant_bytes();
681 let gpr = cols / GROUP_SIZE;
682 let v = Q4tpView::new(bytes, self.rows(), cols);
683 let mut sc = vec![0f32; gpr];
684 v.scales_into(r, gpr, &mut sc);
685 for gi in 0..gpr {
686 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
687 let s = sc[gi];
688 for (k, &b) in tile.iter().enumerate() {
689 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
690 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
691 }
692 }
693 return;
694 }
695 if *dtype == TensorDtype::Q2TiledP {
696 let bytes = self.quant_bytes();
697 let gpr = cols / GROUP_SIZE;
698 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
699 let mut sc = vec![0f32; gpr];
700 v.scales_into(r, gpr, &mut sc);
701 for gi in 0..gpr {
702 let ch =
703 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
704 let s = sc[gi];
705 for (k, &b) in ch.iter().enumerate() {
706 for j in 0..4 {
707 dst[gi * GROUP_SIZE + k * 4 + j] =
708 (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
709 }
710 }
711 }
712 return;
713 }
714 if *dtype == TensorDtype::Q4Block {
715 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
716 let gpr = cols / GROUP_SIZE;
717 for gi in 0..gpr {
718 let g = r * gpr + gi;
719 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
720 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
721 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
722 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
723 }
724 }
725 return;
726 }
727 if *dtype == TensorDtype::Q1 {
728 let bytes = self.quant_bytes();
729 let gpr = cols / GROUP_SIZE;
730 for gi in 0..gpr {
731 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
732 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
733 for (j, &b) in tile[2..].iter().enumerate() {
734 for k in 0..8 {
735 dst[gi * GROUP_SIZE + j * 8 + k] =
736 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
737 }
738 }
739 }
740 return;
741 }
742 if *dtype == TensorDtype::Q1T {
743 let bytes = self.quant_bytes();
744 let gpr = cols / GROUP_SIZE;
745 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
746 for gi in 0..gpr {
747 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
748 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
749 bytes[off],
750 bytes[off + 1],
751 ]));
752 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
753 for k in 0..GROUP_SIZE {
754 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
755 {
756 1 => s,
757 2 => -s,
758 _ => 0.0,
759 };
760 }
761 }
762 let rows = self.rows();
764 let entries = base_len + (rows + 1) * 4;
765 if entries <= bytes.len() {
766 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
767 let r0 = u32::from_le_bytes([
768 ptrs[r * 4],
769 ptrs[r * 4 + 1],
770 ptrs[r * 4 + 2],
771 ptrs[r * 4 + 3],
772 ]) as usize;
773 let r1 = u32::from_le_bytes([
774 ptrs[(r + 1) * 4],
775 ptrs[(r + 1) * 4 + 1],
776 ptrs[(r + 1) * 4 + 2],
777 ptrs[(r + 1) * 4 + 3],
778 ]) as usize;
779 let off = entries + r0 * 4;
780 for i in 0..r1 - r0 {
781 let item = &bytes[off + i * 4..off + i * 4 + 4];
782 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
783 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
784 item[2], item[3],
785 ]));
786 if c < cols {
787 dst[c] = v;
788 }
789 }
790 }
791 return;
792 }
793 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
794 let bytes = self.quant_bytes();
795 let rows = self.rows();
796 let ng = cols / GROUP_SIZE;
797 let bits = &bytes[..rows];
798 let sc_off = rows;
799 let off = vbit_offsets[r];
802 let b = bits[r] as usize;
803 let l = ((1usize << (b - 1)) - 1) as f32;
804 let data = &bytes[off..];
805 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
806 for (i, d) in dst.iter_mut().enumerate() {
807 while nbits < b {
808 acc = (acc << 8) | data[idx] as u64;
809 idx += 1;
810 nbits += 8;
811 }
812 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
813 nbits -= b;
814 let so = (r * ng + i / GROUP_SIZE) * 2;
815 let sv = f16_to_f32(u16::from_le_bytes([
816 bytes[sc_off + so],
817 bytes[sc_off + so + 1],
818 ]));
819 *d = (u - l) * sv;
820 }
821 return;
822 }
823 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
824 let s = row_scale[r];
825 match dtype {
826 TensorDtype::Q8Row => {
827 for (d, &b) in dst.iter_mut().zip(q) {
828 *d = (b as i8) as f32 * s;
829 }
830 }
831 TensorDtype::Q8_2f => {
832 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
833 *d = (b as i8) as f32 * s * col_field[i];
834 }
835 }
836 _ => unreachable!(),
837 }
838 }
839 }
840 }
841
842 pub fn sparse_col_ok(&self) -> bool {
847 match self {
848 Self::F32 { .. } => true,
849 Self::Mapped { dtype, .. } => {
850 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
851 }
852 }
853 }
854
855 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
859 let inter = self.cols();
860 let hidden = self.rows();
861 debug_assert_eq!(out.len(), hidden);
862 match self {
863 Self::F32 { data, .. } => {
864 for (k, o) in out.iter_mut().enumerate() {
865 *o += w * data[k * inter + c];
866 }
867 }
868 Self::Mapped {
869 dtype,
870 row_scale,
871 col_field,
872 ..
873 } => {
874 let q = self.quant_bytes();
875 let colf = if *dtype == TensorDtype::Q8_2f {
876 col_field[c]
877 } else {
878 1.0
879 };
880 let wc = w * colf;
881 for (k, o) in out.iter_mut().enumerate() {
882 let b = q[k * inter + c] as i8 as f32;
883 *o += wc * b * row_scale[k];
884 }
885 }
886 }
887 }
888
889 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
893 let cols = self.cols();
894 match self {
895 Self::F32 { data, .. } => {
896 let row = &data[r * cols..(r + 1) * cols];
897 row.iter().zip(x).map(|(w, v)| w * v).sum()
898 }
899 Self::Mapped {
900 dtype,
901 row_scale,
902 col_field,
903 ..
904 } => match dtype {
905 TensorDtype::Q8Row => {
906 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
907 dot_i8_f32(q, x) * row_scale[r]
908 }
909 TensorDtype::Q8_2f => {
910 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
911 dot_i8_col_f32(q, x, col_field) * row_scale[r]
912 }
913 _ => {
914 self.row_f32(r, scratch);
915 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
916 }
917 },
918 }
919 }
920
921 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
924 match self {
925 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
929 Self::Mapped {
930 model,
931 idx,
932 dtype,
933 rows,
934 cols,
935 row_scale,
936 col_field,
937 vbit_offsets,
938 repack,
939 } => {
940 let _ = (model, idx);
941 assert!(
951 out.len() >= *rows && x.len() >= *cols,
952 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
953 out.len(),
954 x.len(),
955 );
956 if *dtype == TensorDtype::Q4Block {
957 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
961 let t0 = std::time::Instant::now();
962 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
963 crate::gpu::ProbeArm::Gpu => {
964 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
965 crate::gpu::probe_record(
966 crate::gpu::OpClass::Matvec,
967 true,
968 t0.elapsed(),
969 );
970 return;
971 }
972 }
973 crate::gpu::ProbeArm::CpuTimed => {
974 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
975 crate::gpu::probe_record(
976 crate::gpu::OpClass::Matvec,
977 false,
978 t0.elapsed(),
979 );
980 return;
981 }
982 crate::gpu::ProbeArm::Cpu => {}
983 }
984 }
985 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
986 return;
987 }
988 if *dtype == TensorDtype::Q4Tiled {
989 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
994 let t0 = std::time::Instant::now();
995 let cls = crate::gpu::matvec_class(*rows, *cols);
996 match crate::gpu::probe_arm(cls) {
997 crate::gpu::ProbeArm::Gpu => {
998 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
999 crate::gpu::probe_record(cls, true, t0.elapsed());
1000 return;
1001 }
1002 }
1003 crate::gpu::ProbeArm::CpuTimed => {
1004 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1005 crate::gpu::probe_record(cls, false, t0.elapsed());
1006 return;
1007 }
1008 crate::gpu::ProbeArm::Cpu => {}
1009 }
1010 }
1011 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1012 return;
1013 }
1014 if *dtype == TensorDtype::Q4TiledP {
1015 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1021 let t0 = std::time::Instant::now();
1022 let cls = crate::gpu::matvec_class(*rows, *cols);
1023 match crate::gpu::probe_arm(cls) {
1024 crate::gpu::ProbeArm::Gpu => {
1025 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1026 crate::gpu::probe_record(cls, true, t0.elapsed());
1027 return;
1028 }
1029 }
1030 crate::gpu::ProbeArm::CpuTimed => {
1031 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1032 crate::gpu::probe_record(cls, false, t0.elapsed());
1033 return;
1034 }
1035 crate::gpu::ProbeArm::Cpu => {}
1036 }
1037 }
1038 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1039 return;
1040 }
1041 if *dtype == TensorDtype::Q2TiledP {
1042 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1043 return;
1044 }
1045 if *dtype == TensorDtype::Q1 {
1046 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1051 let t0 = std::time::Instant::now();
1052 let arm = if crate::gpu::q1_force() {
1053 crate::gpu::ProbeArm::Gpu
1054 } else {
1055 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1056 };
1057 match arm {
1058 crate::gpu::ProbeArm::Gpu => {
1059 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1060 crate::gpu::probe_record(
1061 crate::gpu::OpClass::Matvec,
1062 true,
1063 t0.elapsed(),
1064 );
1065 return;
1066 }
1067 }
1068 crate::gpu::ProbeArm::CpuTimed => {
1069 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1070 crate::gpu::probe_record(
1071 crate::gpu::OpClass::Matvec,
1072 false,
1073 t0.elapsed(),
1074 );
1075 return;
1076 }
1077 crate::gpu::ProbeArm::Cpu => {}
1078 }
1079 }
1080 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1081 return;
1082 }
1083 if *dtype == TensorDtype::Q1T {
1084 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1088 let t0 = std::time::Instant::now();
1089 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1090 crate::gpu::ProbeArm::Gpu => {
1091 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1092 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1093 crate::gpu::probe_record(
1094 crate::gpu::OpClass::Matvec,
1095 true,
1096 t0.elapsed(),
1097 );
1098 return;
1099 }
1100 }
1101 crate::gpu::ProbeArm::CpuTimed => {
1102 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1103 crate::gpu::probe_record(
1104 crate::gpu::OpClass::Matvec,
1105 false,
1106 t0.elapsed(),
1107 );
1108 return;
1109 }
1110 crate::gpu::ProbeArm::Cpu => {}
1111 }
1112 }
1113 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1114 return;
1115 }
1116 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1117 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1118 return;
1119 }
1120 let xs = prescale(x, col_field, *dtype);
1121 if *rows >= crate::gpu::min_rows()
1126 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1127 && gpu_lmhead_enabled()
1128 && crate::gpu::enabled_here()
1129 {
1130 let t0 = std::time::Instant::now();
1133 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1134 crate::gpu::ProbeArm::Gpu => {}
1135 crate::gpu::ProbeArm::CpuTimed => {
1136 qmatvec(
1137 self.quant_bytes(),
1138 repack,
1139 row_scale,
1140 x,
1141 col_field,
1142 *dtype,
1143 *rows,
1144 *cols,
1145 out,
1146 pool,
1147 );
1148 crate::gpu::probe_record(
1149 crate::gpu::OpClass::Matvec,
1150 false,
1151 t0.elapsed(),
1152 );
1153 return;
1154 }
1155 crate::gpu::ProbeArm::Cpu => {
1156 qmatvec(
1157 self.quant_bytes(),
1158 repack,
1159 row_scale,
1160 x,
1161 col_field,
1162 *dtype,
1163 *rows,
1164 *cols,
1165 out,
1166 pool,
1167 );
1168 return;
1169 }
1170 }
1171 let frac = gpu_split_frac();
1172 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1173 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1174 let bytes = self.quant_bytes();
1175 let ok = std::thread::scope(|sc| {
1176 let g = sc.spawn(|| {
1177 crate::gpu::q8_matvec_range(
1178 model,
1179 *idx,
1180 cpu_rows,
1181 &row_scale[cpu_rows..],
1182 &xs,
1183 *rows - cpu_rows,
1184 *cols,
1185 out_gpu,
1186 )
1187 });
1188 if cpu_rows > 0 {
1189 let rep_cpu = if repack.is_empty() {
1192 &[][..]
1193 } else {
1194 &repack[..(cpu_rows / 4) * 4 * *cols]
1195 };
1196 qmatvec(
1197 &bytes[..cpu_rows * *cols],
1198 rep_cpu,
1199 &row_scale[..cpu_rows],
1200 x,
1201 col_field,
1202 *dtype,
1203 cpu_rows,
1204 *cols,
1205 out_cpu,
1206 pool,
1207 );
1208 }
1209 g.join().unwrap_or(false)
1210 });
1211 if ok {
1212 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1213 return;
1214 }
1215 qmatvec(
1218 &bytes[cpu_rows * *cols..(*rows) * *cols],
1219 &[],
1220 &row_scale[cpu_rows..],
1221 x,
1222 col_field,
1223 *dtype,
1224 *rows - cpu_rows,
1225 *cols,
1226 out_gpu,
1227 pool,
1228 );
1229 return;
1230 }
1231 qmatvec(
1232 self.quant_bytes(),
1233 repack,
1234 row_scale,
1235 x,
1236 col_field,
1237 *dtype,
1238 *rows,
1239 *cols,
1240 out,
1241 pool,
1242 );
1243 }
1244 }
1245 }
1246
1247 pub fn matvec2(
1249 &self,
1250 x1: &[f32],
1251 x2: &[f32],
1252 o1: &mut [f32],
1253 o2: &mut [f32],
1254 pool: Option<&Pool>,
1255 ) {
1256 match self {
1257 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1258 Self::Mapped {
1259 dtype,
1260 rows,
1261 cols,
1262 row_scale,
1263 col_field,
1264 vbit_offsets,
1265 ..
1266 } => {
1267 if *dtype == TensorDtype::Q4Block {
1268 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1269 return;
1270 }
1271 if *dtype == TensorDtype::Q4Tiled {
1272 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1273 return;
1274 }
1275 if *dtype == TensorDtype::Q4TiledP {
1276 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1277 return;
1278 }
1279 if *dtype == TensorDtype::Q2TiledP {
1280 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1281 return;
1282 }
1283 if *dtype == TensorDtype::Q1 {
1284 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1285 return;
1286 }
1287 if *dtype == TensorDtype::Q1T {
1288 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1294 return;
1295 }
1296 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1297 vbitmatvec2(
1298 self.quant_bytes(),
1299 vbit_offsets,
1300 x1,
1301 x2,
1302 *rows,
1303 *cols,
1304 o1,
1305 o2,
1306 pool,
1307 );
1308 return;
1309 }
1310 qmatvec2(
1311 self.quant_bytes(),
1312 row_scale,
1313 x1,
1314 x2,
1315 col_field,
1316 *dtype,
1317 *rows,
1318 *cols,
1319 o1,
1320 o2,
1321 pool,
1322 );
1323 }
1324 }
1325 }
1326}
1327
1328impl QTensor {
1329 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1337 match self {
1338 Self::Mapped {
1339 model, idx, dtype, ..
1340 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1341 _ => None,
1342 }
1343 }
1344
1345 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1346 let cols = self.cols();
1347 let rows = self.rows();
1348 debug_assert_eq!(xs_all.len(), b * cols);
1349 debug_assert_eq!(out.len(), b * rows);
1350 if crate::gptq_capture::capturing() {
1354 if let Self::Mapped { model, idx, .. } = self {
1355 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1356 }
1357 }
1358 match self {
1359 Self::F32 { data, .. } => {
1360 let out_addr = SendMut(out.as_mut_ptr());
1361 let run = |start: usize, end: usize| {
1362 for o in start..end {
1363 let row = &data[o * cols..(o + 1) * cols];
1364 for bi in 0..b {
1365 let x = &xs_all[bi * cols..(bi + 1) * cols];
1366 let mut acc = 0f32;
1367 for j in 0..cols {
1368 acc += row[j] * x[j];
1369 }
1370 unsafe { *out_addr.at(bi * rows + o) = acc };
1371 }
1372 }
1373 };
1374 dispatch_rows(pool, rows, &run);
1375 }
1376 Self::Mapped {
1377 dtype,
1378 row_scale,
1379 col_field,
1380 vbit_offsets,
1381 ..
1382 } => {
1383 if *dtype == TensorDtype::Q4Block {
1384 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1385 return;
1386 }
1387 if *dtype == TensorDtype::Q4TiledP {
1388 if b >= 32
1400 && b * rows * cols >= 128_000_000
1401 && cols % 32 == 0
1402 && !crate::gpu::mm_killed()
1403 && crate::gpu::enabled_here()
1404 {
1405 let class = if b >= 128 {
1406 crate::gpu::OpClass::MatmatWide
1407 } else {
1408 crate::gpu::OpClass::Matmat
1409 };
1410 if let Self::Mapped { model, idx, .. } = self {
1411 if crate::mm_ab::on() {
1423 let mut g = vec![0f32; b * rows];
1424 let t = std::time::Instant::now();
1425 let took = crate::gpu::q4tp_matmat(
1426 model, *idx, xs_all, b, rows, cols, &mut g,
1427 );
1428 let dg = t.elapsed();
1429 let t = std::time::Instant::now();
1430 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1431 let dc = t.elapsed();
1432 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1433 return;
1434 }
1435 let t0 = std::time::Instant::now();
1436 let resident = crate::gpu::weight_is_resident(model, *idx);
1440 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1441 crate::gpu::ProbeArm::Gpu => {
1442 if crate::gpu::q4tp_matmat(
1443 model, *idx, xs_all, b, rows, cols, out,
1444 ) {
1445 let el = t0.elapsed();
1446 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1456 let budget = std::time::Duration::from_secs_f64(
1457 flops / 1.5e12 * 8.0 + 0.020,
1458 );
1459 crate::gpu::mm_budget_check(
1460 "q4tp matmat",
1461 el,
1462 budget,
1463 crate::gpu::probe_was_cold() || !resident,
1464 );
1465 crate::gpu::probe_record(class, true, el);
1466 return;
1467 }
1468 }
1469 crate::gpu::ProbeArm::CpuTimed => {
1470 q4tp_matmat(
1471 self.quant_bytes(),
1472 xs_all,
1473 b,
1474 rows,
1475 cols,
1476 out,
1477 pool,
1478 );
1479 crate::gpu::probe_record(class, false, t0.elapsed());
1480 return;
1481 }
1482 crate::gpu::ProbeArm::Cpu => {}
1483 }
1484 }
1485 }
1486 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1487 return;
1488 }
1489 if *dtype == TensorDtype::Q2TiledP {
1490 if b >= 32
1496 && b * rows * cols >= 128_000_000
1497 && cols % 32 == 0
1498 && !crate::gpu::mm_killed()
1499 && crate::gpu::enabled_here()
1500 {
1501 let class = if b >= 128 {
1502 crate::gpu::OpClass::MatmatWide
1503 } else {
1504 crate::gpu::OpClass::Matmat
1505 };
1506 if let Self::Mapped { model, idx, .. } = self {
1507 let t0 = std::time::Instant::now();
1508 match crate::gpu::probe_arm(class) {
1509 crate::gpu::ProbeArm::Gpu => {
1510 if crate::gpu::q2tp_matmat(
1511 model, *idx, xs_all, b, rows, cols, out,
1512 ) {
1513 crate::gpu::probe_record(class, true, t0.elapsed());
1514 return;
1515 }
1516 }
1517 crate::gpu::ProbeArm::CpuTimed => {
1518 q2tp_matmat(
1519 self.quant_bytes(),
1520 xs_all,
1521 b,
1522 rows,
1523 cols,
1524 out,
1525 pool,
1526 );
1527 crate::gpu::probe_record(class, false, t0.elapsed());
1528 return;
1529 }
1530 crate::gpu::ProbeArm::Cpu => {}
1531 }
1532 }
1533 }
1534 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1539 return;
1540 }
1541 if *dtype == TensorDtype::Q4Tiled {
1542 if b >= 32
1554 && b * rows * cols >= 128_000_000
1555 && cols % 32 == 0
1556 && !crate::gpu::mm_killed()
1557 && crate::gpu::enabled_here()
1558 {
1559 let class = if b >= 128 {
1560 crate::gpu::OpClass::MatmatWide
1561 } else {
1562 crate::gpu::OpClass::Matmat
1563 };
1564 if let Self::Mapped { model, idx, .. } = self {
1565 let t0 = std::time::Instant::now();
1566 match crate::gpu::probe_arm(class) {
1567 crate::gpu::ProbeArm::Gpu => {
1568 if crate::gpu::q4t_matmat(
1569 model, *idx, xs_all, b, rows, cols, out,
1570 ) {
1571 let el = t0.elapsed();
1572 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1582 let budget = std::time::Duration::from_secs_f64(
1583 flops / 1.5e12 * 8.0 + 0.020,
1584 );
1585 crate::gpu::mm_budget_check(
1586 "q4t matmat",
1587 el,
1588 budget,
1589 crate::gpu::probe_was_cold(),
1590 );
1591 crate::gpu::probe_record(class, true, el);
1592 return;
1593 }
1594 }
1595 crate::gpu::ProbeArm::CpuTimed => {
1596 q4t_matmat(
1597 self.quant_bytes(),
1598 xs_all,
1599 b,
1600 rows,
1601 cols,
1602 out,
1603 pool,
1604 );
1605 crate::gpu::probe_record(class, false, t0.elapsed());
1606 return;
1607 }
1608 crate::gpu::ProbeArm::Cpu => {}
1609 }
1610 }
1611 }
1612 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1613 return;
1614 }
1615 if *dtype == TensorDtype::Q1 {
1616 if b >= 32
1619 && b * rows * cols >= 128_000_000
1620 && cols % 64 == 0
1621 && crate::gpu::enabled_here()
1622 {
1623 if let Self::Mapped { model, idx, .. } = self {
1624 let t0 = std::time::Instant::now();
1625 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1626 crate::gpu::ProbeArm::Gpu => {
1627 if crate::gpu::q1_matmat(
1628 model, *idx, xs_all, b, rows, cols, out,
1629 ) {
1630 crate::gpu::probe_record(
1631 crate::gpu::OpClass::Matmat,
1632 true,
1633 t0.elapsed(),
1634 );
1635 return;
1636 }
1637 }
1638 crate::gpu::ProbeArm::CpuTimed => {
1639 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1640 crate::gpu::probe_record(
1641 crate::gpu::OpClass::Matmat,
1642 false,
1643 t0.elapsed(),
1644 );
1645 return;
1646 }
1647 crate::gpu::ProbeArm::Cpu => {}
1648 }
1649 }
1650 }
1651 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1652 return;
1653 }
1654 if *dtype == TensorDtype::Q1T {
1655 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1658 if let Self::Mapped { model, idx, .. } = self {
1659 let t0 = std::time::Instant::now();
1660 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1661 crate::gpu::ProbeArm::Gpu => {
1662 if crate::gpu::q1t_matmat(
1663 model, *idx, xs_all, b, rows, cols, out,
1664 ) {
1665 crate::gpu::probe_record(
1666 crate::gpu::OpClass::Matmat,
1667 true,
1668 t0.elapsed(),
1669 );
1670 return;
1671 }
1672 }
1673 crate::gpu::ProbeArm::CpuTimed => {
1674 q1t_matmat(
1675 self.quant_bytes(),
1676 xs_all,
1677 b,
1678 rows,
1679 cols,
1680 out,
1681 pool,
1682 );
1683 crate::gpu::probe_record(
1684 crate::gpu::OpClass::Matmat,
1685 false,
1686 t0.elapsed(),
1687 );
1688 return;
1689 }
1690 crate::gpu::ProbeArm::Cpu => {}
1691 }
1692 }
1693 }
1694 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1695 return;
1696 }
1697 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1698 vbitmatmat(
1699 self.quant_bytes(),
1700 vbit_offsets,
1701 xs_all,
1702 b,
1703 rows,
1704 cols,
1705 out,
1706 pool,
1707 );
1708 return;
1709 }
1710 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1711 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1712 .collect();
1713 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1719 if let Self::Mapped { model, idx, .. } = self {
1720 let t0 = std::time::Instant::now();
1721 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1722 crate::gpu::ProbeArm::Gpu
1723 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1724 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1725 {
1726 let q = self.quant_bytes();
1730 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1731 return;
1732 }
1733 crate::gpu::ProbeArm::Gpu => {
1734 let flat: Vec<f32> =
1735 pre.iter().flat_map(|v| v.iter().copied()).collect();
1736 if crate::gpu::q8_matmat(
1737 model, *idx, row_scale, &flat, b, rows, cols, out,
1738 ) {
1739 crate::gpu::probe_record(
1740 crate::gpu::OpClass::Matmat,
1741 true,
1742 t0.elapsed(),
1743 );
1744 return;
1745 }
1746 }
1747 crate::gpu::ProbeArm::CpuTimed => {
1748 let q = self.quant_bytes();
1749 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1750 crate::gpu::probe_record(
1751 crate::gpu::OpClass::Matmat,
1752 false,
1753 t0.elapsed(),
1754 );
1755 return;
1756 }
1757 crate::gpu::ProbeArm::Cpu => {}
1758 }
1759 }
1760 }
1761 let q = self.quant_bytes();
1762 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1763 }
1764 }
1765 }
1766}
1767
1768impl QTensor {
1769 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
1779 let (rows, cols) = (self.rows(), self.cols());
1780 let Self::Mapped { model, idx, dtype, row_scale, col_field, .. } = self else {
1781 return false;
1782 };
1783 match *dtype {
1784 TensorDtype::Q4TiledP => {
1785 crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out)
1786 }
1787 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1791 let flat: Vec<f32> = (0..b)
1792 .flat_map(|bi| {
1793 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
1794 })
1795 .collect();
1796 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
1797 }
1798 _ => false,
1799 }
1800 }
1801
1802 pub fn matvec_many<const N: usize>(
1809 ts: [&QTensor; N],
1810 x: &[f32],
1811 mut outs: [&mut [f32]; N],
1812 pool: Option<&Pool>,
1813 ) {
1814 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1815 let uniform_q8 = ts.iter().all(|t| {
1816 matches!(
1817 t,
1818 Self::Mapped {
1819 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1820 ..
1821 }
1822 )
1823 });
1824 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1825 let uniform_q4 = ts.iter().all(|t| {
1826 matches!(
1827 t,
1828 Self::Mapped {
1829 dtype: TensorDtype::Q4Block,
1830 ..
1831 }
1832 )
1833 });
1834 let uniform_vbit = ts.iter().all(|t| {
1835 matches!(
1836 t,
1837 Self::Mapped {
1838 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1839 ..
1840 }
1841 )
1842 });
1843 let uniform_q1 = ts.iter().all(|t| {
1844 matches!(
1845 t,
1846 Self::Mapped {
1847 dtype: TensorDtype::Q1,
1848 ..
1849 }
1850 )
1851 });
1852 let uniform_q1t = ts.iter().all(|t| {
1853 matches!(
1854 t,
1855 Self::Mapped {
1856 dtype: TensorDtype::Q1T,
1857 ..
1858 }
1859 )
1860 });
1861 let uniform_q4tp = ts.iter().all(|t| {
1866 matches!(
1867 t,
1868 Self::Mapped {
1869 dtype: TensorDtype::Q4TiledP,
1870 ..
1871 }
1872 )
1873 }) && ts
1874 .iter()
1875 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1876 let Some(pool) = pool else {
1877 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1878 t.matvec(x, o, None);
1879 }
1880 return;
1881 };
1882 if total_rows < 256
1883 || !(uniform_q8
1884 || uniform_f32
1885 || uniform_q4
1886 || uniform_vbit
1887 || uniform_q1
1888 || uniform_q1t
1889 || uniform_q4tp)
1890 {
1891 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1892 t.matvec(x, o, Some(pool));
1893 }
1894 return;
1895 }
1896
1897 if uniform_q4tp {
1898 let cols = ts[0].cols();
1904 let gpr = cols / GROUP_SIZE;
1905 let views: Vec<Q4tpView> = ts
1906 .iter()
1907 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1908 .collect();
1909 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1910 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1911 let locate = |flat: usize| -> (usize, usize) {
1913 let mut acc = 0;
1914 for (i, &r) in rows_of.iter().enumerate() {
1915 if flat < acc + r {
1916 return (i, flat - acc);
1917 }
1918 acc += r;
1919 }
1920 (rows_of.len() - 1, 0)
1921 };
1922 let (views, outs_addr) = (&views, &outs_addr);
1923 if a8w8_enabled() {
1924 let act = split_act(x);
1925 let act = &act;
1926 let run = |start: usize, end: usize| {
1927 let mut sc = vec![0f32; gpr];
1928 for flat in start..end {
1929 let (t, r) = locate(flat);
1930 let v = &views[t];
1931 v.scales_into(r, gpr, &mut sc);
1932 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1933 for &(j, xv) in &act.outliers {
1934 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1935 acc += w * s * xv;
1936 }
1937 unsafe { *outs_addr[t].at(r) = acc };
1939 }
1940 };
1941 pool.run_rows(total_rows, &run);
1942 } else {
1943 let run = |start: usize, end: usize| {
1944 let mut sc = vec![0f32; gpr];
1945 for flat in start..end {
1946 let (t, r) = locate(flat);
1947 let v = &views[t];
1948 v.scales_into(r, gpr, &mut sc);
1949 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1951 }
1952 };
1953 pool.run_rows(total_rows, &run);
1954 }
1955 return;
1956 }
1957
1958 if uniform_q1 {
1959 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1962 if a8w8_enabled() {
1963 let act = split_act(x);
1964 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1965 let (act, gsum) = (&act, &gsum);
1966 let closures: [_; N] = std::array::from_fn(|i| {
1967 let (bytes, gpr, out) =
1968 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1969 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1970 });
1971 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1972 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1973 pool.run_many(&parts);
1974 } else {
1975 let closures: [_; N] = std::array::from_fn(|i| {
1976 let (bytes, gpr, out) =
1977 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1978 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1979 });
1980 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1981 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1982 pool.run_many(&parts);
1983 }
1984 return;
1985 }
1986
1987 if uniform_q1t {
1988 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1992 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1993 if a8w8_enabled() {
1994 let act = split_act(x);
1995 let act = &act;
1996 let x_ref = x;
1997 let closures: [_; N] = std::array::from_fn(|i| {
1998 let bytes = ts[i].quant_bytes();
1999 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2000 let gpr = cols / GROUP_SIZE;
2001 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2002 let out = outs_addr[i];
2003 move |s: usize, e: usize| {
2004 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2005 }
2006 });
2007 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2008 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2009 pool.run_many(&parts);
2010 } else {
2011 let x_ref = x;
2012 let closures: [_; N] = std::array::from_fn(|i| {
2013 let bytes = ts[i].quant_bytes();
2014 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2015 let gpr = cols / GROUP_SIZE;
2016 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2017 let out = outs_addr[i];
2018 move |s: usize, e: usize| {
2019 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2020 }
2021 });
2022 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2023 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2024 pool.run_many(&parts);
2025 }
2026 return;
2027 }
2028
2029 if uniform_q4 || uniform_vbit {
2030 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2031 if a8w8_enabled() {
2033 let act = split_act(x);
2034 let act = &act;
2035 if uniform_q4 {
2036 let closures: [_; N] = std::array::from_fn(|i| {
2037 let (packed, scales) =
2038 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2039 let (gpr, cols, out) =
2040 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2041 move |s: usize, e: usize| {
2042 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2043 }
2044 });
2045 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2046 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2047 pool.run_many(&parts);
2048 } else {
2049 let closures: [_; N] = std::array::from_fn(|i| {
2050 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2051 unreachable!()
2052 };
2053 let (bytes, rows, cols, out) = (
2054 ts[i].quant_bytes(),
2055 ts[i].rows(),
2056 ts[i].cols(),
2057 outs_addr[i],
2058 );
2059 move |s: usize, e: usize| {
2060 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2061 }
2062 });
2063 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2064 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2065 pool.run_many(&parts);
2066 }
2067 return;
2068 }
2069 if uniform_q4 {
2070 let closures: [_; N] = std::array::from_fn(|i| {
2071 let (packed, scales) =
2072 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2073 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2074 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2075 });
2076 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2077 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2078 pool.run_many(&parts);
2079 } else {
2080 let closures: [_; N] = std::array::from_fn(|i| {
2081 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2082 unreachable!()
2083 };
2084 let (bytes, rows, cols, out) = (
2085 ts[i].quant_bytes(),
2086 ts[i].rows(),
2087 ts[i].cols(),
2088 outs_addr[i],
2089 );
2090 move |s: usize, e: usize| {
2091 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2092 }
2093 });
2094 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2095 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2096 pool.run_many(&parts);
2097 }
2098 return;
2099 }
2100
2101 if uniform_f32 {
2102 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2103 let closures: [_; N] = std::array::from_fn(|i| {
2104 let Self::F32 { data, cols, .. } = ts[i] else {
2105 unreachable!()
2106 };
2107 let out = outs_addr[i];
2108 move |start: usize, end: usize| {
2109 for o in start..end {
2110 let row = &data[o * cols..(o + 1) * cols];
2111 let mut sum = 0.0f32;
2112 for j in 0..*cols {
2113 sum += row[j] * x[j];
2114 }
2115 unsafe { *out.at(o) = sum };
2117 }
2118 }
2119 });
2120 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2121 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2122 pool.run_many(&parts);
2123 return;
2124 }
2125
2126 struct Ctx<'a> {
2129 bytes: &'a [u8],
2130 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2131 rep: &'a [u8],
2132 row_scale: &'a [f32],
2133 cols: usize,
2134 xs: std::borrow::Cow<'a, [f32]>,
2135 }
2136 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2137 let Self::Mapped {
2138 dtype,
2139 cols,
2140 row_scale,
2141 col_field,
2142 repack,
2143 ..
2144 } = ts[i]
2145 else {
2146 unreachable!()
2147 };
2148 Ctx {
2149 bytes: ts[i].quant_bytes(),
2150 rep: repack,
2151 row_scale,
2152 cols: *cols,
2153 xs: prescale(x, col_field, *dtype),
2154 }
2155 });
2156 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2157 #[cfg(target_arch = "aarch64")]
2158 if sdot_enabled() {
2159 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2160 let closures: [_; N] = std::array::from_fn(|i| {
2161 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2162 move |start: usize, end: usize| {
2163 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2164 }
2165 });
2166 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2167 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2168 pool.run_many(&parts);
2169 return;
2170 }
2171 #[cfg(target_arch = "x86_64")]
2172 if avx2_a8w8_enabled() {
2173 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2174 let closures: [_; N] = std::array::from_fn(|i| {
2175 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2176 move |start: usize, end: usize| {
2177 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2178 }
2179 });
2180 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2181 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2182 pool.run_many(&parts);
2183 return;
2184 }
2185 let closures: [_; N] = std::array::from_fn(|i| {
2186 let (c, out) = (&ctxs[i], outs_addr[i]);
2187 move |start: usize, end: usize| {
2188 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2189 }
2190 });
2191 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2192 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2193 pool.run_many(&parts);
2194 }
2195}
2196
2197impl QTensor {
2198 #[allow(clippy::needless_range_loop)]
2203 pub fn matvec2_many<const N: usize>(
2204 ts: [&QTensor; N],
2205 x1: &[f32],
2206 x2: &[f32],
2207 mut o1s: [&mut [f32]; N],
2208 mut o2s: [&mut [f32]; N],
2209 pool: Option<&Pool>,
2210 ) {
2211 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2212 let uniform_q8 = ts.iter().all(|t| {
2213 matches!(
2214 t,
2215 Self::Mapped {
2216 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2217 ..
2218 }
2219 )
2220 });
2221 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2222 let uniform_q4 = ts.iter().all(|t| {
2223 matches!(
2224 t,
2225 Self::Mapped {
2226 dtype: TensorDtype::Q4Block,
2227 ..
2228 }
2229 )
2230 });
2231 let uniform_vbit = ts.iter().all(|t| {
2232 matches!(
2233 t,
2234 Self::Mapped {
2235 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2236 ..
2237 }
2238 )
2239 });
2240 let fusable = pool.is_some()
2241 && total_rows >= 256
2242 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2243 if !fusable {
2244 for i in 0..N {
2245 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2246 }
2247 return;
2248 }
2249 let pool = pool.unwrap();
2250
2251 if uniform_q4 || uniform_vbit {
2252 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2253 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2254 if a8w8_enabled() {
2256 let a1 = split_act(x1);
2257 let a2 = split_act(x2);
2258 let (a1, a2) = (&a1, &a2);
2259 if uniform_q4 {
2260 let closures: [_; N] = std::array::from_fn(|i| {
2261 let (packed, scales) =
2262 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2263 let (gpr, cols, o1, o2) =
2264 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2265 move |s: usize, e: usize| {
2266 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2267 }
2268 });
2269 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2270 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2271 pool.run_many(&parts);
2272 } else {
2273 let closures: [_; N] = std::array::from_fn(|i| {
2274 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2275 unreachable!()
2276 };
2277 let (bytes, rows, cols, o1, o2) = (
2278 ts[i].quant_bytes(),
2279 ts[i].rows(),
2280 ts[i].cols(),
2281 p1[i],
2282 p2[i],
2283 );
2284 move |s: usize, e: usize| {
2285 vbit_range2_a8w8(
2286 bytes,
2287 vbit_offsets,
2288 x1,
2289 x2,
2290 a1,
2291 a2,
2292 rows,
2293 cols,
2294 o1,
2295 o2,
2296 s,
2297 e,
2298 )
2299 }
2300 });
2301 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2302 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2303 pool.run_many(&parts);
2304 }
2305 return;
2306 }
2307 if uniform_q4 {
2308 let closures: [_; N] = std::array::from_fn(|i| {
2309 let (packed, scales) =
2310 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2311 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2312 move |s: usize, e: usize| {
2313 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2314 }
2315 });
2316 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2317 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2318 pool.run_many(&parts);
2319 } else {
2320 let closures: [_; N] = std::array::from_fn(|i| {
2321 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2322 unreachable!()
2323 };
2324 let (bytes, rows, cols, o1, o2) = (
2325 ts[i].quant_bytes(),
2326 ts[i].rows(),
2327 ts[i].cols(),
2328 p1[i],
2329 p2[i],
2330 );
2331 move |s: usize, e: usize| {
2332 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2333 }
2334 });
2335 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2336 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2337 pool.run_many(&parts);
2338 }
2339 return;
2340 }
2341
2342 if uniform_f32 {
2343 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2344 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2345 let closures: [_; N] = std::array::from_fn(|i| {
2346 let Self::F32 { data, cols, .. } = ts[i] else {
2347 unreachable!()
2348 };
2349 let (o1, o2) = (p1[i], p2[i]);
2350 move |start: usize, end: usize| {
2351 for o in start..end {
2352 let row = &data[o * cols..(o + 1) * cols];
2353 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2354 for j in 0..*cols {
2355 s1 += row[j] * x1[j];
2356 s2 += row[j] * x2[j];
2357 }
2358 unsafe {
2360 *o1.at(o) = s1;
2361 *o2.at(o) = s2;
2362 }
2363 }
2364 }
2365 });
2366 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2367 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2368 pool.run_many(&parts);
2369 return;
2370 }
2371
2372 struct Ctx<'a> {
2373 bytes: &'a [u8],
2374 row_scale: &'a [f32],
2375 cols: usize,
2376 xs1: std::borrow::Cow<'a, [f32]>,
2377 xs2: std::borrow::Cow<'a, [f32]>,
2378 }
2379 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2380 let Self::Mapped {
2381 dtype,
2382 cols,
2383 row_scale,
2384 col_field,
2385 ..
2386 } = ts[i]
2387 else {
2388 unreachable!()
2389 };
2390 Ctx {
2391 bytes: ts[i].quant_bytes(),
2392 row_scale,
2393 cols: *cols,
2394 xs1: prescale(x1, col_field, *dtype),
2395 xs2: prescale(x2, col_field, *dtype),
2396 }
2397 });
2398 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2399 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2400 #[cfg(target_arch = "aarch64")]
2401 if sdot_enabled() {
2402 let acts: [(SplitAct, SplitAct); N] =
2403 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2404 let closures: [_; N] = std::array::from_fn(|i| {
2405 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2406 move |start: usize, end: usize| {
2407 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2408 }
2409 });
2410 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2411 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2412 pool.run_many(&parts);
2413 return;
2414 }
2415 #[cfg(target_arch = "x86_64")]
2416 if avx2_a8w8_enabled() {
2417 let acts: [(SplitAct, SplitAct); N] =
2418 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2419 let closures: [_; N] = std::array::from_fn(|i| {
2420 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2421 move |start: usize, end: usize| {
2422 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2423 }
2424 });
2425 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2426 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2427 pool.run_many(&parts);
2428 return;
2429 }
2430 let closures: [_; N] = std::array::from_fn(|i| {
2431 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2432 move |start: usize, end: usize| {
2433 q8_range2_f32(
2434 c.bytes,
2435 c.row_scale,
2436 &c.xs1,
2437 &c.xs2,
2438 c.cols,
2439 o1,
2440 o2,
2441 start,
2442 end,
2443 )
2444 }
2445 });
2446 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2447 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2448 pool.run_many(&parts);
2449 }
2450
2451 pub fn matvec_silu_mul(
2456 gate: &QTensor,
2457 up: &QTensor,
2458 x: &[f32],
2459 out: &mut [f32],
2460 pool: Option<&Pool>,
2461 ) -> bool {
2462 let inter = gate.rows();
2463 debug_assert_eq!(up.rows(), inter);
2464 debug_assert_eq!(out.len(), inter);
2465 debug_assert_eq!(gate.cols(), up.cols());
2466 if !a8w8_enabled() {
2467 return false;
2468 }
2469 let act = split_act(x);
2470 let act = &act;
2471 let x_ref = x;
2472 let out_addr = SendMut(out.as_mut_ptr());
2473
2474 match (gate, up) {
2475 (
2477 Self::Mapped {
2478 dtype: TensorDtype::Q4Block,
2479 ..
2480 },
2481 Self::Mapped {
2482 dtype: TensorDtype::Q4Block,
2483 ..
2484 },
2485 ) => {
2486 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2487 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2488 let gpr = gate.cols() / GROUP_SIZE;
2489 let cols = gate.cols();
2490 let run = move |start: usize, end: usize| {
2491 for r in start..end {
2492 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2493 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2494 for &(j, xv) in &act.outliers {
2495 let flat = r * cols + j;
2496 let gb = gp[flat / 2];
2497 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2498 let gsc = f16_to_f32(u16::from_le_bytes([
2499 gs[(flat / GROUP_SIZE) * 2],
2500 gs[(flat / GROUP_SIZE) * 2 + 1],
2501 ]));
2502 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2503 let ub = up_p[flat / 2];
2504 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2505 let usc = f16_to_f32(u16::from_le_bytes([
2506 up_s[(flat / GROUP_SIZE) * 2],
2507 up_s[(flat / GROUP_SIZE) * 2 + 1],
2508 ]));
2509 uv += ((un as i32 - 8) as f32) * usc * xv;
2510 }
2511 let silu_g = gv / (1.0 + (-gv).exp());
2512 unsafe { *out_addr.at(r) = silu_g * uv };
2514 }
2515 };
2516 dispatch_rows(pool, inter, &run);
2517 true
2518 }
2519 (
2523 Self::Mapped {
2524 dtype: TensorDtype::Q4Tiled,
2525 ..
2526 },
2527 Self::Mapped {
2528 dtype: TensorDtype::Q4Tiled,
2529 ..
2530 },
2531 ) => {
2532 let g_bytes = gate.quant_bytes();
2533 let u_bytes = up.quant_bytes();
2534 let gpr = gate.cols() / GROUP_SIZE;
2535 let run = move |start: usize, end: usize| {
2536 for r in start..end {
2537 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2538 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2539 for &(j, xv) in &act.outliers {
2540 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2541 gv += w * s * xv;
2542 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2543 uv += w * s * xv;
2544 }
2545 let silu_g = gv / (1.0 + (-gv).exp());
2546 unsafe { *out_addr.at(r) = silu_g * uv };
2548 }
2549 };
2550 dispatch_rows(pool, inter, &run);
2551 true
2552 }
2553 (
2556 Self::Mapped {
2557 dtype: TensorDtype::Q4TiledP,
2558 ..
2559 },
2560 Self::Mapped {
2561 dtype: TensorDtype::Q4TiledP,
2562 ..
2563 },
2564 ) => {
2565 let cols = gate.cols();
2566 let gpr = cols / GROUP_SIZE;
2567 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2568 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2569 let run = |start: usize, end: usize| {
2570 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2571 for r in start..end {
2572 gv_view.scales_into(r, gpr, &mut gsc);
2573 uv_view.scales_into(r, gpr, &mut usc);
2574 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2575 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2576 for &(j, xv) in &act.outliers {
2577 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2578 gv += w * s * xv;
2579 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2580 uv += w * s * xv;
2581 }
2582 let silu_g = gv / (1.0 + (-gv).exp());
2583 unsafe { *out_addr.at(r) = silu_g * uv };
2585 }
2586 };
2587 dispatch_rows(pool, inter, &run);
2588 true
2589 }
2590 (
2596 Self::Mapped {
2597 dtype: TensorDtype::Q1,
2598 ..
2599 },
2600 Self::Mapped {
2601 dtype: TensorDtype::Q1,
2602 ..
2603 },
2604 ) => {
2605 let g_bytes = gate.quant_bytes();
2606 let u_bytes = up.quant_bytes();
2607 let gpr = gate.cols() / GROUP_SIZE;
2608 let gsum = q1_group_sums(&act.xq, gpr);
2609 let gsum = &gsum;
2610 let run = move |start: usize, end: usize| {
2611 for r in start..end {
2612 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2613 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2614 for &(j, xv) in &act.outliers {
2615 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2616 gv += w * s * xv;
2617 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2618 uv += w * s * xv;
2619 }
2620 let silu_g = gv / (1.0 + (-gv).exp());
2621 unsafe { *out_addr.at(r) = silu_g * uv };
2623 }
2624 };
2625 dispatch_rows(pool, inter, &run);
2626 true
2627 }
2628 (
2632 Self::Mapped {
2633 dtype: TensorDtype::Q2TiledP,
2634 ..
2635 },
2636 Self::Mapped {
2637 dtype: TensorDtype::Q2TiledP,
2638 ..
2639 },
2640 ) => {
2641 let cols = gate.cols();
2642 let gpr = cols / GROUP_SIZE;
2643 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2644 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2645 let gsum = q1_group_sums(&act.xq, gpr);
2646 let gsum = &gsum;
2647 let run = move |start: usize, end: usize| {
2648 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2649 for r in start..end {
2650 gv_view.scales_into(r, gpr, &mut gsc);
2651 uv_view.scales_into(r, gpr, &mut usc);
2652 let mut gv =
2653 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2654 let mut uv =
2655 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2656 for &(j, xv) in &act.outliers {
2657 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2658 gv += w * s * xv;
2659 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2660 uv += w * s * xv;
2661 }
2662 let silu_g = gv / (1.0 + (-gv).exp());
2663 unsafe { *out_addr.at(r) = silu_g * uv };
2665 }
2666 };
2667 dispatch_rows(pool, inter, &run);
2668 true
2669 }
2670 (
2675 Self::Mapped {
2676 dtype: TensorDtype::Q8Row,
2677 row_scale: g_rs,
2678 ..
2679 },
2680 Self::Mapped {
2681 dtype: TensorDtype::Q8Row,
2682 row_scale: u_rs,
2683 ..
2684 },
2685 ) => {
2686 let g_bytes = gate.quant_bytes();
2687 let u_bytes = up.quant_bytes();
2688 let cols = gate.cols();
2689 let run = move |start: usize, end: usize| {
2690 for r in start..end {
2691 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2692 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2693 let silu_g = gv / (1.0 + (-gv).exp());
2694 unsafe { *out_addr.at(r) = silu_g * uv };
2696 }
2697 };
2698 dispatch_rows(pool, inter, &run);
2699 true
2700 }
2701 (
2703 Self::Mapped {
2704 dtype: TensorDtype::Q1T,
2705 ..
2706 },
2707 Self::Mapped {
2708 dtype: TensorDtype::Q1T,
2709 ..
2710 },
2711 ) => {
2712 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2713 let g_bytes = gate.quant_bytes();
2714 let u_bytes = up.quant_bytes();
2715 let gpr = gate.cols() / GROUP_SIZE;
2716 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2717 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2718 let run = move |start: usize, end: usize| {
2719 for r in start..end {
2720 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2721 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2722 for &(j, xv) in &act.outliers {
2723 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2724 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2725 }
2726 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2727 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2728 let silu_g = gv / (1.0 + (-gv).exp());
2729 unsafe { *out_addr.at(r) = silu_g * uv };
2731 }
2732 };
2733 dispatch_rows(pool, inter, &run);
2734 true
2735 }
2736 _ => false,
2737 }
2738 }
2739
2740 pub fn moe_gate_up_many(
2754 pairs: &[(&QTensor, &QTensor)],
2755 x: &[f32],
2756 outs: &mut [Vec<f32>],
2757 pool: Option<&Pool>,
2758 ) -> bool {
2759 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2760 return false;
2761 }
2762 let inter = pairs[0].0.rows();
2763 let cols = pairs[0].0.cols();
2764 if cols % GROUP_SIZE != 0 {
2765 return false;
2766 }
2767 let gpr = cols / GROUP_SIZE;
2768 let q2 = matches!(
2771 pairs[0].0,
2772 Self::Mapped {
2773 dtype: TensorDtype::Q2TiledP,
2774 ..
2775 }
2776 );
2777 let want = if q2 {
2778 TensorDtype::Q2TiledP
2779 } else {
2780 TensorDtype::Q4TiledP
2781 };
2782 let mut views = Vec::with_capacity(pairs.len() * 2);
2783 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2784 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2785 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2786 if !both
2787 || g.rows() != inter
2788 || u.rows() != inter
2789 || g.cols() != cols
2790 || u.cols() != cols
2791 || o.len() != inter
2792 {
2793 return false;
2794 }
2795 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2796 views.push(mk(g.quant_bytes(), inter, cols));
2797 views.push(mk(u.quant_bytes(), inter, cols));
2798 }
2799 let act = split_act(x);
2800 let gsum = if q2 {
2801 q1_group_sums(&act.xq, gpr)
2802 } else {
2803 Vec::new()
2804 };
2805 let (act, gsum) = (&act, &gsum);
2806 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2807 let (views, ptrs) = (&views, &ptrs);
2808 let run = |start: usize, end: usize| {
2809 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2810 for flat in start..end {
2811 let (e, r) = (flat / inter, flat % inter);
2812 let gv_view = &views[e * 2];
2813 let uv_view = &views[e * 2 + 1];
2814 gv_view.scales_into(r, gpr, &mut gsc);
2815 uv_view.scales_into(r, gpr, &mut usc);
2816 let (mut gv, mut uv) = if q2 {
2817 (
2818 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2819 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2820 )
2821 } else {
2822 (
2823 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2824 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2825 )
2826 };
2827 for &(j, xv) in &act.outliers {
2828 let (og, ou) = if q2 {
2829 (
2830 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2831 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2832 )
2833 } else {
2834 (
2835 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2836 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2837 )
2838 };
2839 gv += og.0 * og.1 * xv;
2840 uv += ou.0 * ou.1 * xv;
2841 }
2842 let silu_g = gv / (1.0 + (-gv).exp());
2843 unsafe { *ptrs[e].at(r) = silu_g * uv };
2845 }
2846 };
2847 dispatch_rows(pool, pairs.len() * inter, &run);
2848 true
2849 }
2850
2851 pub fn moe_down_many(
2860 downs: &[&QTensor],
2861 gs: &[Vec<f32>],
2862 weights: &[f32],
2863 out: &mut [f32],
2864 pool: Option<&Pool>,
2865 ) -> bool {
2866 if downs.is_empty()
2867 || downs.len() != gs.len()
2868 || downs.len() != weights.len()
2869 || !a8w8_enabled()
2870 {
2871 return false;
2872 }
2873 let rows = out.len();
2874 let cols = downs[0].cols();
2875 if cols % GROUP_SIZE != 0 {
2876 return false;
2877 }
2878 let gpr = cols / GROUP_SIZE;
2879 let mut views = Vec::with_capacity(downs.len());
2880 for (d, g) in downs.iter().zip(gs.iter()) {
2881 if !matches!(
2882 d,
2883 Self::Mapped {
2884 dtype: TensorDtype::Q4TiledP,
2885 ..
2886 }
2887 ) || d.rows() != rows
2888 || d.cols() != cols
2889 || g.len() != cols
2890 {
2891 return false;
2892 }
2893 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2894 }
2895 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2897 let out_addr = SendMut(out.as_mut_ptr());
2905 let (views, acts, weights) = (&views, &acts, &weights);
2906 let run = |start: usize, end: usize| {
2907 let mut sc = vec![0f32; gpr];
2908 for r in start..end {
2909 let mut acc = 0f32;
2910 for (e, v) in views.iter().enumerate() {
2911 v.scales_into(r, gpr, &mut sc);
2912 let a = &acts[e];
2913 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2914 for &(j, xv) in &a.outliers {
2915 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2916 d += w * s * xv;
2917 }
2918 acc += weights[e] * d;
2919 }
2920 unsafe { *out_addr.at(r) = acc };
2922 }
2923 };
2924 dispatch_rows(pool, rows, &run);
2925 true
2926 }
2927}
2928
2929#[cfg(target_os = "macos")]
2934mod accel_blas {
2935 #[link(name = "Accelerate", kind = "framework")]
2936 unsafe extern "C" {
2937 pub fn cblas_sgemm(
2938 order: i32,
2939 trans_a: i32,
2940 trans_b: i32,
2941 m: i32,
2942 n: i32,
2943 k: i32,
2944 alpha: f32,
2945 a: *const f32,
2946 lda: i32,
2947 b: *const f32,
2948 ldb: i32,
2949 beta: f32,
2950 c: *mut f32,
2951 ldc: i32,
2952 );
2953 }
2954}
2955
2956#[cfg(target_os = "macos")]
2957pub(crate) fn accel_gemm_enabled() -> bool {
2958 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2959 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2960}
2961
2962#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2965pub(crate) fn accel_gemm_enabled() -> bool {
2966 true
2967}
2968
2969#[cfg(target_arch = "aarch64")]
2976#[allow(clippy::too_many_arguments)]
2977pub(crate) fn neon_gemm_rm(
2978 m: usize,
2979 n: usize,
2980 k: usize,
2981 alpha: f32,
2982 a: &[f32],
2983 lda: usize,
2984 b_mat: &[f32],
2985 ldb: usize,
2986 b_rows_are_n: bool,
2987 c: &mut [f32],
2988 ldc: usize,
2989) {
2990 debug_assert!(a.len() >= (m - 1) * lda + k);
2991 debug_assert!(c.len() >= (m - 1) * ldc + n);
2992 unsafe {
2994 use core::arch::aarch64::*;
2995 let mut i = 0usize;
2996 while i < m {
2997 let mi = (m - i).min(4);
2998 let mut j = 0usize;
2999 while j < n {
3000 let nj = (n - j).min(8);
3001 if mi == 4 && nj == 8 {
3002 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3003 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3004 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3005 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3006 for p in 0..k {
3007 let (b0, b1) = if b_rows_are_n {
3008 let base = b_mat.as_ptr().add(j * ldb + p);
3011 let g = |o: usize| *base.add(o * ldb);
3012 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3013 } else {
3014 let base = b_mat.as_ptr().add(p * ldb + j);
3015 (
3016 [*base, *base.add(1), *base.add(2), *base.add(3)],
3017 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3018 )
3019 };
3020 let bv0 = vld1q_f32(b0.as_ptr());
3021 let bv1 = vld1q_f32(b1.as_ptr());
3022 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3023 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3024 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3025 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3026 c0a = vfmaq_f32(c0a, a0, bv0);
3027 c0b = vfmaq_f32(c0b, a0, bv1);
3028 c1a = vfmaq_f32(c1a, a1, bv0);
3029 c1b = vfmaq_f32(c1b, a1, bv1);
3030 c2a = vfmaq_f32(c2a, a2, bv0);
3031 c2b = vfmaq_f32(c2b, a2, bv1);
3032 c3a = vfmaq_f32(c3a, a3, bv0);
3033 c3b = vfmaq_f32(c3b, a3, bv1);
3034 }
3035 let al = vdupq_n_f32(alpha);
3036 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3037 .iter()
3038 .enumerate()
3039 {
3040 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3041 vst1q_f32(dst, vmulq_f32(*ca, al));
3042 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3043 }
3044 } else {
3045 for r in 0..mi {
3046 for q in 0..nj {
3047 let mut acc = 0f32;
3048 for p in 0..k {
3049 let bv = if b_rows_are_n {
3050 b_mat[(j + q) * ldb + p]
3051 } else {
3052 b_mat[p * ldb + j + q]
3053 };
3054 acc += a[(i + r) * lda + p] * bv;
3055 }
3056 c[(i + r) * ldc + j + q] = acc * alpha;
3057 }
3058 }
3059 }
3060 j += nj;
3061 }
3062 i += mi;
3063 }
3064 }
3065}
3066
3067#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3069#[allow(clippy::too_many_arguments)]
3070pub(crate) fn sgemm_rm(
3071 m: usize,
3072 n: usize,
3073 k: usize,
3074 alpha: f32,
3075 a: &[f32],
3076 lda: usize,
3077 b_mat: &[f32],
3078 ldb: usize,
3079 b_rows_are_n: bool,
3080 c: &mut [f32],
3081 ldc: usize,
3082) {
3083 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3084}
3085
3086#[allow(clippy::too_many_arguments)]
3090pub fn sgemm_public(
3091 m: usize,
3092 n: usize,
3093 k: usize,
3094 alpha: f32,
3095 a: &[f32],
3096 lda: usize,
3097 b_mat: &[f32],
3098 ldb: usize,
3099 b_rows_are_n: bool,
3100 c: &mut [f32],
3101 ldc: usize,
3102) {
3103 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3104 {
3105 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3106 }
3107 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3112 {
3113 for i in 0..m {
3114 for j in 0..n {
3115 let mut acc = 0f32;
3116 for p in 0..k {
3117 let bv = if b_rows_are_n {
3118 b_mat[j * ldb + p]
3119 } else {
3120 b_mat[p * ldb + j]
3121 };
3122 acc += a[i * lda + p] * bv;
3123 }
3124 c[i * ldc + j] = alpha * acc;
3125 }
3126 }
3127 }
3128}
3129
3130#[cfg(target_os = "macos")]
3133#[allow(clippy::too_many_arguments)]
3134pub(crate) fn sgemm_rm(
3135 m: usize,
3136 n: usize,
3137 k: usize,
3138 alpha: f32,
3139 a: &[f32],
3140 lda: usize,
3141 b_mat: &[f32],
3142 ldb: usize,
3143 b_rows_are_n: bool,
3144 c: &mut [f32],
3145 ldc: usize,
3146) {
3147 debug_assert!(a.len() >= (m - 1) * lda + k);
3148 debug_assert!(c.len() >= (m - 1) * ldc + n);
3149 #[cfg(target_arch = "aarch64")]
3154 if std::env::var("CMF_FORCE_NEON_GEMM")
3155 .map(|v| v == "1")
3156 .unwrap_or(false)
3157 {
3158 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3159 }
3160 unsafe {
3161 accel_blas::cblas_sgemm(
3162 101, 111, if b_rows_are_n { 112 } else { 111 },
3165 m as i32,
3166 n as i32,
3167 k as i32,
3168 alpha,
3169 a.as_ptr(),
3170 lda as i32,
3171 b_mat.as_ptr(),
3172 ldb as i32,
3173 0.0,
3174 c.as_mut_ptr(),
3175 ldc as i32,
3176 );
3177 }
3178}
3179
3180#[cfg(target_os = "macos")]
3187fn qmatmat_accel(
3188 q: &[u8],
3189 row_scale: &[f32],
3190 pre: &[std::borrow::Cow<'_, [f32]>],
3191 rows: usize,
3192 cols: usize,
3193 out: &mut [f32],
3194 pool: Option<&Pool>,
3195) {
3196 const TR: usize = 2048;
3201 let b = pre.len();
3202 thread_local! {
3203 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3204 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3205 }
3206 XPANEL.with(|xp| {
3207 WTILE.with(|wt| {
3208 let mut xpanel = xp.borrow_mut();
3209 xpanel.clear();
3210 for x in pre {
3211 xpanel.extend_from_slice(x);
3212 }
3213 let mut wtile = wt.borrow_mut();
3214 wtile.resize(TR * cols, 0.0);
3215 let mut r0 = 0usize;
3216 while r0 < rows {
3217 let tr = TR.min(rows - r0);
3218 let wt_addr = SendMut(wtile.as_mut_ptr());
3220 let run = |start: usize, end: usize| {
3221 for r in start..end {
3222 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3223 let s = row_scale[r0 + r];
3224 let dst =
3226 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3227 for (d, &v) in dst.iter_mut().zip(row) {
3228 *d = (v as i8) as f32 * s;
3229 }
3230 }
3231 };
3232 dispatch_rows(pool, tr, &run);
3233 unsafe {
3235 accel_blas::cblas_sgemm(
3236 101, 111, 112, b as i32,
3240 tr as i32,
3241 cols as i32,
3242 1.0,
3243 xpanel.as_ptr(),
3244 cols as i32,
3245 wtile.as_ptr(),
3246 cols as i32,
3247 0.0,
3248 out.as_mut_ptr().add(r0),
3249 rows as i32,
3250 );
3251 }
3252 r0 += tr;
3253 }
3254 })
3255 });
3256}
3257
3258fn qmatmat(
3259 q: &[u8],
3260 row_scale: &[f32],
3261 pre: &[std::borrow::Cow<'_, [f32]>],
3262 rows: usize,
3263 cols: usize,
3264 out: &mut [f32],
3265 pool: Option<&Pool>,
3266) {
3267 let b = pre.len();
3268 debug_assert_eq!(out.len(), b * rows);
3269 #[cfg(target_os = "macos")]
3274 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3275 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3276 return;
3277 }
3278 #[cfg(target_arch = "aarch64")]
3279 if sdot_enabled() {
3280 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3281 let out_addr = SendMut(out.as_mut_ptr());
3282 let blocked_ok = blocked_enabled();
3285 let use_i8mm = i8mm_enabled();
3286 if blocked_ok {
3287 let run = |start: usize, end: usize| {
3288 let mut o = start;
3289 while o < end {
3290 if o + 2 <= end {
3291 let r0 = &q[o * cols..(o + 1) * cols];
3292 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3293 let mut bi = 0usize;
3294 while bi + 4 <= acts.len() {
3295 let xs = [
3296 acts[bi].xq.as_slice(),
3297 acts[bi + 1].xq.as_slice(),
3298 acts[bi + 2].xq.as_slice(),
3299 acts[bi + 3].xq.as_slice(),
3300 ];
3301 let d = if use_i8mm {
3302 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3303 } else {
3304 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3305 };
3306 for (r, row) in [r0, r1].into_iter().enumerate() {
3307 for k in 0..4 {
3308 let act = &acts[bi + k];
3309 let mut v = d[r][k] as f32 * act.sx;
3310 for &(j, xv) in &act.outliers {
3311 v += (row[j] as i8) as f32 * xv;
3312 }
3313 unsafe {
3314 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3315 };
3316 }
3317 }
3318 bi += 4;
3319 }
3320 while bi < acts.len() {
3321 for (r, row) in [r0, r1].into_iter().enumerate() {
3322 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3323 unsafe { *out_addr.at(bi * rows + o + r) = v };
3324 }
3325 bi += 1;
3326 }
3327 o += 2;
3328 } else {
3329 let row = &q[o * cols..(o + 1) * cols];
3330 for (bi, act) in acts.iter().enumerate() {
3331 let v = row_dot_sdot(row, act) * row_scale[o];
3332 unsafe { *out_addr.at(bi * rows + o) = v };
3333 }
3334 o += 1;
3335 }
3336 }
3337 };
3338 dispatch_rows(pool, rows, &run);
3339 return;
3340 }
3341 let run = |start: usize, end: usize| {
3342 for o in start..end {
3343 let row = &q[o * cols..(o + 1) * cols];
3344 for (bi, act) in acts.iter().enumerate() {
3345 let v = row_dot_sdot(row, act) * row_scale[o];
3346 unsafe { *out_addr.at(bi * rows + o) = v };
3347 }
3348 }
3349 };
3350 dispatch_rows(pool, rows, &run);
3351 return;
3352 }
3353 #[cfg(target_arch = "x86_64")]
3358 if avx2_a8w8_enabled() {
3359 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3360 let out_addr = SendMut(out.as_mut_ptr());
3361 let blocked_ok = blocked_enabled();
3364 if !avx512vnni_enabled() && blocked_ok {
3365 let run = |start: usize, end: usize| {
3366 let mut o = start;
3367 while o < end {
3368 if o + 2 <= end {
3369 let r0 = &q[o * cols..(o + 1) * cols];
3370 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3371 let mut bi = 0usize;
3372 while bi + 4 <= acts.len() {
3373 let xs = [
3374 acts[bi].xq.as_slice(),
3375 acts[bi + 1].xq.as_slice(),
3376 acts[bi + 2].xq.as_slice(),
3377 acts[bi + 3].xq.as_slice(),
3378 ];
3379 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3380 for (r, row) in [r0, r1].into_iter().enumerate() {
3381 for k in 0..4 {
3382 let act = &acts[bi + k];
3383 let mut v = d[r][k] as f32 * act.sx;
3384 for &(j, xv) in &act.outliers {
3385 v += (row[j] as i8) as f32 * xv;
3386 }
3387 unsafe {
3388 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3389 };
3390 }
3391 }
3392 bi += 4;
3393 }
3394 while bi < acts.len() {
3395 for (r, row) in [r0, r1].into_iter().enumerate() {
3396 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3397 unsafe { *out_addr.at(bi * rows + o + r) = v };
3398 }
3399 bi += 1;
3400 }
3401 o += 2;
3402 } else {
3403 let row = &q[o * cols..(o + 1) * cols];
3404 for (bi, act) in acts.iter().enumerate() {
3405 let v = row_dot_avx2(row, act) * row_scale[o];
3406 unsafe { *out_addr.at(bi * rows + o) = v };
3407 }
3408 o += 1;
3409 }
3410 }
3411 };
3412 dispatch_rows(pool, rows, &run);
3413 return;
3414 }
3415 let run = |start: usize, end: usize| {
3416 for o in start..end {
3417 let row = &q[o * cols..(o + 1) * cols];
3418 for (bi, act) in acts.iter().enumerate() {
3419 let v = row_dot_avx2(row, act) * row_scale[o];
3420 unsafe { *out_addr.at(bi * rows + o) = v };
3421 }
3422 }
3423 };
3424 dispatch_rows(pool, rows, &run);
3425 return;
3426 }
3427 let out_addr = SendMut(out.as_mut_ptr());
3428 let run = |start: usize, end: usize| {
3429 for o in start..end {
3430 let row = &q[o * cols..(o + 1) * cols];
3431 for (bi, x) in pre.iter().enumerate() {
3432 let mut acc = 0f32;
3433 for j in 0..cols {
3434 acc += (row[j] as i8) as f32 * x[j];
3435 }
3436 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3437 }
3438 }
3439 };
3440 dispatch_rows(pool, rows, &run);
3441}
3442
3443fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3446 match pool {
3447 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3448 _ => run(0, rows),
3449 }
3450}
3451
3452fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3454 let groups = rows * cols / GROUP_SIZE;
3455 bytes.split_at(groups * 16)
3456}
3457
3458#[inline]
3463fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3464 #[cfg(target_arch = "aarch64")]
3465 unsafe {
3466 return vbit_fill4_neon(data, buf);
3467 }
3468 #[cfg(target_arch = "x86_64")]
3469 if avx2_enabled() {
3470 return unsafe { vbit_fill4_avx2(data, buf) };
3471 }
3472 #[allow(unreachable_code)]
3473 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3474 let u = unpack8::<4>(&data[blk * 4..]);
3475 for k in 0..8 {
3476 chunk[k] = (u[k] - 7) as i8 as u8;
3477 }
3478 }
3479}
3480
3481#[cfg(target_arch = "aarch64")]
3482#[target_feature(enable = "neon")]
3483unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3484 unsafe {
3487 use core::arch::aarch64::*;
3488 let n = buf.len();
3489 let mask = vdupq_n_u8(0x0F);
3490 let seven = vdupq_n_s8(7);
3491 let mut g = 0usize;
3492 while g * 32 + 32 <= n {
3493 let b = vld1q_u8(data.as_ptr().add(g * 16));
3494 let hi = vshrq_n_u8::<4>(b);
3495 let lo = vandq_u8(b, mask);
3496 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3497 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3498 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3499 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3500 g += 1;
3501 }
3502 }
3503}
3504
3505#[cfg(target_arch = "x86_64")]
3506#[target_feature(enable = "avx2")]
3507unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3508 unsafe {
3510 use core::arch::x86_64::*;
3511 let n = buf.len();
3512 let mask = _mm_set1_epi8(0x0F);
3513 let seven = _mm256_set1_epi8(7);
3514 let mut g = 0usize;
3515 while g * 32 + 32 <= n {
3516 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3517 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3518 let lo = _mm_and_si128(b, mask);
3519 let z = _mm256_sub_epi8(
3520 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3521 seven,
3522 );
3523 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3524 g += 1;
3525 }
3526 }
3527}
3528
3529#[inline(always)]
3534fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3535 let mut acc = 0u64;
3536 for i in 0..B {
3537 acc = (acc << 8) | data[i] as u64;
3538 }
3539 let mask = (1u64 << B) - 1;
3540 let mut out = [0i32; 8];
3541 for (k, o) in out.iter_mut().enumerate() {
3542 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3543 }
3544 out
3545}
3546
3547#[allow(clippy::too_many_arguments)]
3553fn vbitmatvec(
3554 bytes: &[u8],
3555 offsets: &[usize],
3556 x: &[f32],
3557 rows: usize,
3558 cols: usize,
3559 out: &mut [f32],
3560 pool: Option<&Pool>,
3561) {
3562 debug_assert_eq!(out.len(), rows);
3563 debug_assert_eq!(offsets.len(), rows + 1);
3564
3565 if a8w8_enabled() {
3569 let act = split_act(x);
3570 let out_addr = SendMut(out.as_mut_ptr());
3571 let run = move |start: usize, end: usize| {
3572 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3573 };
3574 dispatch_rows(pool, rows, &run);
3575 return;
3576 }
3577
3578 let out_addr = SendMut(out.as_mut_ptr());
3579 let run = move |start: usize, end: usize| {
3580 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3581 };
3582 dispatch_rows(pool, rows, &run);
3583}
3584
3585#[allow(clippy::too_many_arguments)]
3589fn vbit_range_a8w8(
3590 bytes: &[u8],
3591 offsets: &[usize],
3592 x: &[f32],
3593 act: &SplitAct,
3594 rows: usize,
3595 cols: usize,
3596 out: SendMut,
3597 start: usize,
3598 end: usize,
3599) {
3600 let ng = cols / GROUP_SIZE;
3601 let bits = &bytes[..rows];
3602 let sc_off = rows;
3603 let row_dot = |r: usize| -> f32 {
3604 let b = bits[r] as usize;
3605 let l = (1i32 << (b - 1)) - 1;
3606 let mask = (1u64 << b) - 1;
3607 let data = &bytes[offsets[r]..offsets[r + 1]];
3608 if b == 8 {
3609 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3611 let mut dot = 0f32;
3612 for g in 0..ng {
3613 let so = (r * ng + g) * 2;
3614 let sgf = f16_to_f32(u16::from_le_bytes([
3615 bytes[sc_off + so],
3616 bytes[sc_off + so + 1],
3617 ]));
3618 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3619 let mut gd = 0f32;
3620 for &xv in xg.iter() {
3621 if nbits < 8 {
3622 acc = (acc << 8) | data[idx] as u64;
3623 idx += 1;
3624 nbits += 8;
3625 }
3626 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3627 nbits -= 8;
3628 gd += (u - l) as f32 * xv;
3629 }
3630 dot += gd * sgf;
3631 }
3632 return dot;
3633 }
3634 thread_local! {
3638 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3639 const { std::cell::RefCell::new(Vec::new()) };
3640 }
3641 #[inline(always)]
3642 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3643 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3644 let u = unpack8::<B>(&data[blk * B..]);
3645 for k in 0..8 {
3646 chunk[k] = (u[k] - l) as i8 as u8;
3647 }
3648 }
3649 }
3650 let _ = mask;
3651 VBIT_SCRATCH.with(|scratch| {
3652 let mut buf = scratch.borrow_mut();
3653 buf.resize(cols, 0);
3654 match b {
3655 3 => fill::<3>(data, l, &mut buf),
3656 4 => vbit_fill4(data, &mut buf),
3657 5 => fill::<5>(data, l, &mut buf),
3658 6 => fill::<6>(data, l, &mut buf),
3659 _ => unreachable!(),
3660 }
3661 let mut dot = 0f32;
3662 for g in 0..ng {
3663 let so = (r * ng + g) * 2;
3664 let s = f16_to_f32(u16::from_le_bytes([
3665 bytes[sc_off + so],
3666 bytes[sc_off + so + 1],
3667 ]));
3668 let d = dot_i8_i8(
3669 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3670 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3671 ) as f32
3672 * act.sx;
3673 dot += d * s;
3674 }
3675 for &(j, xv) in &act.outliers {
3676 let so = (r * ng + j / GROUP_SIZE) * 2;
3677 let s = f16_to_f32(u16::from_le_bytes([
3678 bytes[sc_off + so],
3679 bytes[sc_off + so + 1],
3680 ]));
3681 dot += (buf[j] as i8) as f32 * s * xv;
3683 }
3684 dot
3685 })
3686 };
3687 for r in start..end {
3688 unsafe { *out.at(r) = row_dot(r) };
3690 }
3691}
3692
3693#[allow(clippy::too_many_arguments)]
3695fn vbit_range_f32(
3696 bytes: &[u8],
3697 offsets: &[usize],
3698 x: &[f32],
3699 rows: usize,
3700 cols: usize,
3701 out: SendMut,
3702 start: usize,
3703 end: usize,
3704) {
3705 let ng = cols / GROUP_SIZE;
3706 let bits = &bytes[..rows];
3707 let sc_off = rows;
3708 #[inline(always)]
3712 fn dot_row<const B: usize>(
3713 data: &[u8],
3714 bytes: &[u8],
3715 sc_off: usize,
3716 r: usize,
3717 ng: usize,
3718 x: &[f32],
3719 ) -> f32 {
3720 let l = ((1i32 << (B - 1)) - 1) as f32;
3721 let gbytes = GROUP_SIZE * B / 8;
3722 let mut dot = 0f32;
3723 for g in 0..ng {
3724 let so = (r * ng + g) * 2;
3725 let s = f16_to_f32(u16::from_le_bytes([
3726 bytes[sc_off + so],
3727 bytes[sc_off + so + 1],
3728 ]));
3729 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3730 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3731 let mut gd = 0f32;
3732 for blk in 0..GROUP_SIZE / 8 {
3733 let u = unpack8::<B>(&gd0[blk * B..]);
3734 let xb = &xg[blk * 8..blk * 8 + 8];
3735 for k in 0..8 {
3736 gd += (u[k] as f32 - l) * xb[k];
3737 }
3738 }
3739 dot += gd * s;
3740 }
3741 dot
3742 }
3743 for r in start..end {
3744 let data = &bytes[offsets[r]..offsets[r + 1]];
3745 let v = match bits[r] {
3746 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3747 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3748 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3749 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3750 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3751 b => unreachable!("vbit bit-width {b} (validated at load)"),
3752 };
3753 unsafe { *out.at(r) = v };
3755 }
3756}
3757
3758#[allow(clippy::too_many_arguments)]
3763fn vbitmatvec2(
3764 bytes: &[u8],
3765 offsets: &[usize],
3766 x1: &[f32],
3767 x2: &[f32],
3768 rows: usize,
3769 cols: usize,
3770 o1: &mut [f32],
3771 o2: &mut [f32],
3772 pool: Option<&Pool>,
3773) {
3774 debug_assert_eq!(o1.len(), rows);
3775 debug_assert_eq!(o2.len(), rows);
3776
3777 if a8w8_enabled() {
3778 let a1 = split_act(x1);
3779 let a2 = split_act(x2);
3780 let p1 = SendMut(o1.as_mut_ptr());
3781 let p2 = SendMut(o2.as_mut_ptr());
3782 let run = move |start: usize, end: usize| {
3783 vbit_range2_a8w8(
3784 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3785 )
3786 };
3787 dispatch_rows(pool, rows, &run);
3788 return;
3789 }
3790
3791 let p1 = SendMut(o1.as_mut_ptr());
3792 let p2 = SendMut(o2.as_mut_ptr());
3793 let run = move |start: usize, end: usize| {
3794 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3795 };
3796 dispatch_rows(pool, rows, &run);
3797}
3798
3799#[allow(clippy::too_many_arguments)]
3803fn vbit_range2_a8w8(
3804 bytes: &[u8],
3805 offsets: &[usize],
3806 x1: &[f32],
3807 x2: &[f32],
3808 a1: &SplitAct,
3809 a2: &SplitAct,
3810 rows: usize,
3811 cols: usize,
3812 p1: SendMut,
3813 p2: SendMut,
3814 start: usize,
3815 end: usize,
3816) {
3817 let ng = cols / GROUP_SIZE;
3818 let bits = &bytes[..rows];
3819 let sc_off = rows;
3820 let row_dots = |r: usize| -> (f32, f32) {
3821 let b = bits[r] as usize;
3822 let l = (1i32 << (b - 1)) - 1;
3823 let data = &bytes[offsets[r]..offsets[r + 1]];
3824 if b == 8 {
3825 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3828 let (mut d1, mut d2) = (0f32, 0f32);
3829 for g in 0..ng {
3830 let so = (r * ng + g) * 2;
3831 let sgf = f16_to_f32(u16::from_le_bytes([
3832 bytes[sc_off + so],
3833 bytes[sc_off + so + 1],
3834 ]));
3835 let (mut g1, mut g2) = (0f32, 0f32);
3836 for k in 0..GROUP_SIZE {
3837 if nbits < 8 {
3838 acc = (acc << 8) | data[idx] as u64;
3839 idx += 1;
3840 nbits += 8;
3841 }
3842 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3843 nbits -= 8;
3844 let w = (u - l) as f32;
3845 g1 += w * x1[g * GROUP_SIZE + k];
3846 g2 += w * x2[g * GROUP_SIZE + k];
3847 }
3848 d1 += g1 * sgf;
3849 d2 += g2 * sgf;
3850 }
3851 return (d1, d2);
3852 }
3853 thread_local! {
3854 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3855 const { std::cell::RefCell::new(Vec::new()) };
3856 }
3857 #[inline(always)]
3858 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3859 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3860 let u = unpack8::<B>(&data[blk * B..]);
3861 for k in 0..8 {
3862 chunk[k] = (u[k] - l) as i8 as u8;
3863 }
3864 }
3865 }
3866 VBIT_SCRATCH2.with(|scratch| {
3867 let mut buf = scratch.borrow_mut();
3868 buf.resize(cols, 0);
3869 match b {
3870 3 => fill::<3>(data, l, &mut buf),
3871 4 => vbit_fill4(data, &mut buf),
3872 5 => fill::<5>(data, l, &mut buf),
3873 6 => fill::<6>(data, l, &mut buf),
3874 _ => unreachable!(),
3875 }
3876 let (mut d1, mut d2) = (0f32, 0f32);
3877 for g in 0..ng {
3878 let so = (r * ng + g) * 2;
3879 let s = f16_to_f32(u16::from_le_bytes([
3880 bytes[sc_off + so],
3881 bytes[sc_off + so + 1],
3882 ]));
3883 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3884 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3885 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3886 d1 += v1 * s;
3887 d2 += v2 * s;
3888 }
3889 for &(j, xv) in &a1.outliers {
3890 let so = (r * ng + j / GROUP_SIZE) * 2;
3891 let s = f16_to_f32(u16::from_le_bytes([
3892 bytes[sc_off + so],
3893 bytes[sc_off + so + 1],
3894 ]));
3895 d1 += (buf[j] as i8) as f32 * s * xv;
3896 }
3897 for &(j, xv) in &a2.outliers {
3898 let so = (r * ng + j / GROUP_SIZE) * 2;
3899 let s = f16_to_f32(u16::from_le_bytes([
3900 bytes[sc_off + so],
3901 bytes[sc_off + so + 1],
3902 ]));
3903 d2 += (buf[j] as i8) as f32 * s * xv;
3904 }
3905 (d1, d2)
3906 })
3907 };
3908 for r in start..end {
3909 let (v1, v2) = row_dots(r);
3910 unsafe {
3912 *p1.at(r) = v1;
3913 *p2.at(r) = v2;
3914 }
3915 }
3916}
3917
3918#[allow(clippy::too_many_arguments)]
3922fn vbit_range2_f32(
3923 bytes: &[u8],
3924 offsets: &[usize],
3925 x1: &[f32],
3926 x2: &[f32],
3927 rows: usize,
3928 cols: usize,
3929 p1: SendMut,
3930 p2: SendMut,
3931 start: usize,
3932 end: usize,
3933) {
3934 let ng = cols / GROUP_SIZE;
3935 let bits = &bytes[..rows];
3936 let sc_off = rows;
3937 #[inline(always)]
3938 #[allow(clippy::too_many_arguments)]
3939 fn dot_row2<const B: usize>(
3940 data: &[u8],
3941 bytes: &[u8],
3942 sc_off: usize,
3943 r: usize,
3944 ng: usize,
3945 x1: &[f32],
3946 x2: &[f32],
3947 ) -> (f32, f32) {
3948 let l = ((1i32 << (B - 1)) - 1) as f32;
3949 let gbytes = GROUP_SIZE * B / 8;
3950 let (mut d1, mut d2) = (0f32, 0f32);
3951 for g in 0..ng {
3952 let so = (r * ng + g) * 2;
3953 let s = f16_to_f32(u16::from_le_bytes([
3954 bytes[sc_off + so],
3955 bytes[sc_off + so + 1],
3956 ]));
3957 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3958 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3959 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3960 let (mut g1, mut g2) = (0f32, 0f32);
3961 for blk in 0..GROUP_SIZE / 8 {
3962 let u = unpack8::<B>(&gd0[blk * B..]);
3963 for k in 0..8 {
3964 let w = u[k] as f32 - l;
3965 g1 += w * x1g[blk * 8 + k];
3966 g2 += w * x2g[blk * 8 + k];
3967 }
3968 }
3969 d1 += g1 * s;
3970 d2 += g2 * s;
3971 }
3972 (d1, d2)
3973 }
3974 for r in start..end {
3975 let data = &bytes[offsets[r]..offsets[r + 1]];
3976 let (v1, v2) = match bits[r] {
3977 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3978 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3979 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3980 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3981 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3982 b => unreachable!("vbit bit-width {b} (validated at load)"),
3983 };
3984 unsafe {
3986 *p1.at(r) = v1;
3987 *p2.at(r) = v2;
3988 }
3989 }
3990}
3991
3992#[inline]
3999#[allow(unreachable_code)]
4000fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4001 #[cfg(target_arch = "aarch64")]
4002 unsafe {
4003 return dot_q4t_row_sdot(bytes, r, gpr, xq);
4004 }
4005 #[cfg(target_arch = "x86_64")]
4006 unsafe {
4007 if vnni_tiles_enabled() {
4008 return dot_q4t_row_vnni(bytes, r, gpr, xq);
4009 }
4010 return dot_q4t_row_avx2(bytes, r, gpr, xq);
4011 }
4012 let mut acc = 0f32;
4013 for gi in 0..gpr {
4014 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4015 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4016 let mut d = 0i32;
4017 for (k, &b) in tile[2..].iter().enumerate() {
4018 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4019 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4020 }
4021 acc += d as f32 * s;
4022 }
4023 acc
4024}
4025
4026#[cfg(target_arch = "aarch64")]
4027#[target_feature(enable = "neon,dotprod")]
4028unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4029 unsafe {
4032 use core::arch::aarch64::*;
4033 use core::arch::asm;
4034 let lomask = vdupq_n_u8(0x0F);
4035 let eight = vdupq_n_s8(8);
4036 let mut acc = 0f32;
4037 for gi in 0..gpr {
4038 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4039 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4040 let b = vld1q_u8(t.add(2));
4041 let lo = vandq_u8(b, lomask);
4042 let hi = vshrq_n_u8::<4>(b);
4043 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4044 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4045 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4046 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4047 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4048 asm!(
4049 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4050 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4051 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4052 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4053 options(pure, nomem, nostack),
4054 );
4055 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4056 }
4057 acc
4058 }
4059}
4060
4061#[cfg(target_arch = "x86_64")]
4062#[target_feature(enable = "avx2")]
4063unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4064 unsafe {
4066 use core::arch::x86_64::*;
4067 let lomask = _mm_set1_epi8(0x0F);
4068 let eight = _mm256_set1_epi8(8);
4069 let ones = _mm256_set1_epi16(1);
4070 let mut acc = 0f32;
4071 for gi in 0..gpr {
4072 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4073 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4074 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4075 let lo = _mm_and_si128(b, lomask);
4076 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4077 let w = _mm256_sub_epi8(
4078 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4079 eight,
4080 );
4081 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4082 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4083 let d = _mm256_madd_epi16(p16, ones);
4084 let hi128 = _mm256_extracti128_si256::<1>(d);
4085 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4086 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4087 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4088 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4089 }
4090 acc
4091 }
4092}
4093
4094#[cfg(target_arch = "x86_64")]
4098#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4099unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4100 unsafe {
4102 use core::arch::x86_64::*;
4103 let lomask = _mm_set1_epi8(0x0F);
4104 let eight = _mm256_set1_epi8(8);
4105 let mut acc = 0f32;
4106 for gi in 0..gpr {
4107 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4108 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4109 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4110 let lo = _mm_and_si128(b, lomask);
4111 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4112 let w = _mm256_sub_epi8(
4113 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4114 eight,
4115 );
4116 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4117 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4118 acc += d as f32 * s;
4119 }
4120 acc
4121 }
4122}
4123
4124#[cfg(target_arch = "x86_64")]
4129#[target_feature(enable = "avx2,fma")]
4134unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4135 unsafe {
4137 use core::arch::x86_64::*;
4138 let lomask = _mm_set1_epi8(0x0F);
4139 let eight = _mm256_set1_epi8(8);
4140 let ones = _mm256_set1_epi16(1);
4141 let mut f0 = _mm256_setzero_ps();
4154 let mut f1 = _mm256_setzero_ps();
4155 let mut f2 = _mm256_setzero_ps();
4156 let mut f3 = _mm256_setzero_ps();
4157 for gi in 0..gpr {
4158 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4159 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4160 let sv = _mm256_set1_ps(s);
4161 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4162 let lo = _mm_and_si128(bb, lomask);
4163 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4164 let w = _mm256_sub_epi8(
4165 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4166 eight,
4167 );
4168 let aw = _mm256_abs_epi8(w);
4169 let off = gi * GROUP_SIZE;
4170 let dot = |xq: &[i8]| {
4171 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4172 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4173 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4174 };
4175 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4176 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4177 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4178 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4179 }
4180 [
4181 hsum256_ps(f0),
4182 hsum256_ps(f1),
4183 hsum256_ps(f2),
4184 hsum256_ps(f3),
4185 ]
4186 }
4187}
4188
4189#[cfg(target_arch = "x86_64")]
4192#[target_feature(enable = "avx2")]
4193#[inline]
4194unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4195 unsafe {
4197 use core::arch::x86_64::*;
4198 let hi = _mm256_extractf128_ps::<1>(v);
4199 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4200 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4201 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4202 _mm_cvtss_f32(s)
4203 }
4204}
4205
4206#[cfg(target_arch = "x86_64")]
4208#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4209unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4210 unsafe {
4212 use core::arch::x86_64::*;
4213 let lomask = _mm_set1_epi8(0x0F);
4214 let eight = _mm256_set1_epi8(8);
4215 let mut f0 = _mm256_setzero_ps();
4218 let mut f1 = _mm256_setzero_ps();
4219 let mut f2 = _mm256_setzero_ps();
4220 let mut f3 = _mm256_setzero_ps();
4221 for gi in 0..gpr {
4222 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4223 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4224 let sv = _mm256_set1_ps(s);
4225 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4226 let lo = _mm_and_si128(bb, lomask);
4227 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4228 let w = _mm256_sub_epi8(
4229 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4230 eight,
4231 );
4232 let aw = _mm256_abs_epi8(w);
4233 let off = gi * GROUP_SIZE;
4234 let dot = |xq: &[i8]| {
4235 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4236 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4237 _mm256_setzero_si256(),
4238 aw,
4239 _mm256_sign_epi8(x, w),
4240 ))
4241 };
4242 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4243 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4244 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4245 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4246 }
4247 let acc = [
4248 hsum256_ps(f0),
4249 hsum256_ps(f1),
4250 hsum256_ps(f2),
4251 hsum256_ps(f3),
4252 ];
4253 acc
4254 }
4255}
4256
4257#[cfg(target_arch = "aarch64")]
4262#[target_feature(enable = "neon,dotprod")]
4263unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4264 unsafe {
4266 use core::arch::aarch64::*;
4267 use core::arch::asm;
4268 let lomask = vdupq_n_u8(0x0F);
4269 let eight = vdupq_n_s8(8);
4270 let mut acc = [0f32; 4];
4271 for gi in 0..gpr {
4272 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4273 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4274 let b = vld1q_u8(t.add(2));
4275 let lo = vandq_u8(b, lomask);
4276 let hi = vshrq_n_u8::<4>(b);
4277 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4278 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4279 for (k, xq) in xs.iter().enumerate() {
4280 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4281 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4282 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4283 asm!(
4284 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4285 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4286 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4287 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4288 options(pure, nomem, nostack),
4289 );
4290 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4291 }
4292 }
4293 acc
4294 }
4295}
4296
4297#[inline]
4299fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4300 let gi = j / GROUP_SIZE;
4301 let k = j % GROUP_SIZE;
4302 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4303 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4304 let byte = tile[2 + k / 2];
4305 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4306 ((nib as i32 - 8) as f32, s)
4307}
4308
4309#[inline]
4312fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4313 let mut acc = 0f32;
4314 for gi in 0..gpr {
4315 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4316 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4317 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4318 let mut ga = 0f32;
4319 for (k, &b) in tile[2..].iter().enumerate() {
4320 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4321 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4322 }
4323 acc += ga * s;
4324 }
4325 acc
4326}
4327
4328struct Q4tpView<'a> {
4332 nib: &'a [u8],
4333 params: &'a [u8],
4334 codes: &'a [u8],
4335 stride: usize,
4336 zero_rung: bool,
4338}
4339
4340impl<'a> Q4tpView<'a> {
4341 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4342 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4343 Self {
4344 nib: &bytes[..params_off],
4345 params: &bytes[params_off..codes_off],
4346 codes: &bytes[codes_off..],
4347 stride,
4348 zero_rung: false,
4349 }
4350 }
4351
4352 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4354 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4355 Self {
4356 nib: &bytes[..params_off],
4357 params: &bytes[params_off..codes_off],
4358 codes: &bytes[codes_off..],
4359 stride,
4360 zero_rung: true,
4361 }
4362 }
4363
4364 #[inline]
4380 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4381 let tab = if self.zero_rung {
4382 q2tp_ladder(self.params, r)
4383 } else {
4384 q4tp_ladder(self.params, r)
4385 };
4386 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4387 let out = &mut out[..gpr];
4388 let mut chunks = out.chunks_exact_mut(8);
4389 let mut ci = 0usize;
4390 for c in &mut chunks {
4391 let w = u64::from(codes[ci])
4392 | u64::from(codes[ci + 1]) << 8
4393 | u64::from(codes[ci + 2]) << 16
4394 | u64::from(codes[ci + 3]) << 24
4395 | u64::from(codes[ci + 4]) << 32;
4396 for (k, o) in c.iter_mut().enumerate() {
4397 *o = tab[((w >> (5 * k)) & 31) as usize];
4398 }
4399 ci += 5;
4400 }
4401 let tail = &codes[ci..];
4404 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4405 *o = tab[q4tp_code(tail, k)];
4406 }
4407 }
4408}
4409
4410#[inline]
4411fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4412 #[cfg(target_arch = "aarch64")]
4413 unsafe {
4414 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4415 }
4416 #[cfg(target_arch = "x86_64")]
4417 unsafe {
4418 if vnni_tiles_enabled() {
4419 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4420 }
4421 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4422 }
4423 #[allow(unreachable_code)]
4424 {
4425 let mut acc = 0f32;
4426 for gi in 0..gpr {
4427 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4428 let s = scales[gi];
4429 let mut d = 0i32;
4430 for (k, &b) in tile.iter().enumerate() {
4431 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4432 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4433 }
4434 acc += d as f32 * s;
4435 }
4436 acc
4437 }
4438}
4439
4440#[cfg(target_arch = "aarch64")]
4443#[target_feature(enable = "neon,dotprod")]
4444unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4445 unsafe {
4448 use core::arch::aarch64::*;
4449 use core::arch::asm;
4450 let lomask = vdupq_n_u8(0x0F);
4451 let eight = vdupq_n_s8(8);
4452 let mut acc = 0f32;
4453 for gi in 0..gpr {
4454 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4455 let s = *scales.get_unchecked(gi);
4456 let b = vld1q_u8(t);
4457 let lo = vandq_u8(b, lomask);
4458 let hi = vshrq_n_u8::<4>(b);
4459 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4460 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4461 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4462 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4463 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4464 asm!(
4465 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4466 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4467 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4468 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4469 options(pure, nomem, nostack),
4470 );
4471 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4472 }
4473 acc
4474 }
4475}
4476
4477#[cfg(target_arch = "x86_64")]
4478#[target_feature(enable = "avx2")]
4479unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4480 unsafe {
4482 use core::arch::x86_64::*;
4483 let lomask = _mm_set1_epi8(0x0F);
4484 let eight = _mm256_set1_epi8(8);
4485 let ones = _mm256_set1_epi16(1);
4486 let mut acc = 0f32;
4487 for gi in 0..gpr {
4488 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4489 let s = *scales.get_unchecked(gi);
4490 let b = _mm_loadu_si128(t as *const __m128i);
4491 let lo = _mm_and_si128(b, lomask);
4492 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4493 let w = _mm256_sub_epi8(
4494 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4495 eight,
4496 );
4497 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4498 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4499 let d = _mm256_madd_epi16(p16, ones);
4500 let hi128 = _mm256_extracti128_si256::<1>(d);
4501 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4502 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4503 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4504 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4505 }
4506 acc
4507 }
4508}
4509
4510#[cfg(target_arch = "x86_64")]
4513#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4514unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4515 unsafe {
4517 use core::arch::x86_64::*;
4518 let lomask = _mm_set1_epi8(0x0F);
4519 let eight = _mm256_set1_epi8(8);
4520 let mut acc = 0f32;
4521 for gi in 0..gpr {
4522 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4523 let s = *scales.get_unchecked(gi);
4524 let b = _mm_loadu_si128(t as *const __m128i);
4525 let lo = _mm_and_si128(b, lomask);
4526 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4527 let w = _mm256_sub_epi8(
4528 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4529 eight,
4530 );
4531 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4532 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4533 }
4534 acc
4535 }
4536}
4537
4538#[inline]
4541fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4542 let mut acc = 0f32;
4543 for gi in 0..gpr {
4544 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4545 let s = scales[gi];
4546 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4547 let mut ga = 0f32;
4548 for (k, &b) in tile.iter().enumerate() {
4549 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4550 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4551 }
4552 acc += ga * s;
4553 }
4554 acc
4555}
4556
4557#[inline]
4560fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4561 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4562 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4563 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4564 ((n as i32 - 8) as f32, scales[gi])
4565}
4566
4567fn q4tp_matvec(
4569 bytes: &[u8],
4570 x: &[f32],
4571 rows: usize,
4572 cols: usize,
4573 out: &mut [f32],
4574 pool: Option<&Pool>,
4575) {
4576 debug_assert_eq!(out.len(), rows);
4577 let gpr = cols / GROUP_SIZE;
4578 let v = Q4tpView::new(bytes, rows, cols);
4579 let out_addr = SendMut(out.as_mut_ptr());
4580 if a8w8_enabled() {
4581 let act = split_act(x);
4582 let run = |start: usize, end: usize| {
4583 with_krow(gpr, |sc| {
4585 for r in start..end {
4586 v.scales_into(r, gpr, sc);
4587 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4588 for &(j, xv) in &act.outliers {
4589 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4590 acc += w * s * xv;
4591 }
4592 unsafe { *out_addr.at(r) = acc };
4594 }
4595 })
4596 };
4597 dispatch_rows(pool, rows, &run);
4598 return;
4599 }
4600 let run = |start: usize, end: usize| {
4601 with_krow(gpr, |sc| {
4602 for r in start..end {
4603 v.scales_into(r, gpr, sc);
4604 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4606 }
4607 })
4608 };
4609 dispatch_rows(pool, rows, &run);
4610}
4611
4612#[allow(clippy::too_many_arguments)]
4615fn q4tp_matvec2(
4616 bytes: &[u8],
4617 x1: &[f32],
4618 x2: &[f32],
4619 rows: usize,
4620 cols: usize,
4621 o1: &mut [f32],
4622 o2: &mut [f32],
4623 pool: Option<&Pool>,
4624) {
4625 let gpr = cols / GROUP_SIZE;
4626 let v = Q4tpView::new(bytes, rows, cols);
4627 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4628 let run = |start: usize, end: usize| {
4629 let mut sc = vec![0f32; gpr];
4630 for r in start..end {
4631 v.scales_into(r, gpr, &mut sc);
4632 unsafe {
4634 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4635 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4636 }
4637 }
4638 };
4639 dispatch_rows(pool, rows, &run);
4640}
4641
4642#[inline]
4645fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4646 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4647 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4648 let c = (byte >> (2 * (k % 4))) & 3;
4649 (c as f32 - 1.5, scales[gi])
4650}
4651
4652#[inline]
4659fn dot_q2tp_row_i8(
4660 chunks: &[u8],
4661 r: usize,
4662 gpr: usize,
4663 xq: &[i8],
4664 gsum: &[i32],
4665 scales: &[f32],
4666) -> f32 {
4667 let mut acc = 0f32;
4668 let base = r * gpr * Q2TP_CHUNK;
4669 #[cfg(not(target_arch = "aarch64"))]
4670 let mut codes = [0i8; GROUP_SIZE];
4671 for gi in 0..gpr {
4672 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4673 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4674 #[cfg(target_arch = "aarch64")]
4675 let dot = unsafe {
4681 use core::arch::aarch64::*;
4682 let b = vld1_u8(ch.as_ptr());
4683 let three = vdup_n_u8(3);
4684 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4685 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4686 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4687 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4688 let x4 = vld4_s8(xg.as_ptr());
4689 let mut acc4 = vdupq_n_s32(0);
4690 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4691 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4692 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4693 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4694 vaddvq_s32(acc4)
4695 };
4696 #[cfg(not(target_arch = "aarch64"))]
4697 let dot: i32 = {
4698 for (k, &b) in ch.iter().enumerate() {
4699 codes[k * 4] = (b & 3) as i8;
4700 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4701 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4702 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4703 }
4704 codes
4705 .iter()
4706 .zip(xg)
4707 .map(|(&c, &x)| c as i32 * x as i32)
4708 .sum()
4709 };
4710 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4711 }
4712 acc
4713}
4714
4715fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4719 let mut acc = 0f32;
4720 for gi in 0..gpr {
4721 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4722 let s = scales[gi];
4723 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4724 let mut g = 0f32;
4725 for (k, &b) in ch.iter().enumerate() {
4726 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4727 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4728 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4729 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4730 }
4731 acc += s * g;
4732 }
4733 acc
4734}
4735
4736fn q2tp_matvec(
4737 bytes: &[u8],
4738 x: &[f32],
4739 rows: usize,
4740 cols: usize,
4741 out: &mut [f32],
4742 pool: Option<&Pool>,
4743) {
4744 debug_assert_eq!(out.len(), rows);
4745 let gpr = cols / GROUP_SIZE;
4746 let v = Q4tpView::new_q2(bytes, rows, cols);
4747 let out_addr = SendMut(out.as_mut_ptr());
4748 if a8w8_enabled() {
4753 let act = split_act(x);
4754 let gsum = q1_group_sums(&act.xq, gpr);
4755 let (act, gsum) = (&act, &gsum);
4756 let run = move |start: usize, end: usize| {
4757 with_krow(gpr, |sc| {
4758 for r in start..end {
4759 v.scales_into(r, gpr, sc);
4760 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4761 for &(j, xv) in &act.outliers {
4762 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4763 acc += w * s * xv;
4764 }
4765 unsafe { *out_addr.at(r) = acc };
4767 }
4768 })
4769 };
4770 dispatch_rows(pool, rows, &run);
4771 return;
4772 }
4773 let run = |start: usize, end: usize| {
4774 with_krow(gpr, |sc| {
4775 for r in start..end {
4776 v.scales_into(r, gpr, sc);
4777 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4779 }
4780 })
4781 };
4782 dispatch_rows(pool, rows, &run);
4783}
4784
4785#[allow(clippy::too_many_arguments)]
4787fn q2tp_matvec2(
4788 bytes: &[u8],
4789 x1: &[f32],
4790 x2: &[f32],
4791 rows: usize,
4792 cols: usize,
4793 o1: &mut [f32],
4794 o2: &mut [f32],
4795 pool: Option<&Pool>,
4796) {
4797 let gpr = cols / GROUP_SIZE;
4798 let v = Q4tpView::new_q2(bytes, rows, cols);
4799 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4800 let run = |start: usize, end: usize| {
4801 let mut sc = vec![0f32; gpr];
4802 for r in start..end {
4803 v.scales_into(r, gpr, &mut sc);
4804 unsafe {
4806 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4807 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4808 }
4809 }
4810 };
4811 dispatch_rows(pool, rows, &run);
4812}
4813
4814pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4821 let gpr = cols / GROUP_SIZE;
4826 let v = Q4tpView::new_q2(bytes, rows, cols);
4827 with_krow(gpr, |sc| {
4828 for r in 0..rows {
4829 v.scales_into(r, gpr, sc);
4830 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4831 }
4832 });
4833}
4834
4835pub fn q2tp_matmat_for_test(
4836 bytes: &[u8],
4837 xs_all: &[f32],
4838 b: usize,
4839 rows: usize,
4840 cols: usize,
4841 out: &mut [f32],
4842) {
4843 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4844}
4845
4846fn q2tp_matmat(
4847 bytes: &[u8],
4848 xs_all: &[f32],
4849 b: usize,
4850 rows: usize,
4851 cols: usize,
4852 out: &mut [f32],
4853 pool: Option<&Pool>,
4854) {
4855 debug_assert_eq!(out.len(), b * rows);
4856 let gpr = cols / GROUP_SIZE;
4857 let v = Q4tpView::new_q2(bytes, rows, cols);
4858 let out_addr = SendMut(out.as_mut_ptr());
4859 let run = |start: usize, end: usize| {
4860 let mut sc = vec![0f32; gpr];
4861 for r in start..end {
4862 v.scales_into(r, gpr, &mut sc);
4863 for bi in 0..b {
4864 let x = &xs_all[bi * cols..(bi + 1) * cols];
4865 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4867 }
4868 }
4869 };
4870 dispatch_rows(pool, rows, &run);
4871}
4872
4873#[cfg(target_arch = "aarch64")]
4877#[target_feature(enable = "neon,dotprod")]
4878unsafe fn dot_q4tp_row_1x4_sdot_v1(
4879 nib: &[u8],
4880 r: usize,
4881 gpr: usize,
4882 xs: [&[i8]; 4],
4883 scales: &[f32],
4884) -> [f32; 4] {
4885 unsafe {
4886 use core::arch::aarch64::*;
4887 use core::arch::asm;
4888 let lomask = vdupq_n_u8(0x0F);
4889 let eight = vdupq_n_s8(8);
4890 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4891 for gi in 0..gpr {
4892 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4893 let s = *scales.get_unchecked(gi);
4894 let bb = vld1q_u8(t);
4895 let lo = vandq_u8(bb, lomask);
4896 let hi = vshrq_n_u8::<4>(bb);
4897 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4898 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4899 let mut d = [0f32; 4];
4900 for (k, dk) in d.iter_mut().enumerate() {
4901 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4902 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4903 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4904 asm!(
4905 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4906 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4907 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4908 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4909 options(pure, nomem, nostack),
4910 );
4911 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4912 }
4913 f0 += d[0];
4914 f1 += d[1];
4915 f2 += d[2];
4916 f3 += d[3];
4917 }
4918 [f0, f1, f2, f3]
4919 }
4920}
4921
4922#[allow(dead_code)]
4929static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4930
4931#[cfg(target_arch = "x86_64")]
4937fn q4tp_blocked_x86() -> bool {
4938 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4939 1 => false,
4940 2 => avx512vnni_enabled(),
4945 _ => blocked_enabled() && avx512vnni_enabled(),
4949 }
4950}
4951
4952#[cfg(target_arch = "aarch64")]
4954#[allow(dead_code)]
4955fn q4tp_v1() -> bool {
4956 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4957 1 => true,
4958 2 => false,
4959 _ => {
4960 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4961 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4962 }
4963 }
4964}
4965
4966#[cfg(target_arch = "x86_64")]
4978#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4979unsafe fn dot_q4tp_2x8_avx512(
4980 nib: &[u8],
4981 r0: usize,
4982 gpr: usize,
4983 xs: [&[i8]; 8],
4984 sc0: &[f32],
4985 sc1: &[f32],
4986) -> [[f32; 8]; 2] {
4987 unsafe {
4990 use core::arch::x86_64::*;
4991 let lomask = _mm256_set1_epi8(0x0F);
4992 let eight = _mm256_set1_epi8(8);
4993 let zero = _mm512_setzero_si512();
4994 let mut v0 = [_mm512_setzero_ps(); 8];
4995 let mut v1 = [_mm512_setzero_ps(); 8];
4996 let pairs = gpr / 2;
4997 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4998 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4999 let bb = _mm256_loadu_si256(t as *const __m256i);
5000 let lo = _mm256_and_si256(bb, lomask);
5001 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5002 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5003 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5004 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5005 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5006 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5007 };
5008 for gp in 0..pairs {
5009 let gi = gp * 2;
5010 let (wa0, neg0) = unpack(r0, gi);
5011 let (wa1, neg1) = unpack(r0 + 1, gi);
5012 let off = gi * GROUP_SIZE;
5013 let sv = |sc: &[f32]| {
5014 _mm512_insertf32x8::<1>(
5015 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5016 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5017 )
5018 };
5019 let s0 = sv(sc0);
5020 let s1 = sv(sc1);
5021 for k in 0..8 {
5022 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5023 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5024 zero,
5025 wa0,
5026 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5027 ));
5028 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5029 zero,
5030 wa1,
5031 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5032 ));
5033 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5034 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5035 }
5036 }
5037 let mut acc = [[0f32; 8]; 2];
5038 for k in 0..8 {
5039 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5040 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5041 }
5042 if gpr % 2 == 1 {
5043 let off = (gpr - 1) * GROUP_SIZE;
5044 for j in off..off + GROUP_SIZE {
5045 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5046 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5047 for k in 0..8 {
5048 let x = *xs[k].get_unchecked(j) as f32;
5049 acc[0][k] += w0 * sa * x;
5050 acc[1][k] += w1 * sb * x;
5051 }
5052 }
5053 }
5054 acc
5055 }
5056}
5057
5058#[cfg(target_arch = "x86_64")]
5063#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5064unsafe fn dot_q4tp_row_1x8_avx512(
5065 nib: &[u8],
5066 r: usize,
5067 gpr: usize,
5068 xs: [&[i8]; 8],
5069 scales: &[f32],
5070) -> [f32; 8] {
5071 unsafe {
5073 use core::arch::x86_64::*;
5074 let lomask = _mm256_set1_epi8(0x0F);
5075 let eight = _mm256_set1_epi8(8);
5076 let zero = _mm512_setzero_si512();
5077 let (mut v0, mut v1, mut v2, mut v3) = (
5078 _mm512_setzero_ps(),
5079 _mm512_setzero_ps(),
5080 _mm512_setzero_ps(),
5081 _mm512_setzero_ps(),
5082 );
5083 let (mut v4, mut v5, mut v6, mut v7) = (
5084 _mm512_setzero_ps(),
5085 _mm512_setzero_ps(),
5086 _mm512_setzero_ps(),
5087 _mm512_setzero_ps(),
5088 );
5089 let pairs = gpr / 2;
5090 for gp in 0..pairs {
5091 let gi = gp * 2;
5092 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5093 let bb = _mm256_loadu_si256(t as *const __m256i);
5094 let lo = _mm256_and_si256(bb, lomask);
5095 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5096 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5101 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5102 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5103 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5104 let wabs = _mm512_abs_epi8(w);
5105 let neg = _mm512_movepi8_mask(w);
5106 let off = gi * GROUP_SIZE;
5107 let sv = _mm512_insertf32x8::<1>(
5108 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5109 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5110 );
5111 let dot = |x: &[i8]| -> __m512 {
5112 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5113 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5114 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5115 };
5116 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5117 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5118 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5119 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5120 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5121 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5122 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5123 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5124 }
5125 let mut acc = [
5126 _mm512_reduce_add_ps(v0),
5127 _mm512_reduce_add_ps(v1),
5128 _mm512_reduce_add_ps(v2),
5129 _mm512_reduce_add_ps(v3),
5130 _mm512_reduce_add_ps(v4),
5131 _mm512_reduce_add_ps(v5),
5132 _mm512_reduce_add_ps(v6),
5133 _mm512_reduce_add_ps(v7),
5134 ];
5135 if gpr % 2 == 1 {
5138 let off = (gpr - 1) * GROUP_SIZE;
5139 for j in off..off + GROUP_SIZE {
5140 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5141 let ws = w * s;
5142 for k in 0..8 {
5143 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5144 }
5145 }
5146 }
5147 acc
5148 }
5149}
5150
5151#[cfg(target_arch = "x86_64")]
5164#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5165unsafe fn dot_q4tp_row_1x4_avx512(
5166 nib: &[u8],
5167 r: usize,
5168 gpr: usize,
5169 xs: [&[i8]; 4],
5170 scales: &[f32],
5171) -> [f32; 4] {
5172 unsafe {
5174 use core::arch::x86_64::*;
5175 let lomask = _mm256_set1_epi8(0x0F);
5176 let eight = _mm256_set1_epi8(8);
5177 let zero = _mm512_setzero_si512();
5178 let (mut v0, mut v1, mut v2, mut v3) = (
5179 _mm512_setzero_ps(),
5180 _mm512_setzero_ps(),
5181 _mm512_setzero_ps(),
5182 _mm512_setzero_ps(),
5183 );
5184 let pairs = gpr / 2;
5185 for gp in 0..pairs {
5186 let gi = gp * 2;
5187 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5188 let bb = _mm256_loadu_si256(t as *const __m256i);
5189 let lo = _mm256_and_si256(bb, lomask);
5190 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5191 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5196 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5197 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5198 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5199 let wabs = _mm512_abs_epi8(w);
5200 let neg = _mm512_movepi8_mask(w);
5201 let off = gi * GROUP_SIZE;
5202 let sv = _mm512_insertf32x8::<1>(
5203 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5204 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5205 );
5206 let dot = |x: &[i8]| -> __m512 {
5207 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5208 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5209 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5210 };
5211 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5212 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5213 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5214 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5215 }
5216 let mut acc = [
5217 _mm512_reduce_add_ps(v0),
5218 _mm512_reduce_add_ps(v1),
5219 _mm512_reduce_add_ps(v2),
5220 _mm512_reduce_add_ps(v3),
5221 ];
5222 if gpr % 2 == 1 {
5225 let off = (gpr - 1) * GROUP_SIZE;
5226 for j in off..off + GROUP_SIZE {
5227 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5228 let ws = w * s;
5229 for k in 0..4 {
5230 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5231 }
5232 }
5233 }
5234 acc
5235 }
5236}
5237
5238#[cfg(target_arch = "aarch64")]
5242#[target_feature(enable = "neon,dotprod")]
5243unsafe fn dot_q4tp_row_1x4_sdot(
5244 nib: &[u8],
5245 r: usize,
5246 gpr: usize,
5247 xs: [&[i8]; 4],
5248 scales: &[f32],
5249) -> [f32; 4] {
5250 unsafe {
5252 use core::arch::aarch64::*;
5253 use core::arch::asm;
5254 let lomask = vdupq_n_u8(0x0F);
5255 let eight = vdupq_n_s8(8);
5256 let (mut v0, mut v1, mut v2, mut v3) = (
5272 vdupq_n_f32(0.0),
5273 vdupq_n_f32(0.0),
5274 vdupq_n_f32(0.0),
5275 vdupq_n_f32(0.0),
5276 );
5277 for gi in 0..gpr {
5278 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5279 let s = *scales.get_unchecked(gi);
5280 let bb = vld1q_u8(t);
5281 let lo = vandq_u8(bb, lomask);
5282 let hi = vshrq_n_u8::<4>(bb);
5283 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5284 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5285 let off = gi * GROUP_SIZE;
5286 let dot4 = |x: &[i8]| -> int32x4_t {
5287 let x0 = vld1q_s8(x.as_ptr().add(off));
5288 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5289 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5290 asm!(
5291 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5292 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5293 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5294 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5295 options(pure, nomem, nostack),
5296 );
5297 vaddq_s32(a0, a1)
5298 };
5299 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5300 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5301 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5302 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5303 }
5304 [
5305 vaddvq_f32(v0),
5306 vaddvq_f32(v1),
5307 vaddvq_f32(v2),
5308 vaddvq_f32(v3),
5309 ]
5310 }
5311}
5312
5313fn q4tp_matmat(
5317 bytes: &[u8],
5318 xs_all: &[f32],
5319 b: usize,
5320 rows: usize,
5321 cols: usize,
5322 out: &mut [f32],
5323 pool: Option<&Pool>,
5324) {
5325 debug_assert_eq!(out.len(), b * rows);
5326 let gpr = cols / GROUP_SIZE;
5327 let v = Q4tpView::new(bytes, rows, cols);
5328
5329 #[cfg(target_os = "macos")]
5331 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5332 dequant_matmat_accel(
5333 &|r, dst| {
5334 let mut sc = [0f32; 32];
5335 let mut scv;
5336 let s: &[f32] = if gpr <= 32 {
5337 v.scales_into(r, gpr, &mut sc);
5338 &sc[..gpr]
5339 } else {
5340 scv = vec![0f32; gpr];
5341 v.scales_into(r, gpr, &mut scv);
5342 &scv
5343 };
5344 for gi in 0..gpr {
5345 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5346 for (k, &bb) in tile.iter().enumerate() {
5347 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5348 dst[gi * GROUP_SIZE + k * 2 + 1] =
5349 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5350 }
5351 }
5352 },
5353 xs_all,
5354 b,
5355 rows,
5356 cols,
5357 out,
5358 pool,
5359 );
5360 return;
5361 }
5362
5363 let out_addr = SendMut(out.as_mut_ptr());
5364 if a8w8_enabled() {
5365 let acts: Vec<SplitAct> = (0..b)
5366 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5367 .collect();
5368 let acts = &acts;
5369 #[cfg(target_arch = "aarch64")]
5370 let blocked_ok = sdot_enabled() && blocked_enabled();
5371 #[cfg(target_arch = "x86_64")]
5378 let blocked_ok = q4tp_blocked_x86();
5379 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5380 let blocked_ok = false;
5381 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5391 .ok()
5392 .and_then(|v| v.parse().ok())
5393 .filter(|v| *v > 0)
5394 .unwrap_or(256);
5395 let run = |start: usize, end: usize| {
5396 for abase in (0..acts.len()).step_by(panel_cols) {
5397 let alen = (acts.len() - abase).min(panel_cols);
5398 let mut sc = vec![0f32; gpr];
5399 #[cfg(target_arch = "x86_64")]
5400 let mut r_lo = start;
5401 #[cfg(target_arch = "x86_64")]
5402 if blocked_ok && alen >= 8 {
5403 let mut sc1 = vec![0f32; gpr];
5404 while r_lo + 2 <= end {
5405 v.scales_into(r_lo, gpr, &mut sc);
5406 v.scales_into(r_lo + 1, gpr, &mut sc1);
5407 let mut bi = 0usize;
5408 while bi + 8 <= alen {
5409 let xs = [
5410 acts[abase + bi].xq.as_slice(),
5411 acts[abase + bi + 1].xq.as_slice(),
5412 acts[abase + bi + 2].xq.as_slice(),
5413 acts[abase + bi + 3].xq.as_slice(),
5414 acts[abase + bi + 4].xq.as_slice(),
5415 acts[abase + bi + 5].xq.as_slice(),
5416 acts[abase + bi + 6].xq.as_slice(),
5417 acts[abase + bi + 7].xq.as_slice(),
5418 ];
5419 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5420 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
5421 for k in 0..8 {
5422 let act = &acts[abase + bi + k];
5423 let mut acc = dr[k] * act.sx;
5424 for &(j, xv) in &act.outliers {
5425 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5426 acc += w * s * xv;
5427 }
5428 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5430 }
5431 }
5432 bi += 8;
5433 }
5434 for row in [r_lo, r_lo + 1] {
5437 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5438 for b2 in bi..alen {
5439 let act = &acts[abase + b2];
5440 let xs4 = [
5441 act.xq.as_slice(),
5442 act.xq.as_slice(),
5443 act.xq.as_slice(),
5444 act.xq.as_slice(),
5445 ];
5446 let d =
5447 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5448 let mut acc = d[0] * act.sx;
5449 for &(j, xv) in &act.outliers {
5450 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5451 acc += w * s * xv;
5452 }
5453 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5455 }
5456 }
5457 r_lo += 2;
5458 }
5459 }
5460 #[cfg(target_arch = "x86_64")]
5461 let row_start = r_lo;
5462 #[cfg(not(target_arch = "x86_64"))]
5463 let row_start = start;
5464 for r in row_start..end {
5465 v.scales_into(r, gpr, &mut sc);
5466 let mut bi = 0usize;
5467 #[cfg(target_arch = "x86_64")]
5468 if blocked_ok {
5469 while bi + 8 <= alen {
5470 let xs = [
5471 acts[abase + bi].xq.as_slice(),
5472 acts[abase + bi + 1].xq.as_slice(),
5473 acts[abase + bi + 2].xq.as_slice(),
5474 acts[abase + bi + 3].xq.as_slice(),
5475 acts[abase + bi + 4].xq.as_slice(),
5476 acts[abase + bi + 5].xq.as_slice(),
5477 acts[abase + bi + 6].xq.as_slice(),
5478 acts[abase + bi + 7].xq.as_slice(),
5479 ];
5480 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5481 for k in 0..8 {
5482 let act = &acts[abase + bi + k];
5483 let mut acc = d[k] * act.sx;
5484 for &(j, xv) in &act.outliers {
5485 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5486 acc += w * s * xv;
5487 }
5488 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5490 }
5491 bi += 8;
5492 }
5493 while bi + 4 <= alen {
5494 let xs = [
5495 acts[abase + bi].xq.as_slice(),
5496 acts[abase + bi + 1].xq.as_slice(),
5497 acts[abase + bi + 2].xq.as_slice(),
5498 acts[abase + bi + 3].xq.as_slice(),
5499 ];
5500 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5501 for k in 0..4 {
5502 let act = &acts[abase + bi + k];
5503 let mut acc = d[k] * act.sx;
5504 for &(j, xv) in &act.outliers {
5505 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5506 acc += w * s * xv;
5507 }
5508 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5510 }
5511 bi += 4;
5512 }
5513 }
5514 #[cfg(target_arch = "aarch64")]
5515 if blocked_ok {
5516 while bi + 4 <= alen {
5517 let xs = [
5518 acts[abase + bi].xq.as_slice(),
5519 acts[abase + bi + 1].xq.as_slice(),
5520 acts[abase + bi + 2].xq.as_slice(),
5521 acts[abase + bi + 3].xq.as_slice(),
5522 ];
5523 let d = unsafe {
5524 if q4tp_v1() {
5525 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5526 } else {
5527 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5528 }
5529 };
5530 for k in 0..4 {
5531 let act = &acts[abase + bi + k];
5532 let mut acc = d[k] * act.sx;
5533 for &(j, xv) in &act.outliers {
5534 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5535 acc += w * s * xv;
5536 }
5537 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5539 }
5540 bi += 4;
5541 }
5542 }
5543 let _ = blocked_ok;
5544 while bi < alen {
5545 let act = &acts[abase + bi];
5546 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5547 for &(j, xv) in &act.outliers {
5548 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5549 acc += w * s * xv;
5550 }
5551 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5553 bi += 1;
5554 }
5555 }
5556 }
5557 };
5558 dispatch_rows(pool, rows, &run);
5559 return;
5560 }
5561
5562 let run = |start: usize, end: usize| {
5563 let mut sc = vec![0f32; gpr];
5564 for r in start..end {
5565 v.scales_into(r, gpr, &mut sc);
5566 for bi in 0..b {
5567 let x = &xs_all[bi * cols..(bi + 1) * cols];
5568 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5570 }
5571 }
5572 };
5573 dispatch_rows(pool, rows, &run);
5574}
5575
5576fn q4t_matvec(
5578 bytes: &[u8],
5579 x: &[f32],
5580 rows: usize,
5581 cols: usize,
5582 out: &mut [f32],
5583 pool: Option<&Pool>,
5584) {
5585 debug_assert_eq!(out.len(), rows);
5586 let gpr = cols / GROUP_SIZE;
5587 let out_addr = SendMut(out.as_mut_ptr());
5588 if a8w8_enabled() {
5589 let act = split_act(x);
5590 let run = move |start: usize, end: usize| {
5591 for r in start..end {
5592 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5593 for &(j, xv) in &act.outliers {
5594 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5595 acc += w * s * xv;
5596 }
5597 unsafe { *out_addr.at(r) = acc };
5599 }
5600 };
5601 dispatch_rows(pool, rows, &run);
5602 return;
5603 }
5604 let run = move |start: usize, end: usize| {
5605 for r in start..end {
5606 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5608 }
5609 };
5610 dispatch_rows(pool, rows, &run);
5611}
5612
5613#[allow(clippy::too_many_arguments)]
5615fn q4t_matvec2(
5616 bytes: &[u8],
5617 x1: &[f32],
5618 x2: &[f32],
5619 rows: usize,
5620 cols: usize,
5621 o1: &mut [f32],
5622 o2: &mut [f32],
5623 pool: Option<&Pool>,
5624) {
5625 let gpr = cols / GROUP_SIZE;
5626 let p1 = SendMut(o1.as_mut_ptr());
5627 let p2 = SendMut(o2.as_mut_ptr());
5628 if a8w8_enabled() {
5629 let a1 = split_act(x1);
5630 let a2 = split_act(x2);
5631 let run = move |start: usize, end: usize| {
5632 for r in start..end {
5633 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5634 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5635 for &(j, xv) in &a1.outliers {
5636 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5637 v1 += w * s * xv;
5638 }
5639 for &(j, xv) in &a2.outliers {
5640 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5641 v2 += w * s * xv;
5642 }
5643 unsafe {
5645 *p1.at(r) = v1;
5646 *p2.at(r) = v2;
5647 }
5648 }
5649 };
5650 dispatch_rows(pool, rows, &run);
5651 return;
5652 }
5653 let run = move |start: usize, end: usize| {
5654 for r in start..end {
5655 unsafe {
5657 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5658 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5659 }
5660 }
5661 };
5662 dispatch_rows(pool, rows, &run);
5663}
5664
5665#[allow(clippy::too_many_arguments)]
5667#[cfg(target_os = "macos")]
5673fn dequant_matmat_accel(
5674 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5675 xs_all: &[f32],
5676 b: usize,
5677 rows: usize,
5678 cols: usize,
5679 out: &mut [f32],
5680 pool: Option<&Pool>,
5681) {
5682 const TR: usize = 2048;
5683 thread_local! {
5684 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5685 }
5686 WTILE.with(|wt| {
5687 let mut wtile = wt.borrow_mut();
5688 wtile.resize(TR * cols, 0.0);
5689 let mut r0 = 0usize;
5690 while r0 < rows {
5691 let tr = TR.min(rows - r0);
5692 let wt_addr = SendMut(wtile.as_mut_ptr());
5693 let run = |start: usize, end: usize| {
5694 for r in start..end {
5695 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5697 dequant_row(r0 + r, dst);
5698 }
5699 };
5700 dispatch_rows(pool, tr, &run);
5701 unsafe {
5702 accel_blas::cblas_sgemm(
5703 101, 111, 112, b as i32,
5707 tr as i32,
5708 cols as i32,
5709 1.0,
5710 xs_all.as_ptr(),
5711 cols as i32,
5712 wtile.as_ptr(),
5713 cols as i32,
5714 0.0,
5715 out.as_mut_ptr().add(r0),
5716 rows as i32,
5717 );
5718 }
5719 r0 += tr;
5720 }
5721 });
5722}
5723
5724fn q4t_matmat(
5725 bytes: &[u8],
5726 xs_all: &[f32],
5727 b: usize,
5728 rows: usize,
5729 cols: usize,
5730 out: &mut [f32],
5731 pool: Option<&Pool>,
5732) {
5733 debug_assert_eq!(out.len(), b * rows);
5734 let gpr = cols / GROUP_SIZE;
5735 #[cfg(target_os = "macos")]
5739 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5740 dequant_matmat_accel(
5741 &|r, dst| {
5742 for gi in 0..gpr {
5743 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5744 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5745 for (k, &bb) in tile[2..].iter().enumerate() {
5746 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5747 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5748 }
5749 }
5750 },
5751 xs_all,
5752 b,
5753 rows,
5754 cols,
5755 out,
5756 pool,
5757 );
5758 return;
5759 }
5760 let out_addr = SendMut(out.as_mut_ptr());
5761 if a8w8_enabled() {
5762 let acts: Vec<SplitAct> = (0..b)
5763 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5764 .collect();
5765 let acts = &acts;
5766 #[cfg(target_arch = "x86_64")]
5767 let blocked_ok = avx2_enabled() && blocked_enabled();
5768 #[cfg(target_arch = "aarch64")]
5769 let blocked_ok = sdot_enabled() && blocked_enabled();
5770 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5771 let blocked_ok = false;
5772 let run = move |start: usize, end: usize| {
5773 for r in start..end {
5774 let mut bi = 0usize;
5775 #[cfg(target_arch = "aarch64")]
5776 if blocked_ok {
5777 while bi + 4 <= acts.len() {
5778 let xs = [
5779 acts[bi].xq.as_slice(),
5780 acts[bi + 1].xq.as_slice(),
5781 acts[bi + 2].xq.as_slice(),
5782 acts[bi + 3].xq.as_slice(),
5783 ];
5784 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5785 for k in 0..4 {
5786 let act = &acts[bi + k];
5787 let mut acc = d[k] * act.sx;
5788 for &(j, xv) in &act.outliers {
5789 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5790 acc += w * sc * xv;
5791 }
5792 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5794 }
5795 bi += 4;
5796 }
5797 }
5798 #[cfg(target_arch = "x86_64")]
5799 if blocked_ok {
5800 while bi + 4 <= acts.len() {
5801 let xs = [
5802 acts[bi].xq.as_slice(),
5803 acts[bi + 1].xq.as_slice(),
5804 acts[bi + 2].xq.as_slice(),
5805 acts[bi + 3].xq.as_slice(),
5806 ];
5807 let d = unsafe {
5808 if vnni_tiles_enabled() {
5809 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5810 } else {
5811 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5812 }
5813 };
5814 for k in 0..4 {
5815 let act = &acts[bi + k];
5816 let mut acc = d[k] * act.sx;
5817 for &(j, xv) in &act.outliers {
5818 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5819 acc += w * sc * xv;
5820 }
5821 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5823 }
5824 bi += 4;
5825 }
5826 }
5827 let _ = blocked_ok;
5828 while bi < acts.len() {
5829 let act = &acts[bi];
5830 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5831 for &(j, xv) in &act.outliers {
5832 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5833 acc += w * s * xv;
5834 }
5835 unsafe { *out_addr.at(bi * rows + r) = acc };
5837 bi += 1;
5838 }
5839 }
5840 };
5841 dispatch_rows(pool, rows, &run);
5842 return;
5843 }
5844 let run = move |start: usize, end: usize| {
5845 for r in start..end {
5846 for bi in 0..b {
5847 let x = &xs_all[bi * cols..(bi + 1) * cols];
5848 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5850 }
5851 }
5852 };
5853 dispatch_rows(pool, rows, &run);
5854}
5855
5856fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5865 (0..gpr)
5866 .map(|gi| {
5867 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5868 .iter()
5869 .map(|&v| v as i32)
5870 .sum()
5871 })
5872 .collect()
5873}
5874
5875#[inline]
5879#[allow(unreachable_code)]
5880#[cfg(target_arch = "x86_64")]
5885#[target_feature(enable = "avx2")]
5886unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5887 unsafe {
5889 use core::arch::x86_64::*;
5890 let expand = _mm256_setr_epi8(
5892 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,
5893 3, 3, 3,
5894 );
5895 let bitsel = _mm256_setr_epi8(
5896 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5897 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5898 );
5899 let ones8 = _mm256_set1_epi8(1);
5900 let ones16 = _mm256_set1_epi16(1);
5901 let mut acc = 0f32;
5902 for gi in 0..gpr {
5903 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5904 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5905 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5906 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5907 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5908 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5909 let sel = _mm256_and_si256(x, mask);
5910 let p16 = _mm256_maddubs_epi16(ones8, sel);
5912 let d32 = _mm256_madd_epi16(p16, ones16);
5913 let hi128 = _mm256_extracti128_si256::<1>(d32);
5914 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5915 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5916 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5917 let msum = _mm_cvtsi128_si32(s32);
5918 let d = 2 * msum - gsum[gi];
5921 acc += d as f32 * s;
5922 }
5923 acc
5924 }
5925}
5926
5927#[cfg(target_arch = "x86_64")]
5930#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5931unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5932 unsafe {
5934 use core::arch::x86_64::*;
5935 let expand = _mm256_setr_epi8(
5936 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,
5937 3, 3, 3,
5938 );
5939 let bitsel = _mm256_setr_epi8(
5940 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5941 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5942 );
5943 let ones8 = _mm256_set1_epi8(1);
5944 let mut acc = 0f32;
5945 for gi in 0..gpr {
5946 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5947 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5948 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5949 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5950 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5951 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5952 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5953 let d = 2 * msum - gsum[gi];
5954 acc += d as f32 * s;
5955 }
5956 acc
5957 }
5958}
5959
5960#[cfg(target_arch = "x86_64")]
5962#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5963unsafe fn dot_q1_row_1x4_vnni(
5964 bytes: &[u8],
5965 r: usize,
5966 gpr: usize,
5967 xs: [&[i8]; 4],
5968 gsums: [&[i32]; 4],
5969) -> [f32; 4] {
5970 unsafe {
5972 use core::arch::x86_64::*;
5973 let expand = _mm256_setr_epi8(
5974 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,
5975 3, 3, 3,
5976 );
5977 let bitsel = _mm256_setr_epi8(
5978 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5979 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5980 );
5981 let ones8 = _mm256_set1_epi8(1);
5982 let mut acc = [0f32; 4];
5983 for gi in 0..gpr {
5984 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5985 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5986 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5987 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5988 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5989 for (k, xq) in xs.iter().enumerate() {
5990 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5991 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5992 let d = 2 * msum - gsums[k][gi];
5993 acc[k] += d as f32 * s;
5994 }
5995 }
5996 acc
5997 }
5998}
5999
6000#[cfg(target_arch = "x86_64")]
6003#[target_feature(enable = "avx2")]
6004unsafe fn dot_q1_row_1x4_avx2(
6005 bytes: &[u8],
6006 r: usize,
6007 gpr: usize,
6008 xs: [&[i8]; 4],
6009 gsums: [&[i32]; 4],
6010) -> [f32; 4] {
6011 unsafe {
6013 use core::arch::x86_64::*;
6014 let expand = _mm256_setr_epi8(
6015 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,
6016 3, 3, 3,
6017 );
6018 let bitsel = _mm256_setr_epi8(
6019 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6020 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6021 );
6022 let ones8 = _mm256_set1_epi8(1);
6023 let ones16 = _mm256_set1_epi16(1);
6024 let mut acc = [0f32; 4];
6025 for gi in 0..gpr {
6026 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6027 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6028 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6029 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6030 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6031 for (k, xq) in xs.iter().enumerate() {
6032 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6033 let sel = _mm256_and_si256(x, mask);
6034 let p16 = _mm256_maddubs_epi16(ones8, sel);
6035 let d32 = _mm256_madd_epi16(p16, ones16);
6036 let hi128 = _mm256_extracti128_si256::<1>(d32);
6037 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6038 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6039 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6040 let msum = _mm_cvtsi128_si32(s32);
6041 let d = 2 * msum - gsums[k][gi];
6042 acc[k] += d as f32 * s;
6043 }
6044 }
6045 acc
6046 }
6047}
6048
6049#[allow(unreachable_code)]
6050fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6051 #[cfg(target_arch = "aarch64")]
6052 unsafe {
6053 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6054 }
6055 #[cfg(target_arch = "x86_64")]
6056 if avx2_enabled() {
6057 unsafe {
6058 if vnni_tiles_enabled() {
6059 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6060 }
6061 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6062 }
6063 }
6064 let _ = gsum;
6065 let mut acc = 0f32;
6066 for gi in 0..gpr {
6067 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6068 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6069 let mut d = 0i32;
6070 for (j, &b) in tile[2..].iter().enumerate() {
6071 for k in 0..8 {
6072 let w = ((b >> k) & 1) as i32 * 2 - 1;
6073 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6074 }
6075 }
6076 acc += d as f32 * s;
6077 }
6078 acc
6079}
6080
6081#[cfg(target_arch = "aarch64")]
6090#[target_feature(enable = "neon,dotprod")]
6091unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6092 unsafe {
6095 use core::arch::aarch64::*;
6096 use core::arch::asm;
6097 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6098 let m = vld1q_u8(MASKS.as_ptr());
6099 macro_rules! tile_dot {
6101 ($t:expr, $x:expr) => {{
6102 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6103 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6104 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6105 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6106 let x0 = vld1q_s8($x);
6107 let x1 = vld1q_s8($x.add(16));
6108 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6109 asm!(
6110 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6111 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6112 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6113 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6114 options(pure, nomem, nostack),
6115 );
6116 vaddq_s32(a0, a1)
6117 }};
6118 }
6119 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6128 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6129 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6130 const IW11: [u8; 16] = [
6131 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6132 ];
6133 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6134 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6135 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6136 let isc = vld1_u8(ISC.as_ptr());
6137 macro_rules! tile_dot_tbl {
6139 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6140 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6141 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6142 let x0 = vld1q_s8($x);
6143 let x1 = vld1q_s8($x.add(16));
6144 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6145 asm!(
6146 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6147 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6148 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6149 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6150 options(pure, nomem, nostack),
6151 );
6152 vaddq_s32(a0, a1)
6153 }};
6154 }
6155 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6156 let row_base = r * gpr * Q1_TILE;
6157 let abs_end = bytes.len();
6158 let xp = xq.as_ptr();
6159 let gp = gsum.as_ptr();
6160 let mut accv = vdupq_n_f32(0.0);
6161 let mut gi = 0;
6162 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6165 let t0 = base.add(gi * Q1_TILE);
6166 let ld_a = vld1q_u8(t0);
6167 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6168 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6169 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6170 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6171 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6172 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6174 let g = vld1q_s32(gp.add(gi));
6175 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6176 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6177 let scf: float32x4_t;
6178 asm!(
6179 "fcvtl {o:v}.4s, {i:v}.4h",
6180 o = out(vreg) scf, i = in(vreg) sc16,
6181 options(pure, nomem, nostack),
6182 );
6183 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6184 gi += 4;
6185 }
6186 let mut acc = vaddvq_f32(accv);
6187 while gi < gpr {
6188 let t = base.add(gi * Q1_TILE);
6189 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6190 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6191 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6192 gi += 1;
6193 }
6194 acc
6195 }
6196}
6197
6198#[cfg(target_arch = "aarch64")]
6203#[target_feature(enable = "neon,dotprod")]
6204unsafe fn dot_q1_row_1x4_sdot(
6205 bytes: &[u8],
6206 r: usize,
6207 gpr: usize,
6208 xs: [&[i8]; 4],
6209 gs: [&[i32]; 4],
6210) -> [f32; 4] {
6211 unsafe {
6213 use core::arch::aarch64::*;
6214 use core::arch::asm;
6215 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6216 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6217 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6218 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6219 const IW11: [u8; 16] = [
6220 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6221 ];
6222 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6223 let m = vld1q_u8(MASKS.as_ptr());
6224 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6225 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6226 let isc = vld1_u8(ISC.as_ptr());
6227 macro_rules! sdot2 {
6228 ($w0:expr, $w1:expr, $x:expr) => {{
6229 let x0 = vld1q_s8($x);
6230 let x1 = vld1q_s8($x.add(16));
6231 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6232 asm!(
6233 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6234 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6235 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6236 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6237 options(pure, nomem, nostack),
6238 );
6239 vaddq_s32(a0, a1)
6240 }};
6241 }
6242 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6243 let row_base = r * gpr * Q1_TILE;
6244 let abs_end = bytes.len();
6245 let mut accv = [vdupq_n_f32(0.0); 4];
6246 let mut gi = 0;
6247 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6248 let t0 = base.add(gi * Q1_TILE);
6249 let ld_a = vld1q_u8(t0);
6250 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6251 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6253 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6254 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6255 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6256 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6257 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6258 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6259 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6260 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6261 let scf: float32x4_t;
6262 asm!(
6263 "fcvtl {o:v}.4s, {i:v}.4h",
6264 o = out(vreg) scf, i = in(vreg) sc16,
6265 options(pure, nomem, nostack),
6266 );
6267 for k in 0..4 {
6268 let xp = xs[k].as_ptr();
6269 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6270 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6271 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6272 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6273 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6274 let g = vld1q_s32(gs[k].as_ptr().add(gi));
6275 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6276 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6277 }
6278 gi += 4;
6279 }
6280 let mut acc = [
6281 vaddvq_f32(accv[0]),
6282 vaddvq_f32(accv[1]),
6283 vaddvq_f32(accv[2]),
6284 vaddvq_f32(accv[3]),
6285 ];
6286 while gi < gpr {
6287 let t = base.add(gi * Q1_TILE);
6288 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6289 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6290 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6291 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6292 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6293 for k in 0..4 {
6294 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6295 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6296 }
6297 gi += 1;
6298 }
6299 acc
6300 }
6301}
6302
6303#[inline]
6305fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6306 let gi = j / GROUP_SIZE;
6307 let k = j % GROUP_SIZE;
6308 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6309 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6310 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6311 ((bit as i32 * 2 - 1) as f32, s)
6312}
6313
6314#[inline]
6316fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6317 let mut acc = 0f32;
6318 for gi in 0..gpr {
6319 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6320 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6321 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6322 let mut ga = 0f32;
6323 for (j, &b) in tile[2..].iter().enumerate() {
6324 for k in 0..8 {
6325 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6326 }
6327 }
6328 acc += ga * s;
6329 }
6330 acc
6331}
6332
6333#[allow(clippy::too_many_arguments)]
6336fn q1_range_a8w8(
6337 bytes: &[u8],
6338 gpr: usize,
6339 act: &SplitAct,
6340 gsum: &[i32],
6341 out: SendMut,
6342 start: usize,
6343 end: usize,
6344) {
6345 for r in start..end {
6346 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6347 for &(j, xv) in &act.outliers {
6348 let (w, s) = q1_outlier(bytes, r, gpr, j);
6349 acc += w * s * xv;
6350 }
6351 unsafe { *out.at(r) = acc };
6353 }
6354}
6355
6356fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6358 for r in start..end {
6359 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6361 }
6362}
6363
6364fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6369 let entries = base_len + (rows + 1) * 4;
6370 (base_len, entries, entries <= bytes.len())
6371}
6372
6373#[inline]
6375fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6376 let o = rp_off + r * 4;
6377 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6378}
6379
6380const SIGN5: [[f32; 5]; 256] = {
6384 let mut lut = [[0.0f32; 5]; 256];
6385 let pow3 = [1u16, 3, 9, 27, 81];
6386 let mut byte = 0usize;
6387 while byte < 256 {
6388 let mut i = 0usize;
6389 while i < 5 {
6390 let code = (byte as u16 / pow3[i]) % 3;
6391 lut[byte][i] = if code == 1 {
6392 1.0
6393 } else if code == 2 {
6394 -1.0
6395 } else {
6396 0.0
6397 };
6398 i += 1;
6399 }
6400 byte += 1;
6401 }
6402 lut
6403};
6404
6405const SIGN5_I8: [[i8; 5]; 256] = {
6407 let mut lut = [[0i8; 5]; 256];
6408 let pow3 = [1u16, 3, 9, 27, 81];
6409 let mut byte = 0usize;
6410 while byte < 256 {
6411 let mut i = 0usize;
6412 while i < 5 {
6413 let code = (byte as u16 / pow3[i]) % 3;
6414 lut[byte][i] = if code == 1 {
6415 1
6416 } else if code == 2 {
6417 -1
6418 } else {
6419 0
6420 };
6421 i += 1;
6422 }
6423 byte += 1;
6424 }
6425 lut
6426};
6427
6428const SIGN5_U64: [u64; 256] = {
6434 let mut lut = [0u64; 256];
6435 let pow3 = [1u16, 3, 9, 27, 81];
6436 let mut byte = 0usize;
6437 while byte < 256 {
6438 let mut v = 0u64;
6439 let mut i = 0usize;
6440 while i < 5 {
6441 let code = (byte as u16 / pow3[i]) % 3;
6442 let s: u8 = if code == 1 {
6443 1
6444 } else if code == 2 {
6445 0xFF
6446 } else {
6447 0
6448 };
6449 v |= (s as u64) << (i * 8);
6450 i += 1;
6451 }
6452 lut[byte] = v;
6453 byte += 1;
6454 }
6455 lut
6456};
6457
6458#[inline]
6463fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6464 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6465 let off = (r * gpr + j / GROUP_SIZE) * TILE;
6466 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6467 let within = j % GROUP_SIZE;
6468 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6469}
6470
6471#[cfg(target_arch = "aarch64")]
6474#[target_feature(enable = "neon,dotprod")]
6475#[inline]
6476unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6477 unsafe {
6479 use core::arch::aarch64::*;
6480 use core::arch::asm;
6481 let w0 = vld1q_s8(w);
6482 let w1 = vld1q_s8(w.add(16));
6483 let x0 = vld1q_s8(x);
6484 let x1 = vld1q_s8(x.add(16));
6485 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6486 asm!(
6487 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6488 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6489 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6490 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6491 options(pure, nomem, nostack),
6492 );
6493 vaddvq_s32(vaddq_s32(a0, a1))
6494 }
6495}
6496
6497#[cfg(target_arch = "x86_64")]
6500#[target_feature(enable = "avx2")]
6501#[inline]
6502unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6503 unsafe {
6505 use core::arch::x86_64::*;
6506 let wv = _mm256_loadu_si256(w as *const __m256i);
6507 let xv = _mm256_loadu_si256(x as *const __m256i);
6508 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6509 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6510 let hi128 = _mm256_extracti128_si256::<1>(d);
6511 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6512 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6513 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6514 _mm_cvtsi128_si32(s32)
6515 }
6516}
6517
6518#[inline]
6523fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6524 debug_assert!(dst.len() >= 40);
6525 unsafe {
6528 let p = dst.as_mut_ptr();
6529 for bi in 0..7 {
6530 core::ptr::write_unaligned(
6531 p.add(bi * 5) as *mut u64,
6532 SIGN5_U64[*codes.add(bi) as usize],
6533 );
6534 }
6535 }
6536}
6537
6538#[inline]
6543fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6544 #[cfg(target_arch = "aarch64")]
6545 unsafe {
6546 return sdot32_i8(w, x);
6547 }
6548 #[cfg(target_arch = "x86_64")]
6549 unsafe {
6550 return i8dot32_avx2(w, x);
6551 }
6552 #[allow(unreachable_code)]
6553 unsafe {
6554 let mut s = 0i32;
6555 for k in 0..GROUP_SIZE {
6556 s += *w.add(k) as i32 * *x.add(k) as i32;
6557 }
6558 s
6559 }
6560}
6561
6562#[inline]
6563unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6564 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6565 (
6566 SIGN5_U64[*codes as usize],
6567 SIGN5_U64[*codes.add(1) as usize],
6568 SIGN5_U64[*codes.add(2) as usize],
6569 SIGN5_U64[*codes.add(3) as usize],
6570 SIGN5_U64[*codes.add(4) as usize],
6571 SIGN5_U64[*codes.add(5) as usize],
6572 SIGN5_U64[*codes.add(6) as usize],
6573 )
6574 };
6575
6576 let u0 = s0 | (s1 << 40);
6577 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6578 let u2 = (s3 >> 8) | (s4 << 32);
6579 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6580
6581 (u0, u1, u2, u3)
6582}
6583
6584#[cfg(target_arch = "aarch64")]
6588#[target_feature(enable = "neon,dotprod")]
6589unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6590 use core::arch::aarch64::*;
6591 use core::arch::asm;
6592 unsafe {
6593 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6594 let mut acc = 0f32;
6595 let bytes_ptr = bytes.as_ptr();
6596 let xq_ptr = xq.as_ptr();
6597 let row_off = r * gpr * TILE;
6598
6599 let gpr2 = gpr & !1;
6600 let mut gi = 0;
6601 while gi < gpr2 {
6602 let off0 = row_off + gi * TILE;
6603 let off1 = off0 + TILE;
6604 let s0 = f16_to_f32(u16::from_le_bytes([
6605 *bytes_ptr.add(off0),
6606 *bytes_ptr.add(off0 + 1),
6607 ]));
6608 let s1 = f16_to_f32(u16::from_le_bytes([
6609 *bytes_ptr.add(off1),
6610 *bytes_ptr.add(off1 + 1),
6611 ]));
6612
6613 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6614 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6615
6616 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6617 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6618 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6619 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6620
6621 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6622 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6623 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6624 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6625
6626 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6627 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6628 asm!(
6629 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6630 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6631 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6632 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6633 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6634 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6635 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6636 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6637 options(pure, nomem, nostack),
6638 );
6639 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6640 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6641 acc += d0 as f32 * s0 + d1 as f32 * s1;
6642 gi += 2;
6643 }
6644
6645 if gi < gpr {
6646 let off = row_off + gi * TILE;
6647 let s = f16_to_f32(u16::from_le_bytes([
6648 *bytes_ptr.add(off),
6649 *bytes_ptr.add(off + 1),
6650 ]));
6651 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6652 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6653 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6654 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6655 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6656 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6657 asm!(
6658 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6659 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6660 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6661 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6662 options(pure, nomem, nostack),
6663 );
6664 let d = vaddvq_s32(vaddq_s32(a0, a1));
6665 acc += d as f32 * s;
6666 }
6667 acc
6668 }
6669}
6670
6671#[cfg(target_arch = "x86_64")]
6673#[target_feature(enable = "avx2")]
6674unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6675 use core::arch::x86_64::*;
6676 unsafe {
6677 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6678 let mut acc = 0f32;
6679 let bytes_ptr = bytes.as_ptr();
6680 let xq_ptr = xq.as_ptr();
6681 let row_off = r * gpr * TILE;
6682
6683 let ones = _mm256_set1_epi16(1);
6684 for gi in 0..gpr {
6685 let off = row_off + gi * TILE;
6686 let s = f16_to_f32(u16::from_le_bytes([
6687 *bytes_ptr.add(off),
6688 *bytes_ptr.add(off + 1),
6689 ]));
6690 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6691 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6692 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6693 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6694 let d256 = _mm256_madd_epi16(p16, ones);
6695 let d128 = _mm_add_epi32(
6696 _mm256_castsi256_si128(d256),
6697 _mm256_extracti128_si256(d256, 1),
6698 );
6699 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6700 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6701 acc += d32 as f32 * s;
6702 }
6703 acc
6704 }
6705}
6706
6707#[cfg(target_arch = "x86_64")]
6709#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6710unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6711 use core::arch::x86_64::*;
6712 unsafe {
6714 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6715 let mut acc = 0f32;
6716 let bytes_ptr = bytes.as_ptr();
6717 let xq_ptr = xq.as_ptr();
6718 let row_off = r * gpr * TILE;
6719 for gi in 0..gpr {
6720 let off = row_off + gi * TILE;
6721 let s = f16_to_f32(u16::from_le_bytes([
6722 *bytes_ptr.add(off),
6723 *bytes_ptr.add(off + 1),
6724 ]));
6725 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6726 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6727 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6728 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6729 acc += d as f32 * s;
6730 }
6731 acc
6732 }
6733}
6734
6735#[inline]
6739fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6740 #[cfg(target_arch = "aarch64")]
6741 unsafe {
6742 return q1t_dot_row_sdot(bytes, r, gpr, xq);
6743 }
6744 #[cfg(target_arch = "x86_64")]
6745 unsafe {
6746 if vnni_tiles_enabled() {
6747 return q1t_dot_row_vnni(bytes, r, gpr, xq);
6748 }
6749 return q1t_dot_row_avx2(bytes, r, gpr, xq);
6750 }
6751 #[allow(unreachable_code)]
6752 {
6753 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6754 let mut acc = 0f32;
6755 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
6757 let off = (r * gpr + gi) * TILE;
6758 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6759 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6760 let mut d = 0i32;
6761 for k in 0..GROUP_SIZE {
6762 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6763 }
6764 acc += d as f32 * s;
6765 }
6766 acc
6767 }
6768}
6769
6770fn q1t_row_outlier_correction(
6777 bytes: &[u8],
6778 r: usize,
6779 rp_off: usize,
6780 entries_off: usize,
6781 has_ov: bool,
6782 x: &[f32],
6783) -> f32 {
6784 if !has_ov {
6785 return 0.0;
6786 }
6787 let (c0, c1) = (
6788 q1t_rowptr(bytes, rp_off, r),
6789 q1t_rowptr(bytes, rp_off, r + 1),
6790 );
6791 let mut corr = 0f32;
6792 for p in c0..c1 {
6793 let e = entries_off + p * 4;
6794 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6795 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6796 corr += val * x[col];
6797 }
6798 corr
6799}
6800
6801fn q1t_dequant_row(
6805 bytes: &[u8],
6806 r: usize,
6807 gpr: usize,
6808 rp_off: usize,
6809 entries_off: usize,
6810 has_ov: bool,
6811 buf: &mut [f32],
6812) {
6813 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6814 for g in 0..gpr {
6815 let off = (r * gpr + g) * TILE;
6816 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6817 let codes = &bytes[off + 2..off + TILE];
6818 let bc = g * GROUP_SIZE;
6819 for bi in 0..6 {
6821 let lut = &SIGN5[codes[bi] as usize];
6822 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6823 for i in 0..5 {
6824 d[i] = lut[i] * s;
6825 }
6826 }
6827 let lut = &SIGN5[codes[6] as usize];
6828 buf[bc + 30] = lut[0] * s;
6829 buf[bc + 31] = lut[1] * s;
6830 }
6831 if !has_ov {
6832 return;
6833 }
6834 let (c0, c1) = (
6835 q1t_rowptr(bytes, rp_off, r),
6836 q1t_rowptr(bytes, rp_off, r + 1),
6837 );
6838 for p in c0..c1 {
6839 let e = entries_off + p * 4;
6840 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6841 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6842 }
6843}
6844
6845fn q1t_add_overlay(
6849 bytes: &[u8],
6850 x: &[f32],
6851 rows: usize,
6852 cols: usize,
6853 out: &mut [f32],
6854 pool: Option<&Pool>,
6855) {
6856 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6857 let gpr = cols / GROUP_SIZE;
6858 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6859 if !has_ov {
6860 return;
6861 }
6862 let out_addr = SendMut(out.as_mut_ptr());
6863 let run = move |start: usize, end: usize| {
6864 for r in start..end {
6865 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6866 unsafe { *out_addr.at(r) += corr };
6868 }
6869 };
6870 dispatch_rows(pool, rows, &run);
6871}
6872
6873#[allow(clippy::too_many_arguments)]
6876fn q1t_range_a8w8(
6877 bytes: &[u8],
6878 gpr: usize,
6879 rp_off: usize,
6880 ent_off: usize,
6881 has_ov: bool,
6882 act: &SplitAct,
6883 x: &[f32],
6884 out: SendMut,
6885 start: usize,
6886 end: usize,
6887) {
6888 for r in start..end {
6889 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6890 for &(j, xv) in &act.outliers {
6891 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6892 }
6893 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6894 unsafe { *out.at(r) = acc };
6896 }
6897}
6898
6899#[allow(clippy::too_many_arguments)]
6902fn q1t_range_f32_batch(
6903 bytes: &[u8],
6904 gpr: usize,
6905 rp_off: usize,
6906 ent_off: usize,
6907 has_ov: bool,
6908 x: &[f32],
6909 out: SendMut,
6910 start: usize,
6911 end: usize,
6912) {
6913 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6914 let mut sg = [0f32; GROUP_SIZE];
6915 for r in start..end {
6916 let mut acc = 0f32;
6917 for g in 0..gpr {
6918 let off = (r * gpr + g) * TILE;
6919 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6920 let codes = &bytes[off + 2..off + TILE];
6921 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6922 for bi in 0..6 {
6923 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6924 }
6925 let lut = &SIGN5[codes[6] as usize];
6926 sg[30] = lut[0];
6927 sg[31] = lut[1];
6928 let mut gsum = 0f32;
6929 for k in 0..GROUP_SIZE {
6930 gsum += sg[k] * xg[k];
6931 }
6932 acc += s * gsum;
6933 }
6934 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6935 unsafe { *out.at(r) = acc };
6937 }
6938}
6939
6940fn q1t_matvec(
6944 bytes: &[u8],
6945 x: &[f32],
6946 rows: usize,
6947 cols: usize,
6948 out: &mut [f32],
6949 pool: Option<&Pool>,
6950) {
6951 debug_assert_eq!(out.len(), rows);
6952 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6953 let gpr = cols / GROUP_SIZE;
6954 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6955 let out_addr = SendMut(out.as_mut_ptr());
6956 if a8w8_enabled() {
6960 let act = split_act(x);
6961 let act = &act;
6962 let run = move |start: usize, end: usize| {
6963 for r in start..end {
6964 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6965 for &(j, xv) in &act.outliers {
6966 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6967 }
6968 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6969 unsafe { *out_addr.at(r) = acc };
6971 }
6972 };
6973 dispatch_rows(pool, rows, &run);
6974 return;
6975 }
6976 let run = move |start: usize, end: usize| {
6977 let mut sg = [0f32; GROUP_SIZE];
6981 for r in start..end {
6982 let mut acc = 0f32;
6983 for g in 0..gpr {
6984 let off = (r * gpr + g) * TILE;
6985 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6986 let codes = &bytes[off + 2..off + TILE];
6987 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6988 for bi in 0..6 {
6989 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6990 }
6991 let lut = &SIGN5[codes[6] as usize];
6992 sg[30] = lut[0];
6993 sg[31] = lut[1];
6994 let mut gsum = 0f32;
6995 for k in 0..GROUP_SIZE {
6996 gsum += sg[k] * xg[k];
6997 }
6998 acc += s * gsum;
6999 }
7000 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7001 unsafe { *out_addr.at(r) = acc };
7002 }
7003 };
7004 dispatch_rows(pool, rows, &run);
7005}
7006
7007#[cfg(target_arch = "aarch64")]
7013#[target_feature(enable = "neon,dotprod")]
7014unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7015 use core::arch::aarch64::*;
7016 use core::arch::asm;
7017 unsafe {
7019 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7020 let bytes_ptr = bytes.as_ptr();
7021 let row_off = r * gpr * TILE;
7022 let xp = [xa.as_ptr(), xb.as_ptr()];
7023 let mut acc = [0f32; 2];
7024 macro_rules! sdot2 {
7025 ($w0:expr, $w1:expr, $x:expr) => {{
7026 let x0 = vld1q_s8($x);
7027 let x1 = vld1q_s8($x.add(16));
7028 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7029 asm!(
7030 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7031 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7032 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7033 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7034 options(pure, nomem, nostack),
7035 );
7036 vaddvq_s32(vaddq_s32(a0, a1))
7037 }};
7038 }
7039 let gpr2 = gpr & !1;
7040 let mut gi = 0;
7041 while gi < gpr2 {
7042 let off0 = row_off + gi * TILE;
7043 let off1 = off0 + TILE;
7044 let s0 = f16_to_f32(u16::from_le_bytes([
7045 *bytes_ptr.add(off0),
7046 *bytes_ptr.add(off0 + 1),
7047 ]));
7048 let s1 = f16_to_f32(u16::from_le_bytes([
7049 *bytes_ptr.add(off1),
7050 *bytes_ptr.add(off1 + 1),
7051 ]));
7052 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7053 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7054 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7055 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7056 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7057 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7058 for k in 0..2 {
7059 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7060 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7061 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7062 }
7063 gi += 2;
7064 }
7065 if gi < gpr {
7066 let off = row_off + gi * TILE;
7067 let s = f16_to_f32(u16::from_le_bytes([
7068 *bytes_ptr.add(off),
7069 *bytes_ptr.add(off + 1),
7070 ]));
7071 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7072 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7073 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7074 for k in 0..2 {
7075 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7076 acc[k] += d as f32 * s;
7077 }
7078 }
7079 acc
7080 }
7081}
7082
7083fn q1t_matvec2(
7089 bytes: &[u8],
7090 x1: &[f32],
7091 x2: &[f32],
7092 rows: usize,
7093 cols: usize,
7094 o1: &mut [f32],
7095 o2: &mut [f32],
7096 pool: Option<&Pool>,
7097) {
7098 debug_assert_eq!(o1.len(), rows);
7099 debug_assert_eq!(o2.len(), rows);
7100 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7101 let gpr = cols / GROUP_SIZE;
7102 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7103 let out1 = SendMut(o1.as_mut_ptr());
7104 let out2 = SendMut(o2.as_mut_ptr());
7105 if a8w8_enabled() {
7106 let a1 = split_act(x1);
7107 let a2 = split_act(x2);
7108 let (a1, a2) = (&a1, &a2);
7109 let run = move |start: usize, end: usize| {
7110 for r in start..end {
7111 #[cfg(target_arch = "aarch64")]
7112 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7115 #[cfg(not(target_arch = "aarch64"))]
7116 let ds = [
7117 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7118 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7119 ];
7120 let mut acc1 = ds[0] * a1.sx;
7121 for &(j, xv) in &a1.outliers {
7122 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7123 }
7124 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7125 let mut acc2 = ds[1] * a2.sx;
7126 for &(j, xv) in &a2.outliers {
7127 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7128 }
7129 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7130 unsafe {
7132 *out1.at(r) = acc1;
7133 *out2.at(r) = acc2;
7134 }
7135 }
7136 };
7137 dispatch_rows(pool, rows, &run);
7138 return;
7139 }
7140 let run = move |start: usize, end: usize| {
7141 let mut sg = [0f32; GROUP_SIZE];
7144 for r in start..end {
7145 let mut acc1 = 0f32;
7146 let mut acc2 = 0f32;
7147 for g in 0..gpr {
7148 let off = (r * gpr + g) * TILE;
7149 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7150 let codes = &bytes[off + 2..off + TILE];
7151 for bi in 0..6 {
7152 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7153 }
7154 let lut = &SIGN5[codes[6] as usize];
7155 sg[30] = lut[0];
7156 sg[31] = lut[1];
7157 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7158 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7159 let mut gsum1 = 0f32;
7160 for k in 0..GROUP_SIZE {
7161 gsum1 += sg[k] * xg1[k];
7162 }
7163 acc1 += s * gsum1;
7164 let mut gsum2 = 0f32;
7165 for k in 0..GROUP_SIZE {
7166 gsum2 += sg[k] * xg2[k];
7167 }
7168 acc2 += s * gsum2;
7169 }
7170 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7171 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7172 unsafe {
7174 *out1.at(r) = acc1;
7175 *out2.at(r) = acc2;
7176 }
7177 }
7178 };
7179 dispatch_rows(pool, rows, &run);
7180}
7181
7182fn q1t_matmat(
7185 bytes: &[u8],
7186 xs: &[f32],
7187 b: usize,
7188 rows: usize,
7189 cols: usize,
7190 out: &mut [f32],
7191 pool: Option<&Pool>,
7192) {
7193 debug_assert_eq!(out.len(), b * rows);
7194 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7195 let gpr = cols / GROUP_SIZE;
7196 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7197 let out_addr = SendMut(out.as_mut_ptr());
7198 if a8w8_enabled() {
7202 let acts: Vec<SplitAct> = (0..b)
7203 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7204 .collect();
7205 let acts = &acts;
7206 let run = move |start: usize, end: usize| {
7207 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
7211 for g in 0..gpr {
7212 let off = (r * gpr + g) * TILE;
7213 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7214 q1t_unpack_group_i8(
7215 bytes.as_ptr().wrapping_add(off + 2),
7216 &mut sg[g * GROUP_SIZE..],
7217 );
7218 }
7219 for bi in 0..b {
7220 let act = &acts[bi];
7221 let mut isum = 0f32;
7222 for g in 0..gpr {
7223 let d = q1t_i8dot32(
7224 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7225 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7226 );
7227 isum += d as f32 * sc[g];
7228 }
7229 let mut acc = isum * act.sx;
7230 for &(j, xv) in &act.outliers {
7231 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7232 }
7233 accs[bi] = acc;
7234 }
7235 if has_ov {
7239 let (c0, c1) = (
7240 q1t_rowptr(bytes, rp_off, r),
7241 q1t_rowptr(bytes, rp_off, r + 1),
7242 );
7243 for p in c0..c1 {
7244 let e = ent_off + p * 4;
7245 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7246 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7247 for bi in 0..b {
7248 accs[bi] += val * xs[bi * cols + col];
7249 }
7250 }
7251 }
7252 for bi in 0..b {
7253 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7254 }
7255 }
7256 };
7257 dispatch_rows(pool, rows, &run);
7258 return;
7259 }
7260 let run = move |start: usize, end: usize| {
7261 let mut buf = vec![0f32; cols];
7262 for r in start..end {
7263 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7264 for bi in 0..b {
7265 let xr = &xs[bi * cols..(bi + 1) * cols];
7266 let mut acc = 0f32;
7267 for j in 0..cols {
7268 acc += buf[j] * xr[j];
7269 }
7270 unsafe { *out_addr.at(bi * rows + r) = acc };
7271 }
7272 }
7273 };
7274 dispatch_rows(pool, rows, &run);
7275}
7276
7277fn q1_matvec(
7278 bytes: &[u8],
7279 x: &[f32],
7280 rows: usize,
7281 cols: usize,
7282 out: &mut [f32],
7283 pool: Option<&Pool>,
7284) {
7285 debug_assert_eq!(out.len(), rows);
7286 let gpr = cols / GROUP_SIZE;
7287 let out_addr = SendMut(out.as_mut_ptr());
7288 if a8w8_enabled() {
7289 let act = split_act(x);
7290 let gsum = q1_group_sums(&act.xq, gpr);
7291 let (act, gsum) = (&act, &gsum);
7292 let run = move |start: usize, end: usize| {
7293 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7294 };
7295 dispatch_rows(pool, rows, &run);
7296 return;
7297 }
7298 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7299 dispatch_rows(pool, rows, &run);
7300}
7301
7302#[allow(clippy::too_many_arguments)]
7304fn q1_matvec2(
7305 bytes: &[u8],
7306 x1: &[f32],
7307 x2: &[f32],
7308 rows: usize,
7309 cols: usize,
7310 o1: &mut [f32],
7311 o2: &mut [f32],
7312 pool: Option<&Pool>,
7313) {
7314 let gpr = cols / GROUP_SIZE;
7315 let p1 = SendMut(o1.as_mut_ptr());
7316 let p2 = SendMut(o2.as_mut_ptr());
7317 if a8w8_enabled() {
7318 let a1 = split_act(x1);
7319 let a2 = split_act(x2);
7320 let g1 = q1_group_sums(&a1.xq, gpr);
7321 let g2 = q1_group_sums(&a2.xq, gpr);
7322 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7323 let run = move |start: usize, end: usize| {
7324 for r in start..end {
7325 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7326 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7327 for &(j, xv) in &a1.outliers {
7328 let (w, s) = q1_outlier(bytes, r, gpr, j);
7329 v1 += w * s * xv;
7330 }
7331 for &(j, xv) in &a2.outliers {
7332 let (w, s) = q1_outlier(bytes, r, gpr, j);
7333 v2 += w * s * xv;
7334 }
7335 unsafe {
7337 *p1.at(r) = v1;
7338 *p2.at(r) = v2;
7339 }
7340 }
7341 };
7342 dispatch_rows(pool, rows, &run);
7343 return;
7344 }
7345 let run = move |start: usize, end: usize| {
7346 for r in start..end {
7347 unsafe {
7349 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7350 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7351 }
7352 }
7353 };
7354 dispatch_rows(pool, rows, &run);
7355}
7356
7357#[allow(clippy::too_many_arguments)]
7359fn q1_matmat(
7360 bytes: &[u8],
7361 xs_all: &[f32],
7362 b: usize,
7363 rows: usize,
7364 cols: usize,
7365 out: &mut [f32],
7366 pool: Option<&Pool>,
7367) {
7368 debug_assert_eq!(out.len(), b * rows);
7369 let gpr = cols / GROUP_SIZE;
7370 let out_addr = SendMut(out.as_mut_ptr());
7371 if a8w8_enabled() {
7372 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7373 .map(|bi| {
7374 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7375 let gsum = q1_group_sums(&act.xq, gpr);
7376 (act, gsum)
7377 })
7378 .collect();
7379 let acts = &acts;
7380 #[cfg(target_arch = "x86_64")]
7381 let blocked_ok = avx2_enabled() && blocked_enabled();
7382 #[cfg(target_arch = "aarch64")]
7383 let blocked_ok = sdot_enabled() && blocked_enabled();
7384 let run = move |start: usize, end: usize| {
7385 for r in start..end {
7386 let mut bi = 0usize;
7387 #[cfg(target_arch = "aarch64")]
7390 if blocked_ok {
7391 while bi + 4 <= acts.len() {
7392 let xs = [
7393 acts[bi].0.xq.as_slice(),
7394 acts[bi + 1].0.xq.as_slice(),
7395 acts[bi + 2].0.xq.as_slice(),
7396 acts[bi + 3].0.xq.as_slice(),
7397 ];
7398 let gs = [
7399 acts[bi].1.as_slice(),
7400 acts[bi + 1].1.as_slice(),
7401 acts[bi + 2].1.as_slice(),
7402 acts[bi + 3].1.as_slice(),
7403 ];
7404 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7405 for k in 0..4 {
7406 let (act, _) = &acts[bi + k];
7407 let mut acc = d[k] * act.sx;
7408 for &(j, xv) in &act.outliers {
7409 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7410 acc += w * sc * xv;
7411 }
7412 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7414 }
7415 bi += 4;
7416 }
7417 }
7418 #[cfg(target_arch = "x86_64")]
7419 if blocked_ok {
7420 while bi + 4 <= acts.len() {
7421 let xs = [
7422 acts[bi].0.xq.as_slice(),
7423 acts[bi + 1].0.xq.as_slice(),
7424 acts[bi + 2].0.xq.as_slice(),
7425 acts[bi + 3].0.xq.as_slice(),
7426 ];
7427 let gs = [
7428 acts[bi].1.as_slice(),
7429 acts[bi + 1].1.as_slice(),
7430 acts[bi + 2].1.as_slice(),
7431 acts[bi + 3].1.as_slice(),
7432 ];
7433 let d = unsafe {
7434 if vnni_tiles_enabled() {
7435 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7436 } else {
7437 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7438 }
7439 };
7440 for k in 0..4 {
7441 let (act, _) = &acts[bi + k];
7442 let mut acc = d[k] * act.sx;
7443 for &(j, xv) in &act.outliers {
7444 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7445 acc += w * sc * xv;
7446 }
7447 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7449 }
7450 bi += 4;
7451 }
7452 }
7453 while bi < acts.len() {
7454 let (act, gsum) = &acts[bi];
7455 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7456 for &(j, xv) in &act.outliers {
7457 let (w, s) = q1_outlier(bytes, r, gpr, j);
7458 acc += w * s * xv;
7459 }
7460 unsafe { *out_addr.at(bi * rows + r) = acc };
7462 bi += 1;
7463 }
7464 }
7465 };
7466 dispatch_rows(pool, rows, &run);
7467 return;
7468 }
7469 let run = move |start: usize, end: usize| {
7470 for r in start..end {
7471 for bi in 0..b {
7472 let x = &xs_all[bi * cols..(bi + 1) * cols];
7473 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7475 }
7476 }
7477 };
7478 dispatch_rows(pool, rows, &run);
7479}
7480
7481fn q4matvec(
7487 bytes: &[u8],
7488 x: &[f32],
7489 rows: usize,
7490 cols: usize,
7491 out: &mut [f32],
7492 pool: Option<&Pool>,
7493) {
7494 debug_assert_eq!(out.len(), rows);
7495 let (packed, scales) = q4_split(bytes, rows, cols);
7496 let gpr = cols / GROUP_SIZE;
7497 let out_addr = SendMut(out.as_mut_ptr());
7498
7499 if a8w8_enabled() {
7500 let act = split_act(x);
7501 let run = move |start: usize, end: usize| {
7502 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7503 };
7504 dispatch_rows(pool, rows, &run);
7505 return;
7506 }
7507
7508 let run =
7509 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7510 dispatch_rows(pool, rows, &run);
7511}
7512
7513#[inline]
7516#[allow(unreachable_code)]
7517#[cfg(target_arch = "x86_64")]
7522#[target_feature(enable = "avx2")]
7523unsafe fn dot_q4b_row_1x4_avx2(
7524 buf: &[u8],
7525 scales: &[u8],
7526 g0: usize,
7527 gpr: usize,
7528 xs: [&[i8]; 4],
7529) -> [f32; 4] {
7530 unsafe {
7532 use core::arch::x86_64::*;
7533 let ones = _mm256_set1_epi16(1);
7534 let mut acc = [0f32; 4];
7535 for gi in 0..gpr {
7536 let s = f16_to_f32(u16::from_le_bytes([
7537 scales[(g0 + gi) * 2],
7538 scales[(g0 + gi) * 2 + 1],
7539 ]));
7540 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7541 let aw = _mm256_abs_epi8(w);
7542 for (k, xq) in xs.iter().enumerate() {
7543 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7544 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7545 let d = _mm256_madd_epi16(p16, ones);
7546 let hi128 = _mm256_extracti128_si256::<1>(d);
7547 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7548 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7549 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7550 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7551 }
7552 }
7553 acc
7554 }
7555}
7556
7557#[cfg(target_arch = "x86_64")]
7559#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7560unsafe fn dot_q4b_row_1x4_vnni(
7561 buf: &[u8],
7562 scales: &[u8],
7563 g0: usize,
7564 gpr: usize,
7565 xs: [&[i8]; 4],
7566) -> [f32; 4] {
7567 unsafe {
7569 use core::arch::x86_64::*;
7570 let mut acc = [0f32; 4];
7571 for gi in 0..gpr {
7572 let s = f16_to_f32(u16::from_le_bytes([
7573 scales[(g0 + gi) * 2],
7574 scales[(g0 + gi) * 2 + 1],
7575 ]));
7576 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7577 let aw = _mm256_abs_epi8(w);
7578 for (k, xq) in xs.iter().enumerate() {
7579 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7580 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7581 acc[k] += d as f32 * s;
7582 }
7583 }
7584 acc
7585 }
7586}
7587
7588#[cfg(target_arch = "x86_64")]
7594#[target_feature(enable = "avx2")]
7595unsafe fn dot_q4b_row_1x4_sx_avx2(
7596 buf: &[u8],
7597 scales: &[u8],
7598 g0: usize,
7599 gpr: usize,
7600 xs: [&[i8]; 4],
7601 sxs: [f32; 4],
7602) -> [f32; 4] {
7603 unsafe {
7605 use core::arch::x86_64::*;
7606 let ones = _mm256_set1_epi16(1);
7607 let mut acc = [0f32; 4];
7608 for gi in 0..gpr {
7609 let s = f16_to_f32(u16::from_le_bytes([
7610 scales[(g0 + gi) * 2],
7611 scales[(g0 + gi) * 2 + 1],
7612 ]));
7613 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7614 let aw = _mm256_abs_epi8(w);
7615 for (k, xq) in xs.iter().enumerate() {
7616 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7617 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7618 let d = _mm256_madd_epi16(p16, ones);
7619 let hi128 = _mm256_extracti128_si256::<1>(d);
7620 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7621 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7622 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7623 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7624 }
7625 }
7626 acc
7627 }
7628}
7629
7630#[cfg(target_arch = "x86_64")]
7633#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7634unsafe fn dot_q4b_row_1x4_sx_vnni(
7635 buf: &[u8],
7636 scales: &[u8],
7637 g0: usize,
7638 gpr: usize,
7639 xs: [&[i8]; 4],
7640 sxs: [f32; 4],
7641) -> [f32; 4] {
7642 unsafe {
7644 use core::arch::x86_64::*;
7645 let mut acc = [0f32; 4];
7646 for gi in 0..gpr {
7647 let s = f16_to_f32(u16::from_le_bytes([
7648 scales[(g0 + gi) * 2],
7649 scales[(g0 + gi) * 2 + 1],
7650 ]));
7651 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7652 let aw = _mm256_abs_epi8(w);
7653 for (k, xq) in xs.iter().enumerate() {
7654 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7655 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7656 acc[k] += (d as f32 * sxs[k]) * s;
7657 }
7658 }
7659 acc
7660 }
7661}
7662
7663#[allow(unreachable_code)]
7664fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7665 #[cfg(target_arch = "aarch64")]
7666 unsafe {
7667 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7668 }
7669 #[cfg(target_arch = "x86_64")]
7670 unsafe {
7671 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7672 }
7673 let mut acc = 0f32;
7674 for gi in 0..gpr {
7675 let g = g0 + gi;
7676 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7677 let mut d = 0i32;
7678 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7679 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7680 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7681 }
7682 acc += d as f32 * s;
7683 }
7684 acc
7685}
7686
7687#[inline]
7689#[allow(unreachable_code)]
7690fn dot_q4_row_i8_2(
7691 packed: &[u8],
7692 scales: &[u8],
7693 g0: usize,
7694 gpr: usize,
7695 xq1: &[i8],
7696 xq2: &[i8],
7697) -> (f32, f32) {
7698 #[cfg(target_arch = "aarch64")]
7699 unsafe {
7700 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7701 }
7702 #[cfg(target_arch = "x86_64")]
7703 unsafe {
7704 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7705 }
7706 (
7707 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7708 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7709 )
7710}
7711
7712#[allow(clippy::too_many_arguments)]
7715fn q4_range_a8w8(
7716 packed: &[u8],
7717 scales: &[u8],
7718 gpr: usize,
7719 cols: usize,
7720 act: &SplitAct,
7721 out: SendMut,
7722 start: usize,
7723 end: usize,
7724) {
7725 for r in start..end {
7726 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7727 for &(j, xv) in &act.outliers {
7729 let flat = r * cols + j;
7730 let byte = packed[flat / 2];
7731 let nib = if flat & 1 == 0 {
7732 byte & 0x0F
7733 } else {
7734 byte >> 4
7735 };
7736 let s = f16_to_f32(u16::from_le_bytes([
7737 scales[(flat / GROUP_SIZE) * 2],
7738 scales[(flat / GROUP_SIZE) * 2 + 1],
7739 ]));
7740 acc += ((nib as i32 - 8) as f32) * s * xv;
7741 }
7742 unsafe { *out.at(r) = acc };
7744 }
7745}
7746
7747#[allow(clippy::too_many_arguments)]
7750fn q4_range2_a8w8(
7751 packed: &[u8],
7752 scales: &[u8],
7753 gpr: usize,
7754 cols: usize,
7755 a1: &SplitAct,
7756 a2: &SplitAct,
7757 p1: SendMut,
7758 p2: SendMut,
7759 start: usize,
7760 end: usize,
7761) {
7762 for r in start..end {
7763 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7764 let mut acc1 = s1 * a1.sx;
7765 let mut acc2 = s2 * a2.sx;
7766 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7768 for &(j, xv) in outliers {
7769 let flat = r * cols + j;
7770 let byte = packed[flat / 2];
7771 let nib = if flat & 1 == 0 {
7772 byte & 0x0F
7773 } else {
7774 byte >> 4
7775 };
7776 let s = f16_to_f32(u16::from_le_bytes([
7777 scales[(flat / GROUP_SIZE) * 2],
7778 scales[(flat / GROUP_SIZE) * 2 + 1],
7779 ]));
7780 *acc += ((nib as i32 - 8) as f32) * s * xv;
7781 }
7782 };
7783 fix(&a1.outliers, &mut acc1);
7784 fix(&a2.outliers, &mut acc2);
7785 unsafe {
7787 *p1.at(r) = acc1;
7788 *p2.at(r) = acc2;
7789 }
7790 }
7791}
7792
7793fn q4_range_f32(
7795 packed: &[u8],
7796 scales: &[u8],
7797 gpr: usize,
7798 x: &[f32],
7799 out: SendMut,
7800 start: usize,
7801 end: usize,
7802) {
7803 for r in start..end {
7804 let mut acc = 0f32;
7805 for gi in 0..gpr {
7806 let g = r * gpr + gi;
7807 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7808 let pk = &packed[g * 16..(g + 1) * 16];
7809 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7810 let mut ga = 0f32;
7811 for (k, &b) in pk.iter().enumerate() {
7812 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7813 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7814 }
7815 acc += ga * s;
7816 }
7817 unsafe { *out.at(r) = acc };
7819 }
7820}
7821
7822#[allow(clippy::too_many_arguments)]
7826fn q4matvec2(
7827 bytes: &[u8],
7828 x1: &[f32],
7829 x2: &[f32],
7830 rows: usize,
7831 cols: usize,
7832 o1: &mut [f32],
7833 o2: &mut [f32],
7834 pool: Option<&Pool>,
7835) {
7836 debug_assert_eq!(o1.len(), rows);
7837 debug_assert_eq!(o2.len(), rows);
7838 let (packed, scales) = q4_split(bytes, rows, cols);
7839 let gpr = cols / GROUP_SIZE;
7840
7841 if a8w8_enabled() {
7842 let a1 = split_act(x1);
7843 let a2 = split_act(x2);
7844 let p1 = SendMut(o1.as_mut_ptr());
7845 let p2 = SendMut(o2.as_mut_ptr());
7846 let run = move |start: usize, end: usize| {
7847 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7848 };
7849 dispatch_rows(pool, rows, &run);
7850 return;
7851 }
7852
7853 let p1 = SendMut(o1.as_mut_ptr());
7854 let p2 = SendMut(o2.as_mut_ptr());
7855 let run = move |start: usize, end: usize| {
7856 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7857 };
7858 dispatch_rows(pool, rows, &run);
7859}
7860
7861#[allow(clippy::too_many_arguments)]
7863fn q4_range2_f32(
7864 packed: &[u8],
7865 scales: &[u8],
7866 gpr: usize,
7867 x1: &[f32],
7868 x2: &[f32],
7869 p1: SendMut,
7870 p2: SendMut,
7871 start: usize,
7872 end: usize,
7873) {
7874 for r in start..end {
7875 let (mut acc1, mut acc2) = (0f32, 0f32);
7876 for gi in 0..gpr {
7877 let g = r * gpr + gi;
7878 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7879 let pk = &packed[g * 16..(g + 1) * 16];
7880 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7881 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7882 let (mut g1, mut g2) = (0f32, 0f32);
7883 for (k, &b) in pk.iter().enumerate() {
7884 let wl = (b & 0x0F) as f32 - 8.0;
7885 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7886 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7887 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7888 }
7889 acc1 += g1 * s;
7890 acc2 += g2 * s;
7891 }
7892 unsafe {
7894 *p1.at(r) = acc1;
7895 *p2.at(r) = acc2;
7896 }
7897 }
7898}
7899
7900thread_local! {
7901 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7904 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7905}
7906
7907#[allow(clippy::too_many_arguments)]
7913fn q4matmat(
7914 bytes: &[u8],
7915 xs_all: &[f32],
7916 b: usize,
7917 rows: usize,
7918 cols: usize,
7919 out: &mut [f32],
7920 pool: Option<&Pool>,
7921) {
7922 debug_assert_eq!(xs_all.len(), b * cols);
7923 debug_assert_eq!(out.len(), b * rows);
7924 let (packed, scales) = q4_split(bytes, rows, cols);
7925 let gpr = cols / GROUP_SIZE;
7926 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7927
7928 if a8w8_enabled() {
7929 let acts: Vec<SplitAct> = (0..b)
7930 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7931 .collect();
7932 let acts = &acts;
7933 let out_addr = SendMut(out.as_mut_ptr());
7934 let run = move |start: usize, end: usize| {
7935 ROW_I8.with(|rb| {
7936 let mut buf = rb.borrow_mut();
7937 buf.resize(cols, 0);
7938 for r in start..end {
7939 for gi in 0..gpr {
7943 let g = r * gpr + gi;
7944 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7945 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7946 buf[gi * GROUP_SIZE + k * 2 + 1] =
7947 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7948 }
7949 }
7950 let mut bi = 0usize;
7951 #[cfg(target_arch = "x86_64")]
7952 if avx2_enabled() && blocked_enabled() {
7953 while bi + 4 <= acts.len() {
7954 let xs = [
7955 acts[bi].xq.as_slice(),
7956 acts[bi + 1].xq.as_slice(),
7957 acts[bi + 2].xq.as_slice(),
7958 acts[bi + 3].xq.as_slice(),
7959 ];
7960 let d = unsafe {
7961 if vnni_tiles_enabled() {
7962 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7963 } else {
7964 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7965 }
7966 };
7967 for k in 0..4 {
7968 let act = &acts[bi + k];
7969 let mut acc = d[k] * act.sx;
7970 for &(j, xv) in &act.outliers {
7971 acc += (buf[j] as i8) as f32
7972 * gscale((r * cols + j) / GROUP_SIZE)
7973 * xv;
7974 }
7975 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7977 }
7978 bi += 4;
7979 }
7980 }
7981 while bi < acts.len() {
7982 let act = &acts[bi];
7983 let mut acc = 0f32;
7984 for gi in 0..gpr {
7985 let d = dot_i8_i8(
7986 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7987 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7988 );
7989 acc += d as f32 * gscale(r * gpr + gi);
7990 }
7991 acc *= act.sx;
7992 for &(j, xv) in &act.outliers {
7994 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7995 }
7996 unsafe { *out_addr.at(bi * rows + r) = acc };
7998 bi += 1;
7999 }
8000 }
8001 })
8002 };
8003 dispatch_rows(pool, rows, &run);
8004 return;
8005 }
8006
8007 let out_addr = SendMut(out.as_mut_ptr());
8008 let run = move |start: usize, end: usize| {
8009 ROW_F32.with(|rb| {
8010 let mut buf = rb.borrow_mut();
8011 buf.resize(cols, 0.0);
8012 for r in start..end {
8013 for gi in 0..gpr {
8016 let g = r * gpr + gi;
8017 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8018 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8019 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8020 }
8021 }
8022 for bi in 0..b {
8023 let x = &xs_all[bi * cols..(bi + 1) * cols];
8024 let mut acc = 0f32;
8025 for gi in 0..gpr {
8026 let mut ga = 0f32;
8027 for k in 0..GROUP_SIZE / 2 {
8032 let e = gi * GROUP_SIZE + k * 2;
8033 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8034 }
8035 acc += ga * gscale(r * gpr + gi);
8036 }
8037 unsafe { *out_addr.at(bi * rows + r) = acc };
8039 }
8040 }
8041 })
8042 };
8043 dispatch_rows(pool, rows, &run);
8044}
8045
8046#[allow(clippy::too_many_arguments)]
8051fn vbitmatmat(
8052 bytes: &[u8],
8053 offsets: &[usize],
8054 xs_all: &[f32],
8055 b: usize,
8056 rows: usize,
8057 cols: usize,
8058 out: &mut [f32],
8059 pool: Option<&Pool>,
8060) {
8061 debug_assert_eq!(xs_all.len(), b * cols);
8062 debug_assert_eq!(out.len(), b * rows);
8063 debug_assert_eq!(offsets.len(), rows + 1);
8064 let ng = cols / GROUP_SIZE;
8065 let bits = &bytes[..rows];
8066 let sc_off = rows;
8067 let gscale = |r: usize, g: usize| {
8068 let so = (r * ng + g) * 2;
8069 f16_to_f32(u16::from_le_bytes([
8070 bytes[sc_off + so],
8071 bytes[sc_off + so + 1],
8072 ]))
8073 };
8074
8075 let decode_f32 = |r: usize, dst: &mut [f32]| {
8077 let bw = bits[r] as usize;
8078 let l = ((1i32 << (bw - 1)) - 1) as f32;
8079 let data = &bytes[offsets[r]..offsets[r + 1]];
8080 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8081 for d in dst.iter_mut() {
8082 while nbits < bw {
8083 acc = (acc << 8) | data[idx] as u64;
8084 idx += 1;
8085 nbits += 8;
8086 }
8087 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8088 nbits -= bw;
8089 *d = u - l;
8090 }
8091 };
8092
8093 if a8w8_enabled() {
8094 let acts: Vec<SplitAct> = (0..b)
8095 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8096 .collect();
8097 let acts = &acts;
8098 let out_addr = SendMut(out.as_mut_ptr());
8099 let run = move |start: usize, end: usize| {
8100 for r in start..end {
8101 let bw = bits[r] as usize;
8102 if bw == 8 {
8103 ROW_F32.with(|rb| {
8106 let mut buf = rb.borrow_mut();
8107 buf.resize(cols, 0.0);
8108 decode_f32(r, &mut buf);
8109 for bi in 0..b {
8110 let x = &xs_all[bi * cols..(bi + 1) * cols];
8111 let mut dot = 0f32;
8112 for g in 0..ng {
8113 let mut gd = 0f32;
8114 for k in 0..GROUP_SIZE {
8115 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8116 }
8117 dot += gd * gscale(r, g);
8118 }
8119 unsafe { *out_addr.at(bi * rows + r) = dot };
8121 }
8122 });
8123 continue;
8124 }
8125 let l = (1i32 << (bw - 1)) - 1;
8126 let data = &bytes[offsets[r]..offsets[r + 1]];
8127 ROW_I8.with(|rb| {
8128 let mut buf = rb.borrow_mut();
8129 buf.resize(cols, 0);
8130 #[inline(always)]
8131 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8132 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8133 let u = unpack8::<B>(&data[blk * B..]);
8134 for k in 0..8 {
8135 chunk[k] = (u[k] - l) as i8 as u8;
8136 }
8137 }
8138 }
8139 match bw {
8140 3 => fill::<3>(data, l, &mut buf),
8141 4 => vbit_fill4(data, &mut buf),
8142 5 => fill::<5>(data, l, &mut buf),
8143 6 => fill::<6>(data, l, &mut buf),
8144 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8145 }
8146 let mut bi = 0usize;
8147 #[cfg(target_arch = "x86_64")]
8151 if avx2_enabled() && blocked_enabled() {
8152 while bi + 4 <= acts.len() {
8153 let xs = [
8154 acts[bi].xq.as_slice(),
8155 acts[bi + 1].xq.as_slice(),
8156 acts[bi + 2].xq.as_slice(),
8157 acts[bi + 3].xq.as_slice(),
8158 ];
8159 let sxs = [
8160 acts[bi].sx,
8161 acts[bi + 1].sx,
8162 acts[bi + 2].sx,
8163 acts[bi + 3].sx,
8164 ];
8165 let d = unsafe {
8166 if vnni_tiles_enabled() {
8167 dot_q4b_row_1x4_sx_vnni(
8168 &buf,
8169 &bytes[sc_off..],
8170 r * ng,
8171 ng,
8172 xs,
8173 sxs,
8174 )
8175 } else {
8176 dot_q4b_row_1x4_sx_avx2(
8177 &buf,
8178 &bytes[sc_off..],
8179 r * ng,
8180 ng,
8181 xs,
8182 sxs,
8183 )
8184 }
8185 };
8186 for k in 0..4 {
8187 let act = &acts[bi + k];
8188 let mut dot = d[k];
8189 for &(j, xv) in &act.outliers {
8190 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8191 }
8192 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8194 }
8195 bi += 4;
8196 }
8197 }
8198 while bi < acts.len() {
8199 let act = &acts[bi];
8200 let mut dot = 0f32;
8201 for g in 0..ng {
8202 let d = dot_i8_i8(
8203 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8204 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8205 ) as f32
8206 * act.sx;
8207 dot += d * gscale(r, g);
8208 }
8209 for &(j, xv) in &act.outliers {
8210 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8211 }
8212 unsafe { *out_addr.at(bi * rows + r) = dot };
8214 bi += 1;
8215 }
8216 });
8217 }
8218 };
8219 dispatch_rows(pool, rows, &run);
8220 return;
8221 }
8222
8223 let out_addr = SendMut(out.as_mut_ptr());
8224 let run = move |start: usize, end: usize| {
8225 ROW_F32.with(|rb| {
8226 let mut buf = rb.borrow_mut();
8227 buf.resize(cols, 0.0);
8228 for r in start..end {
8229 decode_f32(r, &mut buf);
8230 for bi in 0..b {
8231 let x = &xs_all[bi * cols..(bi + 1) * cols];
8232 let mut dot = 0f32;
8233 for g in 0..ng {
8234 let mut gd = 0f32;
8235 for k in 0..GROUP_SIZE {
8236 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8237 }
8238 dot += gd * gscale(r, g);
8239 }
8240 unsafe { *out_addr.at(bi * rows + r) = dot };
8242 }
8243 }
8244 })
8245 };
8246 dispatch_rows(pool, rows, &run);
8247}
8248
8249pub(crate) fn gpu_batch_job<'a>(
8253 t: &'a QTensor,
8254 x: &[f32],
8255) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8256 match t {
8257 QTensor::Mapped {
8258 model,
8259 idx,
8260 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8261 rows,
8262 cols,
8263 row_scale,
8264 col_field,
8265 ..
8266 } => Some((
8267 model.clone(),
8268 crate::gpu::BatchJob {
8269 idx: *idx,
8270 rows: *rows,
8271 cols: *cols,
8272 row_scale,
8273 xs: prescale(x, col_field, *dt).into_owned(),
8274 layout: crate::gpu::BatchLayout::Q8,
8275 },
8276 )),
8277 QTensor::Mapped {
8279 model,
8280 idx,
8281 dtype: TensorDtype::Q1,
8282 rows,
8283 cols,
8284 ..
8285 } => Some((
8286 model.clone(),
8287 crate::gpu::BatchJob {
8288 idx: *idx,
8289 rows: *rows,
8290 cols: *cols,
8291 row_scale: &[],
8292 xs: x.to_vec(),
8293 layout: crate::gpu::BatchLayout::Q1,
8294 },
8295 )),
8296 QTensor::Mapped {
8301 model,
8302 idx,
8303 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8304 rows,
8305 cols,
8306 ..
8307 } => Some((
8308 model.clone(),
8309 crate::gpu::BatchJob {
8310 idx: *idx,
8311 rows: *rows,
8312 cols: *cols,
8313 row_scale: &[],
8314 xs: x.to_vec(),
8315 layout: if *dt == TensorDtype::Q4Tiled {
8316 crate::gpu::BatchLayout::Q4t
8317 } else {
8318 crate::gpu::BatchLayout::Q4tp
8319 },
8320 },
8321 )),
8322 _ => None,
8323 }
8324}
8325
8326thread_local! {
8327 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8328 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8329}
8330
8331pub(crate) fn prescale<'a>(
8332 x: &'a [f32],
8333 col_field: &[f32],
8334 dtype: TensorDtype,
8335) -> std::borrow::Cow<'a, [f32]> {
8336 if dtype == TensorDtype::Q8_2f {
8337 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8338 } else {
8339 std::borrow::Cow::Borrowed(x)
8340 }
8341}
8342
8343pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8346 x: &[f32],
8347 col_field: &[f32],
8348 dtype: TensorDtype,
8349 buf_id: u8,
8350 f: F,
8351) -> R {
8352 if dtype == TensorDtype::Q8_2f {
8353 if buf_id == 1 {
8354 PRESCALE_BUF1.with(|b| {
8355 let mut buf = b.borrow_mut();
8356 buf.clear();
8357 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8358 f(&buf)
8359 })
8360 } else {
8361 PRESCALE_BUF2.with(|b| {
8362 let mut buf = b.borrow_mut();
8363 buf.clear();
8364 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8365 f(&buf)
8366 })
8367 }
8368 } else {
8369 f(x)
8370 }
8371}
8372
8373#[cfg(target_arch = "x86_64")]
8378pub(crate) fn avx2_enabled() -> bool {
8379 use std::sync::OnceLock;
8380 static ON: OnceLock<bool> = OnceLock::new();
8381 *ON.get_or_init(|| {
8382 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8383 && std::arch::is_x86_feature_detected!("avx2")
8384 && std::arch::is_x86_feature_detected!("fma")
8385 })
8386}
8387
8388#[cfg(target_arch = "x86_64")]
8393fn avx2_a8w8_enabled() -> bool {
8394 use std::sync::OnceLock;
8395 static ON: OnceLock<bool> = OnceLock::new();
8396 *ON.get_or_init(|| {
8397 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8398 })
8399}
8400
8401#[inline]
8405pub(crate) fn a8w8_enabled() -> bool {
8406 #[cfg(target_arch = "aarch64")]
8407 {
8408 sdot_enabled()
8409 }
8410 #[cfg(target_arch = "x86_64")]
8411 {
8412 avx2_a8w8_enabled()
8413 }
8414 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8415 {
8416 false
8417 }
8418}
8419
8420#[inline]
8423#[allow(unreachable_code)]
8424fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8425 #[cfg(target_arch = "aarch64")]
8426 unsafe {
8427 return dot_i8_sdot(w, xq);
8428 }
8429 #[cfg(target_arch = "x86_64")]
8430 unsafe {
8431 if avx512vnni_enabled() {
8432 return dot_i8_i8_vnni(w, xq);
8433 }
8434 return dot_i8_i8_avx2(w, xq);
8435 }
8436 w.iter()
8437 .zip(xq)
8438 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8439 .sum()
8440}
8441
8442#[cfg(target_arch = "x86_64")]
8446fn avx512vnni_enabled() -> bool {
8447 use std::sync::OnceLock;
8448 static ON: OnceLock<bool> = OnceLock::new();
8449 *ON.get_or_init(|| {
8450 std::env::var("CMF_AVX512")
8451 .map(|v| v != "0")
8452 .unwrap_or(true)
8453 && std::arch::is_x86_feature_detected!("avx512f")
8454 && std::arch::is_x86_feature_detected!("avx512bw")
8455 && std::arch::is_x86_feature_detected!("avx512vl")
8456 && std::arch::is_x86_feature_detected!("avx512vnni")
8457 })
8458}
8459
8460#[cfg(target_arch = "x86_64")]
8468fn vnni_tiles_enabled() -> bool {
8469 use std::sync::OnceLock;
8470 static ON: OnceLock<bool> = OnceLock::new();
8471 *ON.get_or_init(|| {
8472 std::env::var("CMF_VNNI_TILES")
8473 .map(|v| v != "0")
8474 .unwrap_or(true)
8475 && avx512vnni_enabled()
8476 })
8477}
8478
8479#[cfg(target_arch = "x86_64")]
8484#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8485#[inline]
8486unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8487 unsafe {
8489 use core::arch::x86_64::*;
8490 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8491 let hi128 = _mm256_extracti128_si256::<1>(d);
8492 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8493 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8494 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8495 _mm_cvtsi128_si32(s32)
8496 }
8497}
8498
8499#[cfg(target_arch = "x86_64")]
8504#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8505unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8506 unsafe {
8508 use core::arch::x86_64::*;
8509 let n = w.len();
8510 let mut j = 0usize;
8511 let mut total: i32;
8512 {
8517 #[inline(always)]
8518 unsafe fn step(
8519 w: *const u8,
8520 x: *const i8,
8521 acc: core::arch::x86_64::__m512i,
8522 ) -> core::arch::x86_64::__m512i {
8523 unsafe {
8524 use core::arch::x86_64::*;
8525 let wv = _mm512_loadu_si512(w as *const _);
8526 let xv = _mm512_loadu_si512(x as *const _);
8527 let aw = _mm512_abs_epi8(wv);
8528 let neg = _mm512_movepi8_mask(wv);
8529 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8530 _mm512_dpbusd_epi32(acc, aw, sx)
8531 }
8532 }
8533 let (mut a0, mut a1, mut a2, mut a3) = (
8534 _mm512_setzero_si512(),
8535 _mm512_setzero_si512(),
8536 _mm512_setzero_si512(),
8537 _mm512_setzero_si512(),
8538 );
8539 while j + 256 <= n {
8540 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8541 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8542 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8543 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8544 j += 256;
8545 }
8546 while j + 64 <= n {
8547 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8548 j += 64;
8549 }
8550 let s01 = _mm512_add_epi32(a0, a1);
8551 let s23 = _mm512_add_epi32(a2, a3);
8552 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8553 }
8554 if j + 32 <= n {
8556 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8557 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8558 let d = _mm256_dpbusd_epi32(
8559 _mm256_setzero_si256(),
8560 _mm256_abs_epi8(wv),
8561 _mm256_sign_epi8(xv, wv),
8562 );
8563 let hi128 = _mm256_extracti128_si256::<1>(d);
8564 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8565 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8566 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8567 total += _mm_cvtsi128_si32(s32);
8568 j += 32;
8569 }
8570 while j < n {
8571 total += (w[j] as i8) as i32 * xq[j] as i32;
8572 j += 1;
8573 }
8574 total
8575 }
8576}
8577
8578#[cfg(target_arch = "x86_64")]
8580#[target_feature(enable = "avx2,fma")]
8581unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8582 unsafe {
8584 use core::arch::x86_64::*;
8585 let n = x.len();
8586 let wp = w.as_ptr();
8587 let xp = x.as_ptr();
8588 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8589 let mut j = 0usize;
8590 while j + 16 <= n {
8591 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8592 let lo = _mm256_cvtepi8_epi32(wb);
8593 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8594 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8595 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8596 j += 16;
8597 }
8598 let acc = _mm256_add_ps(a0, a1);
8599 let hi128 = _mm256_extractf128_ps::<1>(acc);
8600 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8601 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8602 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8603 let mut sum = _mm_cvtss_f32(s32);
8604 while j < n {
8605 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8606 j += 1;
8607 }
8608 sum
8609 }
8610}
8611
8612#[cfg(target_arch = "x86_64")]
8617#[target_feature(enable = "avx2")]
8618unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8619 unsafe {
8621 use core::arch::x86_64::*;
8622 let n = w.len();
8623 let ones = _mm256_set1_epi16(1);
8624 let mut acc = _mm256_setzero_si256();
8625 let mut j = 0usize;
8626 while j + 32 <= n {
8627 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8628 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8629 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8630 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8631 j += 32;
8632 }
8633 let hi128 = _mm256_extracti128_si256::<1>(acc);
8634 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8635 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8636 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8637 let mut s = _mm_cvtsi128_si32(s32);
8638 while j < n {
8639 s += (w[j] as i8) as i32 * xq[j] as i32;
8640 j += 1;
8641 }
8642 s
8643 }
8644}
8645
8646#[cfg(target_arch = "aarch64")]
8650#[target_feature(enable = "neon,i8mm")]
8651unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8652 unsafe {
8654 use core::arch::aarch64::*;
8655 use core::arch::asm;
8656 let n = w0.len();
8657 let w0p = w0.as_ptr() as *const i8;
8658 let w1p = w1.as_ptr() as *const i8;
8659 let mut acc01 = vdupq_n_s32(0);
8662 let mut acc23 = vdupq_n_s32(0);
8663 let mut i = 0usize;
8664 while i + 8 <= n {
8665 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8666 let xb01 = vcombine_s8(
8667 vld1_s8(xs[0].as_ptr().add(i)),
8668 vld1_s8(xs[1].as_ptr().add(i)),
8669 );
8670 let xb23 = vcombine_s8(
8671 vld1_s8(xs[2].as_ptr().add(i)),
8672 vld1_s8(xs[3].as_ptr().add(i)),
8673 );
8674 asm!(
8675 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8676 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8677 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8678 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8679 options(pure, nomem, nostack),
8680 );
8681 i += 8;
8682 }
8683 let mut out = [[0i32; 4]; 2];
8684 let a01: [i32; 4] = core::mem::transmute(acc01);
8685 let a23: [i32; 4] = core::mem::transmute(acc23);
8686 out[0][0] = a01[0];
8687 out[0][1] = a01[1];
8688 out[1][0] = a01[2];
8689 out[1][1] = a01[3];
8690 out[0][2] = a23[0];
8691 out[0][3] = a23[1];
8692 out[1][2] = a23[2];
8693 out[1][3] = a23[3];
8694 if i < n {
8695 for (k, x) in xs.iter().enumerate() {
8696 for j in i..n {
8697 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8698 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8699 }
8700 }
8701 }
8702 out
8703 }
8704}
8705
8706#[cfg(target_arch = "aarch64")]
8710#[target_feature(enable = "neon,dotprod")]
8711unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8712 unsafe {
8714 use core::arch::aarch64::*;
8715 use core::arch::asm;
8716 let n = w0.len();
8717 let w0p = w0.as_ptr() as *const i8;
8718 let w1p = w1.as_ptr() as *const i8;
8719 let mut acc = [[vdupq_n_s32(0); 4]; 2];
8720 let mut i = 0usize;
8721 while i + 16 <= n {
8722 let wv0 = vld1q_s8(w0p.add(i));
8723 let wv1 = vld1q_s8(w1p.add(i));
8724 for (k, x) in xs.iter().enumerate() {
8725 let xv = vld1q_s8(x.as_ptr().add(i));
8726 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8727 asm!(
8728 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8729 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8730 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8731 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8732 options(pure, nomem, nostack),
8733 );
8734 acc[0][k] = a0;
8735 acc[1][k] = a1;
8736 }
8737 i += 16;
8738 }
8739 let mut out = [[0i32; 4]; 2];
8740 for r in 0..2 {
8741 for k in 0..4 {
8742 out[r][k] = vaddvq_s32(acc[r][k]);
8743 }
8744 }
8745 if i < n {
8746 for (k, x) in xs.iter().enumerate() {
8747 for j in i..n {
8748 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8749 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8750 }
8751 }
8752 }
8753 out
8754 }
8755}
8756
8757#[cfg(target_arch = "x86_64")]
8763#[target_feature(enable = "avx2")]
8764unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8765 unsafe {
8767 use core::arch::x86_64::*;
8768 let n = w0.len();
8769 let ones = _mm256_set1_epi16(1);
8770 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8771 let mut j = 0usize;
8772 while j + 32 <= n {
8773 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8774 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8775 let aw0 = _mm256_abs_epi8(wv0);
8776 let aw1 = _mm256_abs_epi8(wv1);
8777 for (k, x) in xs.iter().enumerate() {
8778 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8779 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8780 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8781 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8782 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8783 }
8784 j += 32;
8785 }
8786 let mut out = [[0i32; 4]; 2];
8787 for r in 0..2 {
8788 for k in 0..4 {
8789 let a = acc[r][k];
8790 let hi128 = _mm256_extracti128_si256::<1>(a);
8791 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8792 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8793 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8794 out[r][k] = _mm_cvtsi128_si32(s32);
8795 }
8796 }
8797 if j < n {
8798 for (k, x) in xs.iter().enumerate() {
8799 for i in j..n {
8800 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8801 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8802 }
8803 }
8804 }
8805 out
8806 }
8807}
8808
8809#[cfg(target_arch = "x86_64")]
8814#[inline]
8815fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8816 let dot = if avx512vnni_enabled() && row.len() >= 64 {
8817 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8818 } else {
8819 unsafe { dot_i8_i8_avx2(row, &act.xq) }
8820 };
8821 let mut acc = dot as f32 * act.sx;
8822 for &(j, xv) in &act.outliers {
8823 acc += (row[j] as i8) as f32 * xv;
8824 }
8825 acc
8826}
8827
8828#[cfg(target_arch = "x86_64")]
8832#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8833unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8834 unsafe {
8836 use core::arch::x86_64::*;
8837 let n = w.len();
8838 let flip = _mm512_set1_epi8(-128); #[inline(always)]
8840 unsafe fn step(
8841 w: *const u8,
8842 x: *const i8,
8843 flip: core::arch::x86_64::__m512i,
8844 acc: core::arch::x86_64::__m512i,
8845 ) -> core::arch::x86_64::__m512i {
8846 unsafe {
8847 use core::arch::x86_64::*;
8848 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8849 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8850 }
8851 }
8852 let (mut a0, mut a1, mut a2, mut a3) = (
8853 _mm512_setzero_si512(),
8854 _mm512_setzero_si512(),
8855 _mm512_setzero_si512(),
8856 _mm512_setzero_si512(),
8857 );
8858 let mut j = 0usize;
8859 while j + 256 <= n {
8860 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8861 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8862 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8863 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8864 j += 256;
8865 }
8866 while j + 64 <= n {
8867 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8868 j += 64;
8869 }
8870 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8871 _mm512_add_epi32(a0, a1),
8872 _mm512_add_epi32(a2, a3),
8873 ));
8874 while j < n {
8876 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8877 j += 1;
8878 }
8879 total
8880 }
8881}
8882
8883#[cfg(target_arch = "x86_64")]
8889#[target_feature(enable = "avx2")]
8890unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8891 unsafe {
8894 use core::arch::x86_64::*;
8895 let lomask = _mm_set1_epi8(0x0F);
8896 let eight = _mm256_set1_epi8(8);
8897 let ones = _mm256_set1_epi16(1);
8898 let mut acc = 0f32;
8899 for gi in 0..gpr {
8900 let g = g0 + gi;
8901 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8902 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8903 let lo = _mm_and_si128(b, lomask);
8904 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8905 let w = _mm256_sub_epi8(
8906 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8907 eight,
8908 );
8909 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8910 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8911 let d = _mm256_madd_epi16(p16, ones);
8912 let hi128 = _mm256_extracti128_si256::<1>(d);
8913 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8914 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8915 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8916 acc += _mm_cvtsi128_si32(s32) as f32 * s;
8917 }
8918 acc
8919 }
8920}
8921
8922#[cfg(target_arch = "x86_64")]
8925#[target_feature(enable = "avx2")]
8926unsafe fn dot_q4_row_avx2_2(
8927 packed: &[u8],
8928 scales: &[u8],
8929 g0: usize,
8930 gpr: usize,
8931 xq1: &[i8],
8932 xq2: &[i8],
8933) -> (f32, f32) {
8934 unsafe {
8936 use core::arch::x86_64::*;
8937 let lomask = _mm_set1_epi8(0x0F);
8938 let eight = _mm256_set1_epi8(8);
8939 let ones = _mm256_set1_epi16(1);
8940 let (mut acc1, mut acc2) = (0f32, 0f32);
8941 #[inline(always)]
8942 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8943 unsafe {
8944 use core::arch::x86_64::*;
8945 let hi128 = _mm256_extracti128_si256::<1>(d);
8946 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8947 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8948 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8949 _mm_cvtsi128_si32(s32)
8950 }
8951 }
8952 for gi in 0..gpr {
8953 let g = g0 + gi;
8954 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8955 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8956 let lo = _mm_and_si128(b, lomask);
8957 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8958 let w = _mm256_sub_epi8(
8959 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8960 eight,
8961 );
8962 let aw = _mm256_abs_epi8(w);
8963 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8964 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8965 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8966 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8967 acc1 += hsum(d1) as f32 * s;
8968 acc2 += hsum(d2) as f32 * s;
8969 }
8970 (acc1, acc2)
8971 }
8972}
8973
8974#[cfg(target_arch = "x86_64")]
8976fn q8_range_avx2(
8977 q: &[u8],
8978 row_scale: &[f32],
8979 act: &SplitAct,
8980 cols: usize,
8981 out_addr: SendMut,
8982 start: usize,
8983 end: usize,
8984) {
8985 for o in start..end {
8986 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8987 unsafe { *out_addr.at(o) = v };
8989 }
8990}
8991
8992#[cfg(target_arch = "x86_64")]
8994#[allow(clippy::too_many_arguments)]
8995fn q8_range2_avx2(
8996 q: &[u8],
8997 row_scale: &[f32],
8998 a1: &SplitAct,
8999 a2: &SplitAct,
9000 cols: usize,
9001 p1: SendMut,
9002 p2: SendMut,
9003 start: usize,
9004 end: usize,
9005) {
9006 for o in start..end {
9007 let row = &q[o * cols..(o + 1) * cols];
9008 unsafe {
9010 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9011 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9012 }
9013 }
9014}
9015
9016#[cfg(target_arch = "aarch64")]
9027fn i8mm_enabled() -> bool {
9028 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9029 *ON.get_or_init(|| {
9030 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9031 && std::arch::is_aarch64_feature_detected!("i8mm")
9032 })
9033}
9034
9035#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9039fn sdot_enabled() -> bool {
9040 use std::sync::OnceLock;
9041 static ON: OnceLock<bool> = OnceLock::new();
9042 *ON.get_or_init(|| {
9043 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9044 if !want {
9045 return false;
9046 }
9047
9048 #[cfg(target_arch = "aarch64")]
9049 {
9050 if std::arch::is_aarch64_feature_detected!("dotprod") {
9051 return true;
9052 }
9053 #[cfg(target_os = "android")]
9054 {
9055 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9056 if cpuinfo.lines().any(|l| {
9057 (l.starts_with("Features") || l.starts_with("features"))
9058 && l.contains("asimddp")
9059 }) {
9060 return true;
9061 }
9062 }
9063 }
9064 false
9065 }
9066 #[cfg(not(target_arch = "aarch64"))]
9067 {
9068 false
9069 }
9070 })
9071}
9072
9073struct SplitAct {
9078 xq: Vec<i8>,
9079 sx: f32,
9080 outliers: Vec<(usize, f32)>,
9081 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9084 xsum: i32,
9085}
9086
9087thread_local! {
9088 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9091 const { std::cell::RefCell::new(Vec::new()) };
9092}
9093
9094impl Drop for SplitAct {
9095 fn drop(&mut self) {
9096 let buf = std::mem::take(&mut self.xq);
9097 if buf.capacity() > 0 {
9098 XQ_FREE.with(|f| {
9099 let mut f = f.borrow_mut();
9100 if f.len() < 16 {
9101 f.push(buf);
9102 }
9103 });
9104 }
9105 }
9106}
9107
9108thread_local! {
9109 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9116}
9117
9118#[inline]
9121fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9122 KROW.with(|s| {
9123 let mut b = s.borrow_mut();
9124 if b.len() < n {
9125 b.resize(n, 0.0);
9126 }
9127 f(&mut b[..n])
9128 })
9129}
9130
9131fn split_act(x: &[f32]) -> SplitAct {
9132 let n = x.len();
9133 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9134 let thr = 8.0 * rms;
9135 let mut outliers: Vec<(usize, f32)> = Vec::new();
9139 let mut amax = 0f32;
9140 for (j, &v) in x.iter().enumerate() {
9141 let a = v.abs();
9142 if a > thr {
9143 outliers.push((j, v));
9144 } else if a > amax {
9145 amax = a;
9146 }
9147 }
9148 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9149 let inv = 1.0 / sx;
9150 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9151 xq.clear();
9152 xq.reserve(n);
9153 if outliers.is_empty() {
9154 xq.extend(
9155 x.iter()
9156 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9157 );
9158 } else {
9159 xq.extend(x.iter().map(|&v| {
9161 if v.abs() > thr {
9162 0
9163 } else {
9164 (v * inv).round().clamp(-127.0, 127.0) as i8
9165 }
9166 }));
9167 }
9168 let xsum = xq.iter().map(|&v| v as i32).sum();
9169 SplitAct {
9170 xq,
9171 sx,
9172 outliers,
9173 xsum,
9174 }
9175}
9176
9177fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9178 let n = x.len();
9179 let rms = (x
9180 .iter()
9181 .zip(col)
9182 .map(|(&a, &c)| {
9183 let v = a * c;
9184 (v * v) as f64
9185 })
9186 .sum::<f64>()
9187 / n.max(1) as f64)
9188 .sqrt() as f32;
9189 let thr = 8.0 * rms;
9190
9191 let mut outliers = Vec::new();
9192 let mut amax = 0f32;
9193 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9194 let v = a * c;
9195 let s = v.abs();
9196 if s > thr {
9197 outliers.push((j, v));
9198 } else if s > amax {
9199 amax = s;
9200 }
9201 }
9202
9203 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9204 let inv = 1.0 / sx;
9205 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9206 xq.clear();
9207 xq.reserve(n);
9208 if outliers.is_empty() {
9209 xq.extend(
9210 x.iter()
9211 .zip(col)
9212 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9213 );
9214 } else {
9215 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9216 let v = a * c;
9217 if v.abs() > thr {
9218 0
9219 } else {
9220 (v * inv).round().clamp(-127.0, 127.0) as i8
9221 }
9222 }));
9223 }
9224 let xsum = xq.iter().map(|&v| v as i32).sum();
9225 SplitAct {
9226 xq,
9227 sx,
9228 outliers,
9229 xsum,
9230 }
9231}
9232
9233#[cfg(target_arch = "aarch64")]
9236#[target_feature(enable = "neon,dotprod")]
9237unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9238 unsafe {
9240 use core::arch::aarch64::*;
9241 use core::arch::asm;
9242 let wp = w.as_ptr() as *const i8;
9243 let n = w.len();
9244 let (mut a0, mut a1, mut a2, mut a3) = (
9245 vdupq_n_s32(0),
9246 vdupq_n_s32(0),
9247 vdupq_n_s32(0),
9248 vdupq_n_s32(0),
9249 );
9250 let mut i = 0;
9251 while i + 64 <= n {
9252 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9253 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9254 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9255 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9256 asm!(
9257 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9258 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9259 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9260 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9261 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9262 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9263 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9264 options(pure, nomem, nostack),
9265 );
9266 i += 64;
9267 }
9268 while i + 16 <= n {
9269 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9270 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9271 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9272 i += 16;
9273 }
9274 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9275 while i < n {
9276 s += (*wp.add(i)) as i32 * xq[i] as i32;
9277 i += 1;
9278 }
9279 s
9280 }
9281}
9282
9283#[cfg(target_arch = "aarch64")]
9287#[target_feature(enable = "neon,dotprod")]
9288unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9289 unsafe {
9291 use core::arch::aarch64::*;
9292 use core::arch::asm;
9293 let n = xq.len();
9294 let px = xq.as_ptr();
9295 let (p0, p1, p2, p3) = (
9296 w0.as_ptr() as *const i8,
9297 w1.as_ptr() as *const i8,
9298 w2.as_ptr() as *const i8,
9299 w3.as_ptr() as *const i8,
9300 );
9301 let (mut a0, mut a1, mut a2, mut a3) = (
9302 vdupq_n_s32(0),
9303 vdupq_n_s32(0),
9304 vdupq_n_s32(0),
9305 vdupq_n_s32(0),
9306 );
9307 let mut i = 0;
9308 while i + 16 <= n {
9309 let x = vld1q_s8(px.add(i));
9310 let v0 = vld1q_s8(p0.add(i));
9311 let v1 = vld1q_s8(p1.add(i));
9312 let v2 = vld1q_s8(p2.add(i));
9313 let v3 = vld1q_s8(p3.add(i));
9314 asm!(
9315 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9316 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9317 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9318 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9319 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9320 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9321 options(pure, nomem, nostack),
9322 );
9323 i += 16;
9324 }
9325 let mut r = [
9326 vaddvq_s32(a0),
9327 vaddvq_s32(a1),
9328 vaddvq_s32(a2),
9329 vaddvq_s32(a3),
9330 ];
9331 while i < n {
9332 let xi = *px.add(i) as i32;
9333 r[0] += (*p0.add(i)) as i32 * xi;
9334 r[1] += (*p1.add(i)) as i32 * xi;
9335 r[2] += (*p2.add(i)) as i32 * xi;
9336 r[3] += (*p3.add(i)) as i32 * xi;
9337 i += 1;
9338 }
9339 r
9340 }
9341}
9342
9343#[cfg(target_arch = "aarch64")]
9350#[target_feature(enable = "neon,dotprod")]
9351unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9352 unsafe {
9355 use core::arch::aarch64::*;
9356 use core::arch::asm;
9357 let n = xq.len();
9358 let px = xq.as_ptr();
9359 let pg = g.as_ptr() as *const i8;
9360 let (mut a0, mut a1, mut a2, mut a3) = (
9361 vdupq_n_s32(0),
9362 vdupq_n_s32(0),
9363 vdupq_n_s32(0),
9364 vdupq_n_s32(0),
9365 );
9366 let mut i = 0;
9367 while i + 16 <= n {
9368 let x = vld1q_s8(px.add(i));
9369 let base = pg.add(4 * i);
9370 let v0 = vld1q_s8(base);
9371 let v1 = vld1q_s8(base.add(16));
9372 let v2 = vld1q_s8(base.add(32));
9373 let v3 = vld1q_s8(base.add(48));
9374 asm!(
9375 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9376 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9377 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9378 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9379 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9380 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9381 options(pure, nomem, nostack),
9382 );
9383 i += 16;
9384 }
9385 [
9386 vaddvq_s32(a0),
9387 vaddvq_s32(a1),
9388 vaddvq_s32(a2),
9389 vaddvq_s32(a3),
9390 ]
9391 }
9392}
9393
9394#[cfg(target_arch = "aarch64")]
9400fn q8_range_sdot(
9401 q: &[u8],
9402 rep: &[u8],
9403 row_scale: &[f32],
9404 act: &SplitAct,
9405 cols: usize,
9406 out_addr: SendMut,
9407 start: usize,
9408 end: usize,
9409) {
9410 let mut o = start;
9411 if !rep.is_empty() {
9414 while o < end && o % 4 != 0 {
9415 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9416 unsafe { *out_addr.at(o) = v };
9417 o += 1;
9418 }
9419 }
9420 while o + 4 <= end {
9421 let r = if rep.is_empty() {
9422 unsafe {
9423 dot_i8_sdot_4rows(
9424 &q[o * cols..(o + 1) * cols],
9425 &q[(o + 1) * cols..(o + 2) * cols],
9426 &q[(o + 2) * cols..(o + 3) * cols],
9427 &q[(o + 3) * cols..(o + 4) * cols],
9428 &act.xq,
9429 )
9430 }
9431 } else {
9432 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9433 };
9434 for k in 0..4 {
9435 let mut acc = r[k] as f32 * act.sx;
9436 for &(j, xv) in &act.outliers {
9437 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9438 }
9439 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9441 }
9442 o += 4;
9443 }
9444 while o < end {
9445 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9446 unsafe { *out_addr.at(o) = v };
9447 o += 1;
9448 }
9449}
9450
9451#[cfg(target_arch = "aarch64")]
9454#[allow(clippy::too_many_arguments)]
9455fn q8_range2_sdot(
9456 q: &[u8],
9457 row_scale: &[f32],
9458 a1: &SplitAct,
9459 a2: &SplitAct,
9460 cols: usize,
9461 p1: SendMut,
9462 p2: SendMut,
9463 start: usize,
9464 end: usize,
9465) {
9466 for o in start..end {
9467 let row = &q[o * cols..(o + 1) * cols];
9468 unsafe {
9470 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9471 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9472 }
9473 }
9474}
9475
9476#[allow(clippy::too_many_arguments)]
9478fn q8_range2_f32(
9479 q: &[u8],
9480 row_scale: &[f32],
9481 x1: &[f32],
9482 x2: &[f32],
9483 cols: usize,
9484 p1: SendMut,
9485 p2: SendMut,
9486 start: usize,
9487 end: usize,
9488) {
9489 for o in start..end {
9490 let row = &q[o * cols..(o + 1) * cols];
9491 unsafe {
9493 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9494 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9495 }
9496 }
9497}
9498
9499fn q8_range_f32(
9501 q: &[u8],
9502 row_scale: &[f32],
9503 xs: &[f32],
9504 cols: usize,
9505 out_addr: SendMut,
9506 start: usize,
9507 end: usize,
9508) {
9509 for o in start..end {
9510 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9511 unsafe { *out_addr.at(o) = v };
9513 }
9514}
9515
9516#[inline]
9520fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9521 #[cfg(target_arch = "aarch64")]
9522 return row_dot_sdot(row, act);
9523 #[cfg(target_arch = "x86_64")]
9524 return row_dot_avx2(row, act);
9525 #[allow(unreachable_code)]
9526 q8_row_dot_scalar(row, act)
9527}
9528
9529#[allow(dead_code)]
9530fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9531 let mut acc = 0i32;
9532 for (k, &b) in row.iter().enumerate() {
9533 acc += (b as i8) as i32 * act.xq[k] as i32;
9534 }
9535 let mut acc = acc as f32 * act.sx;
9536 for &(j, xv) in &act.outliers {
9537 acc += (row[j] as i8) as f32 * xv;
9538 }
9539 acc
9540}
9541
9542#[cfg(target_arch = "aarch64")]
9545#[inline]
9546fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9547 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9548 for &(j, xv) in &act.outliers {
9549 acc += (row[j] as i8) as f32 * xv;
9550 }
9551 acc
9552}
9553
9554#[cfg(target_arch = "aarch64")]
9562#[target_feature(enable = "neon,dotprod")]
9563unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9564 unsafe {
9567 use core::arch::aarch64::*;
9568 use core::arch::asm;
9569 let lomask = vdupq_n_u8(0x0F);
9570 let eight = vdupq_n_s8(8);
9571 let mut acc = 0f32;
9572 for gi in 0..gpr {
9573 let g = g0 + gi;
9574 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9575 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9576 let lo = vandq_u8(b, lomask);
9577 let hi = vshrq_n_u8::<4>(b);
9578 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9579 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9580 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9581 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9582 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9583 asm!(
9584 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9585 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9586 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9587 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9588 options(pure, nomem, nostack),
9589 );
9590 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9591 }
9592 acc
9593 }
9594}
9595
9596#[cfg(target_arch = "aarch64")]
9601#[target_feature(enable = "neon,dotprod")]
9602unsafe fn dot_q4_row_sdot2(
9603 packed: &[u8],
9604 scales: &[u8],
9605 g0: usize,
9606 gpr: usize,
9607 xq1: &[i8],
9608 xq2: &[i8],
9609) -> (f32, f32) {
9610 unsafe {
9613 use core::arch::aarch64::*;
9614 use core::arch::asm;
9615 let lomask = vdupq_n_u8(0x0F);
9616 let eight = vdupq_n_s8(8);
9617 let (mut acc1, mut acc2) = (0f32, 0f32);
9618 for gi in 0..gpr {
9619 let g = g0 + gi;
9620 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9621 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9622 let lo = vandq_u8(b, lomask);
9623 let hi = vshrq_n_u8::<4>(b);
9624 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9625 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9626 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9627 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9628 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9629 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9630 let (mut a0, mut a1, mut b0, mut b1) = (
9631 vdupq_n_s32(0),
9632 vdupq_n_s32(0),
9633 vdupq_n_s32(0),
9634 vdupq_n_s32(0),
9635 );
9636 asm!(
9637 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9638 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9639 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9640 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9641 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9642 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9643 e0 = in(vreg) e0, e1 = in(vreg) e1,
9644 x10 = in(vreg) x10, x11 = in(vreg) x11,
9645 x20 = in(vreg) x20, x21 = in(vreg) x21,
9646 options(pure, nomem, nostack),
9647 );
9648 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9649 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9650 }
9651 (acc1, acc2)
9652 }
9653}
9654
9655#[inline]
9660pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9661 #[cfg(target_arch = "aarch64")]
9662 unsafe {
9663 return axpy_i8_f32_neon(acc, row, w);
9664 }
9665 #[cfg(target_arch = "x86_64")]
9666 if avx2_enabled() {
9667 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9668 }
9669 #[allow(unreachable_code)]
9670 {
9671 for (a, &b) in acc.iter_mut().zip(row) {
9672 *a += w * b as f32;
9673 }
9674 }
9675}
9676
9677#[cfg(target_arch = "x86_64")]
9679#[target_feature(enable = "avx2,fma")]
9680unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9681 unsafe {
9683 use core::arch::x86_64::*;
9684 let n = acc.len().min(row.len());
9685 let ap = acc.as_mut_ptr();
9686 let rp = row.as_ptr();
9687 let wv = _mm256_set1_ps(w);
9688 let mut j = 0usize;
9689 while j + 16 <= n {
9690 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9691 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9692 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9693 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9694 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9695 _mm256_storeu_ps(ap.add(j), v0);
9696 _mm256_storeu_ps(ap.add(j + 8), v1);
9697 j += 16;
9698 }
9699 while j < n {
9700 *ap.add(j) += w * (*rp.add(j)) as f32;
9701 j += 1;
9702 }
9703 }
9704}
9705
9706#[cfg(target_arch = "aarch64")]
9707#[target_feature(enable = "neon")]
9708unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9709 unsafe {
9711 use core::arch::aarch64::*;
9712 let n = acc.len().min(row.len());
9713 let ap = acc.as_mut_ptr();
9714 let rp = row.as_ptr();
9715 let wv = vdupq_n_f32(w);
9716 let mut j = 0usize;
9717 while j + 16 <= n {
9718 let rb = vld1q_s8(rp.add(j));
9719 let lo = vmovl_s8(vget_low_s8(rb));
9720 let hi = vmovl_s8(vget_high_s8(rb));
9721 for (off, half) in [(0, lo), (8, hi)] {
9722 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9723 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9724 let o = j + off;
9725 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9726 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9727 }
9728 j += 16;
9729 }
9730 while j < n {
9731 *ap.add(j) += w * (*rp.add(j)) as f32;
9732 j += 1;
9733 }
9734 }
9735}
9736
9737#[inline]
9740pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9741 #[cfg(target_arch = "aarch64")]
9742 unsafe {
9743 return dot_i8_f32_neon(w, x);
9744 }
9745 #[cfg(target_arch = "x86_64")]
9746 if avx2_enabled() {
9747 return unsafe { dot_i8_f32_avx2(w, x) };
9748 }
9749 #[allow(unreachable_code)]
9750 {
9751 let mut sum = 0.0f32;
9752 for (j, &b) in w.iter().enumerate() {
9753 sum += (b as i8) as f32 * x[j];
9754 }
9755 sum
9756 }
9757}
9758
9759#[inline]
9763fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9764 #[cfg(target_arch = "aarch64")]
9765 unsafe {
9766 return dot_i8_col_f32_neon(w, x, col);
9767 }
9768 #[allow(unreachable_code)]
9769 {
9770 let mut sum = 0.0f32;
9771 for (j, &b) in w.iter().enumerate() {
9772 sum += (b as i8) as f32 * x[j] * col[j];
9773 }
9774 sum
9775 }
9776}
9777
9778#[cfg(target_arch = "aarch64")]
9779#[target_feature(enable = "neon")]
9780unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9781 unsafe {
9783 use core::arch::aarch64::*;
9784 let n = x.len();
9785 let wp = w.as_ptr() as *const i8;
9786 let xp = x.as_ptr();
9787 let cp = col.as_ptr();
9788 let (mut a0, mut a1, mut a2, mut a3) = (
9789 vdupq_n_f32(0.0),
9790 vdupq_n_f32(0.0),
9791 vdupq_n_f32(0.0),
9792 vdupq_n_f32(0.0),
9793 );
9794 let mut j = 0usize;
9795 while j + 16 <= n {
9796 let wb = vld1q_s8(wp.add(j));
9797 let lo = vmovl_s8(vget_low_s8(wb));
9798 let hi = vmovl_s8(vget_high_s8(wb));
9799 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9800 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9801 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9802 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9803 a0 = vfmaq_f32(
9804 a0,
9805 w0,
9806 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9807 );
9808 a1 = vfmaq_f32(
9809 a1,
9810 w1,
9811 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9812 );
9813 a2 = vfmaq_f32(
9814 a2,
9815 w2,
9816 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9817 );
9818 a3 = vfmaq_f32(
9819 a3,
9820 w3,
9821 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9822 );
9823 j += 16;
9824 }
9825 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9826 while j < n {
9827 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9828 j += 1;
9829 }
9830 sum
9831 }
9832}
9833
9834#[cfg(target_arch = "aarch64")]
9835#[target_feature(enable = "neon")]
9836unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9837 unsafe {
9839 use core::arch::aarch64::*;
9840 let n = x.len();
9841 let wp = w.as_ptr() as *const i8;
9842 let xp = x.as_ptr();
9843 let (mut a0, mut a1, mut a2, mut a3) = (
9844 vdupq_n_f32(0.0),
9845 vdupq_n_f32(0.0),
9846 vdupq_n_f32(0.0),
9847 vdupq_n_f32(0.0),
9848 );
9849 let mut j = 0usize;
9850 while j + 16 <= n {
9851 let wb = vld1q_s8(wp.add(j));
9852 let lo = vmovl_s8(vget_low_s8(wb));
9853 let hi = vmovl_s8(vget_high_s8(wb));
9854 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9855 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9856 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9857 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9858 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9859 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9860 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9861 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9862 j += 16;
9863 }
9864 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9865 while j < n {
9866 sum += (*wp.add(j)) as f32 * *xp.add(j);
9867 j += 1;
9868 }
9869 sum
9870 }
9871}
9872
9873#[allow(clippy::too_many_arguments)]
9874fn qmatvec(
9875 q: &[u8],
9876 rep: &[u8],
9877 row_scale: &[f32],
9878 x: &[f32],
9879 col_field: &[f32],
9880 dtype: TensorDtype,
9881 rows: usize,
9882 cols: usize,
9883 out: &mut [f32],
9884 pool: Option<&Pool>,
9885) {
9886 debug_assert_eq!(out.len(), rows);
9887 #[cfg(not(target_arch = "aarch64"))]
9888 let _ = rep;
9889
9890 #[cfg(target_arch = "aarch64")]
9891 if sdot_enabled() {
9892 let act = if dtype == TensorDtype::Q8_2f {
9893 split_act_q8_2f(x, col_field)
9894 } else {
9895 split_act(x)
9896 };
9897 let out_addr = SendMut(out.as_mut_ptr());
9898 let run_range = |start: usize, end: usize| {
9899 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9900 };
9901 match pool {
9902 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9903 _ => run_range(0, rows),
9904 }
9905 return;
9906 }
9907 #[cfg(target_arch = "x86_64")]
9910 if avx2_a8w8_enabled() {
9911 let act = if dtype == TensorDtype::Q8_2f {
9912 split_act_q8_2f(x, col_field)
9913 } else {
9914 split_act(x)
9915 };
9916 let out_addr = SendMut(out.as_mut_ptr());
9917 let run_range = |start: usize, end: usize| {
9918 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9919 };
9920 match pool {
9921 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9922 _ => run_range(0, rows),
9923 }
9924 return;
9925 }
9926
9927 prescale_with(x, col_field, dtype, 1, |xs| {
9928 let out_addr = SendMut(out.as_mut_ptr());
9929 let run_range = move |start: usize, end: usize| {
9930 for o in start..end {
9931 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9932 unsafe { *out_addr.at(o) = v };
9934 }
9935 };
9936 match pool {
9937 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9938 _ => run_range(0, rows),
9939 }
9940 });
9941}
9942
9943#[allow(clippy::too_many_arguments)]
9944fn qmatvec2(
9945 q: &[u8],
9946 row_scale: &[f32],
9947 x1: &[f32],
9948 x2: &[f32],
9949 col_field: &[f32],
9950 dtype: TensorDtype,
9951 rows: usize,
9952 cols: usize,
9953 o1: &mut [f32],
9954 o2: &mut [f32],
9955 pool: Option<&Pool>,
9956) {
9957 #[cfg(target_arch = "aarch64")]
9958 if sdot_enabled() {
9959 let a1s = if dtype == TensorDtype::Q8_2f {
9960 split_act_q8_2f(x1, col_field)
9961 } else {
9962 split_act(x1)
9963 };
9964 let a2s = if dtype == TensorDtype::Q8_2f {
9965 split_act_q8_2f(x2, col_field)
9966 } else {
9967 split_act(x2)
9968 };
9969 let p1 = SendMut(o1.as_mut_ptr());
9970 let p2 = SendMut(o2.as_mut_ptr());
9971 let run_range = |start: usize, end: usize| {
9972 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9973 };
9974 match pool {
9975 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9976 _ => run_range(0, rows),
9977 }
9978 return;
9979 }
9980 #[cfg(target_arch = "x86_64")]
9981 if avx2_a8w8_enabled() {
9982 let a1s = if dtype == TensorDtype::Q8_2f {
9983 split_act_q8_2f(x1, col_field)
9984 } else {
9985 split_act(x1)
9986 };
9987 let a2s = if dtype == TensorDtype::Q8_2f {
9988 split_act_q8_2f(x2, col_field)
9989 } else {
9990 split_act(x2)
9991 };
9992 let p1 = SendMut(o1.as_mut_ptr());
9993 let p2 = SendMut(o2.as_mut_ptr());
9994 let run_range = |start: usize, end: usize| {
9995 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9996 };
9997 match pool {
9998 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9999 _ => run_range(0, rows),
10000 }
10001 return;
10002 }
10003
10004 prescale_with(x1, col_field, dtype, 1, |x1s| {
10005 prescale_with(x2, col_field, dtype, 2, |x2s| {
10006 let p1 = SendMut(o1.as_mut_ptr());
10007 let p2 = SendMut(o2.as_mut_ptr());
10008 let run_range = move |start: usize, end: usize| {
10009 for o in start..end {
10010 let row = &q[o * cols..(o + 1) * cols];
10011 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10012 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10013 unsafe {
10015 *p1.at(o) = s1;
10016 *p2.at(o) = s2;
10017 }
10018 }
10019 };
10020 match pool {
10021 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10022 _ => run_range(0, rows),
10023 }
10024 });
10025 });
10026}
10027
10028#[derive(Clone, Copy)]
10029struct SendMut(*mut f32);
10030unsafe impl Send for SendMut {}
10031unsafe impl Sync for SendMut {}
10032
10033impl SendMut {
10034 #[inline]
10035 fn at(self, i: usize) -> *mut f32 {
10036 unsafe { self.0.add(i) }
10037 }
10038}
10039
10040#[cfg(test)]
10041mod tests {
10042 use super::*;
10043
10044 #[test]
10045 fn q2tp_i8_dot_matches_exact_on_grid() {
10046 let (rows, cols) = (5, 64);
10050 let gpr = cols / GROUP_SIZE;
10051 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10055 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10056 .collect();
10057 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10058 let x: Vec<f32> = (0..cols)
10059 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10060 .collect();
10061 let act = split_act(&x);
10062 assert!(
10063 act.outliers.is_empty(),
10064 "on-grid input must have no outliers"
10065 );
10066 let gsum = q1_group_sums(&act.xq, gpr);
10067 for r in 0..rows {
10068 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10069 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10070 assert!(
10071 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10072 "row {r}: exact {exact} vs i8 {fast}"
10073 );
10074 }
10075 }
10076
10077 #[test]
10078 fn q8_row_dot_fast_matches_scalar() {
10079 let cols = 96;
10082 let row: Vec<u8> = (0..cols)
10083 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10084 .collect();
10085 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10086 let act = split_act(&x);
10087 let fast = q8_row_dot(&row, &act);
10088 let scalar = q8_row_dot_scalar(&row, &act);
10089 assert!(
10090 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10091 "fast {fast} vs scalar {scalar}"
10092 );
10093 }
10094
10095 #[test]
10096 fn f32_matvec_matches_matvec_rows_bitexact() {
10097 let (rows, cols) = (300, 40);
10098 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10099 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10100 let qt = QTensor::from_f32(w.clone(), rows, cols);
10101
10102 let mut a = vec![0.0f32; rows];
10103 matvec_rows(None, &w, &x, &mut a);
10104 let mut b = vec![0.0f32; rows];
10105 qt.matvec(&x, &mut b, None);
10106 assert_eq!(a, b);
10107 }
10108
10109 #[test]
10110 fn sdot_kernel_exact_on_grid() {
10111 eprintln!("sdot_enabled = {}", sdot_enabled());
10116 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
10118 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10119 .collect();
10120 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10121 let x: Vec<f32> = (0..cols)
10122 .map(|i| match i % 3 {
10123 0 => 1.0,
10124 1 => -1.0,
10125 _ => 0.0,
10126 })
10127 .collect();
10128 let mut a = vec![0.0f32; rows];
10129 qmatvec(
10130 &w,
10131 &[],
10132 &scales,
10133 &x,
10134 &[],
10135 TensorDtype::Q8Row,
10136 rows,
10137 cols,
10138 &mut a,
10139 None,
10140 );
10141 for o in 0..rows {
10142 let mut acc = 0.0f32;
10143 for j in 0..cols {
10144 acc += (w[o * cols + j] as i8) as f32 * x[j];
10145 }
10146 let expect = acc * scales[o];
10147 assert!(
10148 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10149 "row {o}: {} vs {expect}",
10150 a[o]
10151 );
10152 }
10153 }
10154
10155 #[test]
10156 fn q1_tbl_fast_path_matches_reference() {
10157 let (rows, cols) = (5, 256);
10162 let gpr = cols / GROUP_SIZE;
10163 let mut bytes = Vec::new();
10164 for t in 0..rows * gpr {
10165 let s = 0.007 + (t % 11) as f32 * 0.004;
10166 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10167 for j in 0..4 {
10168 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10169 }
10170 }
10171 let x: Vec<f32> = (0..cols)
10172 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10173 .collect();
10174 let mut w = vec![0.0f32; rows * cols];
10175 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10176 let mut got = vec![0.0f32; rows];
10177 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10178 for o in 0..rows {
10179 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10180 assert!(
10181 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10182 "row {o}: {} vs {expect}",
10183 got[o]
10184 );
10185 }
10186 let b = 5usize;
10189 let mut xs_all = Vec::new();
10190 for bi in 0..b {
10191 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10192 }
10193 let mut mm = vec![0.0f32; b * rows];
10194 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10195 for bi in 0..b {
10196 let mut single = vec![0.0f32; rows];
10197 q1_matvec(
10198 &bytes,
10199 &xs_all[bi * cols..(bi + 1) * cols],
10200 rows,
10201 cols,
10202 &mut single,
10203 None,
10204 );
10205 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10206 }
10207 }
10208
10209 #[test]
10210 fn q1_kernels_match_exact_reference() {
10211 let (rows, cols) = (7, 96);
10213 let gpr = cols / GROUP_SIZE;
10214 let mut bytes = Vec::new();
10215 for t in 0..rows * gpr {
10216 let s = 0.01 + (t % 13) as f32 * 0.003;
10217 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10218 for j in 0..4 {
10219 bytes.push(((t * 31 + j * 97) % 251) as u8);
10220 }
10221 }
10222 let x: Vec<f32> = (0..cols)
10224 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10225 .collect();
10226 let mut w = vec![0.0f32; rows * cols];
10228 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10229 let mut expect = vec![0.0f32; rows];
10230 for o in 0..rows {
10231 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10232 }
10233 let mut got = vec![0.0f32; rows];
10234 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10235 for o in 0..rows {
10236 assert!(
10237 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10238 "row {o}: {} vs {}",
10239 got[o],
10240 expect[o]
10241 );
10242 }
10243 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10245 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10246 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10247 assert_eq!(a1, got);
10248 let mut xs = x.clone();
10249 xs.extend_from_slice(&x2);
10250 let mut mm = vec![0.0f32; 2 * rows];
10251 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10252 assert_eq!(&mm[..rows], got.as_slice());
10253 assert_eq!(&mm[rows..], a2.as_slice());
10254 }
10255
10256 #[test]
10257 fn repack_is_bit_identical() {
10258 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
10264 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10265 .collect();
10266 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10267 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10268 let rep = q8_repack_layout(&w, rows, cols);
10269 for g in 0..rows / 4 {
10271 for c in 0..cols / 16 {
10272 for lane in 0..4 {
10273 assert_eq!(
10274 &rep[g * 4 * cols + c * 64 + lane * 16
10275 ..g * 4 * cols + c * 64 + lane * 16 + 16],
10276 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10277 );
10278 }
10279 }
10280 }
10281 let mut a = vec![0.0f32; rows];
10282 qmatvec(
10283 &w,
10284 &[],
10285 &scales,
10286 &x,
10287 &[],
10288 TensorDtype::Q8Row,
10289 rows,
10290 cols,
10291 &mut a,
10292 None,
10293 );
10294 let mut b = vec![0.0f32; rows];
10295 qmatvec(
10296 &w,
10297 &rep,
10298 &scales,
10299 &x,
10300 &[],
10301 TensorDtype::Q8Row,
10302 rows,
10303 cols,
10304 &mut b,
10305 None,
10306 );
10307 assert_eq!(a, b, "full-range repack output diverged");
10308
10309 #[cfg(target_arch = "aarch64")]
10310 if sdot_enabled() {
10311 let act = split_act(&x);
10313 let mut c1 = vec![0.0f32; rows];
10314 let mut c2 = vec![0.0f32; rows];
10315 q8_range_sdot(
10316 &w,
10317 &[],
10318 &scales,
10319 &act,
10320 cols,
10321 SendMut(c1.as_mut_ptr()),
10322 3,
10323 rows - 2,
10324 );
10325 q8_range_sdot(
10326 &w,
10327 &rep,
10328 &scales,
10329 &act,
10330 cols,
10331 SendMut(c2.as_mut_ptr()),
10332 3,
10333 rows - 2,
10334 );
10335 assert_eq!(c1, c2, "unaligned-range repack output diverged");
10336 }
10337 }
10338
10339 #[test]
10340 fn sdot_a8w8_noise_is_bounded() {
10341 let (rows, cols) = (16, 512);
10345 let w: Vec<u8> = (0..rows * cols)
10346 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10347 .collect();
10348 let scales = vec![0.01f32; rows];
10349 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10350 let mut a = vec![0.0f32; rows];
10351 qmatvec(
10352 &w,
10353 &[],
10354 &scales,
10355 &x,
10356 &[],
10357 TensorDtype::Q8Row,
10358 rows,
10359 cols,
10360 &mut a,
10361 None,
10362 );
10363 let (mut num, mut den) = (0f64, 0f64);
10364 for o in 0..rows {
10365 let mut acc = 0.0f32;
10366 for j in 0..cols {
10367 acc += (w[o * cols + j] as i8) as f32 * x[j];
10368 }
10369 let expect = acc * scales[o];
10370 num += ((a[o] - expect) as f64).powi(2);
10371 den += (expect as f64).powi(2);
10372 }
10373 let rel = (num / den.max(1e-12)).sqrt();
10374 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10375 }
10376
10377 #[test]
10378 fn i8_dot_neon_matches_scalar() {
10379 let n = 100;
10380 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10381 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10382 let mut scalar = 0.0f32;
10383 for j in 0..n {
10384 scalar += (w[j] as i8) as f32 * x[j];
10385 }
10386 let fast = dot_i8_f32(&w, &x);
10387 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10388 }
10389
10390 #[test]
10392 fn vbitmatvec_matches_full_dequant() {
10393 let (rows, cols) = (6, 64);
10394 let ng = cols / GROUP_SIZE;
10395 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10397 let mut bytes = bits.clone();
10398 for g in 0..rows * ng {
10399 let s = 0.02 + 0.001 * g as f32;
10400 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10401 }
10402 for r in 0..rows {
10403 let b = bits[r] as usize;
10404 let (mut acc, mut nb) = (0u64, 0usize);
10405 let mut rowbytes = Vec::new();
10406 for i in 0..cols {
10407 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10408 acc = (acc << b) | v;
10409 nb += b;
10410 while nb >= 8 {
10411 nb -= 8;
10412 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10413 }
10414 }
10415 if nb > 0 {
10416 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10417 }
10418 bytes.extend_from_slice(&rowbytes);
10419 }
10420 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10421
10422 let mut reference = vec![0f32; rows * cols];
10423 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10424 let mut expect = vec![0f32; rows];
10425 for r in 0..rows {
10426 expect[r] = reference[r * cols..(r + 1) * cols]
10427 .iter()
10428 .zip(&x)
10429 .map(|(w, xv)| w * xv)
10430 .sum();
10431 }
10432 let mut got = vec![0f32; rows];
10433 let offsets = vbit_row_offsets(&bytes, rows, cols);
10434 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10435 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10439 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10440 for r in 0..rows {
10441 assert!(
10442 (got[r] - expect[r]).abs() < tol * scale,
10443 "row {r}: {} vs {}",
10444 got[r],
10445 expect[r]
10446 );
10447 }
10448 }
10449
10450 #[test]
10455 #[cfg(target_arch = "x86_64")]
10456 fn vbit_matmat_blocked_matches_per_row() {
10457 let (rows, cols, b) = (64usize, 128usize, 9usize);
10458 let ng = cols / GROUP_SIZE;
10459 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10460 let mut bytes = bits.clone();
10461 for g in 0..rows * ng {
10462 let sc = 0.02 + 0.0005 * g as f32;
10463 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10464 }
10465 for r in 0..rows {
10466 let bw = bits[r] as usize;
10467 let (mut acc, mut nb) = (0u64, 0usize);
10468 let mut rowbytes = Vec::new();
10469 for i in 0..cols {
10470 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10471 acc = (acc << bw) | v;
10472 nb += bw;
10473 while nb >= 8 {
10474 nb -= 8;
10475 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10476 }
10477 }
10478 if nb > 0 {
10479 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10480 }
10481 bytes.extend_from_slice(&rowbytes);
10482 }
10483 let x: Vec<f32> = (0..b * cols)
10484 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10485 .collect();
10486 let offsets = vbit_row_offsets(&bytes, rows, cols);
10487 let mut y_a = vec![0f32; b * rows];
10488 let mut y_b = vec![0f32; b * rows];
10489 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10490 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10491 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10492 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10493 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10494 let max_d = y_a
10495 .iter()
10496 .zip(&y_b)
10497 .map(|(p, q)| (p - q).abs())
10498 .fold(0.0f32, f32::max);
10499 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10500 }
10501
10502 #[test]
10510 fn q4t_matmat_blocked_matches_per_row() {
10511 let (rows, cols, b) = (16usize, 64usize, 9usize);
10512 let gpr = cols / GROUP_SIZE;
10513 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10514 for r in 0..rows {
10515 for g in 0..gpr {
10516 let t = (r * gpr + g) * Q4_TILE;
10517 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10518 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10519 for k in 0..16 {
10520 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10521 }
10522 }
10523 }
10524 let x: Vec<f32> = (0..b * cols)
10525 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10526 .collect();
10527 let mut y_blk = vec![0f32; b * rows];
10528 let mut y_row = vec![0f32; b * rows];
10529 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10530 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10531 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10532 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10533 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10534 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10535 }
10536
10537 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10544 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10545 let gpr = cols / GROUP_SIZE;
10546 let stride = q4tp_code_stride(gpr);
10547 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10548 let mut b = vec![0u8; codes_off + rows * stride];
10549 for r in 0..rows {
10550 for g in 0..gpr {
10551 let t = (r * gpr + g) * Q4TP_NIB;
10552 for k in 0..16 {
10553 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10554 }
10555 }
10556 let lo = -6.0 - 0.03 * (r % 17) as f32;
10557 let step = 0.01 + 0.004 * (r % 11) as f32;
10558 let p = params_off + r * 4;
10559 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10560 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10561 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10562 for g in 0..gpr {
10563 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10564 }
10565 }
10566 b
10567 }
10568
10569 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10573 let gpr = cols / GROUP_SIZE;
10574 let v = Q4tpView::new(bytes, rows, cols);
10575 let mut out = vec![0u8; rows * gpr * Q4_TILE];
10576 let mut sc = vec![0f32; gpr];
10577 for r in 0..rows {
10578 v.scales_into(r, gpr, &mut sc);
10579 for g in 0..gpr {
10580 let t = (r * gpr + g) * Q4_TILE;
10581 let s = sc[g];
10582 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10583 let src = (r * gpr + g) * Q4TP_NIB;
10584 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10585 }
10586 }
10587 out
10588 }
10589
10590 #[test]
10596 fn q4tp_exact_path_matches_dequant_reference() {
10597 let (rows, cols) = (256usize, 512usize);
10598 let gpr = cols / GROUP_SIZE;
10599 let bytes = synth_q4tp(rows, cols);
10600 let mut w = vec![0f32; rows * cols];
10601 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10602
10603 let x: Vec<f32> = (0..cols)
10604 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10605 .collect();
10606 let v = Q4tpView::new(&bytes, rows, cols);
10607 let mut sc = vec![0f32; gpr];
10608 for r in 0..rows {
10609 v.scales_into(r, gpr, &mut sc);
10610 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10611 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10612 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10616 assert!(
10617 (got - want).abs() <= 1e-5 * mag,
10618 "row {r}: kernel {got} vs dequant {want}"
10619 );
10620 }
10621 }
10622
10623 #[test]
10629 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10630 let (rows, cols) = (256usize, 512usize);
10631 let bytes = synth_q4tp(rows, cols);
10632 let twin = q4tp_as_q4t(&bytes, rows, cols);
10633 let x: Vec<f32> = (0..cols)
10634 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10635 .collect();
10636
10637 let mut got = vec![0f32; rows];
10638 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10639 let mut want = vec![0f32; rows];
10640 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10641
10642 let mut w = vec![0f32; rows * cols];
10645 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10646 for r in 0..rows {
10647 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10648 assert!(
10649 (got[r] - want[r]).abs() <= 1e-3 * mag,
10650 "row {r}: q4tp {} vs q4t {}",
10651 got[r],
10652 want[r]
10653 );
10654 }
10655 }
10656
10657 #[test]
10662 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10663 let (rows, cols, b) = (256usize, 512usize, 5usize);
10664 let bytes = synth_q4tp(rows, cols);
10665 let twin = q4tp_as_q4t(&bytes, rows, cols);
10666 let xs: Vec<f32> = (0..b * cols)
10667 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10668 .collect();
10669
10670 let mut got = vec![0f32; b * rows];
10671 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10672 let mut want = vec![0f32; b * rows];
10673 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10674
10675 let mut w = vec![0f32; rows * cols];
10676 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10677 for t in 0..b {
10678 for r in 0..rows {
10679 let mag: f32 = (0..cols)
10680 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10681 .sum();
10682 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10683 assert!(
10684 (g - wa).abs() <= 1e-3 * mag,
10685 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10686 );
10687 }
10688 }
10689 }
10690
10691 #[test]
10692 fn q4tp_matvec2_matches_the_single_stream_kernel() {
10693 let (rows, cols) = (128usize, 256usize);
10694 let gpr = cols / GROUP_SIZE;
10695 let bytes = synth_q4tp(rows, cols);
10696 let xs: Vec<f32> = (0..2 * cols)
10697 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10698 .collect();
10699
10700 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10701 q4tp_matvec2(
10702 &bytes,
10703 &xs[..cols],
10704 &xs[cols..],
10705 rows,
10706 cols,
10707 &mut o1,
10708 &mut o2,
10709 None,
10710 );
10711
10712 let v = Q4tpView::new(&bytes, rows, cols);
10715 let mut sc = vec![0f32; gpr];
10716 for r in 0..rows {
10717 v.scales_into(r, gpr, &mut sc);
10718 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10719 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10720 }
10721 }
10722
10723 #[test]
10730 fn q4tp_matvec_keeps_pace_with_q4t() {
10731 let (rows, cols) = (4096usize, 3072usize);
10732 let bytes = synth_q4tp(rows, cols);
10733 let twin = q4tp_as_q4t(&bytes, rows, cols);
10734 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10735 let mut o = vec![0f32; rows];
10736 let n = 12;
10737 let mut best = (f64::MAX, f64::MAX);
10738 for _ in 0..3 {
10741 let t0 = std::time::Instant::now();
10742 for _ in 0..n {
10743 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10744 }
10745 best.0 = best.0.min(t0.elapsed().as_secs_f64());
10746 let t0 = std::time::Instant::now();
10747 for _ in 0..n {
10748 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10749 }
10750 best.1 = best.1.min(t0.elapsed().as_secs_f64());
10751 }
10752 let ratio = best.1 / best.0;
10753 println!(
10754 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10755 best.0 * 1e3 / n as f64,
10756 best.1 * 1e3 / n as f64
10757 );
10758 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10759 }
10760
10761 #[cfg(target_os = "macos")]
10762 #[test]
10763 fn q4t_matmat_accel_matches_dequant_reference() {
10764 if !accel_gemm_enabled() {
10765 return; }
10767 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
10769 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10770 for r in 0..rows {
10771 for g in 0..gpr {
10772 let t = (r * gpr + g) * Q4_TILE;
10773 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10774 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10775 for k in 0..16 {
10776 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10777 }
10778 }
10779 }
10780 let x: Vec<f32> = (0..b * cols)
10781 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10782 .collect();
10783 let mut got = vec![0f32; b * rows];
10784 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10785 let mut w = vec![0f32; rows * cols];
10787 for r in 0..rows {
10788 for g in 0..gpr {
10789 let t = (r * gpr + g) * Q4_TILE;
10790 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10791 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10792 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10793 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10794 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10795 }
10796 }
10797 }
10798 for bi in 0..b {
10799 for r in 0..rows {
10800 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10801 let d = (got[bi * rows + r] - want).abs();
10802 assert!(
10803 d <= want.abs().max(1.0) * 1e-4,
10804 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10805 got[bi * rows + r]
10806 );
10807 }
10808 }
10809 }
10810
10811 #[test]
10812 fn q4matvec_matches_full_dequant() {
10813 let (rows, cols) = (8, 64);
10814 let groups = rows * cols / GROUP_SIZE;
10815 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10817 for i in 0..groups * 16 {
10818 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10819 }
10820 for g in 0..groups {
10821 let s = 0.01 + 0.003 * g as f32;
10822 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10823 }
10824 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10825
10826 let mut reference = vec![0.0f32; rows * cols];
10827 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10828 let mut expect = vec![0.0f32; rows];
10829 for r in 0..rows {
10830 expect[r] = reference[r * cols..(r + 1) * cols]
10831 .iter()
10832 .zip(&x)
10833 .map(|(w, xv)| w * xv)
10834 .sum();
10835 }
10836
10837 let mut got = vec![0.0f32; rows];
10838 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10839 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10843 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10844 for r in 0..rows {
10845 assert!(
10846 (got[r] - expect[r]).abs() < tol * scale,
10847 "row {r}: {} vs {}",
10848 got[r],
10849 expect[r]
10850 );
10851 }
10852 }
10853
10854 #[test]
10857 fn vbitmatvec2_equals_two_singles() {
10858 let (rows, cols) = (6, 64);
10859 let ng = cols / GROUP_SIZE;
10860 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10861 let mut bytes = bits.clone();
10862 for g in 0..rows * ng {
10863 let s = 0.02 + 0.001 * g as f32;
10864 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10865 }
10866 for r in 0..rows {
10867 let b = bits[r] as usize;
10868 let (mut acc, mut nb) = (0u64, 0usize);
10869 let mut rowbytes = Vec::new();
10870 for i in 0..cols {
10871 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10872 acc = (acc << b) | v;
10873 nb += b;
10874 while nb >= 8 {
10875 nb -= 8;
10876 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10877 }
10878 }
10879 if nb > 0 {
10880 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10881 }
10882 bytes.extend_from_slice(&rowbytes);
10883 }
10884 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10885 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10886 let offsets = vbit_row_offsets(&bytes, rows, cols);
10887
10888 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10889 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10890 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10891 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10892 vbitmatvec2(
10893 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10894 );
10895 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10896 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10897 }
10898
10899 #[test]
10901 fn q4matvec2_equals_two_singles() {
10902 let (rows, cols) = (8, 128);
10903 let groups = rows * cols / GROUP_SIZE;
10904 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10905 for i in 0..groups * 16 {
10906 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10907 }
10908 for g in 0..groups {
10909 let s = 0.01 + 0.003 * g as f32;
10910 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10911 }
10912 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10915 x1[9] = 250.0;
10916 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10917
10918 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10919 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10920 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10921 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10922 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10923 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10924 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10925 }
10926
10927 #[test]
10930 fn matvec_many_equals_separate_matvecs() {
10931 use crate::pool::Pool;
10932 let (r1, r2, cols) = (300, 200, 64);
10933 let mk = |salt: usize, rows: usize| {
10934 QTensor::from_f32(
10935 (0..rows * cols)
10936 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10937 .collect(),
10938 rows,
10939 cols,
10940 )
10941 };
10942 let (a, b) = (mk(1, r1), mk(5, r2));
10943 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10944 let pool = Pool::new(3);
10945
10946 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10947 a.matvec(&x, &mut ea, Some(&pool));
10948 b.matvec(&x, &mut eb, Some(&pool));
10949 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10950 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10951 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10952 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10953 }
10954
10955 #[test]
10958 fn batched_matmat_equals_per_position_matvec() {
10959 let (rows, cols, b) = (8, 64, 5);
10960 let groups = rows * cols / GROUP_SIZE;
10962 let mut q4 = Vec::new();
10963 for i in 0..groups * 16 {
10964 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10965 }
10966 for g in 0..groups {
10967 q4.extend_from_slice(
10968 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10969 );
10970 }
10971 let ng = cols / GROUP_SIZE;
10973 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10974 let mut vb = bits.clone();
10975 for g in 0..rows * ng {
10976 vb.extend_from_slice(
10977 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10978 );
10979 }
10980 for r in 0..rows {
10981 let bw = bits[r] as usize;
10982 let (mut acc, mut nb) = (0u64, 0usize);
10983 let mut rowbytes = Vec::new();
10984 for i in 0..cols {
10985 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10986 acc = (acc << bw) | v;
10987 nb += bw;
10988 while nb >= 8 {
10989 nb -= 8;
10990 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10991 }
10992 }
10993 if nb > 0 {
10994 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10995 }
10996 vb.extend_from_slice(&rowbytes);
10997 }
10998 let offsets = vbit_row_offsets(&vb, rows, cols);
10999
11000 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11001
11002 let mut got = vec![0f32; b * rows];
11004 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
11005 for bi in 0..b {
11006 let mut expect = vec![0f32; rows];
11007 q4matvec(
11008 &q4,
11009 &xs[bi * cols..(bi + 1) * cols],
11010 rows,
11011 cols,
11012 &mut expect,
11013 None,
11014 );
11015 assert_eq!(
11016 &got[bi * rows..(bi + 1) * rows],
11017 &expect[..],
11018 "q4 batch pos {bi}"
11019 );
11020 }
11021
11022 let mut got = vec![0f32; b * rows];
11024 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
11025 for bi in 0..b {
11026 let mut expect = vec![0f32; rows];
11027 vbitmatvec(
11028 &vb,
11029 &offsets,
11030 &xs[bi * cols..(bi + 1) * cols],
11031 rows,
11032 cols,
11033 &mut expect,
11034 None,
11035 );
11036 assert_eq!(
11037 &got[bi * rows..(bi + 1) * rows],
11038 &expect[..],
11039 "vbit batch pos {bi}"
11040 );
11041 }
11042 }
11043
11044 #[test]
11048 fn q4_tiled_matches_q4_block_bitexact() {
11049 let (rows, cols, b) = (8usize, 128usize, 3usize);
11050 let groups = rows * cols / GROUP_SIZE;
11051 let mut split = Vec::with_capacity(groups * 18);
11052 for i in 0..groups * 16 {
11053 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11054 }
11055 for g in 0..groups {
11056 split.extend_from_slice(
11057 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11058 );
11059 }
11060 let (packed, scales) = split.split_at(groups * 16);
11062 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
11063 for g in 0..groups {
11064 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
11065 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
11066 }
11067
11068 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11069 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11071
11072 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
11073 q4matvec(&split, &x1, rows, cols, &mut a, None);
11074 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
11075 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11076
11077 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11078 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11079 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11080 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11081 assert_eq!(a1, t1);
11082 assert_eq!(a2, t2);
11083
11084 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11085 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11086 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11087 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11088 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11089 }
11090
11091 #[test]
11098 fn q4matvec_sdot_outlier_exact() {
11099 let (rows, cols) = (4, 128);
11100 let groups = rows * cols / GROUP_SIZE;
11101 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11102 for i in 0..groups * 16 {
11103 bytes.push(((i * 11 + 5) % 256) as u8);
11104 }
11105 for g in 0..groups {
11106 let s = 0.02 + 0.002 * g as f32;
11107 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11108 }
11109 let mut x: Vec<f32> = (0..cols)
11110 .map(|i| match i % 3 {
11111 0 => 1.0,
11112 1 => -1.0,
11113 _ => 0.0,
11114 })
11115 .collect();
11116 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
11119 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11120 let mut expect = vec![0.0f32; rows];
11121 for r in 0..rows {
11122 expect[r] = reference[r * cols..(r + 1) * cols]
11123 .iter()
11124 .zip(&x)
11125 .map(|(w, xv)| w * xv)
11126 .sum();
11127 }
11128 let mut got = vec![0.0f32; rows];
11129 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11130 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11131 for r in 0..rows {
11132 assert!(
11133 (got[r] - expect[r]).abs() < 2e-3 * scale,
11134 "row {r}: {} vs {} (outlier term must be exact)",
11135 got[r],
11136 expect[r]
11137 );
11138 }
11139 }
11140
11141 #[test]
11145 fn q1t_matvec_matches_reference() {
11146 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11147 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
11149 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11150 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11152 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11153 let mut bytes = Vec::new();
11154 for r in 0..rows {
11155 for g in 0..gpr {
11156 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11157 let mut c = [0u8; 7];
11158 for k in 0..GROUP_SIZE {
11159 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11161 0
11162 } else {
11163 ((k + r * 3 + g) % 3) as u8 };
11165 cortiq_core::quant::q1t_pack(&mut c, k, code);
11166 }
11167 bytes.extend_from_slice(&c);
11168 }
11169 }
11170 let mut row_ptr = vec![0u32; rows + 1];
11173 for &(idx, _) in &outliers {
11174 row_ptr[idx as usize / cols + 1] += 1;
11175 }
11176 for r in 0..rows {
11177 row_ptr[r + 1] += row_ptr[r];
11178 }
11179 for &p in &row_ptr {
11180 bytes.extend_from_slice(&p.to_le_bytes());
11181 }
11182 for &(idx, v) in &outliers {
11183 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11184 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11185 }
11186
11187 let mut refw = vec![0f32; rows * cols];
11188 dequant_q1t(&bytes, rows, cols, &mut refw);
11189 let x: Vec<f32> = (0..cols)
11192 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11193 .collect();
11194 let mut expect = vec![0f32; rows];
11195 for r in 0..rows {
11196 let mut a = 0.0f32;
11197 for j in 0..cols {
11198 a += refw[r * cols + j] * x[j];
11199 }
11200 expect[r] = a;
11201 }
11202 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11203 let mut got = vec![0f32; rows];
11204 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11205 for r in 0..rows {
11206 assert!(
11207 (got[r] - expect[r]).abs() < tol(expect[r]),
11208 "row {r}: {} vs {}",
11209 got[r],
11210 expect[r]
11211 );
11212 }
11213 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
11215 let mut gm = vec![0f32; 2 * rows];
11216 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11217 for r in 0..rows {
11218 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11219 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11220 }
11221 let xb: Vec<f32> = (0..cols)
11225 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11226 .collect();
11227 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11228 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11229 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11230 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11231 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11232 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11233 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11234 }
11235
11236 #[test]
11239 fn q1t_matvec2_odd_gpr_matches_singles() {
11240 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11241 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
11243 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11244 for r in 0..rows {
11245 for g in 0..gpr {
11246 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11247 let mut c = [0u8; 7];
11248 for k in 0..GROUP_SIZE {
11249 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11250 }
11251 bytes.extend_from_slice(&c);
11252 }
11253 }
11254 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11255 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11256 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11257 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11258 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11259 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11260 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11261 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11262 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11263 }
11264
11265 #[test]
11269 #[ignore]
11270 fn q1t_matvec2_speed() {
11271 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11272 use std::time::Instant;
11273 let (rows, cols) = (8192usize, 4096usize);
11274 let gpr = cols / GROUP_SIZE;
11275 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11276 for r in 0..rows {
11277 for g in 0..gpr {
11278 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11279 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11280 let mut c = [0u8; 7];
11281 for k in 0..GROUP_SIZE {
11282 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11283 }
11284 bytes.extend_from_slice(&c);
11285 }
11286 }
11287 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11288 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11289 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11290 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11291 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11293 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11294 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11295 for _ in 0..8 {
11296 let t0 = Instant::now();
11297 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11298 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11299 let t1 = Instant::now();
11300 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11301 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11302 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11303 }
11304 assert_eq!(p1, s1);
11305 assert_eq!(p2, s2);
11306 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11307 }
11308
11309 #[test]
11313 #[ignore]
11314 fn q1t_matvec_speed() {
11315 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11316 use std::time::Instant;
11317 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
11319 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11320 for r in 0..rows {
11321 for g in 0..gpr {
11322 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11323 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11324 let mut c = [0u8; 7];
11325 for k in 0..GROUP_SIZE {
11326 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11327 }
11328 bytes.extend_from_slice(&c);
11329 }
11330 }
11331 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
11333 let mut idx = 0usize;
11334 while idx < n {
11335 row_ptr[idx / cols + 1] += 1;
11336 idx += stride;
11337 }
11338 for r in 0..rows {
11339 row_ptr[r + 1] += row_ptr[r];
11340 }
11341 for &p in &row_ptr {
11342 bytes.extend_from_slice(&p.to_le_bytes());
11343 }
11344 let mut idx = 0usize;
11345 while idx < n {
11346 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11347 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11348 idx += stride;
11349 }
11350 let x: Vec<f32> = (0..cols)
11353 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11354 .collect();
11355 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11356
11357 let slow = |out: &mut [f32]| {
11359 let mut buf = vec![0f32; cols];
11360 for r in 0..rows {
11361 for g in 0..gpr {
11362 let off = (r * gpr + g) * Q1T_TILE;
11363 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11364 let codes = &bytes[off + 2..off + Q1T_TILE];
11365 for k in 0..GROUP_SIZE {
11366 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11367 1 => s,
11368 2 => -s,
11369 _ => 0.0,
11370 };
11371 }
11372 }
11373 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11374 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11375 }
11376 };
11377 let iters = 5;
11378 let mut a = vec![0f32; rows];
11379 slow(&mut a); let t = Instant::now();
11381 for _ in 0..iters {
11382 slow(&mut a);
11383 }
11384 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11385
11386 let mut b = vec![0f32; rows];
11387 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
11389 for _ in 0..iters {
11390 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11391 }
11392 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11393
11394 for r in 0..rows {
11395 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11396 }
11397 println!(
11398 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
11399 slow_ms / fast_ms
11400 );
11401 }
11402}
11403
11404#[cfg(test)]
11405mod gemm_bench {
11406 #[test]
11417 #[ignore]
11418 fn q4tp_matmat_throughput() {
11419 let b: usize = std::env::var("CMF_BENCH_B")
11423 .ok()
11424 .and_then(|v| v.parse().ok())
11425 .unwrap_or(296);
11426 let (rows, cols) = (9216usize, 2304usize);
11427 let (_, _, _) = (rows, cols, b);
11428 let total =
11429 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
11430 .unwrap();
11431 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11436 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11437 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11438 let step = cortiq_core::quant::f32_to_f16(0.1);
11439 for r in 0..rows {
11440 let o = params_off + r * 4;
11441 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11442 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11443 }
11444 let _ = codes_off;
11445 let xs: Vec<f32> = (0..b * cols)
11446 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11447 .collect();
11448 let mut out = vec![0f32; b * rows];
11449 let pool = crate::pool::Pool::from_env();
11450 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11456 let reps: usize = std::env::var("CMF_BENCH_REPS")
11457 .ok()
11458 .and_then(|v| v.parse().ok())
11459 .unwrap_or(10);
11460 let mut best = [f64::MAX; 2];
11461 let mut sums = [0f32; 2];
11462 for _ in 0..reps {
11463 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11464 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11465 let t = std::time::Instant::now();
11466 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11467 best[k] = best[k].min(t.elapsed().as_secs_f64());
11468 sums[k] = out.iter().take(64).sum::<f32>();
11469 }
11470 }
11471 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11472 for (k, name) in ["previous", "tuned "].iter().enumerate() {
11473 println!(
11474 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11475 best[k] * 1e3,
11476 flops / best[k] / 1e9,
11477 sums[k]
11478 );
11479 }
11480 assert!(
11481 (sums[0] - sums[1]).abs() < 1e-2,
11482 "the tuned kernel changed the result: {} vs {}",
11483 sums[0],
11484 sums[1]
11485 );
11486 }
11487
11488 #[test]
11494 fn q4tp_matmat_blocked_matches_scalar() {
11495 use std::sync::atomic::Ordering::Relaxed;
11496 for &(rows, cols, b) in &[
11504 (64usize, 128usize, 7usize),
11505 (33, 96, 4),
11506 (16, 256, 9),
11507 (192, 2304, 37),
11508 ] {
11509 let total = cortiq_core::quant::expected_nbytes(
11510 cortiq_core::TensorDtype::Q4TiledP,
11511 &[rows, cols],
11512 )
11513 .unwrap();
11514 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11515 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11516 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11517 let step = cortiq_core::quant::f32_to_f16(0.1);
11518 for r in 0..rows {
11519 let o = params_off + r * 4;
11520 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11521 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11522 }
11523 let xs: Vec<f32> = (0..b * cols)
11524 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11525 .collect();
11526 let mut got = vec![0f32; b * rows];
11527 let mut want = vec![0f32; b * rows];
11528 let gpr = cols / 32;
11529 let view = super::Q4tpView::new(&bytes, rows, cols);
11530 let pool = crate::pool::Pool::from_env();
11531 super::Q4TP_ALT.store(2, Relaxed);
11532 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11533 super::Q4TP_ALT.store(1, Relaxed);
11534 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11535 super::Q4TP_ALT.store(0, Relaxed);
11536 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11543 let (mut worst, mut at) = (0f32, 0usize);
11544 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11545 if (g - w).abs() > worst {
11546 worst = (g - w).abs();
11547 at = i;
11548 }
11549 }
11550 assert!(
11551 worst <= 1e-4 * scale,
11552 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11553 (scale {scale:.3e}) at cell {at}: {} vs {}",
11554 got[at],
11555 want[at]
11556 );
11557
11558 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11565 for bi in 0..b {
11566 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11567 for r in 0..rows {
11568 let mut sc = vec![0f32; gpr];
11569 view.scales_into(r, gpr, &mut sc);
11570 let mut exact = 0f64;
11571 for j in 0..cols {
11572 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11573 exact += w as f64 * sq as f64 * act.xq[j] as f64;
11574 }
11575 exact *= act.sx as f64;
11576 for &(j, xv) in &act.outliers {
11577 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11578 exact += w as f64 * sq as f64 * xv as f64;
11579 }
11580 let i = bi * rows + r;
11581 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11582 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11583 }
11584 }
11585 println!(
11586 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11587 per-column {e_scalar:.3e}"
11588 );
11589 assert!(
11594 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11595 "{rows}x{cols} b={b}: error against f64 too large — blocked \
11596 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11597 );
11598 }
11599 }
11600
11601 #[test]
11605 #[ignore]
11606 fn q4t_matmat_throughput() {
11607 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11608 let total =
11609 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
11610 .unwrap();
11611 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11614 let sc = cortiq_core::quant::f32_to_f16(0.02);
11615 for t in bytes.chunks_mut(super::Q4_TILE) {
11616 t[..2].copy_from_slice(&sc.to_le_bytes());
11617 }
11618 let xs: Vec<f32> = (0..b * cols)
11619 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11620 .collect();
11621 let mut out = vec![0f32; b * rows];
11622 let pool = crate::pool::Pool::from_env();
11623 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11624 let reps: usize = std::env::var("CMF_BENCH_REPS")
11625 .ok()
11626 .and_then(|v| v.parse().ok())
11627 .unwrap_or(10);
11628 let mut best = f64::MAX;
11629 for _ in 0..reps {
11630 let t = std::time::Instant::now();
11631 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11632 best = best.min(t.elapsed().as_secs_f64());
11633 }
11634 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11635 println!(
11636 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11637 best * 1e3,
11638 flops / best / 1e9,
11639 out.iter().take(64).sum::<f32>()
11640 );
11641 }
11642}