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_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
533 match self {
534 Self::Mapped {
535 model,
536 idx,
537 dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
538 ..
539 } => Some((model, *idx)),
540 _ => None,
541 }
542 }
543
544 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
547 match self {
548 Self::Mapped {
549 model,
550 idx,
551 dtype: TensorDtype::Q2TiledP,
552 ..
553 } => Some((model, *idx)),
554 _ => None,
555 }
556 }
557
558 pub fn cols(&self) -> usize {
559 match self {
560 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
561 }
562 }
563
564 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
567 match self {
568 Self::Mapped {
569 model,
570 idx,
571 dtype: TensorDtype::Q1,
572 ..
573 } => Some((model, *idx)),
574 _ => None,
575 }
576 }
577
578 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
588 match self {
589 Self::Mapped {
590 model,
591 idx,
592 dtype: TensorDtype::Q8Row,
593 row_scale,
594 ..
595 } => Some((model, *idx, 0, row_scale.as_slice())),
596 Self::Mapped {
597 model,
598 idx,
599 dtype: TensorDtype::Q1,
600 ..
601 } => Some((model, *idx, 1, &[])),
602 Self::Mapped {
607 model,
608 idx,
609 dtype: TensorDtype::Q4Tiled,
610 ..
611 } => Some((model, *idx, 5, &[])),
612 Self::Mapped {
616 model,
617 idx,
618 dtype: TensorDtype::Q4TiledP,
619 ..
620 } => Some((model, *idx, 6, &[])),
621 Self::Mapped {
622 model,
623 idx,
624 dtype: TensorDtype::Q4Block,
625 ..
626 } => Some((model, *idx, 2, &[])),
627 Self::Mapped {
632 model,
633 idx,
634 dtype: TensorDtype::Q8_2f,
635 ..
636 } => Some((model, *idx, 7, &[])),
637 Self::Mapped {
638 model,
639 idx,
640 dtype: TensorDtype::Q1T,
641 ..
642 } => Some((model, *idx, 3, &[])),
643 Self::Mapped {
647 model,
648 idx,
649 dtype: TensorDtype::Q2TiledP,
650 ..
651 } => Some((model, *idx, 9, &[])),
652 _ => None,
653 }
654 }
655
656 pub fn as_f32(&self) -> Option<&[f32]> {
659 match self {
660 Self::F32 { data, .. } => Some(data),
661 Self::Mapped { .. } => None,
662 }
663 }
664
665 fn quant_bytes(&self) -> &[u8] {
666 match self {
667 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
668 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
669 }
670 }
671
672 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
674 let cols = self.cols();
675 debug_assert_eq!(dst.len(), cols);
676 match self {
677 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
678 Self::Mapped {
679 dtype,
680 row_scale,
681 col_field,
682 vbit_offsets,
683 ..
684 } => {
685 if *dtype == TensorDtype::Q4Tiled {
686 let bytes = self.quant_bytes();
687 let gpr = cols / GROUP_SIZE;
688 for gi in 0..gpr {
689 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
690 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
691 for (k, &b) in tile[2..].iter().enumerate() {
692 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
693 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
694 }
695 }
696 return;
697 }
698 if *dtype == TensorDtype::Q4TiledP {
699 let bytes = self.quant_bytes();
700 let gpr = cols / GROUP_SIZE;
701 let v = Q4tpView::new(bytes, self.rows(), cols);
702 let mut sc = vec![0f32; gpr];
703 v.scales_into(r, gpr, &mut sc);
704 for gi in 0..gpr {
705 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
706 let s = sc[gi];
707 for (k, &b) in tile.iter().enumerate() {
708 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
709 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
710 }
711 }
712 return;
713 }
714 if *dtype == TensorDtype::Q2TiledP {
715 let bytes = self.quant_bytes();
716 let gpr = cols / GROUP_SIZE;
717 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
718 let mut sc = vec![0f32; gpr];
719 v.scales_into(r, gpr, &mut sc);
720 for gi in 0..gpr {
721 let ch =
722 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
723 let s = sc[gi];
724 for (k, &b) in ch.iter().enumerate() {
725 for j in 0..4 {
726 dst[gi * GROUP_SIZE + k * 4 + j] =
727 (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
728 }
729 }
730 }
731 return;
732 }
733 if *dtype == TensorDtype::Q4Block {
734 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
735 let gpr = cols / GROUP_SIZE;
736 for gi in 0..gpr {
737 let g = r * gpr + gi;
738 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
739 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
740 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
741 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
742 }
743 }
744 return;
745 }
746 if *dtype == TensorDtype::Q1 {
747 let bytes = self.quant_bytes();
748 let gpr = cols / GROUP_SIZE;
749 for gi in 0..gpr {
750 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
751 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
752 for (j, &b) in tile[2..].iter().enumerate() {
753 for k in 0..8 {
754 dst[gi * GROUP_SIZE + j * 8 + k] =
755 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
756 }
757 }
758 }
759 return;
760 }
761 if *dtype == TensorDtype::Q1T {
762 let bytes = self.quant_bytes();
763 let gpr = cols / GROUP_SIZE;
764 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
765 for gi in 0..gpr {
766 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
767 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
768 bytes[off],
769 bytes[off + 1],
770 ]));
771 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
772 for k in 0..GROUP_SIZE {
773 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
774 {
775 1 => s,
776 2 => -s,
777 _ => 0.0,
778 };
779 }
780 }
781 let rows = self.rows();
783 let entries = base_len + (rows + 1) * 4;
784 if entries <= bytes.len() {
785 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
786 let r0 = u32::from_le_bytes([
787 ptrs[r * 4],
788 ptrs[r * 4 + 1],
789 ptrs[r * 4 + 2],
790 ptrs[r * 4 + 3],
791 ]) as usize;
792 let r1 = u32::from_le_bytes([
793 ptrs[(r + 1) * 4],
794 ptrs[(r + 1) * 4 + 1],
795 ptrs[(r + 1) * 4 + 2],
796 ptrs[(r + 1) * 4 + 3],
797 ]) as usize;
798 let off = entries + r0 * 4;
799 for i in 0..r1 - r0 {
800 let item = &bytes[off + i * 4..off + i * 4 + 4];
801 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
802 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
803 item[2], item[3],
804 ]));
805 if c < cols {
806 dst[c] = v;
807 }
808 }
809 }
810 return;
811 }
812 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
813 let bytes = self.quant_bytes();
814 let rows = self.rows();
815 let ng = cols / GROUP_SIZE;
816 let bits = &bytes[..rows];
817 let sc_off = rows;
818 let off = vbit_offsets[r];
821 let b = bits[r] as usize;
822 let l = ((1usize << (b - 1)) - 1) as f32;
823 let data = &bytes[off..];
824 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
825 for (i, d) in dst.iter_mut().enumerate() {
826 while nbits < b {
827 acc = (acc << 8) | data[idx] as u64;
828 idx += 1;
829 nbits += 8;
830 }
831 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
832 nbits -= b;
833 let so = (r * ng + i / GROUP_SIZE) * 2;
834 let sv = f16_to_f32(u16::from_le_bytes([
835 bytes[sc_off + so],
836 bytes[sc_off + so + 1],
837 ]));
838 *d = (u - l) * sv;
839 }
840 return;
841 }
842 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
843 let s = row_scale[r];
844 match dtype {
845 TensorDtype::Q8Row => {
846 for (d, &b) in dst.iter_mut().zip(q) {
847 *d = (b as i8) as f32 * s;
848 }
849 }
850 TensorDtype::Q8_2f => {
851 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
852 *d = (b as i8) as f32 * s * col_field[i];
853 }
854 }
855 _ => unreachable!(),
856 }
857 }
858 }
859 }
860
861 pub fn sparse_col_ok(&self) -> bool {
866 match self {
867 Self::F32 { .. } => true,
868 Self::Mapped { dtype, .. } => {
869 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
870 }
871 }
872 }
873
874 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
878 let inter = self.cols();
879 let hidden = self.rows();
880 debug_assert_eq!(out.len(), hidden);
881 match self {
882 Self::F32 { data, .. } => {
883 for (k, o) in out.iter_mut().enumerate() {
884 *o += w * data[k * inter + c];
885 }
886 }
887 Self::Mapped {
888 dtype,
889 row_scale,
890 col_field,
891 ..
892 } => {
893 let q = self.quant_bytes();
894 let colf = if *dtype == TensorDtype::Q8_2f {
895 col_field[c]
896 } else {
897 1.0
898 };
899 let wc = w * colf;
900 for (k, o) in out.iter_mut().enumerate() {
901 let b = q[k * inter + c] as i8 as f32;
902 *o += wc * b * row_scale[k];
903 }
904 }
905 }
906 }
907
908 #[inline]
917 pub fn prefetch_row(&self, r: usize) {
918 let Self::Mapped { dtype, .. } = self else {
919 return;
920 };
921 if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
922 return;
923 }
924 let cols = self.cols();
925 let q = self.quant_bytes();
926 let (a, b) = (r * cols, (r + 1) * cols);
927 if b > q.len() {
928 return;
929 }
930 let mut j = a;
931 while j < b {
932 unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
933 j += 512;
934 }
935 }
936
937 pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
946 let cols = self.cols();
947 debug_assert_eq!(out.len(), cols);
948 match self {
949 Self::F32 { data, .. } => {
950 let row = &data[r * cols..(r + 1) * cols];
951 for (o, v) in out.iter_mut().zip(row) {
952 *o += w * v;
953 }
954 }
955 Self::Mapped {
956 dtype,
957 row_scale,
958 col_field,
959 ..
960 } => match dtype {
961 TensorDtype::Q8Row => {
962 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
963 let ws = w * row_scale[r];
964 let row: &[i8] =
965 unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
966 axpy_i8_f32(out, row, ws);
967 }
968 TensorDtype::Q8_2f => {
969 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
970 let ws = w * row_scale[r];
971 for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
972 *o += ws * c * (*b as i8 as f32);
973 }
974 }
975 _ => {
976 self.row_f32(r, scratch);
977 for (o, v) in out.iter_mut().zip(scratch.iter()) {
978 *o += w * v;
979 }
980 }
981 },
982 }
983 }
984
985 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
989 let cols = self.cols();
990 match self {
991 Self::F32 { data, .. } => {
992 let row = &data[r * cols..(r + 1) * cols];
993 row.iter().zip(x).map(|(w, v)| w * v).sum()
994 }
995 Self::Mapped {
996 dtype,
997 row_scale,
998 col_field,
999 ..
1000 } => match dtype {
1001 TensorDtype::Q8Row => {
1002 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1003 dot_i8_f32(q, x) * row_scale[r]
1004 }
1005 TensorDtype::Q8_2f => {
1006 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1007 dot_i8_col_f32(q, x, col_field) * row_scale[r]
1008 }
1009 _ => {
1010 self.row_f32(r, scratch);
1011 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1012 }
1013 },
1014 }
1015 }
1016
1017 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1020 match self {
1021 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
1025 Self::Mapped {
1026 model,
1027 idx,
1028 dtype,
1029 rows,
1030 cols,
1031 row_scale,
1032 col_field,
1033 vbit_offsets,
1034 repack,
1035 } => {
1036 let _ = (model, idx);
1037 assert!(
1047 out.len() >= *rows && x.len() >= *cols,
1048 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1049 out.len(),
1050 x.len(),
1051 );
1052 if *dtype == TensorDtype::Q4Block {
1053 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1057 let t0 = std::time::Instant::now();
1058 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1059 crate::gpu::ProbeArm::Gpu => {
1060 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1061 crate::gpu::probe_record(
1062 crate::gpu::OpClass::Matvec,
1063 true,
1064 t0.elapsed(),
1065 );
1066 return;
1067 }
1068 }
1069 crate::gpu::ProbeArm::CpuTimed => {
1070 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1071 crate::gpu::probe_record(
1072 crate::gpu::OpClass::Matvec,
1073 false,
1074 t0.elapsed(),
1075 );
1076 return;
1077 }
1078 crate::gpu::ProbeArm::Cpu => {}
1079 }
1080 }
1081 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1082 return;
1083 }
1084 if *dtype == TensorDtype::Q4Tiled {
1085 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1090 let t0 = std::time::Instant::now();
1091 let cls = crate::gpu::matvec_class(*rows, *cols);
1092 match crate::gpu::probe_arm(cls) {
1093 crate::gpu::ProbeArm::Gpu => {
1094 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1095 crate::gpu::probe_record(cls, true, t0.elapsed());
1096 return;
1097 }
1098 }
1099 crate::gpu::ProbeArm::CpuTimed => {
1100 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1101 crate::gpu::probe_record(cls, false, t0.elapsed());
1102 return;
1103 }
1104 crate::gpu::ProbeArm::Cpu => {}
1105 }
1106 }
1107 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1108 return;
1109 }
1110 if *dtype == TensorDtype::Q4TiledP {
1111 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1117 let t0 = std::time::Instant::now();
1118 let cls = crate::gpu::matvec_class(*rows, *cols);
1119 match crate::gpu::probe_arm(cls) {
1120 crate::gpu::ProbeArm::Gpu => {
1121 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1122 crate::gpu::probe_record(cls, true, t0.elapsed());
1123 return;
1124 }
1125 }
1126 crate::gpu::ProbeArm::CpuTimed => {
1127 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1128 crate::gpu::probe_record(cls, false, t0.elapsed());
1129 return;
1130 }
1131 crate::gpu::ProbeArm::Cpu => {}
1132 }
1133 }
1134 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1135 return;
1136 }
1137 if *dtype == TensorDtype::Q2TiledP {
1138 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1139 return;
1140 }
1141 if *dtype == TensorDtype::Q1 {
1142 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1147 let t0 = std::time::Instant::now();
1148 let arm = if crate::gpu::q1_force() {
1149 crate::gpu::ProbeArm::Gpu
1150 } else {
1151 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1152 };
1153 match arm {
1154 crate::gpu::ProbeArm::Gpu => {
1155 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1156 crate::gpu::probe_record(
1157 crate::gpu::OpClass::Matvec,
1158 true,
1159 t0.elapsed(),
1160 );
1161 return;
1162 }
1163 }
1164 crate::gpu::ProbeArm::CpuTimed => {
1165 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1166 crate::gpu::probe_record(
1167 crate::gpu::OpClass::Matvec,
1168 false,
1169 t0.elapsed(),
1170 );
1171 return;
1172 }
1173 crate::gpu::ProbeArm::Cpu => {}
1174 }
1175 }
1176 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1177 return;
1178 }
1179 if *dtype == TensorDtype::Q1T {
1180 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1184 let t0 = std::time::Instant::now();
1185 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1186 crate::gpu::ProbeArm::Gpu => {
1187 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1188 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1189 crate::gpu::probe_record(
1190 crate::gpu::OpClass::Matvec,
1191 true,
1192 t0.elapsed(),
1193 );
1194 return;
1195 }
1196 }
1197 crate::gpu::ProbeArm::CpuTimed => {
1198 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1199 crate::gpu::probe_record(
1200 crate::gpu::OpClass::Matvec,
1201 false,
1202 t0.elapsed(),
1203 );
1204 return;
1205 }
1206 crate::gpu::ProbeArm::Cpu => {}
1207 }
1208 }
1209 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1210 return;
1211 }
1212 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1213 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1214 return;
1215 }
1216 let xs = prescale(x, col_field, *dtype);
1217 if *rows >= crate::gpu::min_rows()
1222 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1223 && gpu_lmhead_enabled()
1224 && crate::gpu::enabled_here()
1225 {
1226 let t0 = std::time::Instant::now();
1229 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1230 crate::gpu::ProbeArm::Gpu => {}
1231 crate::gpu::ProbeArm::CpuTimed => {
1232 qmatvec(
1233 self.quant_bytes(),
1234 repack,
1235 row_scale,
1236 x,
1237 col_field,
1238 *dtype,
1239 *rows,
1240 *cols,
1241 out,
1242 pool,
1243 );
1244 crate::gpu::probe_record(
1245 crate::gpu::OpClass::Matvec,
1246 false,
1247 t0.elapsed(),
1248 );
1249 return;
1250 }
1251 crate::gpu::ProbeArm::Cpu => {
1252 qmatvec(
1253 self.quant_bytes(),
1254 repack,
1255 row_scale,
1256 x,
1257 col_field,
1258 *dtype,
1259 *rows,
1260 *cols,
1261 out,
1262 pool,
1263 );
1264 return;
1265 }
1266 }
1267 let frac = gpu_split_frac();
1268 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1269 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1270 let bytes = self.quant_bytes();
1271 let ok = std::thread::scope(|sc| {
1272 let g = sc.spawn(|| {
1273 crate::gpu::q8_matvec_range(
1274 model,
1275 *idx,
1276 cpu_rows,
1277 &row_scale[cpu_rows..],
1278 &xs,
1279 *rows - cpu_rows,
1280 *cols,
1281 out_gpu,
1282 )
1283 });
1284 if cpu_rows > 0 {
1285 let rep_cpu = if repack.is_empty() {
1288 &[][..]
1289 } else {
1290 &repack[..(cpu_rows / 4) * 4 * *cols]
1291 };
1292 qmatvec(
1293 &bytes[..cpu_rows * *cols],
1294 rep_cpu,
1295 &row_scale[..cpu_rows],
1296 x,
1297 col_field,
1298 *dtype,
1299 cpu_rows,
1300 *cols,
1301 out_cpu,
1302 pool,
1303 );
1304 }
1305 g.join().unwrap_or(false)
1306 });
1307 if ok {
1308 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1309 return;
1310 }
1311 qmatvec(
1314 &bytes[cpu_rows * *cols..(*rows) * *cols],
1315 &[],
1316 &row_scale[cpu_rows..],
1317 x,
1318 col_field,
1319 *dtype,
1320 *rows - cpu_rows,
1321 *cols,
1322 out_gpu,
1323 pool,
1324 );
1325 return;
1326 }
1327 qmatvec(
1328 self.quant_bytes(),
1329 repack,
1330 row_scale,
1331 x,
1332 col_field,
1333 *dtype,
1334 *rows,
1335 *cols,
1336 out,
1337 pool,
1338 );
1339 }
1340 }
1341 }
1342
1343 pub fn matvec2(
1345 &self,
1346 x1: &[f32],
1347 x2: &[f32],
1348 o1: &mut [f32],
1349 o2: &mut [f32],
1350 pool: Option<&Pool>,
1351 ) {
1352 match self {
1353 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1354 Self::Mapped {
1355 dtype,
1356 rows,
1357 cols,
1358 row_scale,
1359 col_field,
1360 vbit_offsets,
1361 ..
1362 } => {
1363 if *dtype == TensorDtype::Q4Block {
1364 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1365 return;
1366 }
1367 if *dtype == TensorDtype::Q4Tiled {
1368 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1369 return;
1370 }
1371 if *dtype == TensorDtype::Q4TiledP {
1372 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1373 return;
1374 }
1375 if *dtype == TensorDtype::Q2TiledP {
1376 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1377 return;
1378 }
1379 if *dtype == TensorDtype::Q1 {
1380 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1381 return;
1382 }
1383 if *dtype == TensorDtype::Q1T {
1384 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1390 return;
1391 }
1392 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1393 vbitmatvec2(
1394 self.quant_bytes(),
1395 vbit_offsets,
1396 x1,
1397 x2,
1398 *rows,
1399 *cols,
1400 o1,
1401 o2,
1402 pool,
1403 );
1404 return;
1405 }
1406 qmatvec2(
1407 self.quant_bytes(),
1408 row_scale,
1409 x1,
1410 x2,
1411 col_field,
1412 *dtype,
1413 *rows,
1414 *cols,
1415 o1,
1416 o2,
1417 pool,
1418 );
1419 }
1420 }
1421 }
1422}
1423
1424impl QTensor {
1425 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1433 match self {
1434 Self::Mapped {
1435 model, idx, dtype, ..
1436 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1437 _ => None,
1438 }
1439 }
1440
1441 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1442 let cols = self.cols();
1443 let rows = self.rows();
1444 debug_assert_eq!(xs_all.len(), b * cols);
1445 debug_assert_eq!(out.len(), b * rows);
1446 if crate::gptq_capture::capturing() {
1450 if let Self::Mapped { model, idx, .. } = self {
1451 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1452 }
1453 }
1454 match self {
1455 Self::F32 { data, .. } => {
1456 let out_addr = SendMut(out.as_mut_ptr());
1457 let run = |start: usize, end: usize| {
1458 for o in start..end {
1459 let row = &data[o * cols..(o + 1) * cols];
1460 for bi in 0..b {
1461 let x = &xs_all[bi * cols..(bi + 1) * cols];
1462 let mut acc = 0f32;
1463 for j in 0..cols {
1464 acc += row[j] * x[j];
1465 }
1466 unsafe { *out_addr.at(bi * rows + o) = acc };
1467 }
1468 }
1469 };
1470 dispatch_rows(pool, rows, &run);
1471 }
1472 Self::Mapped {
1473 dtype,
1474 row_scale,
1475 col_field,
1476 vbit_offsets,
1477 ..
1478 } => {
1479 if *dtype == TensorDtype::Q4Block {
1480 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1481 return;
1482 }
1483 if *dtype == TensorDtype::Q4TiledP {
1484 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 if crate::mm_ab::on() {
1519 let mut g = vec![0f32; b * rows];
1520 let t = std::time::Instant::now();
1521 let took = crate::gpu::q4tp_matmat(
1522 model, *idx, xs_all, b, rows, cols, &mut g,
1523 );
1524 let dg = t.elapsed();
1525 let t = std::time::Instant::now();
1526 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1527 let dc = t.elapsed();
1528 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1529 return;
1530 }
1531 let t0 = std::time::Instant::now();
1532 let resident = crate::gpu::weight_is_resident(model, *idx);
1536 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1537 crate::gpu::ProbeArm::Gpu => {
1538 if crate::gpu::q4tp_matmat(
1539 model, *idx, xs_all, b, rows, cols, out,
1540 ) {
1541 let el = t0.elapsed();
1542 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1552 let budget = std::time::Duration::from_secs_f64(
1553 flops / 1.5e12 * 8.0 + 0.020,
1554 );
1555 crate::gpu::mm_budget_check(
1556 "q4tp matmat",
1557 el,
1558 budget,
1559 crate::gpu::probe_was_cold() || !resident,
1560 );
1561 crate::gpu::probe_record(class, true, el);
1562 return;
1563 }
1564 }
1565 crate::gpu::ProbeArm::CpuTimed => {
1566 q4tp_matmat(
1567 self.quant_bytes(),
1568 xs_all,
1569 b,
1570 rows,
1571 cols,
1572 out,
1573 pool,
1574 );
1575 crate::gpu::probe_record(class, false, t0.elapsed());
1576 return;
1577 }
1578 crate::gpu::ProbeArm::Cpu => {}
1579 }
1580 }
1581 }
1582 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1583 return;
1584 }
1585 if *dtype == TensorDtype::Q2TiledP {
1586 if b >= 32
1592 && b * rows * cols >= 128_000_000
1593 && cols % 32 == 0
1594 && !crate::gpu::mm_killed()
1595 && crate::gpu::enabled_here()
1596 {
1597 let class = if b >= 128 {
1598 crate::gpu::OpClass::MatmatWide
1599 } else {
1600 crate::gpu::OpClass::Matmat
1601 };
1602 if let Self::Mapped { model, idx, .. } = self {
1603 let t0 = std::time::Instant::now();
1604 match crate::gpu::probe_arm(class) {
1605 crate::gpu::ProbeArm::Gpu => {
1606 if crate::gpu::q2tp_matmat(
1607 model, *idx, xs_all, b, rows, cols, out,
1608 ) {
1609 crate::gpu::probe_record(class, true, t0.elapsed());
1610 return;
1611 }
1612 }
1613 crate::gpu::ProbeArm::CpuTimed => {
1614 q2tp_matmat(
1615 self.quant_bytes(),
1616 xs_all,
1617 b,
1618 rows,
1619 cols,
1620 out,
1621 pool,
1622 );
1623 crate::gpu::probe_record(class, false, t0.elapsed());
1624 return;
1625 }
1626 crate::gpu::ProbeArm::Cpu => {}
1627 }
1628 }
1629 }
1630 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1635 return;
1636 }
1637 if *dtype == TensorDtype::Q4Tiled {
1638 if b >= 32
1650 && b * rows * cols >= 128_000_000
1651 && cols % 32 == 0
1652 && !crate::gpu::mm_killed()
1653 && crate::gpu::enabled_here()
1654 {
1655 let class = if b >= 128 {
1656 crate::gpu::OpClass::MatmatWide
1657 } else {
1658 crate::gpu::OpClass::Matmat
1659 };
1660 if let Self::Mapped { model, idx, .. } = self {
1661 let t0 = std::time::Instant::now();
1662 match crate::gpu::probe_arm(class) {
1663 crate::gpu::ProbeArm::Gpu => {
1664 if crate::gpu::q4t_matmat(
1665 model, *idx, xs_all, b, rows, cols, out,
1666 ) {
1667 let el = t0.elapsed();
1668 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1678 let budget = std::time::Duration::from_secs_f64(
1679 flops / 1.5e12 * 8.0 + 0.020,
1680 );
1681 crate::gpu::mm_budget_check(
1682 "q4t matmat",
1683 el,
1684 budget,
1685 crate::gpu::probe_was_cold(),
1686 );
1687 crate::gpu::probe_record(class, true, el);
1688 return;
1689 }
1690 }
1691 crate::gpu::ProbeArm::CpuTimed => {
1692 q4t_matmat(
1693 self.quant_bytes(),
1694 xs_all,
1695 b,
1696 rows,
1697 cols,
1698 out,
1699 pool,
1700 );
1701 crate::gpu::probe_record(class, false, t0.elapsed());
1702 return;
1703 }
1704 crate::gpu::ProbeArm::Cpu => {}
1705 }
1706 }
1707 }
1708 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1709 return;
1710 }
1711 if *dtype == TensorDtype::Q1 {
1712 if b >= 32
1715 && b * rows * cols >= 128_000_000
1716 && cols % 64 == 0
1717 && crate::gpu::enabled_here()
1718 {
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::q1_matmat(
1724 model, *idx, xs_all, b, rows, cols, out,
1725 ) {
1726 crate::gpu::probe_record(
1727 crate::gpu::OpClass::Matmat,
1728 true,
1729 t0.elapsed(),
1730 );
1731 return;
1732 }
1733 }
1734 crate::gpu::ProbeArm::CpuTimed => {
1735 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1736 crate::gpu::probe_record(
1737 crate::gpu::OpClass::Matmat,
1738 false,
1739 t0.elapsed(),
1740 );
1741 return;
1742 }
1743 crate::gpu::ProbeArm::Cpu => {}
1744 }
1745 }
1746 }
1747 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1748 return;
1749 }
1750 if *dtype == TensorDtype::Q1T {
1751 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1754 if let Self::Mapped { model, idx, .. } = self {
1755 let t0 = std::time::Instant::now();
1756 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1757 crate::gpu::ProbeArm::Gpu => {
1758 if crate::gpu::q1t_matmat(
1759 model, *idx, xs_all, b, rows, cols, out,
1760 ) {
1761 crate::gpu::probe_record(
1762 crate::gpu::OpClass::Matmat,
1763 true,
1764 t0.elapsed(),
1765 );
1766 return;
1767 }
1768 }
1769 crate::gpu::ProbeArm::CpuTimed => {
1770 q1t_matmat(
1771 self.quant_bytes(),
1772 xs_all,
1773 b,
1774 rows,
1775 cols,
1776 out,
1777 pool,
1778 );
1779 crate::gpu::probe_record(
1780 crate::gpu::OpClass::Matmat,
1781 false,
1782 t0.elapsed(),
1783 );
1784 return;
1785 }
1786 crate::gpu::ProbeArm::Cpu => {}
1787 }
1788 }
1789 }
1790 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1791 return;
1792 }
1793 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1794 vbitmatmat(
1795 self.quant_bytes(),
1796 vbit_offsets,
1797 xs_all,
1798 b,
1799 rows,
1800 cols,
1801 out,
1802 pool,
1803 );
1804 return;
1805 }
1806 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1807 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1808 .collect();
1809 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1815 if let Self::Mapped { model, idx, .. } = self {
1816 let t0 = std::time::Instant::now();
1817 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1818 crate::gpu::ProbeArm::Gpu
1819 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1820 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1821 {
1822 let q = self.quant_bytes();
1826 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1827 return;
1828 }
1829 crate::gpu::ProbeArm::Gpu => {
1830 let flat: Vec<f32> =
1831 pre.iter().flat_map(|v| v.iter().copied()).collect();
1832 if crate::gpu::q8_matmat(
1833 model, *idx, row_scale, &flat, b, rows, cols, out,
1834 ) {
1835 crate::gpu::probe_record(
1836 crate::gpu::OpClass::Matmat,
1837 true,
1838 t0.elapsed(),
1839 );
1840 return;
1841 }
1842 }
1843 crate::gpu::ProbeArm::CpuTimed => {
1844 let q = self.quant_bytes();
1845 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1846 crate::gpu::probe_record(
1847 crate::gpu::OpClass::Matmat,
1848 false,
1849 t0.elapsed(),
1850 );
1851 return;
1852 }
1853 crate::gpu::ProbeArm::Cpu => {}
1854 }
1855 }
1856 }
1857 let q = self.quant_bytes();
1858 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1859 }
1860 }
1861 }
1862}
1863
1864impl QTensor {
1865 pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
1875 let (rows, cols) = (self.rows(), self.cols());
1876 let Self::Mapped { model, idx, dtype, row_scale, col_field, .. } = self else {
1877 return false;
1878 };
1879 match *dtype {
1880 TensorDtype::Q4TiledP => {
1881 crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out)
1882 }
1883 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1887 if *dtype == TensorDtype::Q8_2f
1890 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
1891 && crate::gpu::q8_matmat_2f(
1892 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
1893 )
1894 {
1895 return true;
1896 }
1897 let flat: Vec<f32> = (0..b)
1898 .flat_map(|bi| {
1899 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
1900 })
1901 .collect();
1902 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
1903 }
1904 _ => false,
1905 }
1906 }
1907
1908 pub fn matvec_many<const N: usize>(
1915 ts: [&QTensor; N],
1916 x: &[f32],
1917 mut outs: [&mut [f32]; N],
1918 pool: Option<&Pool>,
1919 ) {
1920 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1921 let uniform_q8 = ts.iter().all(|t| {
1922 matches!(
1923 t,
1924 Self::Mapped {
1925 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1926 ..
1927 }
1928 )
1929 });
1930 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1931 let uniform_q4 = ts.iter().all(|t| {
1932 matches!(
1933 t,
1934 Self::Mapped {
1935 dtype: TensorDtype::Q4Block,
1936 ..
1937 }
1938 )
1939 });
1940 let uniform_vbit = ts.iter().all(|t| {
1941 matches!(
1942 t,
1943 Self::Mapped {
1944 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1945 ..
1946 }
1947 )
1948 });
1949 let uniform_q1 = ts.iter().all(|t| {
1950 matches!(
1951 t,
1952 Self::Mapped {
1953 dtype: TensorDtype::Q1,
1954 ..
1955 }
1956 )
1957 });
1958 let uniform_q1t = ts.iter().all(|t| {
1959 matches!(
1960 t,
1961 Self::Mapped {
1962 dtype: TensorDtype::Q1T,
1963 ..
1964 }
1965 )
1966 });
1967 let uniform_q4tp = ts.iter().all(|t| {
1972 matches!(
1973 t,
1974 Self::Mapped {
1975 dtype: TensorDtype::Q4TiledP,
1976 ..
1977 }
1978 )
1979 }) && ts
1980 .iter()
1981 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1982 let Some(pool) = pool else {
1983 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1984 t.matvec(x, o, None);
1985 }
1986 return;
1987 };
1988 if total_rows < 256
1989 || !(uniform_q8
1990 || uniform_f32
1991 || uniform_q4
1992 || uniform_vbit
1993 || uniform_q1
1994 || uniform_q1t
1995 || uniform_q4tp)
1996 {
1997 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1998 t.matvec(x, o, Some(pool));
1999 }
2000 return;
2001 }
2002
2003 if uniform_q4tp {
2004 let cols = ts[0].cols();
2010 let gpr = cols / GROUP_SIZE;
2011 let views: Vec<Q4tpView> = ts
2012 .iter()
2013 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2014 .collect();
2015 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2016 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2017 let locate = |flat: usize| -> (usize, usize) {
2019 let mut acc = 0;
2020 for (i, &r) in rows_of.iter().enumerate() {
2021 if flat < acc + r {
2022 return (i, flat - acc);
2023 }
2024 acc += r;
2025 }
2026 (rows_of.len() - 1, 0)
2027 };
2028 let (views, outs_addr) = (&views, &outs_addr);
2029 if a8w8_enabled() {
2030 let act = split_act(x);
2031 let act = &act;
2032 let run = |start: usize, end: usize| {
2033 let mut sc = vec![0f32; gpr];
2034 for flat in start..end {
2035 let (t, r) = locate(flat);
2036 let v = &views[t];
2037 v.scales_into(r, gpr, &mut sc);
2038 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2039 for &(j, xv) in &act.outliers {
2040 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2041 acc += w * s * xv;
2042 }
2043 unsafe { *outs_addr[t].at(r) = acc };
2045 }
2046 };
2047 pool.run_rows(total_rows, &run);
2048 } else {
2049 let run = |start: usize, end: usize| {
2050 let mut sc = vec![0f32; gpr];
2051 for flat in start..end {
2052 let (t, r) = locate(flat);
2053 let v = &views[t];
2054 v.scales_into(r, gpr, &mut sc);
2055 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2057 }
2058 };
2059 pool.run_rows(total_rows, &run);
2060 }
2061 return;
2062 }
2063
2064 if uniform_q1 {
2065 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2068 if a8w8_enabled() {
2069 let act = split_act(x);
2070 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2071 let (act, gsum) = (&act, &gsum);
2072 let closures: [_; N] = std::array::from_fn(|i| {
2073 let (bytes, gpr, out) =
2074 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2075 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2076 });
2077 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2078 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2079 pool.run_many(&parts);
2080 } else {
2081 let closures: [_; N] = std::array::from_fn(|i| {
2082 let (bytes, gpr, out) =
2083 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2084 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2085 });
2086 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2087 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2088 pool.run_many(&parts);
2089 }
2090 return;
2091 }
2092
2093 if uniform_q1t {
2094 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2098 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2099 if a8w8_enabled() {
2100 let act = split_act(x);
2101 let act = &act;
2102 let x_ref = x;
2103 let closures: [_; N] = std::array::from_fn(|i| {
2104 let bytes = ts[i].quant_bytes();
2105 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2106 let gpr = cols / GROUP_SIZE;
2107 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2108 let out = outs_addr[i];
2109 move |s: usize, e: usize| {
2110 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2111 }
2112 });
2113 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2114 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2115 pool.run_many(&parts);
2116 } else {
2117 let x_ref = x;
2118 let closures: [_; N] = std::array::from_fn(|i| {
2119 let bytes = ts[i].quant_bytes();
2120 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2121 let gpr = cols / GROUP_SIZE;
2122 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2123 let out = outs_addr[i];
2124 move |s: usize, e: usize| {
2125 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2126 }
2127 });
2128 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2129 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2130 pool.run_many(&parts);
2131 }
2132 return;
2133 }
2134
2135 if uniform_q4 || uniform_vbit {
2136 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2137 if a8w8_enabled() {
2139 let act = split_act(x);
2140 let act = &act;
2141 if uniform_q4 {
2142 let closures: [_; N] = std::array::from_fn(|i| {
2143 let (packed, scales) =
2144 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2145 let (gpr, cols, out) =
2146 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2147 move |s: usize, e: usize| {
2148 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2149 }
2150 });
2151 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2152 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2153 pool.run_many(&parts);
2154 } else {
2155 let closures: [_; N] = std::array::from_fn(|i| {
2156 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2157 unreachable!()
2158 };
2159 let (bytes, rows, cols, out) = (
2160 ts[i].quant_bytes(),
2161 ts[i].rows(),
2162 ts[i].cols(),
2163 outs_addr[i],
2164 );
2165 move |s: usize, e: usize| {
2166 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2167 }
2168 });
2169 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2170 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2171 pool.run_many(&parts);
2172 }
2173 return;
2174 }
2175 if uniform_q4 {
2176 let closures: [_; N] = std::array::from_fn(|i| {
2177 let (packed, scales) =
2178 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2179 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2180 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2181 });
2182 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2183 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2184 pool.run_many(&parts);
2185 } else {
2186 let closures: [_; N] = std::array::from_fn(|i| {
2187 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2188 unreachable!()
2189 };
2190 let (bytes, rows, cols, out) = (
2191 ts[i].quant_bytes(),
2192 ts[i].rows(),
2193 ts[i].cols(),
2194 outs_addr[i],
2195 );
2196 move |s: usize, e: usize| {
2197 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2198 }
2199 });
2200 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2201 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2202 pool.run_many(&parts);
2203 }
2204 return;
2205 }
2206
2207 if uniform_f32 {
2208 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2209 let closures: [_; N] = std::array::from_fn(|i| {
2210 let Self::F32 { data, cols, .. } = ts[i] else {
2211 unreachable!()
2212 };
2213 let out = outs_addr[i];
2214 move |start: usize, end: usize| {
2215 for o in start..end {
2216 let row = &data[o * cols..(o + 1) * cols];
2217 let mut sum = 0.0f32;
2218 for j in 0..*cols {
2219 sum += row[j] * x[j];
2220 }
2221 unsafe { *out.at(o) = sum };
2223 }
2224 }
2225 });
2226 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2227 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2228 pool.run_many(&parts);
2229 return;
2230 }
2231
2232 struct Ctx<'a> {
2235 bytes: &'a [u8],
2236 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2237 rep: &'a [u8],
2238 row_scale: &'a [f32],
2239 cols: usize,
2240 xs: std::borrow::Cow<'a, [f32]>,
2241 }
2242 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2243 let Self::Mapped {
2244 dtype,
2245 cols,
2246 row_scale,
2247 col_field,
2248 repack,
2249 ..
2250 } = ts[i]
2251 else {
2252 unreachable!()
2253 };
2254 Ctx {
2255 bytes: ts[i].quant_bytes(),
2256 rep: repack,
2257 row_scale,
2258 cols: *cols,
2259 xs: prescale(x, col_field, *dtype),
2260 }
2261 });
2262 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2263 #[cfg(target_arch = "aarch64")]
2264 if sdot_enabled() {
2265 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2266 let closures: [_; N] = std::array::from_fn(|i| {
2267 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2268 move |start: usize, end: usize| {
2269 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2270 }
2271 });
2272 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2273 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2274 pool.run_many(&parts);
2275 return;
2276 }
2277 #[cfg(target_arch = "x86_64")]
2278 if avx2_a8w8_enabled() {
2279 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2280 let closures: [_; N] = std::array::from_fn(|i| {
2281 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2282 move |start: usize, end: usize| {
2283 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2284 }
2285 });
2286 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2287 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2288 pool.run_many(&parts);
2289 return;
2290 }
2291 let closures: [_; N] = std::array::from_fn(|i| {
2292 let (c, out) = (&ctxs[i], outs_addr[i]);
2293 move |start: usize, end: usize| {
2294 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2295 }
2296 });
2297 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2298 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2299 pool.run_many(&parts);
2300 }
2301}
2302
2303impl QTensor {
2304 #[allow(clippy::needless_range_loop)]
2309 pub fn matvec2_many<const N: usize>(
2310 ts: [&QTensor; N],
2311 x1: &[f32],
2312 x2: &[f32],
2313 mut o1s: [&mut [f32]; N],
2314 mut o2s: [&mut [f32]; N],
2315 pool: Option<&Pool>,
2316 ) {
2317 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2318 let uniform_q8 = ts.iter().all(|t| {
2319 matches!(
2320 t,
2321 Self::Mapped {
2322 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2323 ..
2324 }
2325 )
2326 });
2327 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2328 let uniform_q4 = ts.iter().all(|t| {
2329 matches!(
2330 t,
2331 Self::Mapped {
2332 dtype: TensorDtype::Q4Block,
2333 ..
2334 }
2335 )
2336 });
2337 let uniform_vbit = ts.iter().all(|t| {
2338 matches!(
2339 t,
2340 Self::Mapped {
2341 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2342 ..
2343 }
2344 )
2345 });
2346 let fusable = pool.is_some()
2347 && total_rows >= 256
2348 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2349 if !fusable {
2350 for i in 0..N {
2351 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2352 }
2353 return;
2354 }
2355 let pool = pool.unwrap();
2356
2357 if uniform_q4 || uniform_vbit {
2358 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2359 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2360 if a8w8_enabled() {
2362 let a1 = split_act(x1);
2363 let a2 = split_act(x2);
2364 let (a1, a2) = (&a1, &a2);
2365 if uniform_q4 {
2366 let closures: [_; N] = std::array::from_fn(|i| {
2367 let (packed, scales) =
2368 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2369 let (gpr, cols, o1, o2) =
2370 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2371 move |s: usize, e: usize| {
2372 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2373 }
2374 });
2375 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2376 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2377 pool.run_many(&parts);
2378 } else {
2379 let closures: [_; N] = std::array::from_fn(|i| {
2380 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2381 unreachable!()
2382 };
2383 let (bytes, rows, cols, o1, o2) = (
2384 ts[i].quant_bytes(),
2385 ts[i].rows(),
2386 ts[i].cols(),
2387 p1[i],
2388 p2[i],
2389 );
2390 move |s: usize, e: usize| {
2391 vbit_range2_a8w8(
2392 bytes,
2393 vbit_offsets,
2394 x1,
2395 x2,
2396 a1,
2397 a2,
2398 rows,
2399 cols,
2400 o1,
2401 o2,
2402 s,
2403 e,
2404 )
2405 }
2406 });
2407 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2408 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2409 pool.run_many(&parts);
2410 }
2411 return;
2412 }
2413 if uniform_q4 {
2414 let closures: [_; N] = std::array::from_fn(|i| {
2415 let (packed, scales) =
2416 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2417 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2418 move |s: usize, e: usize| {
2419 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2420 }
2421 });
2422 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2423 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2424 pool.run_many(&parts);
2425 } else {
2426 let closures: [_; N] = std::array::from_fn(|i| {
2427 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2428 unreachable!()
2429 };
2430 let (bytes, rows, cols, o1, o2) = (
2431 ts[i].quant_bytes(),
2432 ts[i].rows(),
2433 ts[i].cols(),
2434 p1[i],
2435 p2[i],
2436 );
2437 move |s: usize, e: usize| {
2438 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2439 }
2440 });
2441 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2442 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2443 pool.run_many(&parts);
2444 }
2445 return;
2446 }
2447
2448 if uniform_f32 {
2449 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2450 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2451 let closures: [_; N] = std::array::from_fn(|i| {
2452 let Self::F32 { data, cols, .. } = ts[i] else {
2453 unreachable!()
2454 };
2455 let (o1, o2) = (p1[i], p2[i]);
2456 move |start: usize, end: usize| {
2457 for o in start..end {
2458 let row = &data[o * cols..(o + 1) * cols];
2459 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2460 for j in 0..*cols {
2461 s1 += row[j] * x1[j];
2462 s2 += row[j] * x2[j];
2463 }
2464 unsafe {
2466 *o1.at(o) = s1;
2467 *o2.at(o) = s2;
2468 }
2469 }
2470 }
2471 });
2472 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2473 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2474 pool.run_many(&parts);
2475 return;
2476 }
2477
2478 struct Ctx<'a> {
2479 bytes: &'a [u8],
2480 row_scale: &'a [f32],
2481 cols: usize,
2482 xs1: std::borrow::Cow<'a, [f32]>,
2483 xs2: std::borrow::Cow<'a, [f32]>,
2484 }
2485 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2486 let Self::Mapped {
2487 dtype,
2488 cols,
2489 row_scale,
2490 col_field,
2491 ..
2492 } = ts[i]
2493 else {
2494 unreachable!()
2495 };
2496 Ctx {
2497 bytes: ts[i].quant_bytes(),
2498 row_scale,
2499 cols: *cols,
2500 xs1: prescale(x1, col_field, *dtype),
2501 xs2: prescale(x2, col_field, *dtype),
2502 }
2503 });
2504 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2505 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2506 #[cfg(target_arch = "aarch64")]
2507 if sdot_enabled() {
2508 let acts: [(SplitAct, SplitAct); N] =
2509 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2510 let closures: [_; N] = std::array::from_fn(|i| {
2511 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2512 move |start: usize, end: usize| {
2513 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2514 }
2515 });
2516 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2517 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2518 pool.run_many(&parts);
2519 return;
2520 }
2521 #[cfg(target_arch = "x86_64")]
2522 if avx2_a8w8_enabled() {
2523 let acts: [(SplitAct, SplitAct); N] =
2524 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2525 let closures: [_; N] = std::array::from_fn(|i| {
2526 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2527 move |start: usize, end: usize| {
2528 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2529 }
2530 });
2531 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2532 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2533 pool.run_many(&parts);
2534 return;
2535 }
2536 let closures: [_; N] = std::array::from_fn(|i| {
2537 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2538 move |start: usize, end: usize| {
2539 q8_range2_f32(
2540 c.bytes,
2541 c.row_scale,
2542 &c.xs1,
2543 &c.xs2,
2544 c.cols,
2545 o1,
2546 o2,
2547 start,
2548 end,
2549 )
2550 }
2551 });
2552 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2553 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2554 pool.run_many(&parts);
2555 }
2556
2557 pub fn matvec_silu_mul(
2562 gate: &QTensor,
2563 up: &QTensor,
2564 x: &[f32],
2565 out: &mut [f32],
2566 pool: Option<&Pool>,
2567 ) -> bool {
2568 let inter = gate.rows();
2569 debug_assert_eq!(up.rows(), inter);
2570 debug_assert_eq!(out.len(), inter);
2571 debug_assert_eq!(gate.cols(), up.cols());
2572 if !a8w8_enabled() {
2573 return false;
2574 }
2575 let act = split_act(x);
2576 let act = &act;
2577 let x_ref = x;
2578 let out_addr = SendMut(out.as_mut_ptr());
2579
2580 match (gate, up) {
2581 (
2583 Self::Mapped {
2584 dtype: TensorDtype::Q4Block,
2585 ..
2586 },
2587 Self::Mapped {
2588 dtype: TensorDtype::Q4Block,
2589 ..
2590 },
2591 ) => {
2592 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2593 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2594 let gpr = gate.cols() / GROUP_SIZE;
2595 let cols = gate.cols();
2596 let run = move |start: usize, end: usize| {
2597 for r in start..end {
2598 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2599 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2600 for &(j, xv) in &act.outliers {
2601 let flat = r * cols + j;
2602 let gb = gp[flat / 2];
2603 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2604 let gsc = f16_to_f32(u16::from_le_bytes([
2605 gs[(flat / GROUP_SIZE) * 2],
2606 gs[(flat / GROUP_SIZE) * 2 + 1],
2607 ]));
2608 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2609 let ub = up_p[flat / 2];
2610 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2611 let usc = f16_to_f32(u16::from_le_bytes([
2612 up_s[(flat / GROUP_SIZE) * 2],
2613 up_s[(flat / GROUP_SIZE) * 2 + 1],
2614 ]));
2615 uv += ((un as i32 - 8) as f32) * usc * xv;
2616 }
2617 let silu_g = gv / (1.0 + (-gv).exp());
2618 unsafe { *out_addr.at(r) = silu_g * uv };
2620 }
2621 };
2622 dispatch_rows(pool, inter, &run);
2623 true
2624 }
2625 (
2629 Self::Mapped {
2630 dtype: TensorDtype::Q4Tiled,
2631 ..
2632 },
2633 Self::Mapped {
2634 dtype: TensorDtype::Q4Tiled,
2635 ..
2636 },
2637 ) => {
2638 let g_bytes = gate.quant_bytes();
2639 let u_bytes = up.quant_bytes();
2640 let gpr = gate.cols() / GROUP_SIZE;
2641 let run = move |start: usize, end: usize| {
2642 for r in start..end {
2643 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2644 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2645 for &(j, xv) in &act.outliers {
2646 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2647 gv += w * s * xv;
2648 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2649 uv += w * s * xv;
2650 }
2651 let silu_g = gv / (1.0 + (-gv).exp());
2652 unsafe { *out_addr.at(r) = silu_g * uv };
2654 }
2655 };
2656 dispatch_rows(pool, inter, &run);
2657 true
2658 }
2659 (
2662 Self::Mapped {
2663 dtype: TensorDtype::Q4TiledP,
2664 ..
2665 },
2666 Self::Mapped {
2667 dtype: TensorDtype::Q4TiledP,
2668 ..
2669 },
2670 ) => {
2671 let cols = gate.cols();
2672 let gpr = cols / GROUP_SIZE;
2673 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2674 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2675 let run = |start: usize, end: usize| {
2676 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2677 for r in start..end {
2678 gv_view.scales_into(r, gpr, &mut gsc);
2679 uv_view.scales_into(r, gpr, &mut usc);
2680 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2681 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2682 for &(j, xv) in &act.outliers {
2683 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2684 gv += w * s * xv;
2685 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2686 uv += w * s * xv;
2687 }
2688 let silu_g = gv / (1.0 + (-gv).exp());
2689 unsafe { *out_addr.at(r) = silu_g * uv };
2691 }
2692 };
2693 dispatch_rows(pool, inter, &run);
2694 true
2695 }
2696 (
2702 Self::Mapped {
2703 dtype: TensorDtype::Q1,
2704 ..
2705 },
2706 Self::Mapped {
2707 dtype: TensorDtype::Q1,
2708 ..
2709 },
2710 ) => {
2711 let g_bytes = gate.quant_bytes();
2712 let u_bytes = up.quant_bytes();
2713 let gpr = gate.cols() / GROUP_SIZE;
2714 let gsum = q1_group_sums(&act.xq, gpr);
2715 let gsum = &gsum;
2716 let run = move |start: usize, end: usize| {
2717 for r in start..end {
2718 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2719 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2720 for &(j, xv) in &act.outliers {
2721 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2722 gv += w * s * xv;
2723 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2724 uv += w * s * xv;
2725 }
2726 let silu_g = gv / (1.0 + (-gv).exp());
2727 unsafe { *out_addr.at(r) = silu_g * uv };
2729 }
2730 };
2731 dispatch_rows(pool, inter, &run);
2732 true
2733 }
2734 (
2738 Self::Mapped {
2739 dtype: TensorDtype::Q2TiledP,
2740 ..
2741 },
2742 Self::Mapped {
2743 dtype: TensorDtype::Q2TiledP,
2744 ..
2745 },
2746 ) => {
2747 let cols = gate.cols();
2748 let gpr = cols / GROUP_SIZE;
2749 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2750 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2751 let gsum = q1_group_sums(&act.xq, gpr);
2752 let gsum = &gsum;
2753 let run = move |start: usize, end: usize| {
2754 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2755 for r in start..end {
2756 gv_view.scales_into(r, gpr, &mut gsc);
2757 uv_view.scales_into(r, gpr, &mut usc);
2758 let mut gv =
2759 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2760 let mut uv =
2761 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2762 for &(j, xv) in &act.outliers {
2763 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2764 gv += w * s * xv;
2765 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2766 uv += w * s * xv;
2767 }
2768 let silu_g = gv / (1.0 + (-gv).exp());
2769 unsafe { *out_addr.at(r) = silu_g * uv };
2771 }
2772 };
2773 dispatch_rows(pool, inter, &run);
2774 true
2775 }
2776 (
2781 Self::Mapped {
2782 dtype: TensorDtype::Q8Row,
2783 row_scale: g_rs,
2784 ..
2785 },
2786 Self::Mapped {
2787 dtype: TensorDtype::Q8Row,
2788 row_scale: u_rs,
2789 ..
2790 },
2791 ) => {
2792 let g_bytes = gate.quant_bytes();
2793 let u_bytes = up.quant_bytes();
2794 let cols = gate.cols();
2795 let run = move |start: usize, end: usize| {
2796 for r in start..end {
2797 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2798 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2799 let silu_g = gv / (1.0 + (-gv).exp());
2800 unsafe { *out_addr.at(r) = silu_g * uv };
2802 }
2803 };
2804 dispatch_rows(pool, inter, &run);
2805 true
2806 }
2807 (
2809 Self::Mapped {
2810 dtype: TensorDtype::Q1T,
2811 ..
2812 },
2813 Self::Mapped {
2814 dtype: TensorDtype::Q1T,
2815 ..
2816 },
2817 ) => {
2818 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2819 let g_bytes = gate.quant_bytes();
2820 let u_bytes = up.quant_bytes();
2821 let gpr = gate.cols() / GROUP_SIZE;
2822 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2823 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2824 let run = move |start: usize, end: usize| {
2825 for r in start..end {
2826 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2827 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2828 for &(j, xv) in &act.outliers {
2829 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2830 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2831 }
2832 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2833 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2834 let silu_g = gv / (1.0 + (-gv).exp());
2835 unsafe { *out_addr.at(r) = silu_g * uv };
2837 }
2838 };
2839 dispatch_rows(pool, inter, &run);
2840 true
2841 }
2842 _ => false,
2843 }
2844 }
2845
2846 pub fn moe_gate_up_many(
2860 pairs: &[(&QTensor, &QTensor)],
2861 x: &[f32],
2862 outs: &mut [Vec<f32>],
2863 pool: Option<&Pool>,
2864 ) -> bool {
2865 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2866 return false;
2867 }
2868 let inter = pairs[0].0.rows();
2869 let cols = pairs[0].0.cols();
2870 if cols % GROUP_SIZE != 0 {
2871 return false;
2872 }
2873 let gpr = cols / GROUP_SIZE;
2874 let q2 = matches!(
2877 pairs[0].0,
2878 Self::Mapped {
2879 dtype: TensorDtype::Q2TiledP,
2880 ..
2881 }
2882 );
2883 let want = if q2 {
2884 TensorDtype::Q2TiledP
2885 } else {
2886 TensorDtype::Q4TiledP
2887 };
2888 let mut views = Vec::with_capacity(pairs.len() * 2);
2889 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2890 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2891 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2892 if !both
2893 || g.rows() != inter
2894 || u.rows() != inter
2895 || g.cols() != cols
2896 || u.cols() != cols
2897 || o.len() != inter
2898 {
2899 return false;
2900 }
2901 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2902 views.push(mk(g.quant_bytes(), inter, cols));
2903 views.push(mk(u.quant_bytes(), inter, cols));
2904 }
2905 let act = split_act(x);
2906 let gsum = if q2 {
2907 q1_group_sums(&act.xq, gpr)
2908 } else {
2909 Vec::new()
2910 };
2911 let (act, gsum) = (&act, &gsum);
2912 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2913 let (views, ptrs) = (&views, &ptrs);
2914 let run = |start: usize, end: usize| {
2915 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2916 for flat in start..end {
2917 let (e, r) = (flat / inter, flat % inter);
2918 let gv_view = &views[e * 2];
2919 let uv_view = &views[e * 2 + 1];
2920 gv_view.scales_into(r, gpr, &mut gsc);
2921 uv_view.scales_into(r, gpr, &mut usc);
2922 let (mut gv, mut uv) = if q2 {
2923 (
2924 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2925 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2926 )
2927 } else {
2928 (
2929 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2930 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2931 )
2932 };
2933 for &(j, xv) in &act.outliers {
2934 let (og, ou) = if q2 {
2935 (
2936 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2937 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2938 )
2939 } else {
2940 (
2941 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2942 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2943 )
2944 };
2945 gv += og.0 * og.1 * xv;
2946 uv += ou.0 * ou.1 * xv;
2947 }
2948 let silu_g = gv / (1.0 + (-gv).exp());
2949 unsafe { *ptrs[e].at(r) = silu_g * uv };
2951 }
2952 };
2953 dispatch_rows(pool, pairs.len() * inter, &run);
2954 true
2955 }
2956
2957 pub fn moe_down_many(
2966 downs: &[&QTensor],
2967 gs: &[Vec<f32>],
2968 weights: &[f32],
2969 out: &mut [f32],
2970 pool: Option<&Pool>,
2971 ) -> bool {
2972 if downs.is_empty()
2973 || downs.len() != gs.len()
2974 || downs.len() != weights.len()
2975 || !a8w8_enabled()
2976 {
2977 return false;
2978 }
2979 let rows = out.len();
2980 let cols = downs[0].cols();
2981 if cols % GROUP_SIZE != 0 {
2982 return false;
2983 }
2984 let gpr = cols / GROUP_SIZE;
2985 let mut views = Vec::with_capacity(downs.len());
2986 for (d, g) in downs.iter().zip(gs.iter()) {
2987 if !matches!(
2988 d,
2989 Self::Mapped {
2990 dtype: TensorDtype::Q4TiledP,
2991 ..
2992 }
2993 ) || d.rows() != rows
2994 || d.cols() != cols
2995 || g.len() != cols
2996 {
2997 return false;
2998 }
2999 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3000 }
3001 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3003 let out_addr = SendMut(out.as_mut_ptr());
3011 let (views, acts, weights) = (&views, &acts, &weights);
3012 let run = |start: usize, end: usize| {
3013 let mut sc = vec![0f32; gpr];
3014 for r in start..end {
3015 let mut acc = 0f32;
3016 for (e, v) in views.iter().enumerate() {
3017 v.scales_into(r, gpr, &mut sc);
3018 let a = &acts[e];
3019 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3020 for &(j, xv) in &a.outliers {
3021 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3022 d += w * s * xv;
3023 }
3024 acc += weights[e] * d;
3025 }
3026 unsafe { *out_addr.at(r) = acc };
3028 }
3029 };
3030 dispatch_rows(pool, rows, &run);
3031 true
3032 }
3033}
3034
3035#[cfg(target_os = "macos")]
3040mod accel_blas {
3041 #[link(name = "Accelerate", kind = "framework")]
3042 unsafe extern "C" {
3043 pub fn cblas_sgemm(
3044 order: i32,
3045 trans_a: i32,
3046 trans_b: i32,
3047 m: i32,
3048 n: i32,
3049 k: i32,
3050 alpha: f32,
3051 a: *const f32,
3052 lda: i32,
3053 b: *const f32,
3054 ldb: i32,
3055 beta: f32,
3056 c: *mut f32,
3057 ldc: i32,
3058 );
3059 }
3060}
3061
3062#[cfg(target_os = "macos")]
3063pub(crate) fn accel_gemm_enabled() -> bool {
3064 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3065 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3066}
3067
3068#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3071pub(crate) fn accel_gemm_enabled() -> bool {
3072 true
3073}
3074
3075#[cfg(target_arch = "aarch64")]
3082#[allow(clippy::too_many_arguments)]
3083pub(crate) fn neon_gemm_rm(
3084 m: usize,
3085 n: usize,
3086 k: usize,
3087 alpha: f32,
3088 a: &[f32],
3089 lda: usize,
3090 b_mat: &[f32],
3091 ldb: usize,
3092 b_rows_are_n: bool,
3093 c: &mut [f32],
3094 ldc: usize,
3095) {
3096 debug_assert!(a.len() >= (m - 1) * lda + k);
3097 debug_assert!(c.len() >= (m - 1) * ldc + n);
3098 unsafe {
3100 use core::arch::aarch64::*;
3101 let mut i = 0usize;
3102 while i < m {
3103 let mi = (m - i).min(4);
3104 let mut j = 0usize;
3105 while j < n {
3106 let nj = (n - j).min(8);
3107 if mi == 4 && nj == 8 {
3108 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3109 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3110 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3111 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3112 for p in 0..k {
3113 let (b0, b1) = if b_rows_are_n {
3114 let base = b_mat.as_ptr().add(j * ldb + p);
3117 let g = |o: usize| *base.add(o * ldb);
3118 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3119 } else {
3120 let base = b_mat.as_ptr().add(p * ldb + j);
3121 (
3122 [*base, *base.add(1), *base.add(2), *base.add(3)],
3123 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3124 )
3125 };
3126 let bv0 = vld1q_f32(b0.as_ptr());
3127 let bv1 = vld1q_f32(b1.as_ptr());
3128 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3129 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3130 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3131 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3132 c0a = vfmaq_f32(c0a, a0, bv0);
3133 c0b = vfmaq_f32(c0b, a0, bv1);
3134 c1a = vfmaq_f32(c1a, a1, bv0);
3135 c1b = vfmaq_f32(c1b, a1, bv1);
3136 c2a = vfmaq_f32(c2a, a2, bv0);
3137 c2b = vfmaq_f32(c2b, a2, bv1);
3138 c3a = vfmaq_f32(c3a, a3, bv0);
3139 c3b = vfmaq_f32(c3b, a3, bv1);
3140 }
3141 let al = vdupq_n_f32(alpha);
3142 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3143 .iter()
3144 .enumerate()
3145 {
3146 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3147 vst1q_f32(dst, vmulq_f32(*ca, al));
3148 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3149 }
3150 } else {
3151 for r in 0..mi {
3152 for q in 0..nj {
3153 let mut acc = 0f32;
3154 for p in 0..k {
3155 let bv = if b_rows_are_n {
3156 b_mat[(j + q) * ldb + p]
3157 } else {
3158 b_mat[p * ldb + j + q]
3159 };
3160 acc += a[(i + r) * lda + p] * bv;
3161 }
3162 c[(i + r) * ldc + j + q] = acc * alpha;
3163 }
3164 }
3165 }
3166 j += nj;
3167 }
3168 i += mi;
3169 }
3170 }
3171}
3172
3173#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3175#[allow(clippy::too_many_arguments)]
3176pub(crate) fn sgemm_rm(
3177 m: usize,
3178 n: usize,
3179 k: usize,
3180 alpha: f32,
3181 a: &[f32],
3182 lda: usize,
3183 b_mat: &[f32],
3184 ldb: usize,
3185 b_rows_are_n: bool,
3186 c: &mut [f32],
3187 ldc: usize,
3188) {
3189 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3190}
3191
3192#[allow(clippy::too_many_arguments)]
3196pub fn sgemm_public(
3197 m: usize,
3198 n: usize,
3199 k: usize,
3200 alpha: f32,
3201 a: &[f32],
3202 lda: usize,
3203 b_mat: &[f32],
3204 ldb: usize,
3205 b_rows_are_n: bool,
3206 c: &mut [f32],
3207 ldc: usize,
3208) {
3209 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3210 {
3211 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3212 }
3213 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3218 {
3219 for i in 0..m {
3220 for j in 0..n {
3221 let mut acc = 0f32;
3222 for p in 0..k {
3223 let bv = if b_rows_are_n {
3224 b_mat[j * ldb + p]
3225 } else {
3226 b_mat[p * ldb + j]
3227 };
3228 acc += a[i * lda + p] * bv;
3229 }
3230 c[i * ldc + j] = alpha * acc;
3231 }
3232 }
3233 }
3234}
3235
3236#[cfg(target_os = "macos")]
3239#[allow(clippy::too_many_arguments)]
3240pub(crate) fn sgemm_rm(
3241 m: usize,
3242 n: usize,
3243 k: usize,
3244 alpha: f32,
3245 a: &[f32],
3246 lda: usize,
3247 b_mat: &[f32],
3248 ldb: usize,
3249 b_rows_are_n: bool,
3250 c: &mut [f32],
3251 ldc: usize,
3252) {
3253 debug_assert!(a.len() >= (m - 1) * lda + k);
3254 debug_assert!(c.len() >= (m - 1) * ldc + n);
3255 #[cfg(target_arch = "aarch64")]
3260 if std::env::var("CMF_FORCE_NEON_GEMM")
3261 .map(|v| v == "1")
3262 .unwrap_or(false)
3263 {
3264 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3265 }
3266 unsafe {
3267 accel_blas::cblas_sgemm(
3268 101, 111, if b_rows_are_n { 112 } else { 111 },
3271 m as i32,
3272 n as i32,
3273 k as i32,
3274 alpha,
3275 a.as_ptr(),
3276 lda as i32,
3277 b_mat.as_ptr(),
3278 ldb as i32,
3279 0.0,
3280 c.as_mut_ptr(),
3281 ldc as i32,
3282 );
3283 }
3284}
3285
3286#[cfg(target_os = "macos")]
3293fn qmatmat_accel(
3294 q: &[u8],
3295 row_scale: &[f32],
3296 pre: &[std::borrow::Cow<'_, [f32]>],
3297 rows: usize,
3298 cols: usize,
3299 out: &mut [f32],
3300 pool: Option<&Pool>,
3301) {
3302 const TR: usize = 2048;
3307 let b = pre.len();
3308 thread_local! {
3309 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3310 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3311 }
3312 XPANEL.with(|xp| {
3313 WTILE.with(|wt| {
3314 let mut xpanel = xp.borrow_mut();
3315 xpanel.clear();
3316 for x in pre {
3317 xpanel.extend_from_slice(x);
3318 }
3319 let mut wtile = wt.borrow_mut();
3320 wtile.resize(TR * cols, 0.0);
3321 let mut r0 = 0usize;
3322 while r0 < rows {
3323 let tr = TR.min(rows - r0);
3324 let wt_addr = SendMut(wtile.as_mut_ptr());
3326 let run = |start: usize, end: usize| {
3327 for r in start..end {
3328 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3329 let s = row_scale[r0 + r];
3330 let dst =
3332 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3333 for (d, &v) in dst.iter_mut().zip(row) {
3334 *d = (v as i8) as f32 * s;
3335 }
3336 }
3337 };
3338 dispatch_rows(pool, tr, &run);
3339 unsafe {
3341 accel_blas::cblas_sgemm(
3342 101, 111, 112, b as i32,
3346 tr as i32,
3347 cols as i32,
3348 1.0,
3349 xpanel.as_ptr(),
3350 cols as i32,
3351 wtile.as_ptr(),
3352 cols as i32,
3353 0.0,
3354 out.as_mut_ptr().add(r0),
3355 rows as i32,
3356 );
3357 }
3358 r0 += tr;
3359 }
3360 })
3361 });
3362}
3363
3364fn qmatmat(
3365 q: &[u8],
3366 row_scale: &[f32],
3367 pre: &[std::borrow::Cow<'_, [f32]>],
3368 rows: usize,
3369 cols: usize,
3370 out: &mut [f32],
3371 pool: Option<&Pool>,
3372) {
3373 let b = pre.len();
3374 debug_assert_eq!(out.len(), b * rows);
3375 #[cfg(target_os = "macos")]
3380 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3381 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3382 return;
3383 }
3384 #[cfg(target_arch = "aarch64")]
3385 if sdot_enabled() {
3386 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3387 let out_addr = SendMut(out.as_mut_ptr());
3388 let blocked_ok = blocked_enabled();
3391 let use_i8mm = i8mm_enabled();
3392 if blocked_ok {
3393 let run = |start: usize, end: usize| {
3394 let mut o = start;
3395 while o < end {
3396 if o + 2 <= end {
3397 let r0 = &q[o * cols..(o + 1) * cols];
3398 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3399 let mut bi = 0usize;
3400 while bi + 4 <= acts.len() {
3401 let xs = [
3402 acts[bi].xq.as_slice(),
3403 acts[bi + 1].xq.as_slice(),
3404 acts[bi + 2].xq.as_slice(),
3405 acts[bi + 3].xq.as_slice(),
3406 ];
3407 let d = if use_i8mm {
3408 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3409 } else {
3410 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3411 };
3412 for (r, row) in [r0, r1].into_iter().enumerate() {
3413 for k in 0..4 {
3414 let act = &acts[bi + k];
3415 let mut v = d[r][k] as f32 * act.sx;
3416 for &(j, xv) in &act.outliers {
3417 v += (row[j] as i8) as f32 * xv;
3418 }
3419 unsafe {
3420 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3421 };
3422 }
3423 }
3424 bi += 4;
3425 }
3426 while bi < acts.len() {
3427 for (r, row) in [r0, r1].into_iter().enumerate() {
3428 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3429 unsafe { *out_addr.at(bi * rows + o + r) = v };
3430 }
3431 bi += 1;
3432 }
3433 o += 2;
3434 } else {
3435 let row = &q[o * cols..(o + 1) * cols];
3436 for (bi, act) in acts.iter().enumerate() {
3437 let v = row_dot_sdot(row, act) * row_scale[o];
3438 unsafe { *out_addr.at(bi * rows + o) = v };
3439 }
3440 o += 1;
3441 }
3442 }
3443 };
3444 dispatch_rows(pool, rows, &run);
3445 return;
3446 }
3447 let run = |start: usize, end: usize| {
3448 for o in start..end {
3449 let row = &q[o * cols..(o + 1) * cols];
3450 for (bi, act) in acts.iter().enumerate() {
3451 let v = row_dot_sdot(row, act) * row_scale[o];
3452 unsafe { *out_addr.at(bi * rows + o) = v };
3453 }
3454 }
3455 };
3456 dispatch_rows(pool, rows, &run);
3457 return;
3458 }
3459 #[cfg(target_arch = "x86_64")]
3464 if avx2_a8w8_enabled() {
3465 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3466 let out_addr = SendMut(out.as_mut_ptr());
3467 let blocked_ok = blocked_enabled();
3470 if !avx512vnni_enabled() && blocked_ok {
3471 let run = |start: usize, end: usize| {
3472 let mut o = start;
3473 while o < end {
3474 if o + 2 <= end {
3475 let r0 = &q[o * cols..(o + 1) * cols];
3476 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3477 let mut bi = 0usize;
3478 while bi + 4 <= acts.len() {
3479 let xs = [
3480 acts[bi].xq.as_slice(),
3481 acts[bi + 1].xq.as_slice(),
3482 acts[bi + 2].xq.as_slice(),
3483 acts[bi + 3].xq.as_slice(),
3484 ];
3485 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3486 for (r, row) in [r0, r1].into_iter().enumerate() {
3487 for k in 0..4 {
3488 let act = &acts[bi + k];
3489 let mut v = d[r][k] as f32 * act.sx;
3490 for &(j, xv) in &act.outliers {
3491 v += (row[j] as i8) as f32 * xv;
3492 }
3493 unsafe {
3494 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3495 };
3496 }
3497 }
3498 bi += 4;
3499 }
3500 while bi < acts.len() {
3501 for (r, row) in [r0, r1].into_iter().enumerate() {
3502 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3503 unsafe { *out_addr.at(bi * rows + o + r) = v };
3504 }
3505 bi += 1;
3506 }
3507 o += 2;
3508 } else {
3509 let row = &q[o * cols..(o + 1) * cols];
3510 for (bi, act) in acts.iter().enumerate() {
3511 let v = row_dot_avx2(row, act) * row_scale[o];
3512 unsafe { *out_addr.at(bi * rows + o) = v };
3513 }
3514 o += 1;
3515 }
3516 }
3517 };
3518 dispatch_rows(pool, rows, &run);
3519 return;
3520 }
3521 let run = |start: usize, end: usize| {
3522 for o in start..end {
3523 let row = &q[o * cols..(o + 1) * cols];
3524 for (bi, act) in acts.iter().enumerate() {
3525 let v = row_dot_avx2(row, act) * row_scale[o];
3526 unsafe { *out_addr.at(bi * rows + o) = v };
3527 }
3528 }
3529 };
3530 dispatch_rows(pool, rows, &run);
3531 return;
3532 }
3533 let out_addr = SendMut(out.as_mut_ptr());
3534 let run = |start: usize, end: usize| {
3535 for o in start..end {
3536 let row = &q[o * cols..(o + 1) * cols];
3537 for (bi, x) in pre.iter().enumerate() {
3538 let mut acc = 0f32;
3539 for j in 0..cols {
3540 acc += (row[j] as i8) as f32 * x[j];
3541 }
3542 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3543 }
3544 }
3545 };
3546 dispatch_rows(pool, rows, &run);
3547}
3548
3549fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3552 match pool {
3553 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3554 _ => run(0, rows),
3555 }
3556}
3557
3558fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3560 let groups = rows * cols / GROUP_SIZE;
3561 bytes.split_at(groups * 16)
3562}
3563
3564#[inline]
3569fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3570 #[cfg(target_arch = "aarch64")]
3571 unsafe {
3572 return vbit_fill4_neon(data, buf);
3573 }
3574 #[cfg(target_arch = "x86_64")]
3575 if avx2_enabled() {
3576 return unsafe { vbit_fill4_avx2(data, buf) };
3577 }
3578 #[allow(unreachable_code)]
3579 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3580 let u = unpack8::<4>(&data[blk * 4..]);
3581 for k in 0..8 {
3582 chunk[k] = (u[k] - 7) as i8 as u8;
3583 }
3584 }
3585}
3586
3587#[cfg(target_arch = "aarch64")]
3588#[target_feature(enable = "neon")]
3589unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3590 unsafe {
3593 use core::arch::aarch64::*;
3594 let n = buf.len();
3595 let mask = vdupq_n_u8(0x0F);
3596 let seven = vdupq_n_s8(7);
3597 let mut g = 0usize;
3598 while g * 32 + 32 <= n {
3599 let b = vld1q_u8(data.as_ptr().add(g * 16));
3600 let hi = vshrq_n_u8::<4>(b);
3601 let lo = vandq_u8(b, mask);
3602 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3603 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3604 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3605 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3606 g += 1;
3607 }
3608 }
3609}
3610
3611#[cfg(target_arch = "x86_64")]
3612#[target_feature(enable = "avx2")]
3613unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3614 unsafe {
3616 use core::arch::x86_64::*;
3617 let n = buf.len();
3618 let mask = _mm_set1_epi8(0x0F);
3619 let seven = _mm256_set1_epi8(7);
3620 let mut g = 0usize;
3621 while g * 32 + 32 <= n {
3622 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3623 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3624 let lo = _mm_and_si128(b, mask);
3625 let z = _mm256_sub_epi8(
3626 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3627 seven,
3628 );
3629 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3630 g += 1;
3631 }
3632 }
3633}
3634
3635#[inline(always)]
3640fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3641 let mut acc = 0u64;
3642 for i in 0..B {
3643 acc = (acc << 8) | data[i] as u64;
3644 }
3645 let mask = (1u64 << B) - 1;
3646 let mut out = [0i32; 8];
3647 for (k, o) in out.iter_mut().enumerate() {
3648 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3649 }
3650 out
3651}
3652
3653#[allow(clippy::too_many_arguments)]
3659fn vbitmatvec(
3660 bytes: &[u8],
3661 offsets: &[usize],
3662 x: &[f32],
3663 rows: usize,
3664 cols: usize,
3665 out: &mut [f32],
3666 pool: Option<&Pool>,
3667) {
3668 debug_assert_eq!(out.len(), rows);
3669 debug_assert_eq!(offsets.len(), rows + 1);
3670
3671 if a8w8_enabled() {
3675 let act = split_act(x);
3676 let out_addr = SendMut(out.as_mut_ptr());
3677 let run = move |start: usize, end: usize| {
3678 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3679 };
3680 dispatch_rows(pool, rows, &run);
3681 return;
3682 }
3683
3684 let out_addr = SendMut(out.as_mut_ptr());
3685 let run = move |start: usize, end: usize| {
3686 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3687 };
3688 dispatch_rows(pool, rows, &run);
3689}
3690
3691#[allow(clippy::too_many_arguments)]
3695fn vbit_range_a8w8(
3696 bytes: &[u8],
3697 offsets: &[usize],
3698 x: &[f32],
3699 act: &SplitAct,
3700 rows: usize,
3701 cols: usize,
3702 out: SendMut,
3703 start: usize,
3704 end: usize,
3705) {
3706 let ng = cols / GROUP_SIZE;
3707 let bits = &bytes[..rows];
3708 let sc_off = rows;
3709 let row_dot = |r: usize| -> f32 {
3710 let b = bits[r] as usize;
3711 let l = (1i32 << (b - 1)) - 1;
3712 let mask = (1u64 << b) - 1;
3713 let data = &bytes[offsets[r]..offsets[r + 1]];
3714 if b == 8 {
3715 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3717 let mut dot = 0f32;
3718 for g in 0..ng {
3719 let so = (r * ng + g) * 2;
3720 let sgf = f16_to_f32(u16::from_le_bytes([
3721 bytes[sc_off + so],
3722 bytes[sc_off + so + 1],
3723 ]));
3724 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3725 let mut gd = 0f32;
3726 for &xv in xg.iter() {
3727 if nbits < 8 {
3728 acc = (acc << 8) | data[idx] as u64;
3729 idx += 1;
3730 nbits += 8;
3731 }
3732 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3733 nbits -= 8;
3734 gd += (u - l) as f32 * xv;
3735 }
3736 dot += gd * sgf;
3737 }
3738 return dot;
3739 }
3740 thread_local! {
3744 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3745 const { std::cell::RefCell::new(Vec::new()) };
3746 }
3747 #[inline(always)]
3748 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3749 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3750 let u = unpack8::<B>(&data[blk * B..]);
3751 for k in 0..8 {
3752 chunk[k] = (u[k] - l) as i8 as u8;
3753 }
3754 }
3755 }
3756 let _ = mask;
3757 VBIT_SCRATCH.with(|scratch| {
3758 let mut buf = scratch.borrow_mut();
3759 buf.resize(cols, 0);
3760 match b {
3761 3 => fill::<3>(data, l, &mut buf),
3762 4 => vbit_fill4(data, &mut buf),
3763 5 => fill::<5>(data, l, &mut buf),
3764 6 => fill::<6>(data, l, &mut buf),
3765 _ => unreachable!(),
3766 }
3767 let mut dot = 0f32;
3768 for g in 0..ng {
3769 let so = (r * ng + g) * 2;
3770 let s = f16_to_f32(u16::from_le_bytes([
3771 bytes[sc_off + so],
3772 bytes[sc_off + so + 1],
3773 ]));
3774 let d = dot_i8_i8(
3775 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3776 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3777 ) as f32
3778 * act.sx;
3779 dot += d * s;
3780 }
3781 for &(j, xv) in &act.outliers {
3782 let so = (r * ng + j / GROUP_SIZE) * 2;
3783 let s = f16_to_f32(u16::from_le_bytes([
3784 bytes[sc_off + so],
3785 bytes[sc_off + so + 1],
3786 ]));
3787 dot += (buf[j] as i8) as f32 * s * xv;
3789 }
3790 dot
3791 })
3792 };
3793 for r in start..end {
3794 unsafe { *out.at(r) = row_dot(r) };
3796 }
3797}
3798
3799#[allow(clippy::too_many_arguments)]
3801fn vbit_range_f32(
3802 bytes: &[u8],
3803 offsets: &[usize],
3804 x: &[f32],
3805 rows: usize,
3806 cols: usize,
3807 out: SendMut,
3808 start: usize,
3809 end: usize,
3810) {
3811 let ng = cols / GROUP_SIZE;
3812 let bits = &bytes[..rows];
3813 let sc_off = rows;
3814 #[inline(always)]
3818 fn dot_row<const B: usize>(
3819 data: &[u8],
3820 bytes: &[u8],
3821 sc_off: usize,
3822 r: usize,
3823 ng: usize,
3824 x: &[f32],
3825 ) -> f32 {
3826 let l = ((1i32 << (B - 1)) - 1) as f32;
3827 let gbytes = GROUP_SIZE * B / 8;
3828 let mut dot = 0f32;
3829 for g in 0..ng {
3830 let so = (r * ng + g) * 2;
3831 let s = f16_to_f32(u16::from_le_bytes([
3832 bytes[sc_off + so],
3833 bytes[sc_off + so + 1],
3834 ]));
3835 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3836 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3837 let mut gd = 0f32;
3838 for blk in 0..GROUP_SIZE / 8 {
3839 let u = unpack8::<B>(&gd0[blk * B..]);
3840 let xb = &xg[blk * 8..blk * 8 + 8];
3841 for k in 0..8 {
3842 gd += (u[k] as f32 - l) * xb[k];
3843 }
3844 }
3845 dot += gd * s;
3846 }
3847 dot
3848 }
3849 for r in start..end {
3850 let data = &bytes[offsets[r]..offsets[r + 1]];
3851 let v = match bits[r] {
3852 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3853 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3854 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3855 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3856 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3857 b => unreachable!("vbit bit-width {b} (validated at load)"),
3858 };
3859 unsafe { *out.at(r) = v };
3861 }
3862}
3863
3864#[allow(clippy::too_many_arguments)]
3869fn vbitmatvec2(
3870 bytes: &[u8],
3871 offsets: &[usize],
3872 x1: &[f32],
3873 x2: &[f32],
3874 rows: usize,
3875 cols: usize,
3876 o1: &mut [f32],
3877 o2: &mut [f32],
3878 pool: Option<&Pool>,
3879) {
3880 debug_assert_eq!(o1.len(), rows);
3881 debug_assert_eq!(o2.len(), rows);
3882
3883 if a8w8_enabled() {
3884 let a1 = split_act(x1);
3885 let a2 = split_act(x2);
3886 let p1 = SendMut(o1.as_mut_ptr());
3887 let p2 = SendMut(o2.as_mut_ptr());
3888 let run = move |start: usize, end: usize| {
3889 vbit_range2_a8w8(
3890 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3891 )
3892 };
3893 dispatch_rows(pool, rows, &run);
3894 return;
3895 }
3896
3897 let p1 = SendMut(o1.as_mut_ptr());
3898 let p2 = SendMut(o2.as_mut_ptr());
3899 let run = move |start: usize, end: usize| {
3900 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3901 };
3902 dispatch_rows(pool, rows, &run);
3903}
3904
3905#[allow(clippy::too_many_arguments)]
3909fn vbit_range2_a8w8(
3910 bytes: &[u8],
3911 offsets: &[usize],
3912 x1: &[f32],
3913 x2: &[f32],
3914 a1: &SplitAct,
3915 a2: &SplitAct,
3916 rows: usize,
3917 cols: usize,
3918 p1: SendMut,
3919 p2: SendMut,
3920 start: usize,
3921 end: usize,
3922) {
3923 let ng = cols / GROUP_SIZE;
3924 let bits = &bytes[..rows];
3925 let sc_off = rows;
3926 let row_dots = |r: usize| -> (f32, f32) {
3927 let b = bits[r] as usize;
3928 let l = (1i32 << (b - 1)) - 1;
3929 let data = &bytes[offsets[r]..offsets[r + 1]];
3930 if b == 8 {
3931 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3934 let (mut d1, mut d2) = (0f32, 0f32);
3935 for g in 0..ng {
3936 let so = (r * ng + g) * 2;
3937 let sgf = f16_to_f32(u16::from_le_bytes([
3938 bytes[sc_off + so],
3939 bytes[sc_off + so + 1],
3940 ]));
3941 let (mut g1, mut g2) = (0f32, 0f32);
3942 for k in 0..GROUP_SIZE {
3943 if nbits < 8 {
3944 acc = (acc << 8) | data[idx] as u64;
3945 idx += 1;
3946 nbits += 8;
3947 }
3948 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3949 nbits -= 8;
3950 let w = (u - l) as f32;
3951 g1 += w * x1[g * GROUP_SIZE + k];
3952 g2 += w * x2[g * GROUP_SIZE + k];
3953 }
3954 d1 += g1 * sgf;
3955 d2 += g2 * sgf;
3956 }
3957 return (d1, d2);
3958 }
3959 thread_local! {
3960 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3961 const { std::cell::RefCell::new(Vec::new()) };
3962 }
3963 #[inline(always)]
3964 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3965 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3966 let u = unpack8::<B>(&data[blk * B..]);
3967 for k in 0..8 {
3968 chunk[k] = (u[k] - l) as i8 as u8;
3969 }
3970 }
3971 }
3972 VBIT_SCRATCH2.with(|scratch| {
3973 let mut buf = scratch.borrow_mut();
3974 buf.resize(cols, 0);
3975 match b {
3976 3 => fill::<3>(data, l, &mut buf),
3977 4 => vbit_fill4(data, &mut buf),
3978 5 => fill::<5>(data, l, &mut buf),
3979 6 => fill::<6>(data, l, &mut buf),
3980 _ => unreachable!(),
3981 }
3982 let (mut d1, mut d2) = (0f32, 0f32);
3983 for g in 0..ng {
3984 let so = (r * ng + g) * 2;
3985 let s = f16_to_f32(u16::from_le_bytes([
3986 bytes[sc_off + so],
3987 bytes[sc_off + so + 1],
3988 ]));
3989 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3990 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3991 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3992 d1 += v1 * s;
3993 d2 += v2 * s;
3994 }
3995 for &(j, xv) in &a1.outliers {
3996 let so = (r * ng + j / GROUP_SIZE) * 2;
3997 let s = f16_to_f32(u16::from_le_bytes([
3998 bytes[sc_off + so],
3999 bytes[sc_off + so + 1],
4000 ]));
4001 d1 += (buf[j] as i8) as f32 * s * xv;
4002 }
4003 for &(j, xv) in &a2.outliers {
4004 let so = (r * ng + j / GROUP_SIZE) * 2;
4005 let s = f16_to_f32(u16::from_le_bytes([
4006 bytes[sc_off + so],
4007 bytes[sc_off + so + 1],
4008 ]));
4009 d2 += (buf[j] as i8) as f32 * s * xv;
4010 }
4011 (d1, d2)
4012 })
4013 };
4014 for r in start..end {
4015 let (v1, v2) = row_dots(r);
4016 unsafe {
4018 *p1.at(r) = v1;
4019 *p2.at(r) = v2;
4020 }
4021 }
4022}
4023
4024#[allow(clippy::too_many_arguments)]
4028fn vbit_range2_f32(
4029 bytes: &[u8],
4030 offsets: &[usize],
4031 x1: &[f32],
4032 x2: &[f32],
4033 rows: usize,
4034 cols: usize,
4035 p1: SendMut,
4036 p2: SendMut,
4037 start: usize,
4038 end: usize,
4039) {
4040 let ng = cols / GROUP_SIZE;
4041 let bits = &bytes[..rows];
4042 let sc_off = rows;
4043 #[inline(always)]
4044 #[allow(clippy::too_many_arguments)]
4045 fn dot_row2<const B: usize>(
4046 data: &[u8],
4047 bytes: &[u8],
4048 sc_off: usize,
4049 r: usize,
4050 ng: usize,
4051 x1: &[f32],
4052 x2: &[f32],
4053 ) -> (f32, f32) {
4054 let l = ((1i32 << (B - 1)) - 1) as f32;
4055 let gbytes = GROUP_SIZE * B / 8;
4056 let (mut d1, mut d2) = (0f32, 0f32);
4057 for g in 0..ng {
4058 let so = (r * ng + g) * 2;
4059 let s = f16_to_f32(u16::from_le_bytes([
4060 bytes[sc_off + so],
4061 bytes[sc_off + so + 1],
4062 ]));
4063 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4064 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4065 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4066 let (mut g1, mut g2) = (0f32, 0f32);
4067 for blk in 0..GROUP_SIZE / 8 {
4068 let u = unpack8::<B>(&gd0[blk * B..]);
4069 for k in 0..8 {
4070 let w = u[k] as f32 - l;
4071 g1 += w * x1g[blk * 8 + k];
4072 g2 += w * x2g[blk * 8 + k];
4073 }
4074 }
4075 d1 += g1 * s;
4076 d2 += g2 * s;
4077 }
4078 (d1, d2)
4079 }
4080 for r in start..end {
4081 let data = &bytes[offsets[r]..offsets[r + 1]];
4082 let (v1, v2) = match bits[r] {
4083 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4084 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4085 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4086 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4087 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4088 b => unreachable!("vbit bit-width {b} (validated at load)"),
4089 };
4090 unsafe {
4092 *p1.at(r) = v1;
4093 *p2.at(r) = v2;
4094 }
4095 }
4096}
4097
4098#[inline]
4105#[allow(unreachable_code)]
4106fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4107 #[cfg(target_arch = "aarch64")]
4108 unsafe {
4109 return dot_q4t_row_sdot(bytes, r, gpr, xq);
4110 }
4111 #[cfg(target_arch = "x86_64")]
4112 unsafe {
4113 if vnni_tiles_enabled() {
4114 return dot_q4t_row_vnni(bytes, r, gpr, xq);
4115 }
4116 return dot_q4t_row_avx2(bytes, r, gpr, xq);
4117 }
4118 let mut acc = 0f32;
4119 for gi in 0..gpr {
4120 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4121 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4122 let mut d = 0i32;
4123 for (k, &b) in tile[2..].iter().enumerate() {
4124 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4125 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4126 }
4127 acc += d as f32 * s;
4128 }
4129 acc
4130}
4131
4132#[cfg(target_arch = "aarch64")]
4133#[target_feature(enable = "neon,dotprod")]
4134unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4135 unsafe {
4138 use core::arch::aarch64::*;
4139 use core::arch::asm;
4140 let lomask = vdupq_n_u8(0x0F);
4141 let eight = vdupq_n_s8(8);
4142 let mut acc = 0f32;
4143 for gi in 0..gpr {
4144 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4145 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4146 let b = vld1q_u8(t.add(2));
4147 let lo = vandq_u8(b, lomask);
4148 let hi = vshrq_n_u8::<4>(b);
4149 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4150 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4151 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4152 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4153 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4154 asm!(
4155 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4156 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4157 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4158 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4159 options(pure, nomem, nostack),
4160 );
4161 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4162 }
4163 acc
4164 }
4165}
4166
4167#[cfg(target_arch = "x86_64")]
4168#[target_feature(enable = "avx2")]
4169unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4170 unsafe {
4172 use core::arch::x86_64::*;
4173 let lomask = _mm_set1_epi8(0x0F);
4174 let eight = _mm256_set1_epi8(8);
4175 let ones = _mm256_set1_epi16(1);
4176 let mut acc = 0f32;
4177 for gi in 0..gpr {
4178 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4179 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4180 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4181 let lo = _mm_and_si128(b, lomask);
4182 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4183 let w = _mm256_sub_epi8(
4184 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4185 eight,
4186 );
4187 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4188 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4189 let d = _mm256_madd_epi16(p16, ones);
4190 let hi128 = _mm256_extracti128_si256::<1>(d);
4191 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4192 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4193 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4194 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4195 }
4196 acc
4197 }
4198}
4199
4200#[cfg(target_arch = "x86_64")]
4204#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4205unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4206 unsafe {
4208 use core::arch::x86_64::*;
4209 let lomask = _mm_set1_epi8(0x0F);
4210 let eight = _mm256_set1_epi8(8);
4211 let mut acc = 0f32;
4212 for gi in 0..gpr {
4213 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4214 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4215 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4216 let lo = _mm_and_si128(b, lomask);
4217 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4218 let w = _mm256_sub_epi8(
4219 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4220 eight,
4221 );
4222 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4223 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4224 acc += d as f32 * s;
4225 }
4226 acc
4227 }
4228}
4229
4230#[cfg(target_arch = "x86_64")]
4235#[target_feature(enable = "avx2,fma")]
4240unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4241 unsafe {
4243 use core::arch::x86_64::*;
4244 let lomask = _mm_set1_epi8(0x0F);
4245 let eight = _mm256_set1_epi8(8);
4246 let ones = _mm256_set1_epi16(1);
4247 let mut f0 = _mm256_setzero_ps();
4260 let mut f1 = _mm256_setzero_ps();
4261 let mut f2 = _mm256_setzero_ps();
4262 let mut f3 = _mm256_setzero_ps();
4263 for gi in 0..gpr {
4264 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4265 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4266 let sv = _mm256_set1_ps(s);
4267 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4268 let lo = _mm_and_si128(bb, lomask);
4269 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4270 let w = _mm256_sub_epi8(
4271 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4272 eight,
4273 );
4274 let aw = _mm256_abs_epi8(w);
4275 let off = gi * GROUP_SIZE;
4276 let dot = |xq: &[i8]| {
4277 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4278 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4279 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4280 };
4281 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4282 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4283 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4284 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4285 }
4286 [
4287 hsum256_ps(f0),
4288 hsum256_ps(f1),
4289 hsum256_ps(f2),
4290 hsum256_ps(f3),
4291 ]
4292 }
4293}
4294
4295#[cfg(target_arch = "x86_64")]
4298#[target_feature(enable = "avx2")]
4299#[inline]
4300unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4301 unsafe {
4303 use core::arch::x86_64::*;
4304 let hi = _mm256_extractf128_ps::<1>(v);
4305 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4306 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4307 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4308 _mm_cvtss_f32(s)
4309 }
4310}
4311
4312#[cfg(target_arch = "x86_64")]
4314#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4315unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4316 unsafe {
4318 use core::arch::x86_64::*;
4319 let lomask = _mm_set1_epi8(0x0F);
4320 let eight = _mm256_set1_epi8(8);
4321 let mut f0 = _mm256_setzero_ps();
4324 let mut f1 = _mm256_setzero_ps();
4325 let mut f2 = _mm256_setzero_ps();
4326 let mut f3 = _mm256_setzero_ps();
4327 for gi in 0..gpr {
4328 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4329 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4330 let sv = _mm256_set1_ps(s);
4331 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4332 let lo = _mm_and_si128(bb, lomask);
4333 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4334 let w = _mm256_sub_epi8(
4335 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4336 eight,
4337 );
4338 let aw = _mm256_abs_epi8(w);
4339 let off = gi * GROUP_SIZE;
4340 let dot = |xq: &[i8]| {
4341 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4342 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4343 _mm256_setzero_si256(),
4344 aw,
4345 _mm256_sign_epi8(x, w),
4346 ))
4347 };
4348 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4349 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4350 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4351 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4352 }
4353 let acc = [
4354 hsum256_ps(f0),
4355 hsum256_ps(f1),
4356 hsum256_ps(f2),
4357 hsum256_ps(f3),
4358 ];
4359 acc
4360 }
4361}
4362
4363#[cfg(target_arch = "aarch64")]
4368#[target_feature(enable = "neon,dotprod")]
4369unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4370 unsafe {
4372 use core::arch::aarch64::*;
4373 use core::arch::asm;
4374 let lomask = vdupq_n_u8(0x0F);
4375 let eight = vdupq_n_s8(8);
4376 let mut acc = [0f32; 4];
4377 for gi in 0..gpr {
4378 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4379 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4380 let b = vld1q_u8(t.add(2));
4381 let lo = vandq_u8(b, lomask);
4382 let hi = vshrq_n_u8::<4>(b);
4383 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4384 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4385 for (k, xq) in xs.iter().enumerate() {
4386 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4387 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4388 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4389 asm!(
4390 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4391 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4392 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4393 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4394 options(pure, nomem, nostack),
4395 );
4396 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4397 }
4398 }
4399 acc
4400 }
4401}
4402
4403#[inline]
4405fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4406 let gi = j / GROUP_SIZE;
4407 let k = j % GROUP_SIZE;
4408 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4409 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4410 let byte = tile[2 + k / 2];
4411 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4412 ((nib as i32 - 8) as f32, s)
4413}
4414
4415#[inline]
4418fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4419 let mut acc = 0f32;
4420 for gi in 0..gpr {
4421 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4422 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4423 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4424 let mut ga = 0f32;
4425 for (k, &b) in tile[2..].iter().enumerate() {
4426 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4427 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4428 }
4429 acc += ga * s;
4430 }
4431 acc
4432}
4433
4434struct Q4tpView<'a> {
4438 nib: &'a [u8],
4439 params: &'a [u8],
4440 codes: &'a [u8],
4441 stride: usize,
4442 zero_rung: bool,
4444}
4445
4446impl<'a> Q4tpView<'a> {
4447 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4448 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4449 Self {
4450 nib: &bytes[..params_off],
4451 params: &bytes[params_off..codes_off],
4452 codes: &bytes[codes_off..],
4453 stride,
4454 zero_rung: false,
4455 }
4456 }
4457
4458 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4460 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4461 Self {
4462 nib: &bytes[..params_off],
4463 params: &bytes[params_off..codes_off],
4464 codes: &bytes[codes_off..],
4465 stride,
4466 zero_rung: true,
4467 }
4468 }
4469
4470 #[inline]
4486 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4487 let tab = if self.zero_rung {
4488 q2tp_ladder(self.params, r)
4489 } else {
4490 q4tp_ladder(self.params, r)
4491 };
4492 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4493 let out = &mut out[..gpr];
4494 let mut chunks = out.chunks_exact_mut(8);
4495 let mut ci = 0usize;
4496 for c in &mut chunks {
4497 let w = u64::from(codes[ci])
4498 | u64::from(codes[ci + 1]) << 8
4499 | u64::from(codes[ci + 2]) << 16
4500 | u64::from(codes[ci + 3]) << 24
4501 | u64::from(codes[ci + 4]) << 32;
4502 for (k, o) in c.iter_mut().enumerate() {
4503 *o = tab[((w >> (5 * k)) & 31) as usize];
4504 }
4505 ci += 5;
4506 }
4507 let tail = &codes[ci..];
4510 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4511 *o = tab[q4tp_code(tail, k)];
4512 }
4513 }
4514}
4515
4516#[inline]
4517fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4518 #[cfg(target_arch = "aarch64")]
4519 unsafe {
4520 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4521 }
4522 #[cfg(target_arch = "x86_64")]
4523 unsafe {
4524 if vnni_tiles_enabled() {
4525 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4526 }
4527 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4528 }
4529 #[allow(unreachable_code)]
4530 {
4531 let mut acc = 0f32;
4532 for gi in 0..gpr {
4533 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4534 let s = scales[gi];
4535 let mut d = 0i32;
4536 for (k, &b) in tile.iter().enumerate() {
4537 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4538 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4539 }
4540 acc += d as f32 * s;
4541 }
4542 acc
4543 }
4544}
4545
4546#[cfg(target_arch = "aarch64")]
4549#[target_feature(enable = "neon,dotprod")]
4550unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4551 unsafe {
4554 use core::arch::aarch64::*;
4555 use core::arch::asm;
4556 let lomask = vdupq_n_u8(0x0F);
4557 let eight = vdupq_n_s8(8);
4558 let mut acc = 0f32;
4559 for gi in 0..gpr {
4560 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4561 let s = *scales.get_unchecked(gi);
4562 let b = vld1q_u8(t);
4563 let lo = vandq_u8(b, lomask);
4564 let hi = vshrq_n_u8::<4>(b);
4565 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4566 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4567 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4568 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4569 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4570 asm!(
4571 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4572 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4573 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4574 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4575 options(pure, nomem, nostack),
4576 );
4577 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4578 }
4579 acc
4580 }
4581}
4582
4583#[cfg(target_arch = "x86_64")]
4584#[target_feature(enable = "avx2")]
4585unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4586 unsafe {
4588 use core::arch::x86_64::*;
4589 let lomask = _mm_set1_epi8(0x0F);
4590 let eight = _mm256_set1_epi8(8);
4591 let ones = _mm256_set1_epi16(1);
4592 let mut acc = 0f32;
4593 for gi in 0..gpr {
4594 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4595 let s = *scales.get_unchecked(gi);
4596 let b = _mm_loadu_si128(t as *const __m128i);
4597 let lo = _mm_and_si128(b, lomask);
4598 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4599 let w = _mm256_sub_epi8(
4600 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4601 eight,
4602 );
4603 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4604 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4605 let d = _mm256_madd_epi16(p16, ones);
4606 let hi128 = _mm256_extracti128_si256::<1>(d);
4607 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4608 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4609 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4610 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4611 }
4612 acc
4613 }
4614}
4615
4616#[cfg(target_arch = "x86_64")]
4619#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4620unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4621 unsafe {
4623 use core::arch::x86_64::*;
4624 let lomask = _mm_set1_epi8(0x0F);
4625 let eight = _mm256_set1_epi8(8);
4626 let mut acc = 0f32;
4627 for gi in 0..gpr {
4628 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4629 let s = *scales.get_unchecked(gi);
4630 let b = _mm_loadu_si128(t as *const __m128i);
4631 let lo = _mm_and_si128(b, lomask);
4632 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4633 let w = _mm256_sub_epi8(
4634 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4635 eight,
4636 );
4637 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4638 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4639 }
4640 acc
4641 }
4642}
4643
4644#[inline]
4647fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4648 let mut acc = 0f32;
4649 for gi in 0..gpr {
4650 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4651 let s = scales[gi];
4652 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4653 let mut ga = 0f32;
4654 for (k, &b) in tile.iter().enumerate() {
4655 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4656 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4657 }
4658 acc += ga * s;
4659 }
4660 acc
4661}
4662
4663#[inline]
4666fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4667 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4668 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4669 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4670 ((n as i32 - 8) as f32, scales[gi])
4671}
4672
4673fn q4tp_matvec(
4675 bytes: &[u8],
4676 x: &[f32],
4677 rows: usize,
4678 cols: usize,
4679 out: &mut [f32],
4680 pool: Option<&Pool>,
4681) {
4682 debug_assert_eq!(out.len(), rows);
4683 let gpr = cols / GROUP_SIZE;
4684 let v = Q4tpView::new(bytes, rows, cols);
4685 let out_addr = SendMut(out.as_mut_ptr());
4686 if a8w8_enabled() {
4687 let act = split_act(x);
4688 let run = |start: usize, end: usize| {
4689 with_krow(gpr, |sc| {
4691 for r in start..end {
4692 v.scales_into(r, gpr, sc);
4693 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4694 for &(j, xv) in &act.outliers {
4695 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4696 acc += w * s * xv;
4697 }
4698 unsafe { *out_addr.at(r) = acc };
4700 }
4701 })
4702 };
4703 dispatch_rows(pool, rows, &run);
4704 return;
4705 }
4706 let run = |start: usize, end: usize| {
4707 with_krow(gpr, |sc| {
4708 for r in start..end {
4709 v.scales_into(r, gpr, sc);
4710 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4712 }
4713 })
4714 };
4715 dispatch_rows(pool, rows, &run);
4716}
4717
4718#[allow(clippy::too_many_arguments)]
4721fn q4tp_matvec2(
4722 bytes: &[u8],
4723 x1: &[f32],
4724 x2: &[f32],
4725 rows: usize,
4726 cols: usize,
4727 o1: &mut [f32],
4728 o2: &mut [f32],
4729 pool: Option<&Pool>,
4730) {
4731 let gpr = cols / GROUP_SIZE;
4732 let v = Q4tpView::new(bytes, rows, cols);
4733 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4734 let run = |start: usize, end: usize| {
4735 let mut sc = vec![0f32; gpr];
4736 for r in start..end {
4737 v.scales_into(r, gpr, &mut sc);
4738 unsafe {
4740 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4741 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4742 }
4743 }
4744 };
4745 dispatch_rows(pool, rows, &run);
4746}
4747
4748#[inline]
4751fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4752 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4753 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4754 let c = (byte >> (2 * (k % 4))) & 3;
4755 (c as f32 - 1.5, scales[gi])
4756}
4757
4758#[inline]
4765fn dot_q2tp_row_i8(
4766 chunks: &[u8],
4767 r: usize,
4768 gpr: usize,
4769 xq: &[i8],
4770 gsum: &[i32],
4771 scales: &[f32],
4772) -> f32 {
4773 let mut acc = 0f32;
4774 let base = r * gpr * Q2TP_CHUNK;
4775 #[cfg(not(target_arch = "aarch64"))]
4776 let mut codes = [0i8; GROUP_SIZE];
4777 for gi in 0..gpr {
4778 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4779 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4780 #[cfg(target_arch = "aarch64")]
4781 let dot = unsafe {
4787 use core::arch::aarch64::*;
4788 let b = vld1_u8(ch.as_ptr());
4789 let three = vdup_n_u8(3);
4790 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4791 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4792 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4793 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4794 let x4 = vld4_s8(xg.as_ptr());
4795 let mut acc4 = vdupq_n_s32(0);
4796 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4797 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4798 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4799 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4800 vaddvq_s32(acc4)
4801 };
4802 #[cfg(not(target_arch = "aarch64"))]
4803 let dot: i32 = {
4804 for (k, &b) in ch.iter().enumerate() {
4805 codes[k * 4] = (b & 3) as i8;
4806 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4807 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4808 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4809 }
4810 codes
4811 .iter()
4812 .zip(xg)
4813 .map(|(&c, &x)| c as i32 * x as i32)
4814 .sum()
4815 };
4816 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4817 }
4818 acc
4819}
4820
4821fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4825 let mut acc = 0f32;
4826 for gi in 0..gpr {
4827 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4828 let s = scales[gi];
4829 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4830 let mut g = 0f32;
4831 for (k, &b) in ch.iter().enumerate() {
4832 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4833 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4834 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4835 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4836 }
4837 acc += s * g;
4838 }
4839 acc
4840}
4841
4842fn q2tp_matvec(
4843 bytes: &[u8],
4844 x: &[f32],
4845 rows: usize,
4846 cols: usize,
4847 out: &mut [f32],
4848 pool: Option<&Pool>,
4849) {
4850 debug_assert_eq!(out.len(), rows);
4851 let gpr = cols / GROUP_SIZE;
4852 let v = Q4tpView::new_q2(bytes, rows, cols);
4853 let out_addr = SendMut(out.as_mut_ptr());
4854 if a8w8_enabled() {
4859 let act = split_act(x);
4860 let gsum = q1_group_sums(&act.xq, gpr);
4861 let (act, gsum) = (&act, &gsum);
4862 let run = move |start: usize, end: usize| {
4863 with_krow(gpr, |sc| {
4864 for r in start..end {
4865 v.scales_into(r, gpr, sc);
4866 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4867 for &(j, xv) in &act.outliers {
4868 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4869 acc += w * s * xv;
4870 }
4871 unsafe { *out_addr.at(r) = acc };
4873 }
4874 })
4875 };
4876 dispatch_rows(pool, rows, &run);
4877 return;
4878 }
4879 let run = |start: usize, end: usize| {
4880 with_krow(gpr, |sc| {
4881 for r in start..end {
4882 v.scales_into(r, gpr, sc);
4883 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4885 }
4886 })
4887 };
4888 dispatch_rows(pool, rows, &run);
4889}
4890
4891#[allow(clippy::too_many_arguments)]
4893fn q2tp_matvec2(
4894 bytes: &[u8],
4895 x1: &[f32],
4896 x2: &[f32],
4897 rows: usize,
4898 cols: usize,
4899 o1: &mut [f32],
4900 o2: &mut [f32],
4901 pool: Option<&Pool>,
4902) {
4903 let gpr = cols / GROUP_SIZE;
4904 let v = Q4tpView::new_q2(bytes, rows, cols);
4905 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4906 let run = |start: usize, end: usize| {
4907 let mut sc = vec![0f32; gpr];
4908 for r in start..end {
4909 v.scales_into(r, gpr, &mut sc);
4910 unsafe {
4912 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4913 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4914 }
4915 }
4916 };
4917 dispatch_rows(pool, rows, &run);
4918}
4919
4920pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4927 let gpr = cols / GROUP_SIZE;
4932 let v = Q4tpView::new_q2(bytes, rows, cols);
4933 with_krow(gpr, |sc| {
4934 for r in 0..rows {
4935 v.scales_into(r, gpr, sc);
4936 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4937 }
4938 });
4939}
4940
4941pub fn q2tp_matmat_for_test(
4942 bytes: &[u8],
4943 xs_all: &[f32],
4944 b: usize,
4945 rows: usize,
4946 cols: usize,
4947 out: &mut [f32],
4948) {
4949 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4950}
4951
4952fn q2tp_matmat(
4953 bytes: &[u8],
4954 xs_all: &[f32],
4955 b: usize,
4956 rows: usize,
4957 cols: usize,
4958 out: &mut [f32],
4959 pool: Option<&Pool>,
4960) {
4961 debug_assert_eq!(out.len(), b * rows);
4962 let gpr = cols / GROUP_SIZE;
4963 let v = Q4tpView::new_q2(bytes, rows, cols);
4964 let out_addr = SendMut(out.as_mut_ptr());
4965 let run = |start: usize, end: usize| {
4966 let mut sc = vec![0f32; gpr];
4967 for r in start..end {
4968 v.scales_into(r, gpr, &mut sc);
4969 for bi in 0..b {
4970 let x = &xs_all[bi * cols..(bi + 1) * cols];
4971 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4973 }
4974 }
4975 };
4976 dispatch_rows(pool, rows, &run);
4977}
4978
4979#[cfg(target_arch = "aarch64")]
4983#[target_feature(enable = "neon,dotprod")]
4984unsafe fn dot_q4tp_row_1x4_sdot_v1(
4985 nib: &[u8],
4986 r: usize,
4987 gpr: usize,
4988 xs: [&[i8]; 4],
4989 scales: &[f32],
4990) -> [f32; 4] {
4991 unsafe {
4992 use core::arch::aarch64::*;
4993 use core::arch::asm;
4994 let lomask = vdupq_n_u8(0x0F);
4995 let eight = vdupq_n_s8(8);
4996 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4997 for gi in 0..gpr {
4998 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4999 let s = *scales.get_unchecked(gi);
5000 let bb = vld1q_u8(t);
5001 let lo = vandq_u8(bb, lomask);
5002 let hi = vshrq_n_u8::<4>(bb);
5003 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5004 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5005 let mut d = [0f32; 4];
5006 for (k, dk) in d.iter_mut().enumerate() {
5007 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5008 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5009 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5010 asm!(
5011 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5012 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5013 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5014 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5015 options(pure, nomem, nostack),
5016 );
5017 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5018 }
5019 f0 += d[0];
5020 f1 += d[1];
5021 f2 += d[2];
5022 f3 += d[3];
5023 }
5024 [f0, f1, f2, f3]
5025 }
5026}
5027
5028#[allow(dead_code)]
5035static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
5036
5037#[cfg(target_arch = "x86_64")]
5043fn q4tp_blocked_x86() -> bool {
5044 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5045 1 => false,
5046 2 => avx512vnni_enabled(),
5051 _ => blocked_enabled() && avx512vnni_enabled(),
5055 }
5056}
5057
5058#[cfg(target_arch = "aarch64")]
5060#[allow(dead_code)]
5061fn q4tp_v1() -> bool {
5062 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5063 1 => true,
5064 2 => false,
5065 _ => {
5066 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5067 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
5068 }
5069 }
5070}
5071
5072#[cfg(target_arch = "x86_64")]
5084#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5085unsafe fn dot_q4tp_2x8_avx512(
5086 nib: &[u8],
5087 r0: usize,
5088 gpr: usize,
5089 xs: [&[i8]; 8],
5090 sc0: &[f32],
5091 sc1: &[f32],
5092) -> [[f32; 8]; 2] {
5093 unsafe {
5096 use core::arch::x86_64::*;
5097 let lomask = _mm256_set1_epi8(0x0F);
5098 let eight = _mm256_set1_epi8(8);
5099 let zero = _mm512_setzero_si512();
5100 let mut v0 = [_mm512_setzero_ps(); 8];
5101 let mut v1 = [_mm512_setzero_ps(); 8];
5102 let pairs = gpr / 2;
5103 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
5104 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5105 let bb = _mm256_loadu_si256(t as *const __m256i);
5106 let lo = _mm256_and_si256(bb, lomask);
5107 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5108 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5109 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5110 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5111 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5112 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5113 };
5114 for gp in 0..pairs {
5115 let gi = gp * 2;
5116 let (wa0, neg0) = unpack(r0, gi);
5117 let (wa1, neg1) = unpack(r0 + 1, gi);
5118 let off = gi * GROUP_SIZE;
5119 let sv = |sc: &[f32]| {
5120 _mm512_insertf32x8::<1>(
5121 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5122 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5123 )
5124 };
5125 let s0 = sv(sc0);
5126 let s1 = sv(sc1);
5127 for k in 0..8 {
5128 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5129 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5130 zero,
5131 wa0,
5132 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5133 ));
5134 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5135 zero,
5136 wa1,
5137 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5138 ));
5139 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5140 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5141 }
5142 }
5143 let mut acc = [[0f32; 8]; 2];
5144 for k in 0..8 {
5145 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5146 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5147 }
5148 if gpr % 2 == 1 {
5149 let off = (gpr - 1) * GROUP_SIZE;
5150 for j in off..off + GROUP_SIZE {
5151 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5152 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5153 for k in 0..8 {
5154 let x = *xs[k].get_unchecked(j) as f32;
5155 acc[0][k] += w0 * sa * x;
5156 acc[1][k] += w1 * sb * x;
5157 }
5158 }
5159 }
5160 acc
5161 }
5162}
5163
5164#[cfg(target_arch = "x86_64")]
5169#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5170unsafe fn dot_q4tp_row_1x8_avx512(
5171 nib: &[u8],
5172 r: usize,
5173 gpr: usize,
5174 xs: [&[i8]; 8],
5175 scales: &[f32],
5176) -> [f32; 8] {
5177 unsafe {
5179 use core::arch::x86_64::*;
5180 let lomask = _mm256_set1_epi8(0x0F);
5181 let eight = _mm256_set1_epi8(8);
5182 let zero = _mm512_setzero_si512();
5183 let (mut v0, mut v1, mut v2, mut v3) = (
5184 _mm512_setzero_ps(),
5185 _mm512_setzero_ps(),
5186 _mm512_setzero_ps(),
5187 _mm512_setzero_ps(),
5188 );
5189 let (mut v4, mut v5, mut v6, mut v7) = (
5190 _mm512_setzero_ps(),
5191 _mm512_setzero_ps(),
5192 _mm512_setzero_ps(),
5193 _mm512_setzero_ps(),
5194 );
5195 let pairs = gpr / 2;
5196 for gp in 0..pairs {
5197 let gi = gp * 2;
5198 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5199 let bb = _mm256_loadu_si256(t as *const __m256i);
5200 let lo = _mm256_and_si256(bb, lomask);
5201 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5202 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5207 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5208 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5209 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5210 let wabs = _mm512_abs_epi8(w);
5211 let neg = _mm512_movepi8_mask(w);
5212 let off = gi * GROUP_SIZE;
5213 let sv = _mm512_insertf32x8::<1>(
5214 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5215 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5216 );
5217 let dot = |x: &[i8]| -> __m512 {
5218 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5219 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5220 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5221 };
5222 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5223 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5224 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5225 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5226 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5227 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5228 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5229 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5230 }
5231 let mut acc = [
5232 _mm512_reduce_add_ps(v0),
5233 _mm512_reduce_add_ps(v1),
5234 _mm512_reduce_add_ps(v2),
5235 _mm512_reduce_add_ps(v3),
5236 _mm512_reduce_add_ps(v4),
5237 _mm512_reduce_add_ps(v5),
5238 _mm512_reduce_add_ps(v6),
5239 _mm512_reduce_add_ps(v7),
5240 ];
5241 if gpr % 2 == 1 {
5244 let off = (gpr - 1) * GROUP_SIZE;
5245 for j in off..off + GROUP_SIZE {
5246 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5247 let ws = w * s;
5248 for k in 0..8 {
5249 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5250 }
5251 }
5252 }
5253 acc
5254 }
5255}
5256
5257#[cfg(target_arch = "x86_64")]
5270#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5271unsafe fn dot_q4tp_row_1x4_avx512(
5272 nib: &[u8],
5273 r: usize,
5274 gpr: usize,
5275 xs: [&[i8]; 4],
5276 scales: &[f32],
5277) -> [f32; 4] {
5278 unsafe {
5280 use core::arch::x86_64::*;
5281 let lomask = _mm256_set1_epi8(0x0F);
5282 let eight = _mm256_set1_epi8(8);
5283 let zero = _mm512_setzero_si512();
5284 let (mut v0, mut v1, mut v2, mut v3) = (
5285 _mm512_setzero_ps(),
5286 _mm512_setzero_ps(),
5287 _mm512_setzero_ps(),
5288 _mm512_setzero_ps(),
5289 );
5290 let pairs = gpr / 2;
5291 for gp in 0..pairs {
5292 let gi = gp * 2;
5293 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5294 let bb = _mm256_loadu_si256(t as *const __m256i);
5295 let lo = _mm256_and_si256(bb, lomask);
5296 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5297 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5302 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5303 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5304 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5305 let wabs = _mm512_abs_epi8(w);
5306 let neg = _mm512_movepi8_mask(w);
5307 let off = gi * GROUP_SIZE;
5308 let sv = _mm512_insertf32x8::<1>(
5309 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5310 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5311 );
5312 let dot = |x: &[i8]| -> __m512 {
5313 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5314 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5315 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5316 };
5317 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5318 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5319 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5320 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5321 }
5322 let mut acc = [
5323 _mm512_reduce_add_ps(v0),
5324 _mm512_reduce_add_ps(v1),
5325 _mm512_reduce_add_ps(v2),
5326 _mm512_reduce_add_ps(v3),
5327 ];
5328 if gpr % 2 == 1 {
5331 let off = (gpr - 1) * GROUP_SIZE;
5332 for j in off..off + GROUP_SIZE {
5333 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5334 let ws = w * s;
5335 for k in 0..4 {
5336 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5337 }
5338 }
5339 }
5340 acc
5341 }
5342}
5343
5344#[cfg(target_arch = "aarch64")]
5348#[target_feature(enable = "neon,dotprod")]
5349unsafe fn dot_q4tp_row_1x4_sdot(
5350 nib: &[u8],
5351 r: usize,
5352 gpr: usize,
5353 xs: [&[i8]; 4],
5354 scales: &[f32],
5355) -> [f32; 4] {
5356 unsafe {
5358 use core::arch::aarch64::*;
5359 use core::arch::asm;
5360 let lomask = vdupq_n_u8(0x0F);
5361 let eight = vdupq_n_s8(8);
5362 let (mut v0, mut v1, mut v2, mut v3) = (
5378 vdupq_n_f32(0.0),
5379 vdupq_n_f32(0.0),
5380 vdupq_n_f32(0.0),
5381 vdupq_n_f32(0.0),
5382 );
5383 for gi in 0..gpr {
5384 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5385 let s = *scales.get_unchecked(gi);
5386 let bb = vld1q_u8(t);
5387 let lo = vandq_u8(bb, lomask);
5388 let hi = vshrq_n_u8::<4>(bb);
5389 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5390 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5391 let off = gi * GROUP_SIZE;
5392 let dot4 = |x: &[i8]| -> int32x4_t {
5393 let x0 = vld1q_s8(x.as_ptr().add(off));
5394 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5395 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5396 asm!(
5397 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5398 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5399 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5400 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5401 options(pure, nomem, nostack),
5402 );
5403 vaddq_s32(a0, a1)
5404 };
5405 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5406 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5407 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5408 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5409 }
5410 [
5411 vaddvq_f32(v0),
5412 vaddvq_f32(v1),
5413 vaddvq_f32(v2),
5414 vaddvq_f32(v3),
5415 ]
5416 }
5417}
5418
5419fn q4tp_matmat(
5423 bytes: &[u8],
5424 xs_all: &[f32],
5425 b: usize,
5426 rows: usize,
5427 cols: usize,
5428 out: &mut [f32],
5429 pool: Option<&Pool>,
5430) {
5431 debug_assert_eq!(out.len(), b * rows);
5432 let gpr = cols / GROUP_SIZE;
5433 let v = Q4tpView::new(bytes, rows, cols);
5434
5435 #[cfg(target_os = "macos")]
5437 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5438 dequant_matmat_accel(
5439 &|r, dst| {
5440 let mut sc = [0f32; 32];
5441 let mut scv;
5442 let s: &[f32] = if gpr <= 32 {
5443 v.scales_into(r, gpr, &mut sc);
5444 &sc[..gpr]
5445 } else {
5446 scv = vec![0f32; gpr];
5447 v.scales_into(r, gpr, &mut scv);
5448 &scv
5449 };
5450 for gi in 0..gpr {
5451 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5452 for (k, &bb) in tile.iter().enumerate() {
5453 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5454 dst[gi * GROUP_SIZE + k * 2 + 1] =
5455 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5456 }
5457 }
5458 },
5459 xs_all,
5460 b,
5461 rows,
5462 cols,
5463 out,
5464 pool,
5465 );
5466 return;
5467 }
5468
5469 let out_addr = SendMut(out.as_mut_ptr());
5470 if a8w8_enabled() {
5471 let acts: Vec<SplitAct> = (0..b)
5472 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5473 .collect();
5474 let acts = &acts;
5475 #[cfg(target_arch = "aarch64")]
5476 let blocked_ok = sdot_enabled() && blocked_enabled();
5477 #[cfg(target_arch = "x86_64")]
5484 let blocked_ok = q4tp_blocked_x86();
5485 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5486 let blocked_ok = false;
5487 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5497 .ok()
5498 .and_then(|v| v.parse().ok())
5499 .filter(|v| *v > 0)
5500 .unwrap_or(256);
5501 let run = |start: usize, end: usize| {
5502 for abase in (0..acts.len()).step_by(panel_cols) {
5503 let alen = (acts.len() - abase).min(panel_cols);
5504 let mut sc = vec![0f32; gpr];
5505 #[cfg(target_arch = "x86_64")]
5506 let mut r_lo = start;
5507 #[cfg(target_arch = "x86_64")]
5508 if blocked_ok && alen >= 8 {
5509 let mut sc1 = vec![0f32; gpr];
5510 while r_lo + 2 <= end {
5511 v.scales_into(r_lo, gpr, &mut sc);
5512 v.scales_into(r_lo + 1, gpr, &mut sc1);
5513 let mut bi = 0usize;
5514 while bi + 8 <= alen {
5515 let xs = [
5516 acts[abase + bi].xq.as_slice(),
5517 acts[abase + bi + 1].xq.as_slice(),
5518 acts[abase + bi + 2].xq.as_slice(),
5519 acts[abase + bi + 3].xq.as_slice(),
5520 acts[abase + bi + 4].xq.as_slice(),
5521 acts[abase + bi + 5].xq.as_slice(),
5522 acts[abase + bi + 6].xq.as_slice(),
5523 acts[abase + bi + 7].xq.as_slice(),
5524 ];
5525 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5526 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
5527 for k in 0..8 {
5528 let act = &acts[abase + bi + k];
5529 let mut acc = dr[k] * act.sx;
5530 for &(j, xv) in &act.outliers {
5531 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5532 acc += w * s * xv;
5533 }
5534 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5536 }
5537 }
5538 bi += 8;
5539 }
5540 for row in [r_lo, r_lo + 1] {
5543 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5544 for b2 in bi..alen {
5545 let act = &acts[abase + b2];
5546 let xs4 = [
5547 act.xq.as_slice(),
5548 act.xq.as_slice(),
5549 act.xq.as_slice(),
5550 act.xq.as_slice(),
5551 ];
5552 let d =
5553 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5554 let mut acc = d[0] * act.sx;
5555 for &(j, xv) in &act.outliers {
5556 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5557 acc += w * s * xv;
5558 }
5559 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5561 }
5562 }
5563 r_lo += 2;
5564 }
5565 }
5566 #[cfg(target_arch = "x86_64")]
5567 let row_start = r_lo;
5568 #[cfg(not(target_arch = "x86_64"))]
5569 let row_start = start;
5570 for r in row_start..end {
5571 v.scales_into(r, gpr, &mut sc);
5572 let mut bi = 0usize;
5573 #[cfg(target_arch = "x86_64")]
5574 if blocked_ok {
5575 while bi + 8 <= alen {
5576 let xs = [
5577 acts[abase + bi].xq.as_slice(),
5578 acts[abase + bi + 1].xq.as_slice(),
5579 acts[abase + bi + 2].xq.as_slice(),
5580 acts[abase + bi + 3].xq.as_slice(),
5581 acts[abase + bi + 4].xq.as_slice(),
5582 acts[abase + bi + 5].xq.as_slice(),
5583 acts[abase + bi + 6].xq.as_slice(),
5584 acts[abase + bi + 7].xq.as_slice(),
5585 ];
5586 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5587 for k in 0..8 {
5588 let act = &acts[abase + bi + k];
5589 let mut acc = d[k] * act.sx;
5590 for &(j, xv) in &act.outliers {
5591 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5592 acc += w * s * xv;
5593 }
5594 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5596 }
5597 bi += 8;
5598 }
5599 while bi + 4 <= alen {
5600 let xs = [
5601 acts[abase + bi].xq.as_slice(),
5602 acts[abase + bi + 1].xq.as_slice(),
5603 acts[abase + bi + 2].xq.as_slice(),
5604 acts[abase + bi + 3].xq.as_slice(),
5605 ];
5606 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5607 for k in 0..4 {
5608 let act = &acts[abase + bi + k];
5609 let mut acc = d[k] * act.sx;
5610 for &(j, xv) in &act.outliers {
5611 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5612 acc += w * s * xv;
5613 }
5614 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5616 }
5617 bi += 4;
5618 }
5619 }
5620 #[cfg(target_arch = "aarch64")]
5621 if blocked_ok {
5622 while bi + 4 <= alen {
5623 let xs = [
5624 acts[abase + bi].xq.as_slice(),
5625 acts[abase + bi + 1].xq.as_slice(),
5626 acts[abase + bi + 2].xq.as_slice(),
5627 acts[abase + bi + 3].xq.as_slice(),
5628 ];
5629 let d = unsafe {
5630 if q4tp_v1() {
5631 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5632 } else {
5633 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5634 }
5635 };
5636 for k in 0..4 {
5637 let act = &acts[abase + bi + k];
5638 let mut acc = d[k] * act.sx;
5639 for &(j, xv) in &act.outliers {
5640 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5641 acc += w * s * xv;
5642 }
5643 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5645 }
5646 bi += 4;
5647 }
5648 }
5649 let _ = blocked_ok;
5650 while bi < alen {
5651 let act = &acts[abase + bi];
5652 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5653 for &(j, xv) in &act.outliers {
5654 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5655 acc += w * s * xv;
5656 }
5657 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5659 bi += 1;
5660 }
5661 }
5662 }
5663 };
5664 dispatch_rows(pool, rows, &run);
5665 return;
5666 }
5667
5668 let run = |start: usize, end: usize| {
5669 let mut sc = vec![0f32; gpr];
5670 for r in start..end {
5671 v.scales_into(r, gpr, &mut sc);
5672 for bi in 0..b {
5673 let x = &xs_all[bi * cols..(bi + 1) * cols];
5674 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5676 }
5677 }
5678 };
5679 dispatch_rows(pool, rows, &run);
5680}
5681
5682fn q4t_matvec(
5684 bytes: &[u8],
5685 x: &[f32],
5686 rows: usize,
5687 cols: usize,
5688 out: &mut [f32],
5689 pool: Option<&Pool>,
5690) {
5691 debug_assert_eq!(out.len(), rows);
5692 let gpr = cols / GROUP_SIZE;
5693 let out_addr = SendMut(out.as_mut_ptr());
5694 if a8w8_enabled() {
5695 let act = split_act(x);
5696 let run = move |start: usize, end: usize| {
5697 for r in start..end {
5698 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5699 for &(j, xv) in &act.outliers {
5700 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5701 acc += w * s * xv;
5702 }
5703 unsafe { *out_addr.at(r) = acc };
5705 }
5706 };
5707 dispatch_rows(pool, rows, &run);
5708 return;
5709 }
5710 let run = move |start: usize, end: usize| {
5711 for r in start..end {
5712 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5714 }
5715 };
5716 dispatch_rows(pool, rows, &run);
5717}
5718
5719#[allow(clippy::too_many_arguments)]
5721fn q4t_matvec2(
5722 bytes: &[u8],
5723 x1: &[f32],
5724 x2: &[f32],
5725 rows: usize,
5726 cols: usize,
5727 o1: &mut [f32],
5728 o2: &mut [f32],
5729 pool: Option<&Pool>,
5730) {
5731 let gpr = cols / GROUP_SIZE;
5732 let p1 = SendMut(o1.as_mut_ptr());
5733 let p2 = SendMut(o2.as_mut_ptr());
5734 if a8w8_enabled() {
5735 let a1 = split_act(x1);
5736 let a2 = split_act(x2);
5737 let run = move |start: usize, end: usize| {
5738 for r in start..end {
5739 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5740 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5741 for &(j, xv) in &a1.outliers {
5742 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5743 v1 += w * s * xv;
5744 }
5745 for &(j, xv) in &a2.outliers {
5746 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5747 v2 += w * s * xv;
5748 }
5749 unsafe {
5751 *p1.at(r) = v1;
5752 *p2.at(r) = v2;
5753 }
5754 }
5755 };
5756 dispatch_rows(pool, rows, &run);
5757 return;
5758 }
5759 let run = move |start: usize, end: usize| {
5760 for r in start..end {
5761 unsafe {
5763 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5764 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5765 }
5766 }
5767 };
5768 dispatch_rows(pool, rows, &run);
5769}
5770
5771#[allow(clippy::too_many_arguments)]
5773#[cfg(target_os = "macos")]
5779fn dequant_matmat_accel(
5780 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5781 xs_all: &[f32],
5782 b: usize,
5783 rows: usize,
5784 cols: usize,
5785 out: &mut [f32],
5786 pool: Option<&Pool>,
5787) {
5788 const TR: usize = 2048;
5789 thread_local! {
5790 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5791 }
5792 WTILE.with(|wt| {
5793 let mut wtile = wt.borrow_mut();
5794 wtile.resize(TR * cols, 0.0);
5795 let mut r0 = 0usize;
5796 while r0 < rows {
5797 let tr = TR.min(rows - r0);
5798 let wt_addr = SendMut(wtile.as_mut_ptr());
5799 let run = |start: usize, end: usize| {
5800 for r in start..end {
5801 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5803 dequant_row(r0 + r, dst);
5804 }
5805 };
5806 dispatch_rows(pool, tr, &run);
5807 unsafe {
5808 accel_blas::cblas_sgemm(
5809 101, 111, 112, b as i32,
5813 tr as i32,
5814 cols as i32,
5815 1.0,
5816 xs_all.as_ptr(),
5817 cols as i32,
5818 wtile.as_ptr(),
5819 cols as i32,
5820 0.0,
5821 out.as_mut_ptr().add(r0),
5822 rows as i32,
5823 );
5824 }
5825 r0 += tr;
5826 }
5827 });
5828}
5829
5830fn q4t_matmat(
5831 bytes: &[u8],
5832 xs_all: &[f32],
5833 b: usize,
5834 rows: usize,
5835 cols: usize,
5836 out: &mut [f32],
5837 pool: Option<&Pool>,
5838) {
5839 debug_assert_eq!(out.len(), b * rows);
5840 let gpr = cols / GROUP_SIZE;
5841 #[cfg(target_os = "macos")]
5845 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5846 dequant_matmat_accel(
5847 &|r, dst| {
5848 for gi in 0..gpr {
5849 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5850 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5851 for (k, &bb) in tile[2..].iter().enumerate() {
5852 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5853 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5854 }
5855 }
5856 },
5857 xs_all,
5858 b,
5859 rows,
5860 cols,
5861 out,
5862 pool,
5863 );
5864 return;
5865 }
5866 let out_addr = SendMut(out.as_mut_ptr());
5867 if a8w8_enabled() {
5868 let acts: Vec<SplitAct> = (0..b)
5869 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5870 .collect();
5871 let acts = &acts;
5872 #[cfg(target_arch = "x86_64")]
5873 let blocked_ok = avx2_enabled() && blocked_enabled();
5874 #[cfg(target_arch = "aarch64")]
5875 let blocked_ok = sdot_enabled() && blocked_enabled();
5876 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5877 let blocked_ok = false;
5878 let run = move |start: usize, end: usize| {
5879 for r in start..end {
5880 let mut bi = 0usize;
5881 #[cfg(target_arch = "aarch64")]
5882 if blocked_ok {
5883 while bi + 4 <= acts.len() {
5884 let xs = [
5885 acts[bi].xq.as_slice(),
5886 acts[bi + 1].xq.as_slice(),
5887 acts[bi + 2].xq.as_slice(),
5888 acts[bi + 3].xq.as_slice(),
5889 ];
5890 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5891 for k in 0..4 {
5892 let act = &acts[bi + k];
5893 let mut acc = d[k] * act.sx;
5894 for &(j, xv) in &act.outliers {
5895 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5896 acc += w * sc * xv;
5897 }
5898 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5900 }
5901 bi += 4;
5902 }
5903 }
5904 #[cfg(target_arch = "x86_64")]
5905 if blocked_ok {
5906 while bi + 4 <= acts.len() {
5907 let xs = [
5908 acts[bi].xq.as_slice(),
5909 acts[bi + 1].xq.as_slice(),
5910 acts[bi + 2].xq.as_slice(),
5911 acts[bi + 3].xq.as_slice(),
5912 ];
5913 let d = unsafe {
5914 if vnni_tiles_enabled() {
5915 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5916 } else {
5917 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5918 }
5919 };
5920 for k in 0..4 {
5921 let act = &acts[bi + k];
5922 let mut acc = d[k] * act.sx;
5923 for &(j, xv) in &act.outliers {
5924 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5925 acc += w * sc * xv;
5926 }
5927 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5929 }
5930 bi += 4;
5931 }
5932 }
5933 let _ = blocked_ok;
5934 while bi < acts.len() {
5935 let act = &acts[bi];
5936 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5937 for &(j, xv) in &act.outliers {
5938 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5939 acc += w * s * xv;
5940 }
5941 unsafe { *out_addr.at(bi * rows + r) = acc };
5943 bi += 1;
5944 }
5945 }
5946 };
5947 dispatch_rows(pool, rows, &run);
5948 return;
5949 }
5950 let run = move |start: usize, end: usize| {
5951 for r in start..end {
5952 for bi in 0..b {
5953 let x = &xs_all[bi * cols..(bi + 1) * cols];
5954 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5956 }
5957 }
5958 };
5959 dispatch_rows(pool, rows, &run);
5960}
5961
5962fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5971 (0..gpr)
5972 .map(|gi| {
5973 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5974 .iter()
5975 .map(|&v| v as i32)
5976 .sum()
5977 })
5978 .collect()
5979}
5980
5981#[inline]
5985#[allow(unreachable_code)]
5986#[cfg(target_arch = "x86_64")]
5991#[target_feature(enable = "avx2")]
5992unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5993 unsafe {
5995 use core::arch::x86_64::*;
5996 let expand = _mm256_setr_epi8(
5998 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,
5999 3, 3, 3,
6000 );
6001 let bitsel = _mm256_setr_epi8(
6002 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6003 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6004 );
6005 let ones8 = _mm256_set1_epi8(1);
6006 let ones16 = _mm256_set1_epi16(1);
6007 let mut acc = 0f32;
6008 for gi in 0..gpr {
6009 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6010 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6011 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6012 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6013 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6014 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6015 let sel = _mm256_and_si256(x, mask);
6016 let p16 = _mm256_maddubs_epi16(ones8, sel);
6018 let d32 = _mm256_madd_epi16(p16, ones16);
6019 let hi128 = _mm256_extracti128_si256::<1>(d32);
6020 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6021 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6022 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6023 let msum = _mm_cvtsi128_si32(s32);
6024 let d = 2 * msum - gsum[gi];
6027 acc += d as f32 * s;
6028 }
6029 acc
6030 }
6031}
6032
6033#[cfg(target_arch = "x86_64")]
6036#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6037unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6038 unsafe {
6040 use core::arch::x86_64::*;
6041 let expand = _mm256_setr_epi8(
6042 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,
6043 3, 3, 3,
6044 );
6045 let bitsel = _mm256_setr_epi8(
6046 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6047 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6048 );
6049 let ones8 = _mm256_set1_epi8(1);
6050 let mut acc = 0f32;
6051 for gi in 0..gpr {
6052 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6053 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6054 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6055 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6056 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6057 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6058 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6059 let d = 2 * msum - gsum[gi];
6060 acc += d as f32 * s;
6061 }
6062 acc
6063 }
6064}
6065
6066#[cfg(target_arch = "x86_64")]
6068#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6069unsafe fn dot_q1_row_1x4_vnni(
6070 bytes: &[u8],
6071 r: usize,
6072 gpr: usize,
6073 xs: [&[i8]; 4],
6074 gsums: [&[i32]; 4],
6075) -> [f32; 4] {
6076 unsafe {
6078 use core::arch::x86_64::*;
6079 let expand = _mm256_setr_epi8(
6080 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,
6081 3, 3, 3,
6082 );
6083 let bitsel = _mm256_setr_epi8(
6084 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6085 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6086 );
6087 let ones8 = _mm256_set1_epi8(1);
6088 let mut acc = [0f32; 4];
6089 for gi in 0..gpr {
6090 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6091 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6092 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6093 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6094 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6095 for (k, xq) in xs.iter().enumerate() {
6096 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6097 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6098 let d = 2 * msum - gsums[k][gi];
6099 acc[k] += d as f32 * s;
6100 }
6101 }
6102 acc
6103 }
6104}
6105
6106#[cfg(target_arch = "x86_64")]
6109#[target_feature(enable = "avx2")]
6110unsafe fn dot_q1_row_1x4_avx2(
6111 bytes: &[u8],
6112 r: usize,
6113 gpr: usize,
6114 xs: [&[i8]; 4],
6115 gsums: [&[i32]; 4],
6116) -> [f32; 4] {
6117 unsafe {
6119 use core::arch::x86_64::*;
6120 let expand = _mm256_setr_epi8(
6121 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,
6122 3, 3, 3,
6123 );
6124 let bitsel = _mm256_setr_epi8(
6125 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6126 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6127 );
6128 let ones8 = _mm256_set1_epi8(1);
6129 let ones16 = _mm256_set1_epi16(1);
6130 let mut acc = [0f32; 4];
6131 for gi in 0..gpr {
6132 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6133 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6134 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6135 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6136 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6137 for (k, xq) in xs.iter().enumerate() {
6138 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6139 let sel = _mm256_and_si256(x, mask);
6140 let p16 = _mm256_maddubs_epi16(ones8, sel);
6141 let d32 = _mm256_madd_epi16(p16, ones16);
6142 let hi128 = _mm256_extracti128_si256::<1>(d32);
6143 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6144 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6145 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6146 let msum = _mm_cvtsi128_si32(s32);
6147 let d = 2 * msum - gsums[k][gi];
6148 acc[k] += d as f32 * s;
6149 }
6150 }
6151 acc
6152 }
6153}
6154
6155#[allow(unreachable_code)]
6156fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6157 #[cfg(target_arch = "aarch64")]
6158 unsafe {
6159 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6160 }
6161 #[cfg(target_arch = "x86_64")]
6162 if avx2_enabled() {
6163 unsafe {
6164 if vnni_tiles_enabled() {
6165 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6166 }
6167 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6168 }
6169 }
6170 let _ = gsum;
6171 let mut acc = 0f32;
6172 for gi in 0..gpr {
6173 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6174 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6175 let mut d = 0i32;
6176 for (j, &b) in tile[2..].iter().enumerate() {
6177 for k in 0..8 {
6178 let w = ((b >> k) & 1) as i32 * 2 - 1;
6179 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6180 }
6181 }
6182 acc += d as f32 * s;
6183 }
6184 acc
6185}
6186
6187#[cfg(target_arch = "aarch64")]
6196#[target_feature(enable = "neon,dotprod")]
6197unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6198 unsafe {
6201 use core::arch::aarch64::*;
6202 use core::arch::asm;
6203 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6204 let m = vld1q_u8(MASKS.as_ptr());
6205 macro_rules! tile_dot {
6207 ($t:expr, $x:expr) => {{
6208 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6209 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6210 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6211 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6212 let x0 = vld1q_s8($x);
6213 let x1 = vld1q_s8($x.add(16));
6214 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6215 asm!(
6216 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6217 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6218 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6219 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6220 options(pure, nomem, nostack),
6221 );
6222 vaddq_s32(a0, a1)
6223 }};
6224 }
6225 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6234 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6235 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6236 const IW11: [u8; 16] = [
6237 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6238 ];
6239 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6240 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6241 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6242 let isc = vld1_u8(ISC.as_ptr());
6243 macro_rules! tile_dot_tbl {
6245 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6246 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6247 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6248 let x0 = vld1q_s8($x);
6249 let x1 = vld1q_s8($x.add(16));
6250 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6251 asm!(
6252 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6253 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6254 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6255 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6256 options(pure, nomem, nostack),
6257 );
6258 vaddq_s32(a0, a1)
6259 }};
6260 }
6261 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6262 let row_base = r * gpr * Q1_TILE;
6263 let abs_end = bytes.len();
6264 let xp = xq.as_ptr();
6265 let gp = gsum.as_ptr();
6266 let mut accv = vdupq_n_f32(0.0);
6267 let mut gi = 0;
6268 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6271 let t0 = base.add(gi * Q1_TILE);
6272 let ld_a = vld1q_u8(t0);
6273 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6274 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6275 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6276 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6277 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6278 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6280 let g = vld1q_s32(gp.add(gi));
6281 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6282 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6283 let scf: float32x4_t;
6284 asm!(
6285 "fcvtl {o:v}.4s, {i:v}.4h",
6286 o = out(vreg) scf, i = in(vreg) sc16,
6287 options(pure, nomem, nostack),
6288 );
6289 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6290 gi += 4;
6291 }
6292 let mut acc = vaddvq_f32(accv);
6293 while gi < gpr {
6294 let t = base.add(gi * Q1_TILE);
6295 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6296 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6297 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6298 gi += 1;
6299 }
6300 acc
6301 }
6302}
6303
6304#[cfg(target_arch = "aarch64")]
6309#[target_feature(enable = "neon,dotprod")]
6310unsafe fn dot_q1_row_1x4_sdot(
6311 bytes: &[u8],
6312 r: usize,
6313 gpr: usize,
6314 xs: [&[i8]; 4],
6315 gs: [&[i32]; 4],
6316) -> [f32; 4] {
6317 unsafe {
6319 use core::arch::aarch64::*;
6320 use core::arch::asm;
6321 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6322 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6323 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6324 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6325 const IW11: [u8; 16] = [
6326 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6327 ];
6328 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6329 let m = vld1q_u8(MASKS.as_ptr());
6330 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6331 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6332 let isc = vld1_u8(ISC.as_ptr());
6333 macro_rules! sdot2 {
6334 ($w0:expr, $w1:expr, $x:expr) => {{
6335 let x0 = vld1q_s8($x);
6336 let x1 = vld1q_s8($x.add(16));
6337 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6338 asm!(
6339 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6340 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6341 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6342 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6343 options(pure, nomem, nostack),
6344 );
6345 vaddq_s32(a0, a1)
6346 }};
6347 }
6348 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6349 let row_base = r * gpr * Q1_TILE;
6350 let abs_end = bytes.len();
6351 let mut accv = [vdupq_n_f32(0.0); 4];
6352 let mut gi = 0;
6353 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6354 let t0 = base.add(gi * Q1_TILE);
6355 let ld_a = vld1q_u8(t0);
6356 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6357 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6359 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6360 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6361 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6362 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6363 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6364 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6365 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6366 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6367 let scf: float32x4_t;
6368 asm!(
6369 "fcvtl {o:v}.4s, {i:v}.4h",
6370 o = out(vreg) scf, i = in(vreg) sc16,
6371 options(pure, nomem, nostack),
6372 );
6373 for k in 0..4 {
6374 let xp = xs[k].as_ptr();
6375 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6376 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6377 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6378 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6379 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6380 let g = vld1q_s32(gs[k].as_ptr().add(gi));
6381 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6382 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6383 }
6384 gi += 4;
6385 }
6386 let mut acc = [
6387 vaddvq_f32(accv[0]),
6388 vaddvq_f32(accv[1]),
6389 vaddvq_f32(accv[2]),
6390 vaddvq_f32(accv[3]),
6391 ];
6392 while gi < gpr {
6393 let t = base.add(gi * Q1_TILE);
6394 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6395 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6396 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6397 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6398 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6399 for k in 0..4 {
6400 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6401 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6402 }
6403 gi += 1;
6404 }
6405 acc
6406 }
6407}
6408
6409#[inline]
6411fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6412 let gi = j / GROUP_SIZE;
6413 let k = j % GROUP_SIZE;
6414 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6415 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6416 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6417 ((bit as i32 * 2 - 1) as f32, s)
6418}
6419
6420#[inline]
6422fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6423 let mut acc = 0f32;
6424 for gi in 0..gpr {
6425 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6426 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6427 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6428 let mut ga = 0f32;
6429 for (j, &b) in tile[2..].iter().enumerate() {
6430 for k in 0..8 {
6431 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6432 }
6433 }
6434 acc += ga * s;
6435 }
6436 acc
6437}
6438
6439#[allow(clippy::too_many_arguments)]
6442fn q1_range_a8w8(
6443 bytes: &[u8],
6444 gpr: usize,
6445 act: &SplitAct,
6446 gsum: &[i32],
6447 out: SendMut,
6448 start: usize,
6449 end: usize,
6450) {
6451 for r in start..end {
6452 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6453 for &(j, xv) in &act.outliers {
6454 let (w, s) = q1_outlier(bytes, r, gpr, j);
6455 acc += w * s * xv;
6456 }
6457 unsafe { *out.at(r) = acc };
6459 }
6460}
6461
6462fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6464 for r in start..end {
6465 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6467 }
6468}
6469
6470fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6475 let entries = base_len + (rows + 1) * 4;
6476 (base_len, entries, entries <= bytes.len())
6477}
6478
6479#[inline]
6481fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6482 let o = rp_off + r * 4;
6483 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6484}
6485
6486const SIGN5: [[f32; 5]; 256] = {
6490 let mut lut = [[0.0f32; 5]; 256];
6491 let pow3 = [1u16, 3, 9, 27, 81];
6492 let mut byte = 0usize;
6493 while byte < 256 {
6494 let mut i = 0usize;
6495 while i < 5 {
6496 let code = (byte as u16 / pow3[i]) % 3;
6497 lut[byte][i] = if code == 1 {
6498 1.0
6499 } else if code == 2 {
6500 -1.0
6501 } else {
6502 0.0
6503 };
6504 i += 1;
6505 }
6506 byte += 1;
6507 }
6508 lut
6509};
6510
6511const SIGN5_I8: [[i8; 5]; 256] = {
6513 let mut lut = [[0i8; 5]; 256];
6514 let pow3 = [1u16, 3, 9, 27, 81];
6515 let mut byte = 0usize;
6516 while byte < 256 {
6517 let mut i = 0usize;
6518 while i < 5 {
6519 let code = (byte as u16 / pow3[i]) % 3;
6520 lut[byte][i] = if code == 1 {
6521 1
6522 } else if code == 2 {
6523 -1
6524 } else {
6525 0
6526 };
6527 i += 1;
6528 }
6529 byte += 1;
6530 }
6531 lut
6532};
6533
6534const SIGN5_U64: [u64; 256] = {
6540 let mut lut = [0u64; 256];
6541 let pow3 = [1u16, 3, 9, 27, 81];
6542 let mut byte = 0usize;
6543 while byte < 256 {
6544 let mut v = 0u64;
6545 let mut i = 0usize;
6546 while i < 5 {
6547 let code = (byte as u16 / pow3[i]) % 3;
6548 let s: u8 = if code == 1 {
6549 1
6550 } else if code == 2 {
6551 0xFF
6552 } else {
6553 0
6554 };
6555 v |= (s as u64) << (i * 8);
6556 i += 1;
6557 }
6558 lut[byte] = v;
6559 byte += 1;
6560 }
6561 lut
6562};
6563
6564#[inline]
6569fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6570 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6571 let off = (r * gpr + j / GROUP_SIZE) * TILE;
6572 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6573 let within = j % GROUP_SIZE;
6574 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6575}
6576
6577#[cfg(target_arch = "aarch64")]
6580#[target_feature(enable = "neon,dotprod")]
6581#[inline]
6582unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6583 unsafe {
6585 use core::arch::aarch64::*;
6586 use core::arch::asm;
6587 let w0 = vld1q_s8(w);
6588 let w1 = vld1q_s8(w.add(16));
6589 let x0 = vld1q_s8(x);
6590 let x1 = vld1q_s8(x.add(16));
6591 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6592 asm!(
6593 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6594 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6595 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6596 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6597 options(pure, nomem, nostack),
6598 );
6599 vaddvq_s32(vaddq_s32(a0, a1))
6600 }
6601}
6602
6603#[cfg(target_arch = "x86_64")]
6606#[target_feature(enable = "avx2")]
6607#[inline]
6608unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6609 unsafe {
6611 use core::arch::x86_64::*;
6612 let wv = _mm256_loadu_si256(w as *const __m256i);
6613 let xv = _mm256_loadu_si256(x as *const __m256i);
6614 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6615 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6616 let hi128 = _mm256_extracti128_si256::<1>(d);
6617 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6618 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6619 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6620 _mm_cvtsi128_si32(s32)
6621 }
6622}
6623
6624#[inline]
6629fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6630 debug_assert!(dst.len() >= 40);
6631 unsafe {
6634 let p = dst.as_mut_ptr();
6635 for bi in 0..7 {
6636 core::ptr::write_unaligned(
6637 p.add(bi * 5) as *mut u64,
6638 SIGN5_U64[*codes.add(bi) as usize],
6639 );
6640 }
6641 }
6642}
6643
6644#[inline]
6649fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6650 #[cfg(target_arch = "aarch64")]
6651 unsafe {
6652 return sdot32_i8(w, x);
6653 }
6654 #[cfg(target_arch = "x86_64")]
6655 unsafe {
6656 return i8dot32_avx2(w, x);
6657 }
6658 #[allow(unreachable_code)]
6659 unsafe {
6660 let mut s = 0i32;
6661 for k in 0..GROUP_SIZE {
6662 s += *w.add(k) as i32 * *x.add(k) as i32;
6663 }
6664 s
6665 }
6666}
6667
6668#[inline]
6669unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6670 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6671 (
6672 SIGN5_U64[*codes as usize],
6673 SIGN5_U64[*codes.add(1) as usize],
6674 SIGN5_U64[*codes.add(2) as usize],
6675 SIGN5_U64[*codes.add(3) as usize],
6676 SIGN5_U64[*codes.add(4) as usize],
6677 SIGN5_U64[*codes.add(5) as usize],
6678 SIGN5_U64[*codes.add(6) as usize],
6679 )
6680 };
6681
6682 let u0 = s0 | (s1 << 40);
6683 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6684 let u2 = (s3 >> 8) | (s4 << 32);
6685 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6686
6687 (u0, u1, u2, u3)
6688}
6689
6690#[cfg(target_arch = "aarch64")]
6694#[target_feature(enable = "neon,dotprod")]
6695unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6696 use core::arch::aarch64::*;
6697 use core::arch::asm;
6698 unsafe {
6699 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6700 let mut acc = 0f32;
6701 let bytes_ptr = bytes.as_ptr();
6702 let xq_ptr = xq.as_ptr();
6703 let row_off = r * gpr * TILE;
6704
6705 let gpr2 = gpr & !1;
6706 let mut gi = 0;
6707 while gi < gpr2 {
6708 let off0 = row_off + gi * TILE;
6709 let off1 = off0 + TILE;
6710 let s0 = f16_to_f32(u16::from_le_bytes([
6711 *bytes_ptr.add(off0),
6712 *bytes_ptr.add(off0 + 1),
6713 ]));
6714 let s1 = f16_to_f32(u16::from_le_bytes([
6715 *bytes_ptr.add(off1),
6716 *bytes_ptr.add(off1 + 1),
6717 ]));
6718
6719 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6720 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6721
6722 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6723 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6724 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6725 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6726
6727 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6728 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6729 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6730 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6731
6732 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6733 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6734 asm!(
6735 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6736 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6737 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6738 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6739 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6740 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6741 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6742 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6743 options(pure, nomem, nostack),
6744 );
6745 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6746 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6747 acc += d0 as f32 * s0 + d1 as f32 * s1;
6748 gi += 2;
6749 }
6750
6751 if gi < gpr {
6752 let off = row_off + gi * TILE;
6753 let s = f16_to_f32(u16::from_le_bytes([
6754 *bytes_ptr.add(off),
6755 *bytes_ptr.add(off + 1),
6756 ]));
6757 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6758 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6759 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6760 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6761 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6762 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6763 asm!(
6764 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6765 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6766 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6767 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6768 options(pure, nomem, nostack),
6769 );
6770 let d = vaddvq_s32(vaddq_s32(a0, a1));
6771 acc += d as f32 * s;
6772 }
6773 acc
6774 }
6775}
6776
6777#[cfg(target_arch = "x86_64")]
6779#[target_feature(enable = "avx2")]
6780unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6781 use core::arch::x86_64::*;
6782 unsafe {
6783 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6784 let mut acc = 0f32;
6785 let bytes_ptr = bytes.as_ptr();
6786 let xq_ptr = xq.as_ptr();
6787 let row_off = r * gpr * TILE;
6788
6789 let ones = _mm256_set1_epi16(1);
6790 for gi in 0..gpr {
6791 let off = row_off + gi * TILE;
6792 let s = f16_to_f32(u16::from_le_bytes([
6793 *bytes_ptr.add(off),
6794 *bytes_ptr.add(off + 1),
6795 ]));
6796 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6797 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6798 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6799 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6800 let d256 = _mm256_madd_epi16(p16, ones);
6801 let d128 = _mm_add_epi32(
6802 _mm256_castsi256_si128(d256),
6803 _mm256_extracti128_si256(d256, 1),
6804 );
6805 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6806 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6807 acc += d32 as f32 * s;
6808 }
6809 acc
6810 }
6811}
6812
6813#[cfg(target_arch = "x86_64")]
6815#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6816unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6817 use core::arch::x86_64::*;
6818 unsafe {
6820 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6821 let mut acc = 0f32;
6822 let bytes_ptr = bytes.as_ptr();
6823 let xq_ptr = xq.as_ptr();
6824 let row_off = r * gpr * TILE;
6825 for gi in 0..gpr {
6826 let off = row_off + gi * TILE;
6827 let s = f16_to_f32(u16::from_le_bytes([
6828 *bytes_ptr.add(off),
6829 *bytes_ptr.add(off + 1),
6830 ]));
6831 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6832 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6833 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6834 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6835 acc += d as f32 * s;
6836 }
6837 acc
6838 }
6839}
6840
6841#[inline]
6845fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6846 #[cfg(target_arch = "aarch64")]
6847 unsafe {
6848 return q1t_dot_row_sdot(bytes, r, gpr, xq);
6849 }
6850 #[cfg(target_arch = "x86_64")]
6851 unsafe {
6852 if vnni_tiles_enabled() {
6853 return q1t_dot_row_vnni(bytes, r, gpr, xq);
6854 }
6855 return q1t_dot_row_avx2(bytes, r, gpr, xq);
6856 }
6857 #[allow(unreachable_code)]
6858 {
6859 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6860 let mut acc = 0f32;
6861 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
6863 let off = (r * gpr + gi) * TILE;
6864 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6865 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6866 let mut d = 0i32;
6867 for k in 0..GROUP_SIZE {
6868 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6869 }
6870 acc += d as f32 * s;
6871 }
6872 acc
6873 }
6874}
6875
6876fn q1t_row_outlier_correction(
6883 bytes: &[u8],
6884 r: usize,
6885 rp_off: usize,
6886 entries_off: usize,
6887 has_ov: bool,
6888 x: &[f32],
6889) -> f32 {
6890 if !has_ov {
6891 return 0.0;
6892 }
6893 let (c0, c1) = (
6894 q1t_rowptr(bytes, rp_off, r),
6895 q1t_rowptr(bytes, rp_off, r + 1),
6896 );
6897 let mut corr = 0f32;
6898 for p in c0..c1 {
6899 let e = entries_off + p * 4;
6900 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6901 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6902 corr += val * x[col];
6903 }
6904 corr
6905}
6906
6907fn q1t_dequant_row(
6911 bytes: &[u8],
6912 r: usize,
6913 gpr: usize,
6914 rp_off: usize,
6915 entries_off: usize,
6916 has_ov: bool,
6917 buf: &mut [f32],
6918) {
6919 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6920 for g in 0..gpr {
6921 let off = (r * gpr + g) * TILE;
6922 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6923 let codes = &bytes[off + 2..off + TILE];
6924 let bc = g * GROUP_SIZE;
6925 for bi in 0..6 {
6927 let lut = &SIGN5[codes[bi] as usize];
6928 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6929 for i in 0..5 {
6930 d[i] = lut[i] * s;
6931 }
6932 }
6933 let lut = &SIGN5[codes[6] as usize];
6934 buf[bc + 30] = lut[0] * s;
6935 buf[bc + 31] = lut[1] * s;
6936 }
6937 if !has_ov {
6938 return;
6939 }
6940 let (c0, c1) = (
6941 q1t_rowptr(bytes, rp_off, r),
6942 q1t_rowptr(bytes, rp_off, r + 1),
6943 );
6944 for p in c0..c1 {
6945 let e = entries_off + p * 4;
6946 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6947 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6948 }
6949}
6950
6951fn q1t_add_overlay(
6955 bytes: &[u8],
6956 x: &[f32],
6957 rows: usize,
6958 cols: usize,
6959 out: &mut [f32],
6960 pool: Option<&Pool>,
6961) {
6962 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6963 let gpr = cols / GROUP_SIZE;
6964 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6965 if !has_ov {
6966 return;
6967 }
6968 let out_addr = SendMut(out.as_mut_ptr());
6969 let run = move |start: usize, end: usize| {
6970 for r in start..end {
6971 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6972 unsafe { *out_addr.at(r) += corr };
6974 }
6975 };
6976 dispatch_rows(pool, rows, &run);
6977}
6978
6979#[allow(clippy::too_many_arguments)]
6982fn q1t_range_a8w8(
6983 bytes: &[u8],
6984 gpr: usize,
6985 rp_off: usize,
6986 ent_off: usize,
6987 has_ov: bool,
6988 act: &SplitAct,
6989 x: &[f32],
6990 out: SendMut,
6991 start: usize,
6992 end: usize,
6993) {
6994 for r in start..end {
6995 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6996 for &(j, xv) in &act.outliers {
6997 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6998 }
6999 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7000 unsafe { *out.at(r) = acc };
7002 }
7003}
7004
7005#[allow(clippy::too_many_arguments)]
7008fn q1t_range_f32_batch(
7009 bytes: &[u8],
7010 gpr: usize,
7011 rp_off: usize,
7012 ent_off: usize,
7013 has_ov: bool,
7014 x: &[f32],
7015 out: SendMut,
7016 start: usize,
7017 end: usize,
7018) {
7019 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7020 let mut sg = [0f32; GROUP_SIZE];
7021 for r in start..end {
7022 let mut acc = 0f32;
7023 for g in 0..gpr {
7024 let off = (r * gpr + g) * TILE;
7025 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7026 let codes = &bytes[off + 2..off + TILE];
7027 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7028 for bi in 0..6 {
7029 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7030 }
7031 let lut = &SIGN5[codes[6] as usize];
7032 sg[30] = lut[0];
7033 sg[31] = lut[1];
7034 let mut gsum = 0f32;
7035 for k in 0..GROUP_SIZE {
7036 gsum += sg[k] * xg[k];
7037 }
7038 acc += s * gsum;
7039 }
7040 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7041 unsafe { *out.at(r) = acc };
7043 }
7044}
7045
7046fn q1t_matvec(
7050 bytes: &[u8],
7051 x: &[f32],
7052 rows: usize,
7053 cols: usize,
7054 out: &mut [f32],
7055 pool: Option<&Pool>,
7056) {
7057 debug_assert_eq!(out.len(), rows);
7058 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7059 let gpr = cols / GROUP_SIZE;
7060 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7061 let out_addr = SendMut(out.as_mut_ptr());
7062 if a8w8_enabled() {
7066 let act = split_act(x);
7067 let act = &act;
7068 let run = move |start: usize, end: usize| {
7069 for r in start..end {
7070 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7071 for &(j, xv) in &act.outliers {
7072 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7073 }
7074 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7075 unsafe { *out_addr.at(r) = acc };
7077 }
7078 };
7079 dispatch_rows(pool, rows, &run);
7080 return;
7081 }
7082 let run = move |start: usize, end: usize| {
7083 let mut sg = [0f32; GROUP_SIZE];
7087 for r in start..end {
7088 let mut acc = 0f32;
7089 for g in 0..gpr {
7090 let off = (r * gpr + g) * TILE;
7091 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7092 let codes = &bytes[off + 2..off + TILE];
7093 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7094 for bi in 0..6 {
7095 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7096 }
7097 let lut = &SIGN5[codes[6] as usize];
7098 sg[30] = lut[0];
7099 sg[31] = lut[1];
7100 let mut gsum = 0f32;
7101 for k in 0..GROUP_SIZE {
7102 gsum += sg[k] * xg[k];
7103 }
7104 acc += s * gsum;
7105 }
7106 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7107 unsafe { *out_addr.at(r) = acc };
7108 }
7109 };
7110 dispatch_rows(pool, rows, &run);
7111}
7112
7113#[cfg(target_arch = "aarch64")]
7119#[target_feature(enable = "neon,dotprod")]
7120unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7121 use core::arch::aarch64::*;
7122 use core::arch::asm;
7123 unsafe {
7125 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7126 let bytes_ptr = bytes.as_ptr();
7127 let row_off = r * gpr * TILE;
7128 let xp = [xa.as_ptr(), xb.as_ptr()];
7129 let mut acc = [0f32; 2];
7130 macro_rules! sdot2 {
7131 ($w0:expr, $w1:expr, $x:expr) => {{
7132 let x0 = vld1q_s8($x);
7133 let x1 = vld1q_s8($x.add(16));
7134 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7135 asm!(
7136 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7137 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7138 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7139 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7140 options(pure, nomem, nostack),
7141 );
7142 vaddvq_s32(vaddq_s32(a0, a1))
7143 }};
7144 }
7145 let gpr2 = gpr & !1;
7146 let mut gi = 0;
7147 while gi < gpr2 {
7148 let off0 = row_off + gi * TILE;
7149 let off1 = off0 + TILE;
7150 let s0 = f16_to_f32(u16::from_le_bytes([
7151 *bytes_ptr.add(off0),
7152 *bytes_ptr.add(off0 + 1),
7153 ]));
7154 let s1 = f16_to_f32(u16::from_le_bytes([
7155 *bytes_ptr.add(off1),
7156 *bytes_ptr.add(off1 + 1),
7157 ]));
7158 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7159 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7160 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7161 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7162 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7163 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7164 for k in 0..2 {
7165 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7166 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7167 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7168 }
7169 gi += 2;
7170 }
7171 if gi < gpr {
7172 let off = row_off + gi * TILE;
7173 let s = f16_to_f32(u16::from_le_bytes([
7174 *bytes_ptr.add(off),
7175 *bytes_ptr.add(off + 1),
7176 ]));
7177 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7178 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7179 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7180 for k in 0..2 {
7181 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7182 acc[k] += d as f32 * s;
7183 }
7184 }
7185 acc
7186 }
7187}
7188
7189fn q1t_matvec2(
7195 bytes: &[u8],
7196 x1: &[f32],
7197 x2: &[f32],
7198 rows: usize,
7199 cols: usize,
7200 o1: &mut [f32],
7201 o2: &mut [f32],
7202 pool: Option<&Pool>,
7203) {
7204 debug_assert_eq!(o1.len(), rows);
7205 debug_assert_eq!(o2.len(), rows);
7206 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7207 let gpr = cols / GROUP_SIZE;
7208 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7209 let out1 = SendMut(o1.as_mut_ptr());
7210 let out2 = SendMut(o2.as_mut_ptr());
7211 if a8w8_enabled() {
7212 let a1 = split_act(x1);
7213 let a2 = split_act(x2);
7214 let (a1, a2) = (&a1, &a2);
7215 let run = move |start: usize, end: usize| {
7216 for r in start..end {
7217 #[cfg(target_arch = "aarch64")]
7218 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7221 #[cfg(not(target_arch = "aarch64"))]
7222 let ds = [
7223 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7224 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7225 ];
7226 let mut acc1 = ds[0] * a1.sx;
7227 for &(j, xv) in &a1.outliers {
7228 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7229 }
7230 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7231 let mut acc2 = ds[1] * a2.sx;
7232 for &(j, xv) in &a2.outliers {
7233 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7234 }
7235 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7236 unsafe {
7238 *out1.at(r) = acc1;
7239 *out2.at(r) = acc2;
7240 }
7241 }
7242 };
7243 dispatch_rows(pool, rows, &run);
7244 return;
7245 }
7246 let run = move |start: usize, end: usize| {
7247 let mut sg = [0f32; GROUP_SIZE];
7250 for r in start..end {
7251 let mut acc1 = 0f32;
7252 let mut acc2 = 0f32;
7253 for g in 0..gpr {
7254 let off = (r * gpr + g) * TILE;
7255 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7256 let codes = &bytes[off + 2..off + TILE];
7257 for bi in 0..6 {
7258 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7259 }
7260 let lut = &SIGN5[codes[6] as usize];
7261 sg[30] = lut[0];
7262 sg[31] = lut[1];
7263 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7264 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7265 let mut gsum1 = 0f32;
7266 for k in 0..GROUP_SIZE {
7267 gsum1 += sg[k] * xg1[k];
7268 }
7269 acc1 += s * gsum1;
7270 let mut gsum2 = 0f32;
7271 for k in 0..GROUP_SIZE {
7272 gsum2 += sg[k] * xg2[k];
7273 }
7274 acc2 += s * gsum2;
7275 }
7276 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7277 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7278 unsafe {
7280 *out1.at(r) = acc1;
7281 *out2.at(r) = acc2;
7282 }
7283 }
7284 };
7285 dispatch_rows(pool, rows, &run);
7286}
7287
7288fn q1t_matmat(
7291 bytes: &[u8],
7292 xs: &[f32],
7293 b: usize,
7294 rows: usize,
7295 cols: usize,
7296 out: &mut [f32],
7297 pool: Option<&Pool>,
7298) {
7299 debug_assert_eq!(out.len(), b * rows);
7300 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7301 let gpr = cols / GROUP_SIZE;
7302 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7303 let out_addr = SendMut(out.as_mut_ptr());
7304 if a8w8_enabled() {
7308 let acts: Vec<SplitAct> = (0..b)
7309 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7310 .collect();
7311 let acts = &acts;
7312 let run = move |start: usize, end: usize| {
7313 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
7317 for g in 0..gpr {
7318 let off = (r * gpr + g) * TILE;
7319 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7320 q1t_unpack_group_i8(
7321 bytes.as_ptr().wrapping_add(off + 2),
7322 &mut sg[g * GROUP_SIZE..],
7323 );
7324 }
7325 for bi in 0..b {
7326 let act = &acts[bi];
7327 let mut isum = 0f32;
7328 for g in 0..gpr {
7329 let d = q1t_i8dot32(
7330 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7331 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7332 );
7333 isum += d as f32 * sc[g];
7334 }
7335 let mut acc = isum * act.sx;
7336 for &(j, xv) in &act.outliers {
7337 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7338 }
7339 accs[bi] = acc;
7340 }
7341 if has_ov {
7345 let (c0, c1) = (
7346 q1t_rowptr(bytes, rp_off, r),
7347 q1t_rowptr(bytes, rp_off, r + 1),
7348 );
7349 for p in c0..c1 {
7350 let e = ent_off + p * 4;
7351 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7352 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7353 for bi in 0..b {
7354 accs[bi] += val * xs[bi * cols + col];
7355 }
7356 }
7357 }
7358 for bi in 0..b {
7359 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7360 }
7361 }
7362 };
7363 dispatch_rows(pool, rows, &run);
7364 return;
7365 }
7366 let run = move |start: usize, end: usize| {
7367 let mut buf = vec![0f32; cols];
7368 for r in start..end {
7369 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7370 for bi in 0..b {
7371 let xr = &xs[bi * cols..(bi + 1) * cols];
7372 let mut acc = 0f32;
7373 for j in 0..cols {
7374 acc += buf[j] * xr[j];
7375 }
7376 unsafe { *out_addr.at(bi * rows + r) = acc };
7377 }
7378 }
7379 };
7380 dispatch_rows(pool, rows, &run);
7381}
7382
7383fn q1_matvec(
7384 bytes: &[u8],
7385 x: &[f32],
7386 rows: usize,
7387 cols: usize,
7388 out: &mut [f32],
7389 pool: Option<&Pool>,
7390) {
7391 debug_assert_eq!(out.len(), rows);
7392 let gpr = cols / GROUP_SIZE;
7393 let out_addr = SendMut(out.as_mut_ptr());
7394 if a8w8_enabled() {
7395 let act = split_act(x);
7396 let gsum = q1_group_sums(&act.xq, gpr);
7397 let (act, gsum) = (&act, &gsum);
7398 let run = move |start: usize, end: usize| {
7399 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7400 };
7401 dispatch_rows(pool, rows, &run);
7402 return;
7403 }
7404 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7405 dispatch_rows(pool, rows, &run);
7406}
7407
7408#[allow(clippy::too_many_arguments)]
7410fn q1_matvec2(
7411 bytes: &[u8],
7412 x1: &[f32],
7413 x2: &[f32],
7414 rows: usize,
7415 cols: usize,
7416 o1: &mut [f32],
7417 o2: &mut [f32],
7418 pool: Option<&Pool>,
7419) {
7420 let gpr = cols / GROUP_SIZE;
7421 let p1 = SendMut(o1.as_mut_ptr());
7422 let p2 = SendMut(o2.as_mut_ptr());
7423 if a8w8_enabled() {
7424 let a1 = split_act(x1);
7425 let a2 = split_act(x2);
7426 let g1 = q1_group_sums(&a1.xq, gpr);
7427 let g2 = q1_group_sums(&a2.xq, gpr);
7428 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7429 let run = move |start: usize, end: usize| {
7430 for r in start..end {
7431 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7432 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7433 for &(j, xv) in &a1.outliers {
7434 let (w, s) = q1_outlier(bytes, r, gpr, j);
7435 v1 += w * s * xv;
7436 }
7437 for &(j, xv) in &a2.outliers {
7438 let (w, s) = q1_outlier(bytes, r, gpr, j);
7439 v2 += w * s * xv;
7440 }
7441 unsafe {
7443 *p1.at(r) = v1;
7444 *p2.at(r) = v2;
7445 }
7446 }
7447 };
7448 dispatch_rows(pool, rows, &run);
7449 return;
7450 }
7451 let run = move |start: usize, end: usize| {
7452 for r in start..end {
7453 unsafe {
7455 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7456 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7457 }
7458 }
7459 };
7460 dispatch_rows(pool, rows, &run);
7461}
7462
7463#[allow(clippy::too_many_arguments)]
7465fn q1_matmat(
7466 bytes: &[u8],
7467 xs_all: &[f32],
7468 b: usize,
7469 rows: usize,
7470 cols: usize,
7471 out: &mut [f32],
7472 pool: Option<&Pool>,
7473) {
7474 debug_assert_eq!(out.len(), b * rows);
7475 let gpr = cols / GROUP_SIZE;
7476 let out_addr = SendMut(out.as_mut_ptr());
7477 if a8w8_enabled() {
7478 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7479 .map(|bi| {
7480 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7481 let gsum = q1_group_sums(&act.xq, gpr);
7482 (act, gsum)
7483 })
7484 .collect();
7485 let acts = &acts;
7486 #[cfg(target_arch = "x86_64")]
7487 let blocked_ok = avx2_enabled() && blocked_enabled();
7488 #[cfg(target_arch = "aarch64")]
7489 let blocked_ok = sdot_enabled() && blocked_enabled();
7490 let run = move |start: usize, end: usize| {
7491 for r in start..end {
7492 let mut bi = 0usize;
7493 #[cfg(target_arch = "aarch64")]
7496 if blocked_ok {
7497 while bi + 4 <= acts.len() {
7498 let xs = [
7499 acts[bi].0.xq.as_slice(),
7500 acts[bi + 1].0.xq.as_slice(),
7501 acts[bi + 2].0.xq.as_slice(),
7502 acts[bi + 3].0.xq.as_slice(),
7503 ];
7504 let gs = [
7505 acts[bi].1.as_slice(),
7506 acts[bi + 1].1.as_slice(),
7507 acts[bi + 2].1.as_slice(),
7508 acts[bi + 3].1.as_slice(),
7509 ];
7510 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7511 for k in 0..4 {
7512 let (act, _) = &acts[bi + k];
7513 let mut acc = d[k] * act.sx;
7514 for &(j, xv) in &act.outliers {
7515 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7516 acc += w * sc * xv;
7517 }
7518 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7520 }
7521 bi += 4;
7522 }
7523 }
7524 #[cfg(target_arch = "x86_64")]
7525 if blocked_ok {
7526 while bi + 4 <= acts.len() {
7527 let xs = [
7528 acts[bi].0.xq.as_slice(),
7529 acts[bi + 1].0.xq.as_slice(),
7530 acts[bi + 2].0.xq.as_slice(),
7531 acts[bi + 3].0.xq.as_slice(),
7532 ];
7533 let gs = [
7534 acts[bi].1.as_slice(),
7535 acts[bi + 1].1.as_slice(),
7536 acts[bi + 2].1.as_slice(),
7537 acts[bi + 3].1.as_slice(),
7538 ];
7539 let d = unsafe {
7540 if vnni_tiles_enabled() {
7541 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7542 } else {
7543 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7544 }
7545 };
7546 for k in 0..4 {
7547 let (act, _) = &acts[bi + k];
7548 let mut acc = d[k] * act.sx;
7549 for &(j, xv) in &act.outliers {
7550 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7551 acc += w * sc * xv;
7552 }
7553 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7555 }
7556 bi += 4;
7557 }
7558 }
7559 while bi < acts.len() {
7560 let (act, gsum) = &acts[bi];
7561 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7562 for &(j, xv) in &act.outliers {
7563 let (w, s) = q1_outlier(bytes, r, gpr, j);
7564 acc += w * s * xv;
7565 }
7566 unsafe { *out_addr.at(bi * rows + r) = acc };
7568 bi += 1;
7569 }
7570 }
7571 };
7572 dispatch_rows(pool, rows, &run);
7573 return;
7574 }
7575 let run = move |start: usize, end: usize| {
7576 for r in start..end {
7577 for bi in 0..b {
7578 let x = &xs_all[bi * cols..(bi + 1) * cols];
7579 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7581 }
7582 }
7583 };
7584 dispatch_rows(pool, rows, &run);
7585}
7586
7587fn q4matvec(
7593 bytes: &[u8],
7594 x: &[f32],
7595 rows: usize,
7596 cols: usize,
7597 out: &mut [f32],
7598 pool: Option<&Pool>,
7599) {
7600 debug_assert_eq!(out.len(), rows);
7601 let (packed, scales) = q4_split(bytes, rows, cols);
7602 let gpr = cols / GROUP_SIZE;
7603 let out_addr = SendMut(out.as_mut_ptr());
7604
7605 if a8w8_enabled() {
7606 let act = split_act(x);
7607 let run = move |start: usize, end: usize| {
7608 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7609 };
7610 dispatch_rows(pool, rows, &run);
7611 return;
7612 }
7613
7614 let run =
7615 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7616 dispatch_rows(pool, rows, &run);
7617}
7618
7619#[inline]
7622#[allow(unreachable_code)]
7623#[cfg(target_arch = "x86_64")]
7628#[target_feature(enable = "avx2")]
7629unsafe fn dot_q4b_row_1x4_avx2(
7630 buf: &[u8],
7631 scales: &[u8],
7632 g0: usize,
7633 gpr: usize,
7634 xs: [&[i8]; 4],
7635) -> [f32; 4] {
7636 unsafe {
7638 use core::arch::x86_64::*;
7639 let ones = _mm256_set1_epi16(1);
7640 let mut acc = [0f32; 4];
7641 for gi in 0..gpr {
7642 let s = f16_to_f32(u16::from_le_bytes([
7643 scales[(g0 + gi) * 2],
7644 scales[(g0 + gi) * 2 + 1],
7645 ]));
7646 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7647 let aw = _mm256_abs_epi8(w);
7648 for (k, xq) in xs.iter().enumerate() {
7649 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7650 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7651 let d = _mm256_madd_epi16(p16, ones);
7652 let hi128 = _mm256_extracti128_si256::<1>(d);
7653 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7654 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7655 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7656 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7657 }
7658 }
7659 acc
7660 }
7661}
7662
7663#[cfg(target_arch = "x86_64")]
7665#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7666unsafe fn dot_q4b_row_1x4_vnni(
7667 buf: &[u8],
7668 scales: &[u8],
7669 g0: usize,
7670 gpr: usize,
7671 xs: [&[i8]; 4],
7672) -> [f32; 4] {
7673 unsafe {
7675 use core::arch::x86_64::*;
7676 let mut acc = [0f32; 4];
7677 for gi in 0..gpr {
7678 let s = f16_to_f32(u16::from_le_bytes([
7679 scales[(g0 + gi) * 2],
7680 scales[(g0 + gi) * 2 + 1],
7681 ]));
7682 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7683 let aw = _mm256_abs_epi8(w);
7684 for (k, xq) in xs.iter().enumerate() {
7685 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7686 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7687 acc[k] += d as f32 * s;
7688 }
7689 }
7690 acc
7691 }
7692}
7693
7694#[cfg(target_arch = "x86_64")]
7700#[target_feature(enable = "avx2")]
7701unsafe fn dot_q4b_row_1x4_sx_avx2(
7702 buf: &[u8],
7703 scales: &[u8],
7704 g0: usize,
7705 gpr: usize,
7706 xs: [&[i8]; 4],
7707 sxs: [f32; 4],
7708) -> [f32; 4] {
7709 unsafe {
7711 use core::arch::x86_64::*;
7712 let ones = _mm256_set1_epi16(1);
7713 let mut acc = [0f32; 4];
7714 for gi in 0..gpr {
7715 let s = f16_to_f32(u16::from_le_bytes([
7716 scales[(g0 + gi) * 2],
7717 scales[(g0 + gi) * 2 + 1],
7718 ]));
7719 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7720 let aw = _mm256_abs_epi8(w);
7721 for (k, xq) in xs.iter().enumerate() {
7722 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7723 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7724 let d = _mm256_madd_epi16(p16, ones);
7725 let hi128 = _mm256_extracti128_si256::<1>(d);
7726 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7727 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7728 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7729 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7730 }
7731 }
7732 acc
7733 }
7734}
7735
7736#[cfg(target_arch = "x86_64")]
7739#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7740unsafe fn dot_q4b_row_1x4_sx_vnni(
7741 buf: &[u8],
7742 scales: &[u8],
7743 g0: usize,
7744 gpr: usize,
7745 xs: [&[i8]; 4],
7746 sxs: [f32; 4],
7747) -> [f32; 4] {
7748 unsafe {
7750 use core::arch::x86_64::*;
7751 let mut acc = [0f32; 4];
7752 for gi in 0..gpr {
7753 let s = f16_to_f32(u16::from_le_bytes([
7754 scales[(g0 + gi) * 2],
7755 scales[(g0 + gi) * 2 + 1],
7756 ]));
7757 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7758 let aw = _mm256_abs_epi8(w);
7759 for (k, xq) in xs.iter().enumerate() {
7760 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7761 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7762 acc[k] += (d as f32 * sxs[k]) * s;
7763 }
7764 }
7765 acc
7766 }
7767}
7768
7769#[allow(unreachable_code)]
7770fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7771 #[cfg(target_arch = "aarch64")]
7772 unsafe {
7773 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7774 }
7775 #[cfg(target_arch = "x86_64")]
7776 unsafe {
7777 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7778 }
7779 let mut acc = 0f32;
7780 for gi in 0..gpr {
7781 let g = g0 + gi;
7782 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7783 let mut d = 0i32;
7784 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7785 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7786 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7787 }
7788 acc += d as f32 * s;
7789 }
7790 acc
7791}
7792
7793#[inline]
7795#[allow(unreachable_code)]
7796fn dot_q4_row_i8_2(
7797 packed: &[u8],
7798 scales: &[u8],
7799 g0: usize,
7800 gpr: usize,
7801 xq1: &[i8],
7802 xq2: &[i8],
7803) -> (f32, f32) {
7804 #[cfg(target_arch = "aarch64")]
7805 unsafe {
7806 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7807 }
7808 #[cfg(target_arch = "x86_64")]
7809 unsafe {
7810 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7811 }
7812 (
7813 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7814 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7815 )
7816}
7817
7818#[allow(clippy::too_many_arguments)]
7821fn q4_range_a8w8(
7822 packed: &[u8],
7823 scales: &[u8],
7824 gpr: usize,
7825 cols: usize,
7826 act: &SplitAct,
7827 out: SendMut,
7828 start: usize,
7829 end: usize,
7830) {
7831 for r in start..end {
7832 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7833 for &(j, xv) in &act.outliers {
7835 let flat = r * cols + j;
7836 let byte = packed[flat / 2];
7837 let nib = if flat & 1 == 0 {
7838 byte & 0x0F
7839 } else {
7840 byte >> 4
7841 };
7842 let s = f16_to_f32(u16::from_le_bytes([
7843 scales[(flat / GROUP_SIZE) * 2],
7844 scales[(flat / GROUP_SIZE) * 2 + 1],
7845 ]));
7846 acc += ((nib as i32 - 8) as f32) * s * xv;
7847 }
7848 unsafe { *out.at(r) = acc };
7850 }
7851}
7852
7853#[allow(clippy::too_many_arguments)]
7856fn q4_range2_a8w8(
7857 packed: &[u8],
7858 scales: &[u8],
7859 gpr: usize,
7860 cols: usize,
7861 a1: &SplitAct,
7862 a2: &SplitAct,
7863 p1: SendMut,
7864 p2: SendMut,
7865 start: usize,
7866 end: usize,
7867) {
7868 for r in start..end {
7869 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7870 let mut acc1 = s1 * a1.sx;
7871 let mut acc2 = s2 * a2.sx;
7872 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7874 for &(j, xv) in outliers {
7875 let flat = r * cols + j;
7876 let byte = packed[flat / 2];
7877 let nib = if flat & 1 == 0 {
7878 byte & 0x0F
7879 } else {
7880 byte >> 4
7881 };
7882 let s = f16_to_f32(u16::from_le_bytes([
7883 scales[(flat / GROUP_SIZE) * 2],
7884 scales[(flat / GROUP_SIZE) * 2 + 1],
7885 ]));
7886 *acc += ((nib as i32 - 8) as f32) * s * xv;
7887 }
7888 };
7889 fix(&a1.outliers, &mut acc1);
7890 fix(&a2.outliers, &mut acc2);
7891 unsafe {
7893 *p1.at(r) = acc1;
7894 *p2.at(r) = acc2;
7895 }
7896 }
7897}
7898
7899fn q4_range_f32(
7901 packed: &[u8],
7902 scales: &[u8],
7903 gpr: usize,
7904 x: &[f32],
7905 out: SendMut,
7906 start: usize,
7907 end: usize,
7908) {
7909 for r in start..end {
7910 let mut acc = 0f32;
7911 for gi in 0..gpr {
7912 let g = r * gpr + gi;
7913 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7914 let pk = &packed[g * 16..(g + 1) * 16];
7915 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7916 let mut ga = 0f32;
7917 for (k, &b) in pk.iter().enumerate() {
7918 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7919 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7920 }
7921 acc += ga * s;
7922 }
7923 unsafe { *out.at(r) = acc };
7925 }
7926}
7927
7928#[allow(clippy::too_many_arguments)]
7932fn q4matvec2(
7933 bytes: &[u8],
7934 x1: &[f32],
7935 x2: &[f32],
7936 rows: usize,
7937 cols: usize,
7938 o1: &mut [f32],
7939 o2: &mut [f32],
7940 pool: Option<&Pool>,
7941) {
7942 debug_assert_eq!(o1.len(), rows);
7943 debug_assert_eq!(o2.len(), rows);
7944 let (packed, scales) = q4_split(bytes, rows, cols);
7945 let gpr = cols / GROUP_SIZE;
7946
7947 if a8w8_enabled() {
7948 let a1 = split_act(x1);
7949 let a2 = split_act(x2);
7950 let p1 = SendMut(o1.as_mut_ptr());
7951 let p2 = SendMut(o2.as_mut_ptr());
7952 let run = move |start: usize, end: usize| {
7953 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7954 };
7955 dispatch_rows(pool, rows, &run);
7956 return;
7957 }
7958
7959 let p1 = SendMut(o1.as_mut_ptr());
7960 let p2 = SendMut(o2.as_mut_ptr());
7961 let run = move |start: usize, end: usize| {
7962 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7963 };
7964 dispatch_rows(pool, rows, &run);
7965}
7966
7967#[allow(clippy::too_many_arguments)]
7969fn q4_range2_f32(
7970 packed: &[u8],
7971 scales: &[u8],
7972 gpr: usize,
7973 x1: &[f32],
7974 x2: &[f32],
7975 p1: SendMut,
7976 p2: SendMut,
7977 start: usize,
7978 end: usize,
7979) {
7980 for r in start..end {
7981 let (mut acc1, mut acc2) = (0f32, 0f32);
7982 for gi in 0..gpr {
7983 let g = r * gpr + gi;
7984 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7985 let pk = &packed[g * 16..(g + 1) * 16];
7986 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7987 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7988 let (mut g1, mut g2) = (0f32, 0f32);
7989 for (k, &b) in pk.iter().enumerate() {
7990 let wl = (b & 0x0F) as f32 - 8.0;
7991 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7992 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7993 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7994 }
7995 acc1 += g1 * s;
7996 acc2 += g2 * s;
7997 }
7998 unsafe {
8000 *p1.at(r) = acc1;
8001 *p2.at(r) = acc2;
8002 }
8003 }
8004}
8005
8006thread_local! {
8007 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8010 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8011}
8012
8013#[allow(clippy::too_many_arguments)]
8019fn q4matmat(
8020 bytes: &[u8],
8021 xs_all: &[f32],
8022 b: usize,
8023 rows: usize,
8024 cols: usize,
8025 out: &mut [f32],
8026 pool: Option<&Pool>,
8027) {
8028 debug_assert_eq!(xs_all.len(), b * cols);
8029 debug_assert_eq!(out.len(), b * rows);
8030 let (packed, scales) = q4_split(bytes, rows, cols);
8031 let gpr = cols / GROUP_SIZE;
8032 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8033
8034 if a8w8_enabled() {
8035 let acts: Vec<SplitAct> = (0..b)
8036 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8037 .collect();
8038 let acts = &acts;
8039 let out_addr = SendMut(out.as_mut_ptr());
8040 let run = move |start: usize, end: usize| {
8041 ROW_I8.with(|rb| {
8042 let mut buf = rb.borrow_mut();
8043 buf.resize(cols, 0);
8044 for r in start..end {
8045 for gi in 0..gpr {
8049 let g = r * gpr + gi;
8050 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8051 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
8052 buf[gi * GROUP_SIZE + k * 2 + 1] =
8053 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
8054 }
8055 }
8056 let mut bi = 0usize;
8057 #[cfg(target_arch = "x86_64")]
8058 if avx2_enabled() && blocked_enabled() {
8059 while bi + 4 <= acts.len() {
8060 let xs = [
8061 acts[bi].xq.as_slice(),
8062 acts[bi + 1].xq.as_slice(),
8063 acts[bi + 2].xq.as_slice(),
8064 acts[bi + 3].xq.as_slice(),
8065 ];
8066 let d = unsafe {
8067 if vnni_tiles_enabled() {
8068 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
8069 } else {
8070 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
8071 }
8072 };
8073 for k in 0..4 {
8074 let act = &acts[bi + k];
8075 let mut acc = d[k] * act.sx;
8076 for &(j, xv) in &act.outliers {
8077 acc += (buf[j] as i8) as f32
8078 * gscale((r * cols + j) / GROUP_SIZE)
8079 * xv;
8080 }
8081 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8083 }
8084 bi += 4;
8085 }
8086 }
8087 while bi < acts.len() {
8088 let act = &acts[bi];
8089 let mut acc = 0f32;
8090 for gi in 0..gpr {
8091 let d = dot_i8_i8(
8092 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8093 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8094 );
8095 acc += d as f32 * gscale(r * gpr + gi);
8096 }
8097 acc *= act.sx;
8098 for &(j, xv) in &act.outliers {
8100 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
8101 }
8102 unsafe { *out_addr.at(bi * rows + r) = acc };
8104 bi += 1;
8105 }
8106 }
8107 })
8108 };
8109 dispatch_rows(pool, rows, &run);
8110 return;
8111 }
8112
8113 let out_addr = SendMut(out.as_mut_ptr());
8114 let run = move |start: usize, end: usize| {
8115 ROW_F32.with(|rb| {
8116 let mut buf = rb.borrow_mut();
8117 buf.resize(cols, 0.0);
8118 for r in start..end {
8119 for gi in 0..gpr {
8122 let g = r * gpr + gi;
8123 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8124 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8125 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8126 }
8127 }
8128 for bi in 0..b {
8129 let x = &xs_all[bi * cols..(bi + 1) * cols];
8130 let mut acc = 0f32;
8131 for gi in 0..gpr {
8132 let mut ga = 0f32;
8133 for k in 0..GROUP_SIZE / 2 {
8138 let e = gi * GROUP_SIZE + k * 2;
8139 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8140 }
8141 acc += ga * gscale(r * gpr + gi);
8142 }
8143 unsafe { *out_addr.at(bi * rows + r) = acc };
8145 }
8146 }
8147 })
8148 };
8149 dispatch_rows(pool, rows, &run);
8150}
8151
8152#[allow(clippy::too_many_arguments)]
8157fn vbitmatmat(
8158 bytes: &[u8],
8159 offsets: &[usize],
8160 xs_all: &[f32],
8161 b: usize,
8162 rows: usize,
8163 cols: usize,
8164 out: &mut [f32],
8165 pool: Option<&Pool>,
8166) {
8167 debug_assert_eq!(xs_all.len(), b * cols);
8168 debug_assert_eq!(out.len(), b * rows);
8169 debug_assert_eq!(offsets.len(), rows + 1);
8170 let ng = cols / GROUP_SIZE;
8171 let bits = &bytes[..rows];
8172 let sc_off = rows;
8173 let gscale = |r: usize, g: usize| {
8174 let so = (r * ng + g) * 2;
8175 f16_to_f32(u16::from_le_bytes([
8176 bytes[sc_off + so],
8177 bytes[sc_off + so + 1],
8178 ]))
8179 };
8180
8181 let decode_f32 = |r: usize, dst: &mut [f32]| {
8183 let bw = bits[r] as usize;
8184 let l = ((1i32 << (bw - 1)) - 1) as f32;
8185 let data = &bytes[offsets[r]..offsets[r + 1]];
8186 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8187 for d in dst.iter_mut() {
8188 while nbits < bw {
8189 acc = (acc << 8) | data[idx] as u64;
8190 idx += 1;
8191 nbits += 8;
8192 }
8193 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8194 nbits -= bw;
8195 *d = u - l;
8196 }
8197 };
8198
8199 if a8w8_enabled() {
8200 let acts: Vec<SplitAct> = (0..b)
8201 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8202 .collect();
8203 let acts = &acts;
8204 let out_addr = SendMut(out.as_mut_ptr());
8205 let run = move |start: usize, end: usize| {
8206 for r in start..end {
8207 let bw = bits[r] as usize;
8208 if bw == 8 {
8209 ROW_F32.with(|rb| {
8212 let mut buf = rb.borrow_mut();
8213 buf.resize(cols, 0.0);
8214 decode_f32(r, &mut buf);
8215 for bi in 0..b {
8216 let x = &xs_all[bi * cols..(bi + 1) * cols];
8217 let mut dot = 0f32;
8218 for g in 0..ng {
8219 let mut gd = 0f32;
8220 for k in 0..GROUP_SIZE {
8221 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8222 }
8223 dot += gd * gscale(r, g);
8224 }
8225 unsafe { *out_addr.at(bi * rows + r) = dot };
8227 }
8228 });
8229 continue;
8230 }
8231 let l = (1i32 << (bw - 1)) - 1;
8232 let data = &bytes[offsets[r]..offsets[r + 1]];
8233 ROW_I8.with(|rb| {
8234 let mut buf = rb.borrow_mut();
8235 buf.resize(cols, 0);
8236 #[inline(always)]
8237 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8238 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8239 let u = unpack8::<B>(&data[blk * B..]);
8240 for k in 0..8 {
8241 chunk[k] = (u[k] - l) as i8 as u8;
8242 }
8243 }
8244 }
8245 match bw {
8246 3 => fill::<3>(data, l, &mut buf),
8247 4 => vbit_fill4(data, &mut buf),
8248 5 => fill::<5>(data, l, &mut buf),
8249 6 => fill::<6>(data, l, &mut buf),
8250 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8251 }
8252 let mut bi = 0usize;
8253 #[cfg(target_arch = "x86_64")]
8257 if avx2_enabled() && blocked_enabled() {
8258 while bi + 4 <= acts.len() {
8259 let xs = [
8260 acts[bi].xq.as_slice(),
8261 acts[bi + 1].xq.as_slice(),
8262 acts[bi + 2].xq.as_slice(),
8263 acts[bi + 3].xq.as_slice(),
8264 ];
8265 let sxs = [
8266 acts[bi].sx,
8267 acts[bi + 1].sx,
8268 acts[bi + 2].sx,
8269 acts[bi + 3].sx,
8270 ];
8271 let d = unsafe {
8272 if vnni_tiles_enabled() {
8273 dot_q4b_row_1x4_sx_vnni(
8274 &buf,
8275 &bytes[sc_off..],
8276 r * ng,
8277 ng,
8278 xs,
8279 sxs,
8280 )
8281 } else {
8282 dot_q4b_row_1x4_sx_avx2(
8283 &buf,
8284 &bytes[sc_off..],
8285 r * ng,
8286 ng,
8287 xs,
8288 sxs,
8289 )
8290 }
8291 };
8292 for k in 0..4 {
8293 let act = &acts[bi + k];
8294 let mut dot = d[k];
8295 for &(j, xv) in &act.outliers {
8296 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8297 }
8298 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8300 }
8301 bi += 4;
8302 }
8303 }
8304 while bi < acts.len() {
8305 let act = &acts[bi];
8306 let mut dot = 0f32;
8307 for g in 0..ng {
8308 let d = dot_i8_i8(
8309 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8310 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8311 ) as f32
8312 * act.sx;
8313 dot += d * gscale(r, g);
8314 }
8315 for &(j, xv) in &act.outliers {
8316 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8317 }
8318 unsafe { *out_addr.at(bi * rows + r) = dot };
8320 bi += 1;
8321 }
8322 });
8323 }
8324 };
8325 dispatch_rows(pool, rows, &run);
8326 return;
8327 }
8328
8329 let out_addr = SendMut(out.as_mut_ptr());
8330 let run = move |start: usize, end: usize| {
8331 ROW_F32.with(|rb| {
8332 let mut buf = rb.borrow_mut();
8333 buf.resize(cols, 0.0);
8334 for r in start..end {
8335 decode_f32(r, &mut buf);
8336 for bi in 0..b {
8337 let x = &xs_all[bi * cols..(bi + 1) * cols];
8338 let mut dot = 0f32;
8339 for g in 0..ng {
8340 let mut gd = 0f32;
8341 for k in 0..GROUP_SIZE {
8342 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8343 }
8344 dot += gd * gscale(r, g);
8345 }
8346 unsafe { *out_addr.at(bi * rows + r) = dot };
8348 }
8349 }
8350 })
8351 };
8352 dispatch_rows(pool, rows, &run);
8353}
8354
8355pub(crate) fn gpu_batch_job<'a>(
8359 t: &'a QTensor,
8360 x: &[f32],
8361) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8362 match t {
8363 QTensor::Mapped {
8364 model,
8365 idx,
8366 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8367 rows,
8368 cols,
8369 row_scale,
8370 col_field,
8371 ..
8372 } => Some((
8373 model.clone(),
8374 crate::gpu::BatchJob {
8375 idx: *idx,
8376 rows: *rows,
8377 cols: *cols,
8378 row_scale,
8379 xs: prescale(x, col_field, *dt).into_owned(),
8380 layout: crate::gpu::BatchLayout::Q8,
8381 },
8382 )),
8383 QTensor::Mapped {
8385 model,
8386 idx,
8387 dtype: TensorDtype::Q1,
8388 rows,
8389 cols,
8390 ..
8391 } => Some((
8392 model.clone(),
8393 crate::gpu::BatchJob {
8394 idx: *idx,
8395 rows: *rows,
8396 cols: *cols,
8397 row_scale: &[],
8398 xs: x.to_vec(),
8399 layout: crate::gpu::BatchLayout::Q1,
8400 },
8401 )),
8402 QTensor::Mapped {
8407 model,
8408 idx,
8409 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8410 rows,
8411 cols,
8412 ..
8413 } => Some((
8414 model.clone(),
8415 crate::gpu::BatchJob {
8416 idx: *idx,
8417 rows: *rows,
8418 cols: *cols,
8419 row_scale: &[],
8420 xs: x.to_vec(),
8421 layout: if *dt == TensorDtype::Q4Tiled {
8422 crate::gpu::BatchLayout::Q4t
8423 } else {
8424 crate::gpu::BatchLayout::Q4tp
8425 },
8426 },
8427 )),
8428 _ => None,
8429 }
8430}
8431
8432thread_local! {
8433 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8434 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8435}
8436
8437pub(crate) fn prescale<'a>(
8438 x: &'a [f32],
8439 col_field: &[f32],
8440 dtype: TensorDtype,
8441) -> std::borrow::Cow<'a, [f32]> {
8442 if dtype == TensorDtype::Q8_2f {
8443 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8444 } else {
8445 std::borrow::Cow::Borrowed(x)
8446 }
8447}
8448
8449pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8452 x: &[f32],
8453 col_field: &[f32],
8454 dtype: TensorDtype,
8455 buf_id: u8,
8456 f: F,
8457) -> R {
8458 if dtype == TensorDtype::Q8_2f {
8459 if buf_id == 1 {
8460 PRESCALE_BUF1.with(|b| {
8461 let mut buf = b.borrow_mut();
8462 buf.clear();
8463 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8464 f(&buf)
8465 })
8466 } else {
8467 PRESCALE_BUF2.with(|b| {
8468 let mut buf = b.borrow_mut();
8469 buf.clear();
8470 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8471 f(&buf)
8472 })
8473 }
8474 } else {
8475 f(x)
8476 }
8477}
8478
8479#[cfg(target_arch = "x86_64")]
8484pub(crate) fn avx2_enabled() -> bool {
8485 use std::sync::OnceLock;
8486 static ON: OnceLock<bool> = OnceLock::new();
8487 *ON.get_or_init(|| {
8488 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8489 && std::arch::is_x86_feature_detected!("avx2")
8490 && std::arch::is_x86_feature_detected!("fma")
8491 })
8492}
8493
8494#[cfg(target_arch = "x86_64")]
8499fn avx2_a8w8_enabled() -> bool {
8500 use std::sync::OnceLock;
8501 static ON: OnceLock<bool> = OnceLock::new();
8502 *ON.get_or_init(|| {
8503 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8504 })
8505}
8506
8507#[inline]
8511pub(crate) fn a8w8_enabled() -> bool {
8512 #[cfg(target_arch = "aarch64")]
8513 {
8514 sdot_enabled()
8515 }
8516 #[cfg(target_arch = "x86_64")]
8517 {
8518 avx2_a8w8_enabled()
8519 }
8520 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8521 {
8522 false
8523 }
8524}
8525
8526#[inline]
8529#[allow(unreachable_code)]
8530fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8531 #[cfg(target_arch = "aarch64")]
8532 unsafe {
8533 return dot_i8_sdot(w, xq);
8534 }
8535 #[cfg(target_arch = "x86_64")]
8536 unsafe {
8537 if avx512vnni_enabled() {
8538 return dot_i8_i8_vnni(w, xq);
8539 }
8540 return dot_i8_i8_avx2(w, xq);
8541 }
8542 w.iter()
8543 .zip(xq)
8544 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8545 .sum()
8546}
8547
8548#[cfg(target_arch = "x86_64")]
8552fn avx512vnni_enabled() -> bool {
8553 use std::sync::OnceLock;
8554 static ON: OnceLock<bool> = OnceLock::new();
8555 *ON.get_or_init(|| {
8556 std::env::var("CMF_AVX512")
8557 .map(|v| v != "0")
8558 .unwrap_or(true)
8559 && std::arch::is_x86_feature_detected!("avx512f")
8560 && std::arch::is_x86_feature_detected!("avx512bw")
8561 && std::arch::is_x86_feature_detected!("avx512vl")
8562 && std::arch::is_x86_feature_detected!("avx512vnni")
8563 })
8564}
8565
8566#[cfg(target_arch = "x86_64")]
8574fn vnni_tiles_enabled() -> bool {
8575 use std::sync::OnceLock;
8576 static ON: OnceLock<bool> = OnceLock::new();
8577 *ON.get_or_init(|| {
8578 std::env::var("CMF_VNNI_TILES")
8579 .map(|v| v != "0")
8580 .unwrap_or(true)
8581 && avx512vnni_enabled()
8582 })
8583}
8584
8585#[cfg(target_arch = "x86_64")]
8590#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8591#[inline]
8592unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8593 unsafe {
8595 use core::arch::x86_64::*;
8596 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8597 let hi128 = _mm256_extracti128_si256::<1>(d);
8598 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8599 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8600 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8601 _mm_cvtsi128_si32(s32)
8602 }
8603}
8604
8605#[cfg(target_arch = "x86_64")]
8610#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8611unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8612 unsafe {
8614 use core::arch::x86_64::*;
8615 let n = w.len();
8616 let mut j = 0usize;
8617 let mut total: i32;
8618 {
8623 #[inline(always)]
8624 unsafe fn step(
8625 w: *const u8,
8626 x: *const i8,
8627 acc: core::arch::x86_64::__m512i,
8628 ) -> core::arch::x86_64::__m512i {
8629 unsafe {
8630 use core::arch::x86_64::*;
8631 let wv = _mm512_loadu_si512(w as *const _);
8632 let xv = _mm512_loadu_si512(x as *const _);
8633 let aw = _mm512_abs_epi8(wv);
8634 let neg = _mm512_movepi8_mask(wv);
8635 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8636 _mm512_dpbusd_epi32(acc, aw, sx)
8637 }
8638 }
8639 let (mut a0, mut a1, mut a2, mut a3) = (
8640 _mm512_setzero_si512(),
8641 _mm512_setzero_si512(),
8642 _mm512_setzero_si512(),
8643 _mm512_setzero_si512(),
8644 );
8645 while j + 256 <= n {
8646 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8647 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8648 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8649 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8650 j += 256;
8651 }
8652 while j + 64 <= n {
8653 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8654 j += 64;
8655 }
8656 let s01 = _mm512_add_epi32(a0, a1);
8657 let s23 = _mm512_add_epi32(a2, a3);
8658 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8659 }
8660 if j + 32 <= n {
8662 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8663 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8664 let d = _mm256_dpbusd_epi32(
8665 _mm256_setzero_si256(),
8666 _mm256_abs_epi8(wv),
8667 _mm256_sign_epi8(xv, wv),
8668 );
8669 let hi128 = _mm256_extracti128_si256::<1>(d);
8670 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8671 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8672 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8673 total += _mm_cvtsi128_si32(s32);
8674 j += 32;
8675 }
8676 while j < n {
8677 total += (w[j] as i8) as i32 * xq[j] as i32;
8678 j += 1;
8679 }
8680 total
8681 }
8682}
8683
8684#[cfg(target_arch = "x86_64")]
8686#[target_feature(enable = "avx2,fma")]
8687unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8688 unsafe {
8690 use core::arch::x86_64::*;
8691 let n = x.len();
8692 let wp = w.as_ptr();
8693 let xp = x.as_ptr();
8694 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8695 let mut j = 0usize;
8696 while j + 16 <= n {
8697 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8698 let lo = _mm256_cvtepi8_epi32(wb);
8699 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8700 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8701 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8702 j += 16;
8703 }
8704 let acc = _mm256_add_ps(a0, a1);
8705 let hi128 = _mm256_extractf128_ps::<1>(acc);
8706 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8707 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8708 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8709 let mut sum = _mm_cvtss_f32(s32);
8710 while j < n {
8711 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8712 j += 1;
8713 }
8714 sum
8715 }
8716}
8717
8718#[cfg(target_arch = "x86_64")]
8723#[target_feature(enable = "avx2")]
8724unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8725 unsafe {
8727 use core::arch::x86_64::*;
8728 let n = w.len();
8729 let ones = _mm256_set1_epi16(1);
8730 let mut acc = _mm256_setzero_si256();
8731 let mut j = 0usize;
8732 while j + 32 <= n {
8733 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8734 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8735 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8736 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8737 j += 32;
8738 }
8739 let hi128 = _mm256_extracti128_si256::<1>(acc);
8740 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8741 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8742 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8743 let mut s = _mm_cvtsi128_si32(s32);
8744 while j < n {
8745 s += (w[j] as i8) as i32 * xq[j] as i32;
8746 j += 1;
8747 }
8748 s
8749 }
8750}
8751
8752#[cfg(target_arch = "aarch64")]
8756#[target_feature(enable = "neon,i8mm")]
8757unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8758 unsafe {
8760 use core::arch::aarch64::*;
8761 use core::arch::asm;
8762 let n = w0.len();
8763 let w0p = w0.as_ptr() as *const i8;
8764 let w1p = w1.as_ptr() as *const i8;
8765 let mut acc01 = vdupq_n_s32(0);
8768 let mut acc23 = vdupq_n_s32(0);
8769 let mut i = 0usize;
8770 while i + 8 <= n {
8771 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8772 let xb01 = vcombine_s8(
8773 vld1_s8(xs[0].as_ptr().add(i)),
8774 vld1_s8(xs[1].as_ptr().add(i)),
8775 );
8776 let xb23 = vcombine_s8(
8777 vld1_s8(xs[2].as_ptr().add(i)),
8778 vld1_s8(xs[3].as_ptr().add(i)),
8779 );
8780 asm!(
8781 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8782 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8783 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8784 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8785 options(pure, nomem, nostack),
8786 );
8787 i += 8;
8788 }
8789 let mut out = [[0i32; 4]; 2];
8790 let a01: [i32; 4] = core::mem::transmute(acc01);
8791 let a23: [i32; 4] = core::mem::transmute(acc23);
8792 out[0][0] = a01[0];
8793 out[0][1] = a01[1];
8794 out[1][0] = a01[2];
8795 out[1][1] = a01[3];
8796 out[0][2] = a23[0];
8797 out[0][3] = a23[1];
8798 out[1][2] = a23[2];
8799 out[1][3] = a23[3];
8800 if i < n {
8801 for (k, x) in xs.iter().enumerate() {
8802 for j in i..n {
8803 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8804 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8805 }
8806 }
8807 }
8808 out
8809 }
8810}
8811
8812#[cfg(target_arch = "aarch64")]
8816#[target_feature(enable = "neon,dotprod")]
8817unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8818 unsafe {
8820 use core::arch::aarch64::*;
8821 use core::arch::asm;
8822 let n = w0.len();
8823 let w0p = w0.as_ptr() as *const i8;
8824 let w1p = w1.as_ptr() as *const i8;
8825 let mut acc = [[vdupq_n_s32(0); 4]; 2];
8826 let mut i = 0usize;
8827 while i + 16 <= n {
8828 let wv0 = vld1q_s8(w0p.add(i));
8829 let wv1 = vld1q_s8(w1p.add(i));
8830 for (k, x) in xs.iter().enumerate() {
8831 let xv = vld1q_s8(x.as_ptr().add(i));
8832 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8833 asm!(
8834 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8835 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8836 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8837 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8838 options(pure, nomem, nostack),
8839 );
8840 acc[0][k] = a0;
8841 acc[1][k] = a1;
8842 }
8843 i += 16;
8844 }
8845 let mut out = [[0i32; 4]; 2];
8846 for r in 0..2 {
8847 for k in 0..4 {
8848 out[r][k] = vaddvq_s32(acc[r][k]);
8849 }
8850 }
8851 if i < n {
8852 for (k, x) in xs.iter().enumerate() {
8853 for j in i..n {
8854 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8855 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8856 }
8857 }
8858 }
8859 out
8860 }
8861}
8862
8863#[cfg(target_arch = "x86_64")]
8869#[target_feature(enable = "avx2")]
8870unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8871 unsafe {
8873 use core::arch::x86_64::*;
8874 let n = w0.len();
8875 let ones = _mm256_set1_epi16(1);
8876 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8877 let mut j = 0usize;
8878 while j + 32 <= n {
8879 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8880 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8881 let aw0 = _mm256_abs_epi8(wv0);
8882 let aw1 = _mm256_abs_epi8(wv1);
8883 for (k, x) in xs.iter().enumerate() {
8884 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8885 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8886 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8887 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8888 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8889 }
8890 j += 32;
8891 }
8892 let mut out = [[0i32; 4]; 2];
8893 for r in 0..2 {
8894 for k in 0..4 {
8895 let a = acc[r][k];
8896 let hi128 = _mm256_extracti128_si256::<1>(a);
8897 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8898 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8899 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8900 out[r][k] = _mm_cvtsi128_si32(s32);
8901 }
8902 }
8903 if j < n {
8904 for (k, x) in xs.iter().enumerate() {
8905 for i in j..n {
8906 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8907 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8908 }
8909 }
8910 }
8911 out
8912 }
8913}
8914
8915#[cfg(target_arch = "x86_64")]
8920#[inline]
8921fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8922 let dot = if avx512vnni_enabled() && row.len() >= 64 {
8923 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8924 } else {
8925 unsafe { dot_i8_i8_avx2(row, &act.xq) }
8926 };
8927 let mut acc = dot as f32 * act.sx;
8928 for &(j, xv) in &act.outliers {
8929 acc += (row[j] as i8) as f32 * xv;
8930 }
8931 acc
8932}
8933
8934#[cfg(target_arch = "x86_64")]
8938#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8939unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8940 unsafe {
8942 use core::arch::x86_64::*;
8943 let n = w.len();
8944 let flip = _mm512_set1_epi8(-128); #[inline(always)]
8946 unsafe fn step(
8947 w: *const u8,
8948 x: *const i8,
8949 flip: core::arch::x86_64::__m512i,
8950 acc: core::arch::x86_64::__m512i,
8951 ) -> core::arch::x86_64::__m512i {
8952 unsafe {
8953 use core::arch::x86_64::*;
8954 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8955 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8956 }
8957 }
8958 let (mut a0, mut a1, mut a2, mut a3) = (
8959 _mm512_setzero_si512(),
8960 _mm512_setzero_si512(),
8961 _mm512_setzero_si512(),
8962 _mm512_setzero_si512(),
8963 );
8964 let mut j = 0usize;
8965 while j + 256 <= n {
8966 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8967 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8968 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8969 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8970 j += 256;
8971 }
8972 while j + 64 <= n {
8973 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8974 j += 64;
8975 }
8976 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8977 _mm512_add_epi32(a0, a1),
8978 _mm512_add_epi32(a2, a3),
8979 ));
8980 while j < n {
8982 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8983 j += 1;
8984 }
8985 total
8986 }
8987}
8988
8989#[cfg(target_arch = "x86_64")]
8995#[target_feature(enable = "avx2")]
8996unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8997 unsafe {
9000 use core::arch::x86_64::*;
9001 let lomask = _mm_set1_epi8(0x0F);
9002 let eight = _mm256_set1_epi8(8);
9003 let ones = _mm256_set1_epi16(1);
9004 let mut acc = 0f32;
9005 for gi in 0..gpr {
9006 let g = g0 + gi;
9007 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9008 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9009 let lo = _mm_and_si128(b, lomask);
9010 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9011 let w = _mm256_sub_epi8(
9012 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9013 eight,
9014 );
9015 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9016 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
9017 let d = _mm256_madd_epi16(p16, ones);
9018 let hi128 = _mm256_extracti128_si256::<1>(d);
9019 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9020 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9021 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9022 acc += _mm_cvtsi128_si32(s32) as f32 * s;
9023 }
9024 acc
9025 }
9026}
9027
9028#[cfg(target_arch = "x86_64")]
9031#[target_feature(enable = "avx2")]
9032unsafe fn dot_q4_row_avx2_2(
9033 packed: &[u8],
9034 scales: &[u8],
9035 g0: usize,
9036 gpr: usize,
9037 xq1: &[i8],
9038 xq2: &[i8],
9039) -> (f32, f32) {
9040 unsafe {
9042 use core::arch::x86_64::*;
9043 let lomask = _mm_set1_epi8(0x0F);
9044 let eight = _mm256_set1_epi8(8);
9045 let ones = _mm256_set1_epi16(1);
9046 let (mut acc1, mut acc2) = (0f32, 0f32);
9047 #[inline(always)]
9048 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
9049 unsafe {
9050 use core::arch::x86_64::*;
9051 let hi128 = _mm256_extracti128_si256::<1>(d);
9052 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9053 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9054 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9055 _mm_cvtsi128_si32(s32)
9056 }
9057 }
9058 for gi in 0..gpr {
9059 let g = g0 + gi;
9060 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9061 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9062 let lo = _mm_and_si128(b, lomask);
9063 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9064 let w = _mm256_sub_epi8(
9065 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9066 eight,
9067 );
9068 let aw = _mm256_abs_epi8(w);
9069 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9070 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9071 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
9072 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
9073 acc1 += hsum(d1) as f32 * s;
9074 acc2 += hsum(d2) as f32 * s;
9075 }
9076 (acc1, acc2)
9077 }
9078}
9079
9080#[cfg(target_arch = "x86_64")]
9082fn q8_range_avx2(
9083 q: &[u8],
9084 row_scale: &[f32],
9085 act: &SplitAct,
9086 cols: usize,
9087 out_addr: SendMut,
9088 start: usize,
9089 end: usize,
9090) {
9091 for o in start..end {
9092 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9093 unsafe { *out_addr.at(o) = v };
9095 }
9096}
9097
9098#[cfg(target_arch = "x86_64")]
9100#[allow(clippy::too_many_arguments)]
9101fn q8_range2_avx2(
9102 q: &[u8],
9103 row_scale: &[f32],
9104 a1: &SplitAct,
9105 a2: &SplitAct,
9106 cols: usize,
9107 p1: SendMut,
9108 p2: SendMut,
9109 start: usize,
9110 end: usize,
9111) {
9112 for o in start..end {
9113 let row = &q[o * cols..(o + 1) * cols];
9114 unsafe {
9116 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9117 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9118 }
9119 }
9120}
9121
9122#[cfg(target_arch = "aarch64")]
9133fn i8mm_enabled() -> bool {
9134 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9135 *ON.get_or_init(|| {
9136 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9137 && std::arch::is_aarch64_feature_detected!("i8mm")
9138 })
9139}
9140
9141#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9145fn sdot_enabled() -> bool {
9146 use std::sync::OnceLock;
9147 static ON: OnceLock<bool> = OnceLock::new();
9148 *ON.get_or_init(|| {
9149 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9150 if !want {
9151 return false;
9152 }
9153
9154 #[cfg(target_arch = "aarch64")]
9155 {
9156 if std::arch::is_aarch64_feature_detected!("dotprod") {
9157 return true;
9158 }
9159 #[cfg(target_os = "android")]
9160 {
9161 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9162 if cpuinfo.lines().any(|l| {
9163 (l.starts_with("Features") || l.starts_with("features"))
9164 && l.contains("asimddp")
9165 }) {
9166 return true;
9167 }
9168 }
9169 }
9170 false
9171 }
9172 #[cfg(not(target_arch = "aarch64"))]
9173 {
9174 false
9175 }
9176 })
9177}
9178
9179struct SplitAct {
9184 xq: Vec<i8>,
9185 sx: f32,
9186 outliers: Vec<(usize, f32)>,
9187 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9190 xsum: i32,
9191}
9192
9193thread_local! {
9194 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9197 const { std::cell::RefCell::new(Vec::new()) };
9198}
9199
9200impl Drop for SplitAct {
9201 fn drop(&mut self) {
9202 let buf = std::mem::take(&mut self.xq);
9203 if buf.capacity() > 0 {
9204 XQ_FREE.with(|f| {
9205 let mut f = f.borrow_mut();
9206 if f.len() < 16 {
9207 f.push(buf);
9208 }
9209 });
9210 }
9211 }
9212}
9213
9214thread_local! {
9215 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9222}
9223
9224#[inline]
9227fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9228 KROW.with(|s| {
9229 let mut b = s.borrow_mut();
9230 if b.len() < n {
9231 b.resize(n, 0.0);
9232 }
9233 f(&mut b[..n])
9234 })
9235}
9236
9237fn split_act(x: &[f32]) -> SplitAct {
9238 let n = x.len();
9239 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9240 let thr = 8.0 * rms;
9241 let mut outliers: Vec<(usize, f32)> = Vec::new();
9245 let mut amax = 0f32;
9246 for (j, &v) in x.iter().enumerate() {
9247 let a = v.abs();
9248 if a > thr {
9249 outliers.push((j, v));
9250 } else if a > amax {
9251 amax = a;
9252 }
9253 }
9254 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9255 let inv = 1.0 / sx;
9256 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9257 xq.clear();
9258 xq.reserve(n);
9259 if outliers.is_empty() {
9260 xq.extend(
9261 x.iter()
9262 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9263 );
9264 } else {
9265 xq.extend(x.iter().map(|&v| {
9267 if v.abs() > thr {
9268 0
9269 } else {
9270 (v * inv).round().clamp(-127.0, 127.0) as i8
9271 }
9272 }));
9273 }
9274 let xsum = xq.iter().map(|&v| v as i32).sum();
9275 SplitAct {
9276 xq,
9277 sx,
9278 outliers,
9279 xsum,
9280 }
9281}
9282
9283fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9284 let n = x.len();
9285 let rms = (x
9286 .iter()
9287 .zip(col)
9288 .map(|(&a, &c)| {
9289 let v = a * c;
9290 (v * v) as f64
9291 })
9292 .sum::<f64>()
9293 / n.max(1) as f64)
9294 .sqrt() as f32;
9295 let thr = 8.0 * rms;
9296
9297 let mut outliers = Vec::new();
9298 let mut amax = 0f32;
9299 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9300 let v = a * c;
9301 let s = v.abs();
9302 if s > thr {
9303 outliers.push((j, v));
9304 } else if s > amax {
9305 amax = s;
9306 }
9307 }
9308
9309 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9310 let inv = 1.0 / sx;
9311 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9312 xq.clear();
9313 xq.reserve(n);
9314 if outliers.is_empty() {
9315 xq.extend(
9316 x.iter()
9317 .zip(col)
9318 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9319 );
9320 } else {
9321 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9322 let v = a * c;
9323 if v.abs() > thr {
9324 0
9325 } else {
9326 (v * inv).round().clamp(-127.0, 127.0) as i8
9327 }
9328 }));
9329 }
9330 let xsum = xq.iter().map(|&v| v as i32).sum();
9331 SplitAct {
9332 xq,
9333 sx,
9334 outliers,
9335 xsum,
9336 }
9337}
9338
9339#[cfg(target_arch = "aarch64")]
9342#[target_feature(enable = "neon,dotprod")]
9343unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9344 unsafe {
9346 use core::arch::aarch64::*;
9347 use core::arch::asm;
9348 let wp = w.as_ptr() as *const i8;
9349 let n = w.len();
9350 let (mut a0, mut a1, mut a2, mut a3) = (
9351 vdupq_n_s32(0),
9352 vdupq_n_s32(0),
9353 vdupq_n_s32(0),
9354 vdupq_n_s32(0),
9355 );
9356 let mut i = 0;
9357 while i + 64 <= n {
9358 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9359 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9360 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9361 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9362 asm!(
9363 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9364 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9365 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9366 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9367 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9368 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9369 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9370 options(pure, nomem, nostack),
9371 );
9372 i += 64;
9373 }
9374 while i + 16 <= n {
9375 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9376 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9377 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9378 i += 16;
9379 }
9380 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9381 while i < n {
9382 s += (*wp.add(i)) as i32 * xq[i] as i32;
9383 i += 1;
9384 }
9385 s
9386 }
9387}
9388
9389#[cfg(target_arch = "aarch64")]
9393#[target_feature(enable = "neon,dotprod")]
9394unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9395 unsafe {
9397 use core::arch::aarch64::*;
9398 use core::arch::asm;
9399 let n = xq.len();
9400 let px = xq.as_ptr();
9401 let (p0, p1, p2, p3) = (
9402 w0.as_ptr() as *const i8,
9403 w1.as_ptr() as *const i8,
9404 w2.as_ptr() as *const i8,
9405 w3.as_ptr() as *const i8,
9406 );
9407 let (mut a0, mut a1, mut a2, mut a3) = (
9408 vdupq_n_s32(0),
9409 vdupq_n_s32(0),
9410 vdupq_n_s32(0),
9411 vdupq_n_s32(0),
9412 );
9413 let mut i = 0;
9414 while i + 16 <= n {
9415 let x = vld1q_s8(px.add(i));
9416 let v0 = vld1q_s8(p0.add(i));
9417 let v1 = vld1q_s8(p1.add(i));
9418 let v2 = vld1q_s8(p2.add(i));
9419 let v3 = vld1q_s8(p3.add(i));
9420 asm!(
9421 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9422 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9423 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9424 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9425 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9426 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9427 options(pure, nomem, nostack),
9428 );
9429 i += 16;
9430 }
9431 let mut r = [
9432 vaddvq_s32(a0),
9433 vaddvq_s32(a1),
9434 vaddvq_s32(a2),
9435 vaddvq_s32(a3),
9436 ];
9437 while i < n {
9438 let xi = *px.add(i) as i32;
9439 r[0] += (*p0.add(i)) as i32 * xi;
9440 r[1] += (*p1.add(i)) as i32 * xi;
9441 r[2] += (*p2.add(i)) as i32 * xi;
9442 r[3] += (*p3.add(i)) as i32 * xi;
9443 i += 1;
9444 }
9445 r
9446 }
9447}
9448
9449#[cfg(target_arch = "aarch64")]
9456#[target_feature(enable = "neon,dotprod")]
9457unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9458 unsafe {
9461 use core::arch::aarch64::*;
9462 use core::arch::asm;
9463 let n = xq.len();
9464 let px = xq.as_ptr();
9465 let pg = g.as_ptr() as *const i8;
9466 let (mut a0, mut a1, mut a2, mut a3) = (
9467 vdupq_n_s32(0),
9468 vdupq_n_s32(0),
9469 vdupq_n_s32(0),
9470 vdupq_n_s32(0),
9471 );
9472 let mut i = 0;
9473 while i + 16 <= n {
9474 let x = vld1q_s8(px.add(i));
9475 let base = pg.add(4 * i);
9476 let v0 = vld1q_s8(base);
9477 let v1 = vld1q_s8(base.add(16));
9478 let v2 = vld1q_s8(base.add(32));
9479 let v3 = vld1q_s8(base.add(48));
9480 asm!(
9481 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9482 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9483 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9484 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9485 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9486 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9487 options(pure, nomem, nostack),
9488 );
9489 i += 16;
9490 }
9491 [
9492 vaddvq_s32(a0),
9493 vaddvq_s32(a1),
9494 vaddvq_s32(a2),
9495 vaddvq_s32(a3),
9496 ]
9497 }
9498}
9499
9500#[cfg(target_arch = "aarch64")]
9506fn q8_range_sdot(
9507 q: &[u8],
9508 rep: &[u8],
9509 row_scale: &[f32],
9510 act: &SplitAct,
9511 cols: usize,
9512 out_addr: SendMut,
9513 start: usize,
9514 end: usize,
9515) {
9516 let mut o = start;
9517 if !rep.is_empty() {
9520 while o < end && o % 4 != 0 {
9521 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9522 unsafe { *out_addr.at(o) = v };
9523 o += 1;
9524 }
9525 }
9526 while o + 4 <= end {
9527 let r = if rep.is_empty() {
9528 unsafe {
9529 dot_i8_sdot_4rows(
9530 &q[o * cols..(o + 1) * cols],
9531 &q[(o + 1) * cols..(o + 2) * cols],
9532 &q[(o + 2) * cols..(o + 3) * cols],
9533 &q[(o + 3) * cols..(o + 4) * cols],
9534 &act.xq,
9535 )
9536 }
9537 } else {
9538 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9539 };
9540 for k in 0..4 {
9541 let mut acc = r[k] as f32 * act.sx;
9542 for &(j, xv) in &act.outliers {
9543 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9544 }
9545 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9547 }
9548 o += 4;
9549 }
9550 while o < end {
9551 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9552 unsafe { *out_addr.at(o) = v };
9553 o += 1;
9554 }
9555}
9556
9557#[cfg(target_arch = "aarch64")]
9560#[allow(clippy::too_many_arguments)]
9561fn q8_range2_sdot(
9562 q: &[u8],
9563 row_scale: &[f32],
9564 a1: &SplitAct,
9565 a2: &SplitAct,
9566 cols: usize,
9567 p1: SendMut,
9568 p2: SendMut,
9569 start: usize,
9570 end: usize,
9571) {
9572 for o in start..end {
9573 let row = &q[o * cols..(o + 1) * cols];
9574 unsafe {
9576 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9577 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9578 }
9579 }
9580}
9581
9582#[allow(clippy::too_many_arguments)]
9584fn q8_range2_f32(
9585 q: &[u8],
9586 row_scale: &[f32],
9587 x1: &[f32],
9588 x2: &[f32],
9589 cols: usize,
9590 p1: SendMut,
9591 p2: SendMut,
9592 start: usize,
9593 end: usize,
9594) {
9595 for o in start..end {
9596 let row = &q[o * cols..(o + 1) * cols];
9597 unsafe {
9599 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9600 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9601 }
9602 }
9603}
9604
9605fn q8_range_f32(
9607 q: &[u8],
9608 row_scale: &[f32],
9609 xs: &[f32],
9610 cols: usize,
9611 out_addr: SendMut,
9612 start: usize,
9613 end: usize,
9614) {
9615 for o in start..end {
9616 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9617 unsafe { *out_addr.at(o) = v };
9619 }
9620}
9621
9622#[inline]
9626fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9627 #[cfg(target_arch = "aarch64")]
9628 return row_dot_sdot(row, act);
9629 #[cfg(target_arch = "x86_64")]
9630 return row_dot_avx2(row, act);
9631 #[allow(unreachable_code)]
9632 q8_row_dot_scalar(row, act)
9633}
9634
9635#[allow(dead_code)]
9636fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9637 let mut acc = 0i32;
9638 for (k, &b) in row.iter().enumerate() {
9639 acc += (b as i8) as i32 * act.xq[k] as i32;
9640 }
9641 let mut acc = acc as f32 * act.sx;
9642 for &(j, xv) in &act.outliers {
9643 acc += (row[j] as i8) as f32 * xv;
9644 }
9645 acc
9646}
9647
9648#[cfg(target_arch = "aarch64")]
9651#[inline]
9652fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9653 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9654 for &(j, xv) in &act.outliers {
9655 acc += (row[j] as i8) as f32 * xv;
9656 }
9657 acc
9658}
9659
9660#[cfg(target_arch = "aarch64")]
9668#[target_feature(enable = "neon,dotprod")]
9669unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9670 unsafe {
9673 use core::arch::aarch64::*;
9674 use core::arch::asm;
9675 let lomask = vdupq_n_u8(0x0F);
9676 let eight = vdupq_n_s8(8);
9677 let mut acc = 0f32;
9678 for gi in 0..gpr {
9679 let g = g0 + gi;
9680 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9681 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9682 let lo = vandq_u8(b, lomask);
9683 let hi = vshrq_n_u8::<4>(b);
9684 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9685 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9686 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9687 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9688 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9689 asm!(
9690 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9691 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9692 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9693 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9694 options(pure, nomem, nostack),
9695 );
9696 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9697 }
9698 acc
9699 }
9700}
9701
9702#[cfg(target_arch = "aarch64")]
9707#[target_feature(enable = "neon,dotprod")]
9708unsafe fn dot_q4_row_sdot2(
9709 packed: &[u8],
9710 scales: &[u8],
9711 g0: usize,
9712 gpr: usize,
9713 xq1: &[i8],
9714 xq2: &[i8],
9715) -> (f32, f32) {
9716 unsafe {
9719 use core::arch::aarch64::*;
9720 use core::arch::asm;
9721 let lomask = vdupq_n_u8(0x0F);
9722 let eight = vdupq_n_s8(8);
9723 let (mut acc1, mut acc2) = (0f32, 0f32);
9724 for gi in 0..gpr {
9725 let g = g0 + gi;
9726 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9727 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9728 let lo = vandq_u8(b, lomask);
9729 let hi = vshrq_n_u8::<4>(b);
9730 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9731 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9732 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9733 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9734 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9735 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9736 let (mut a0, mut a1, mut b0, mut b1) = (
9737 vdupq_n_s32(0),
9738 vdupq_n_s32(0),
9739 vdupq_n_s32(0),
9740 vdupq_n_s32(0),
9741 );
9742 asm!(
9743 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9744 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9745 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9746 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9747 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9748 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9749 e0 = in(vreg) e0, e1 = in(vreg) e1,
9750 x10 = in(vreg) x10, x11 = in(vreg) x11,
9751 x20 = in(vreg) x20, x21 = in(vreg) x21,
9752 options(pure, nomem, nostack),
9753 );
9754 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9755 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9756 }
9757 (acc1, acc2)
9758 }
9759}
9760
9761#[inline]
9766pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9767 #[cfg(target_arch = "aarch64")]
9768 unsafe {
9769 return axpy_i8_f32_neon(acc, row, w);
9770 }
9771 #[cfg(target_arch = "x86_64")]
9772 if avx2_enabled() {
9773 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9774 }
9775 #[allow(unreachable_code)]
9776 {
9777 for (a, &b) in acc.iter_mut().zip(row) {
9778 *a += w * b as f32;
9779 }
9780 }
9781}
9782
9783#[cfg(target_arch = "x86_64")]
9785#[target_feature(enable = "avx2,fma")]
9786unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9787 unsafe {
9789 use core::arch::x86_64::*;
9790 let n = acc.len().min(row.len());
9791 let ap = acc.as_mut_ptr();
9792 let rp = row.as_ptr();
9793 let wv = _mm256_set1_ps(w);
9794 let mut j = 0usize;
9795 while j + 16 <= n {
9796 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9797 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9798 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9799 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9800 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9801 _mm256_storeu_ps(ap.add(j), v0);
9802 _mm256_storeu_ps(ap.add(j + 8), v1);
9803 j += 16;
9804 }
9805 while j < n {
9806 *ap.add(j) += w * (*rp.add(j)) as f32;
9807 j += 1;
9808 }
9809 }
9810}
9811
9812#[cfg(target_arch = "aarch64")]
9813#[target_feature(enable = "neon")]
9814unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9815 unsafe {
9817 use core::arch::aarch64::*;
9818 let n = acc.len().min(row.len());
9819 let ap = acc.as_mut_ptr();
9820 let rp = row.as_ptr();
9821 let wv = vdupq_n_f32(w);
9822 let mut j = 0usize;
9823 while j + 16 <= n {
9824 let rb = vld1q_s8(rp.add(j));
9825 let lo = vmovl_s8(vget_low_s8(rb));
9826 let hi = vmovl_s8(vget_high_s8(rb));
9827 for (off, half) in [(0, lo), (8, hi)] {
9828 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9829 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9830 let o = j + off;
9831 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9832 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9833 }
9834 j += 16;
9835 }
9836 while j < n {
9837 *ap.add(j) += w * (*rp.add(j)) as f32;
9838 j += 1;
9839 }
9840 }
9841}
9842
9843#[inline]
9846pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9847 #[cfg(target_arch = "aarch64")]
9848 unsafe {
9849 return dot_i8_f32_neon(w, x);
9850 }
9851 #[cfg(target_arch = "x86_64")]
9852 if avx2_enabled() {
9853 return unsafe { dot_i8_f32_avx2(w, x) };
9854 }
9855 #[allow(unreachable_code)]
9856 {
9857 let mut sum = 0.0f32;
9858 for (j, &b) in w.iter().enumerate() {
9859 sum += (b as i8) as f32 * x[j];
9860 }
9861 sum
9862 }
9863}
9864
9865#[inline]
9869fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9870 #[cfg(target_arch = "aarch64")]
9871 unsafe {
9872 return dot_i8_col_f32_neon(w, x, col);
9873 }
9874 #[allow(unreachable_code)]
9875 {
9876 let mut sum = 0.0f32;
9877 for (j, &b) in w.iter().enumerate() {
9878 sum += (b as i8) as f32 * x[j] * col[j];
9879 }
9880 sum
9881 }
9882}
9883
9884#[cfg(target_arch = "aarch64")]
9885#[target_feature(enable = "neon")]
9886unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9887 unsafe {
9889 use core::arch::aarch64::*;
9890 let n = x.len();
9891 let wp = w.as_ptr() as *const i8;
9892 let xp = x.as_ptr();
9893 let cp = col.as_ptr();
9894 let (mut a0, mut a1, mut a2, mut a3) = (
9895 vdupq_n_f32(0.0),
9896 vdupq_n_f32(0.0),
9897 vdupq_n_f32(0.0),
9898 vdupq_n_f32(0.0),
9899 );
9900 let mut j = 0usize;
9901 while j + 16 <= n {
9902 let wb = vld1q_s8(wp.add(j));
9903 let lo = vmovl_s8(vget_low_s8(wb));
9904 let hi = vmovl_s8(vget_high_s8(wb));
9905 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9906 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9907 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9908 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9909 a0 = vfmaq_f32(
9910 a0,
9911 w0,
9912 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9913 );
9914 a1 = vfmaq_f32(
9915 a1,
9916 w1,
9917 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9918 );
9919 a2 = vfmaq_f32(
9920 a2,
9921 w2,
9922 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9923 );
9924 a3 = vfmaq_f32(
9925 a3,
9926 w3,
9927 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9928 );
9929 j += 16;
9930 }
9931 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9932 while j < n {
9933 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9934 j += 1;
9935 }
9936 sum
9937 }
9938}
9939
9940#[cfg(target_arch = "aarch64")]
9941#[target_feature(enable = "neon")]
9942unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9943 unsafe {
9945 use core::arch::aarch64::*;
9946 let n = x.len();
9947 let wp = w.as_ptr() as *const i8;
9948 let xp = x.as_ptr();
9949 let (mut a0, mut a1, mut a2, mut a3) = (
9950 vdupq_n_f32(0.0),
9951 vdupq_n_f32(0.0),
9952 vdupq_n_f32(0.0),
9953 vdupq_n_f32(0.0),
9954 );
9955 let mut j = 0usize;
9956 while j + 16 <= n {
9957 let wb = vld1q_s8(wp.add(j));
9958 let lo = vmovl_s8(vget_low_s8(wb));
9959 let hi = vmovl_s8(vget_high_s8(wb));
9960 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9961 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9962 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9963 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9964 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9965 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9966 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9967 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9968 j += 16;
9969 }
9970 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9971 while j < n {
9972 sum += (*wp.add(j)) as f32 * *xp.add(j);
9973 j += 1;
9974 }
9975 sum
9976 }
9977}
9978
9979#[allow(clippy::too_many_arguments)]
9980fn qmatvec(
9981 q: &[u8],
9982 rep: &[u8],
9983 row_scale: &[f32],
9984 x: &[f32],
9985 col_field: &[f32],
9986 dtype: TensorDtype,
9987 rows: usize,
9988 cols: usize,
9989 out: &mut [f32],
9990 pool: Option<&Pool>,
9991) {
9992 debug_assert_eq!(out.len(), rows);
9993 #[cfg(not(target_arch = "aarch64"))]
9994 let _ = rep;
9995
9996 #[cfg(target_arch = "aarch64")]
9997 if sdot_enabled() {
9998 let act = if dtype == TensorDtype::Q8_2f {
9999 split_act_q8_2f(x, col_field)
10000 } else {
10001 split_act(x)
10002 };
10003 let out_addr = SendMut(out.as_mut_ptr());
10004 let run_range = |start: usize, end: usize| {
10005 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
10006 };
10007 match pool {
10008 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10009 _ => run_range(0, rows),
10010 }
10011 return;
10012 }
10013 #[cfg(target_arch = "x86_64")]
10016 if avx2_a8w8_enabled() {
10017 let act = if dtype == TensorDtype::Q8_2f {
10018 split_act_q8_2f(x, col_field)
10019 } else {
10020 split_act(x)
10021 };
10022 let out_addr = SendMut(out.as_mut_ptr());
10023 let run_range = |start: usize, end: usize| {
10024 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
10025 };
10026 match pool {
10027 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10028 _ => run_range(0, rows),
10029 }
10030 return;
10031 }
10032
10033 prescale_with(x, col_field, dtype, 1, |xs| {
10034 let out_addr = SendMut(out.as_mut_ptr());
10035 let run_range = move |start: usize, end: usize| {
10036 for o in start..end {
10037 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10038 unsafe { *out_addr.at(o) = v };
10040 }
10041 };
10042 match pool {
10043 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10044 _ => run_range(0, rows),
10045 }
10046 });
10047}
10048
10049#[allow(clippy::too_many_arguments)]
10050fn qmatvec2(
10051 q: &[u8],
10052 row_scale: &[f32],
10053 x1: &[f32],
10054 x2: &[f32],
10055 col_field: &[f32],
10056 dtype: TensorDtype,
10057 rows: usize,
10058 cols: usize,
10059 o1: &mut [f32],
10060 o2: &mut [f32],
10061 pool: Option<&Pool>,
10062) {
10063 #[cfg(target_arch = "aarch64")]
10064 if sdot_enabled() {
10065 let a1s = if dtype == TensorDtype::Q8_2f {
10066 split_act_q8_2f(x1, col_field)
10067 } else {
10068 split_act(x1)
10069 };
10070 let a2s = if dtype == TensorDtype::Q8_2f {
10071 split_act_q8_2f(x2, col_field)
10072 } else {
10073 split_act(x2)
10074 };
10075 let p1 = SendMut(o1.as_mut_ptr());
10076 let p2 = SendMut(o2.as_mut_ptr());
10077 let run_range = |start: usize, end: usize| {
10078 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10079 };
10080 match pool {
10081 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10082 _ => run_range(0, rows),
10083 }
10084 return;
10085 }
10086 #[cfg(target_arch = "x86_64")]
10087 if avx2_a8w8_enabled() {
10088 let a1s = if dtype == TensorDtype::Q8_2f {
10089 split_act_q8_2f(x1, col_field)
10090 } else {
10091 split_act(x1)
10092 };
10093 let a2s = if dtype == TensorDtype::Q8_2f {
10094 split_act_q8_2f(x2, col_field)
10095 } else {
10096 split_act(x2)
10097 };
10098 let p1 = SendMut(o1.as_mut_ptr());
10099 let p2 = SendMut(o2.as_mut_ptr());
10100 let run_range = |start: usize, end: usize| {
10101 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10102 };
10103 match pool {
10104 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10105 _ => run_range(0, rows),
10106 }
10107 return;
10108 }
10109
10110 prescale_with(x1, col_field, dtype, 1, |x1s| {
10111 prescale_with(x2, col_field, dtype, 2, |x2s| {
10112 let p1 = SendMut(o1.as_mut_ptr());
10113 let p2 = SendMut(o2.as_mut_ptr());
10114 let run_range = move |start: usize, end: usize| {
10115 for o in start..end {
10116 let row = &q[o * cols..(o + 1) * cols];
10117 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10118 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10119 unsafe {
10121 *p1.at(o) = s1;
10122 *p2.at(o) = s2;
10123 }
10124 }
10125 };
10126 match pool {
10127 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10128 _ => run_range(0, rows),
10129 }
10130 });
10131 });
10132}
10133
10134#[derive(Clone, Copy)]
10135struct SendMut(*mut f32);
10136unsafe impl Send for SendMut {}
10137unsafe impl Sync for SendMut {}
10138
10139impl SendMut {
10140 #[inline]
10141 fn at(self, i: usize) -> *mut f32 {
10142 unsafe { self.0.add(i) }
10143 }
10144}
10145
10146#[cfg(test)]
10147mod tests {
10148 use super::*;
10149
10150 #[test]
10151 fn q2tp_i8_dot_matches_exact_on_grid() {
10152 let (rows, cols) = (5, 64);
10156 let gpr = cols / GROUP_SIZE;
10157 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10161 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10162 .collect();
10163 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10164 let x: Vec<f32> = (0..cols)
10165 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10166 .collect();
10167 let act = split_act(&x);
10168 assert!(
10169 act.outliers.is_empty(),
10170 "on-grid input must have no outliers"
10171 );
10172 let gsum = q1_group_sums(&act.xq, gpr);
10173 for r in 0..rows {
10174 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10175 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10176 assert!(
10177 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10178 "row {r}: exact {exact} vs i8 {fast}"
10179 );
10180 }
10181 }
10182
10183 #[test]
10184 fn q8_row_dot_fast_matches_scalar() {
10185 let cols = 96;
10188 let row: Vec<u8> = (0..cols)
10189 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10190 .collect();
10191 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10192 let act = split_act(&x);
10193 let fast = q8_row_dot(&row, &act);
10194 let scalar = q8_row_dot_scalar(&row, &act);
10195 assert!(
10196 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10197 "fast {fast} vs scalar {scalar}"
10198 );
10199 }
10200
10201 #[test]
10202 fn f32_matvec_matches_matvec_rows_bitexact() {
10203 let (rows, cols) = (300, 40);
10204 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10205 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10206 let qt = QTensor::from_f32(w.clone(), rows, cols);
10207
10208 let mut a = vec![0.0f32; rows];
10209 matvec_rows(None, &w, &x, &mut a);
10210 let mut b = vec![0.0f32; rows];
10211 qt.matvec(&x, &mut b, None);
10212 assert_eq!(a, b);
10213 }
10214
10215 #[test]
10216 fn sdot_kernel_exact_on_grid() {
10217 eprintln!("sdot_enabled = {}", sdot_enabled());
10222 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
10224 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10225 .collect();
10226 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10227 let x: Vec<f32> = (0..cols)
10228 .map(|i| match i % 3 {
10229 0 => 1.0,
10230 1 => -1.0,
10231 _ => 0.0,
10232 })
10233 .collect();
10234 let mut a = vec![0.0f32; rows];
10235 qmatvec(
10236 &w,
10237 &[],
10238 &scales,
10239 &x,
10240 &[],
10241 TensorDtype::Q8Row,
10242 rows,
10243 cols,
10244 &mut a,
10245 None,
10246 );
10247 for o in 0..rows {
10248 let mut acc = 0.0f32;
10249 for j in 0..cols {
10250 acc += (w[o * cols + j] as i8) as f32 * x[j];
10251 }
10252 let expect = acc * scales[o];
10253 assert!(
10254 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10255 "row {o}: {} vs {expect}",
10256 a[o]
10257 );
10258 }
10259 }
10260
10261 #[test]
10262 fn q1_tbl_fast_path_matches_reference() {
10263 let (rows, cols) = (5, 256);
10268 let gpr = cols / GROUP_SIZE;
10269 let mut bytes = Vec::new();
10270 for t in 0..rows * gpr {
10271 let s = 0.007 + (t % 11) as f32 * 0.004;
10272 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10273 for j in 0..4 {
10274 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10275 }
10276 }
10277 let x: Vec<f32> = (0..cols)
10278 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10279 .collect();
10280 let mut w = vec![0.0f32; rows * cols];
10281 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10282 let mut got = vec![0.0f32; rows];
10283 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10284 for o in 0..rows {
10285 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10286 assert!(
10287 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10288 "row {o}: {} vs {expect}",
10289 got[o]
10290 );
10291 }
10292 let b = 5usize;
10295 let mut xs_all = Vec::new();
10296 for bi in 0..b {
10297 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10298 }
10299 let mut mm = vec![0.0f32; b * rows];
10300 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10301 for bi in 0..b {
10302 let mut single = vec![0.0f32; rows];
10303 q1_matvec(
10304 &bytes,
10305 &xs_all[bi * cols..(bi + 1) * cols],
10306 rows,
10307 cols,
10308 &mut single,
10309 None,
10310 );
10311 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10312 }
10313 }
10314
10315 #[test]
10316 fn q1_kernels_match_exact_reference() {
10317 let (rows, cols) = (7, 96);
10319 let gpr = cols / GROUP_SIZE;
10320 let mut bytes = Vec::new();
10321 for t in 0..rows * gpr {
10322 let s = 0.01 + (t % 13) as f32 * 0.003;
10323 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10324 for j in 0..4 {
10325 bytes.push(((t * 31 + j * 97) % 251) as u8);
10326 }
10327 }
10328 let x: Vec<f32> = (0..cols)
10330 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10331 .collect();
10332 let mut w = vec![0.0f32; rows * cols];
10334 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10335 let mut expect = vec![0.0f32; rows];
10336 for o in 0..rows {
10337 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10338 }
10339 let mut got = vec![0.0f32; rows];
10340 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10341 for o in 0..rows {
10342 assert!(
10343 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10344 "row {o}: {} vs {}",
10345 got[o],
10346 expect[o]
10347 );
10348 }
10349 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10351 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10352 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10353 assert_eq!(a1, got);
10354 let mut xs = x.clone();
10355 xs.extend_from_slice(&x2);
10356 let mut mm = vec![0.0f32; 2 * rows];
10357 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10358 assert_eq!(&mm[..rows], got.as_slice());
10359 assert_eq!(&mm[rows..], a2.as_slice());
10360 }
10361
10362 #[test]
10363 fn repack_is_bit_identical() {
10364 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
10370 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10371 .collect();
10372 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10373 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10374 let rep = q8_repack_layout(&w, rows, cols);
10375 for g in 0..rows / 4 {
10377 for c in 0..cols / 16 {
10378 for lane in 0..4 {
10379 assert_eq!(
10380 &rep[g * 4 * cols + c * 64 + lane * 16
10381 ..g * 4 * cols + c * 64 + lane * 16 + 16],
10382 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10383 );
10384 }
10385 }
10386 }
10387 let mut a = vec![0.0f32; rows];
10388 qmatvec(
10389 &w,
10390 &[],
10391 &scales,
10392 &x,
10393 &[],
10394 TensorDtype::Q8Row,
10395 rows,
10396 cols,
10397 &mut a,
10398 None,
10399 );
10400 let mut b = vec![0.0f32; rows];
10401 qmatvec(
10402 &w,
10403 &rep,
10404 &scales,
10405 &x,
10406 &[],
10407 TensorDtype::Q8Row,
10408 rows,
10409 cols,
10410 &mut b,
10411 None,
10412 );
10413 assert_eq!(a, b, "full-range repack output diverged");
10414
10415 #[cfg(target_arch = "aarch64")]
10416 if sdot_enabled() {
10417 let act = split_act(&x);
10419 let mut c1 = vec![0.0f32; rows];
10420 let mut c2 = vec![0.0f32; rows];
10421 q8_range_sdot(
10422 &w,
10423 &[],
10424 &scales,
10425 &act,
10426 cols,
10427 SendMut(c1.as_mut_ptr()),
10428 3,
10429 rows - 2,
10430 );
10431 q8_range_sdot(
10432 &w,
10433 &rep,
10434 &scales,
10435 &act,
10436 cols,
10437 SendMut(c2.as_mut_ptr()),
10438 3,
10439 rows - 2,
10440 );
10441 assert_eq!(c1, c2, "unaligned-range repack output diverged");
10442 }
10443 }
10444
10445 #[test]
10446 fn sdot_a8w8_noise_is_bounded() {
10447 let (rows, cols) = (16, 512);
10451 let w: Vec<u8> = (0..rows * cols)
10452 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10453 .collect();
10454 let scales = vec![0.01f32; rows];
10455 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10456 let mut a = vec![0.0f32; rows];
10457 qmatvec(
10458 &w,
10459 &[],
10460 &scales,
10461 &x,
10462 &[],
10463 TensorDtype::Q8Row,
10464 rows,
10465 cols,
10466 &mut a,
10467 None,
10468 );
10469 let (mut num, mut den) = (0f64, 0f64);
10470 for o in 0..rows {
10471 let mut acc = 0.0f32;
10472 for j in 0..cols {
10473 acc += (w[o * cols + j] as i8) as f32 * x[j];
10474 }
10475 let expect = acc * scales[o];
10476 num += ((a[o] - expect) as f64).powi(2);
10477 den += (expect as f64).powi(2);
10478 }
10479 let rel = (num / den.max(1e-12)).sqrt();
10480 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10481 }
10482
10483 #[test]
10484 fn i8_dot_neon_matches_scalar() {
10485 let n = 100;
10486 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10487 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10488 let mut scalar = 0.0f32;
10489 for j in 0..n {
10490 scalar += (w[j] as i8) as f32 * x[j];
10491 }
10492 let fast = dot_i8_f32(&w, &x);
10493 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10494 }
10495
10496 #[test]
10498 fn vbitmatvec_matches_full_dequant() {
10499 let (rows, cols) = (6, 64);
10500 let ng = cols / GROUP_SIZE;
10501 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10503 let mut bytes = bits.clone();
10504 for g in 0..rows * ng {
10505 let s = 0.02 + 0.001 * g as f32;
10506 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10507 }
10508 for r in 0..rows {
10509 let b = bits[r] as usize;
10510 let (mut acc, mut nb) = (0u64, 0usize);
10511 let mut rowbytes = Vec::new();
10512 for i in 0..cols {
10513 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10514 acc = (acc << b) | v;
10515 nb += b;
10516 while nb >= 8 {
10517 nb -= 8;
10518 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10519 }
10520 }
10521 if nb > 0 {
10522 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10523 }
10524 bytes.extend_from_slice(&rowbytes);
10525 }
10526 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10527
10528 let mut reference = vec![0f32; rows * cols];
10529 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10530 let mut expect = vec![0f32; rows];
10531 for r in 0..rows {
10532 expect[r] = reference[r * cols..(r + 1) * cols]
10533 .iter()
10534 .zip(&x)
10535 .map(|(w, xv)| w * xv)
10536 .sum();
10537 }
10538 let mut got = vec![0f32; rows];
10539 let offsets = vbit_row_offsets(&bytes, rows, cols);
10540 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10541 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10545 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10546 for r in 0..rows {
10547 assert!(
10548 (got[r] - expect[r]).abs() < tol * scale,
10549 "row {r}: {} vs {}",
10550 got[r],
10551 expect[r]
10552 );
10553 }
10554 }
10555
10556 #[test]
10561 #[cfg(target_arch = "x86_64")]
10562 fn vbit_matmat_blocked_matches_per_row() {
10563 let (rows, cols, b) = (64usize, 128usize, 9usize);
10564 let ng = cols / GROUP_SIZE;
10565 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10566 let mut bytes = bits.clone();
10567 for g in 0..rows * ng {
10568 let sc = 0.02 + 0.0005 * g as f32;
10569 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10570 }
10571 for r in 0..rows {
10572 let bw = bits[r] as usize;
10573 let (mut acc, mut nb) = (0u64, 0usize);
10574 let mut rowbytes = Vec::new();
10575 for i in 0..cols {
10576 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10577 acc = (acc << bw) | v;
10578 nb += bw;
10579 while nb >= 8 {
10580 nb -= 8;
10581 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10582 }
10583 }
10584 if nb > 0 {
10585 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10586 }
10587 bytes.extend_from_slice(&rowbytes);
10588 }
10589 let x: Vec<f32> = (0..b * cols)
10590 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10591 .collect();
10592 let offsets = vbit_row_offsets(&bytes, rows, cols);
10593 let mut y_a = vec![0f32; b * rows];
10594 let mut y_b = vec![0f32; b * rows];
10595 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10596 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10597 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10598 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10599 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10600 let max_d = y_a
10601 .iter()
10602 .zip(&y_b)
10603 .map(|(p, q)| (p - q).abs())
10604 .fold(0.0f32, f32::max);
10605 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10606 }
10607
10608 #[test]
10616 fn q4t_matmat_blocked_matches_per_row() {
10617 let (rows, cols, b) = (16usize, 64usize, 9usize);
10618 let gpr = cols / GROUP_SIZE;
10619 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10620 for r in 0..rows {
10621 for g in 0..gpr {
10622 let t = (r * gpr + g) * Q4_TILE;
10623 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10624 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10625 for k in 0..16 {
10626 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10627 }
10628 }
10629 }
10630 let x: Vec<f32> = (0..b * cols)
10631 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10632 .collect();
10633 let mut y_blk = vec![0f32; b * rows];
10634 let mut y_row = vec![0f32; b * rows];
10635 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10636 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10637 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10638 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10639 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10640 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10641 }
10642
10643 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10650 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10651 let gpr = cols / GROUP_SIZE;
10652 let stride = q4tp_code_stride(gpr);
10653 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10654 let mut b = vec![0u8; codes_off + rows * stride];
10655 for r in 0..rows {
10656 for g in 0..gpr {
10657 let t = (r * gpr + g) * Q4TP_NIB;
10658 for k in 0..16 {
10659 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10660 }
10661 }
10662 let lo = -6.0 - 0.03 * (r % 17) as f32;
10663 let step = 0.01 + 0.004 * (r % 11) as f32;
10664 let p = params_off + r * 4;
10665 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10666 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10667 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10668 for g in 0..gpr {
10669 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10670 }
10671 }
10672 b
10673 }
10674
10675 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10679 let gpr = cols / GROUP_SIZE;
10680 let v = Q4tpView::new(bytes, rows, cols);
10681 let mut out = vec![0u8; rows * gpr * Q4_TILE];
10682 let mut sc = vec![0f32; gpr];
10683 for r in 0..rows {
10684 v.scales_into(r, gpr, &mut sc);
10685 for g in 0..gpr {
10686 let t = (r * gpr + g) * Q4_TILE;
10687 let s = sc[g];
10688 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10689 let src = (r * gpr + g) * Q4TP_NIB;
10690 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10691 }
10692 }
10693 out
10694 }
10695
10696 #[test]
10702 fn q4tp_exact_path_matches_dequant_reference() {
10703 let (rows, cols) = (256usize, 512usize);
10704 let gpr = cols / GROUP_SIZE;
10705 let bytes = synth_q4tp(rows, cols);
10706 let mut w = vec![0f32; rows * cols];
10707 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10708
10709 let x: Vec<f32> = (0..cols)
10710 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10711 .collect();
10712 let v = Q4tpView::new(&bytes, rows, cols);
10713 let mut sc = vec![0f32; gpr];
10714 for r in 0..rows {
10715 v.scales_into(r, gpr, &mut sc);
10716 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10717 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10718 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10722 assert!(
10723 (got - want).abs() <= 1e-5 * mag,
10724 "row {r}: kernel {got} vs dequant {want}"
10725 );
10726 }
10727 }
10728
10729 #[test]
10735 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10736 let (rows, cols) = (256usize, 512usize);
10737 let bytes = synth_q4tp(rows, cols);
10738 let twin = q4tp_as_q4t(&bytes, rows, cols);
10739 let x: Vec<f32> = (0..cols)
10740 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10741 .collect();
10742
10743 let mut got = vec![0f32; rows];
10744 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10745 let mut want = vec![0f32; rows];
10746 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10747
10748 let mut w = vec![0f32; rows * cols];
10751 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10752 for r in 0..rows {
10753 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10754 assert!(
10755 (got[r] - want[r]).abs() <= 1e-3 * mag,
10756 "row {r}: q4tp {} vs q4t {}",
10757 got[r],
10758 want[r]
10759 );
10760 }
10761 }
10762
10763 #[test]
10768 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10769 let (rows, cols, b) = (256usize, 512usize, 5usize);
10770 let bytes = synth_q4tp(rows, cols);
10771 let twin = q4tp_as_q4t(&bytes, rows, cols);
10772 let xs: Vec<f32> = (0..b * cols)
10773 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10774 .collect();
10775
10776 let mut got = vec![0f32; b * rows];
10777 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10778 let mut want = vec![0f32; b * rows];
10779 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10780
10781 let mut w = vec![0f32; rows * cols];
10782 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10783 for t in 0..b {
10784 for r in 0..rows {
10785 let mag: f32 = (0..cols)
10786 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10787 .sum();
10788 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10789 assert!(
10790 (g - wa).abs() <= 1e-3 * mag,
10791 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10792 );
10793 }
10794 }
10795 }
10796
10797 #[test]
10798 fn q4tp_matvec2_matches_the_single_stream_kernel() {
10799 let (rows, cols) = (128usize, 256usize);
10800 let gpr = cols / GROUP_SIZE;
10801 let bytes = synth_q4tp(rows, cols);
10802 let xs: Vec<f32> = (0..2 * cols)
10803 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10804 .collect();
10805
10806 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10807 q4tp_matvec2(
10808 &bytes,
10809 &xs[..cols],
10810 &xs[cols..],
10811 rows,
10812 cols,
10813 &mut o1,
10814 &mut o2,
10815 None,
10816 );
10817
10818 let v = Q4tpView::new(&bytes, rows, cols);
10821 let mut sc = vec![0f32; gpr];
10822 for r in 0..rows {
10823 v.scales_into(r, gpr, &mut sc);
10824 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10825 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10826 }
10827 }
10828
10829 #[test]
10836 fn q4tp_matvec_keeps_pace_with_q4t() {
10837 let (rows, cols) = (4096usize, 3072usize);
10838 let bytes = synth_q4tp(rows, cols);
10839 let twin = q4tp_as_q4t(&bytes, rows, cols);
10840 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10841 let mut o = vec![0f32; rows];
10842 let n = 12;
10843 let mut best = (f64::MAX, f64::MAX);
10844 for _ in 0..3 {
10847 let t0 = std::time::Instant::now();
10848 for _ in 0..n {
10849 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10850 }
10851 best.0 = best.0.min(t0.elapsed().as_secs_f64());
10852 let t0 = std::time::Instant::now();
10853 for _ in 0..n {
10854 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10855 }
10856 best.1 = best.1.min(t0.elapsed().as_secs_f64());
10857 }
10858 let ratio = best.1 / best.0;
10859 println!(
10860 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10861 best.0 * 1e3 / n as f64,
10862 best.1 * 1e3 / n as f64
10863 );
10864 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10865 }
10866
10867 #[cfg(target_os = "macos")]
10868 #[test]
10869 fn q4t_matmat_accel_matches_dequant_reference() {
10870 if !accel_gemm_enabled() {
10871 return; }
10873 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
10875 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10876 for r in 0..rows {
10877 for g in 0..gpr {
10878 let t = (r * gpr + g) * Q4_TILE;
10879 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10880 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10881 for k in 0..16 {
10882 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10883 }
10884 }
10885 }
10886 let x: Vec<f32> = (0..b * cols)
10887 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10888 .collect();
10889 let mut got = vec![0f32; b * rows];
10890 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10891 let mut w = vec![0f32; rows * cols];
10893 for r in 0..rows {
10894 for g in 0..gpr {
10895 let t = (r * gpr + g) * Q4_TILE;
10896 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10897 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10898 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10899 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10900 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10901 }
10902 }
10903 }
10904 for bi in 0..b {
10905 for r in 0..rows {
10906 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10907 let d = (got[bi * rows + r] - want).abs();
10908 assert!(
10909 d <= want.abs().max(1.0) * 1e-4,
10910 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10911 got[bi * rows + r]
10912 );
10913 }
10914 }
10915 }
10916
10917 #[test]
10918 fn q4matvec_matches_full_dequant() {
10919 let (rows, cols) = (8, 64);
10920 let groups = rows * cols / GROUP_SIZE;
10921 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10923 for i in 0..groups * 16 {
10924 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10925 }
10926 for g in 0..groups {
10927 let s = 0.01 + 0.003 * g as f32;
10928 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10929 }
10930 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10931
10932 let mut reference = vec![0.0f32; rows * cols];
10933 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10934 let mut expect = vec![0.0f32; rows];
10935 for r in 0..rows {
10936 expect[r] = reference[r * cols..(r + 1) * cols]
10937 .iter()
10938 .zip(&x)
10939 .map(|(w, xv)| w * xv)
10940 .sum();
10941 }
10942
10943 let mut got = vec![0.0f32; rows];
10944 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10945 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10949 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10950 for r in 0..rows {
10951 assert!(
10952 (got[r] - expect[r]).abs() < tol * scale,
10953 "row {r}: {} vs {}",
10954 got[r],
10955 expect[r]
10956 );
10957 }
10958 }
10959
10960 #[test]
10963 fn vbitmatvec2_equals_two_singles() {
10964 let (rows, cols) = (6, 64);
10965 let ng = cols / GROUP_SIZE;
10966 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10967 let mut bytes = bits.clone();
10968 for g in 0..rows * ng {
10969 let s = 0.02 + 0.001 * g as f32;
10970 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10971 }
10972 for r in 0..rows {
10973 let b = bits[r] as usize;
10974 let (mut acc, mut nb) = (0u64, 0usize);
10975 let mut rowbytes = Vec::new();
10976 for i in 0..cols {
10977 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10978 acc = (acc << b) | v;
10979 nb += b;
10980 while nb >= 8 {
10981 nb -= 8;
10982 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10983 }
10984 }
10985 if nb > 0 {
10986 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10987 }
10988 bytes.extend_from_slice(&rowbytes);
10989 }
10990 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10991 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10992 let offsets = vbit_row_offsets(&bytes, rows, cols);
10993
10994 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10995 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10996 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10997 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10998 vbitmatvec2(
10999 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
11000 );
11001 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
11002 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
11003 }
11004
11005 #[test]
11007 fn q4matvec2_equals_two_singles() {
11008 let (rows, cols) = (8, 128);
11009 let groups = rows * cols / GROUP_SIZE;
11010 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11011 for i in 0..groups * 16 {
11012 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11013 }
11014 for g in 0..groups {
11015 let s = 0.01 + 0.003 * g as f32;
11016 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11017 }
11018 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11021 x1[9] = 250.0;
11022 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11023
11024 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11025 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
11026 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
11027 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11028 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
11029 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
11030 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
11031 }
11032
11033 #[test]
11036 fn matvec_many_equals_separate_matvecs() {
11037 use crate::pool::Pool;
11038 let (r1, r2, cols) = (300, 200, 64);
11039 let mk = |salt: usize, rows: usize| {
11040 QTensor::from_f32(
11041 (0..rows * cols)
11042 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
11043 .collect(),
11044 rows,
11045 cols,
11046 )
11047 };
11048 let (a, b) = (mk(1, r1), mk(5, r2));
11049 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
11050 let pool = Pool::new(3);
11051
11052 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
11053 a.matvec(&x, &mut ea, Some(&pool));
11054 b.matvec(&x, &mut eb, Some(&pool));
11055 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
11056 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11057 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
11058 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
11059 }
11060
11061 #[test]
11064 fn batched_matmat_equals_per_position_matvec() {
11065 let (rows, cols, b) = (8, 64, 5);
11066 let groups = rows * cols / GROUP_SIZE;
11068 let mut q4 = Vec::new();
11069 for i in 0..groups * 16 {
11070 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11071 }
11072 for g in 0..groups {
11073 q4.extend_from_slice(
11074 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11075 );
11076 }
11077 let ng = cols / GROUP_SIZE;
11079 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
11080 let mut vb = bits.clone();
11081 for g in 0..rows * ng {
11082 vb.extend_from_slice(
11083 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
11084 );
11085 }
11086 for r in 0..rows {
11087 let bw = bits[r] as usize;
11088 let (mut acc, mut nb) = (0u64, 0usize);
11089 let mut rowbytes = Vec::new();
11090 for i in 0..cols {
11091 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11092 acc = (acc << bw) | v;
11093 nb += bw;
11094 while nb >= 8 {
11095 nb -= 8;
11096 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11097 }
11098 }
11099 if nb > 0 {
11100 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11101 }
11102 vb.extend_from_slice(&rowbytes);
11103 }
11104 let offsets = vbit_row_offsets(&vb, rows, cols);
11105
11106 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11107
11108 let mut got = vec![0f32; b * rows];
11110 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
11111 for bi in 0..b {
11112 let mut expect = vec![0f32; rows];
11113 q4matvec(
11114 &q4,
11115 &xs[bi * cols..(bi + 1) * cols],
11116 rows,
11117 cols,
11118 &mut expect,
11119 None,
11120 );
11121 assert_eq!(
11122 &got[bi * rows..(bi + 1) * rows],
11123 &expect[..],
11124 "q4 batch pos {bi}"
11125 );
11126 }
11127
11128 let mut got = vec![0f32; b * rows];
11130 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
11131 for bi in 0..b {
11132 let mut expect = vec![0f32; rows];
11133 vbitmatvec(
11134 &vb,
11135 &offsets,
11136 &xs[bi * cols..(bi + 1) * cols],
11137 rows,
11138 cols,
11139 &mut expect,
11140 None,
11141 );
11142 assert_eq!(
11143 &got[bi * rows..(bi + 1) * rows],
11144 &expect[..],
11145 "vbit batch pos {bi}"
11146 );
11147 }
11148 }
11149
11150 #[test]
11154 fn q4_tiled_matches_q4_block_bitexact() {
11155 let (rows, cols, b) = (8usize, 128usize, 3usize);
11156 let groups = rows * cols / GROUP_SIZE;
11157 let mut split = Vec::with_capacity(groups * 18);
11158 for i in 0..groups * 16 {
11159 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11160 }
11161 for g in 0..groups {
11162 split.extend_from_slice(
11163 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11164 );
11165 }
11166 let (packed, scales) = split.split_at(groups * 16);
11168 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
11169 for g in 0..groups {
11170 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
11171 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
11172 }
11173
11174 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11175 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11177
11178 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
11179 q4matvec(&split, &x1, rows, cols, &mut a, None);
11180 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
11181 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11182
11183 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11184 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11185 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11186 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11187 assert_eq!(a1, t1);
11188 assert_eq!(a2, t2);
11189
11190 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11191 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11192 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11193 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11194 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11195 }
11196
11197 #[test]
11204 fn q4matvec_sdot_outlier_exact() {
11205 let (rows, cols) = (4, 128);
11206 let groups = rows * cols / GROUP_SIZE;
11207 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11208 for i in 0..groups * 16 {
11209 bytes.push(((i * 11 + 5) % 256) as u8);
11210 }
11211 for g in 0..groups {
11212 let s = 0.02 + 0.002 * g as f32;
11213 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11214 }
11215 let mut x: Vec<f32> = (0..cols)
11216 .map(|i| match i % 3 {
11217 0 => 1.0,
11218 1 => -1.0,
11219 _ => 0.0,
11220 })
11221 .collect();
11222 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
11225 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11226 let mut expect = vec![0.0f32; rows];
11227 for r in 0..rows {
11228 expect[r] = reference[r * cols..(r + 1) * cols]
11229 .iter()
11230 .zip(&x)
11231 .map(|(w, xv)| w * xv)
11232 .sum();
11233 }
11234 let mut got = vec![0.0f32; rows];
11235 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11236 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11237 for r in 0..rows {
11238 assert!(
11239 (got[r] - expect[r]).abs() < 2e-3 * scale,
11240 "row {r}: {} vs {} (outlier term must be exact)",
11241 got[r],
11242 expect[r]
11243 );
11244 }
11245 }
11246
11247 #[test]
11251 fn q1t_matvec_matches_reference() {
11252 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11253 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
11255 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11256 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11258 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11259 let mut bytes = Vec::new();
11260 for r in 0..rows {
11261 for g in 0..gpr {
11262 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11263 let mut c = [0u8; 7];
11264 for k in 0..GROUP_SIZE {
11265 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11267 0
11268 } else {
11269 ((k + r * 3 + g) % 3) as u8 };
11271 cortiq_core::quant::q1t_pack(&mut c, k, code);
11272 }
11273 bytes.extend_from_slice(&c);
11274 }
11275 }
11276 let mut row_ptr = vec![0u32; rows + 1];
11279 for &(idx, _) in &outliers {
11280 row_ptr[idx as usize / cols + 1] += 1;
11281 }
11282 for r in 0..rows {
11283 row_ptr[r + 1] += row_ptr[r];
11284 }
11285 for &p in &row_ptr {
11286 bytes.extend_from_slice(&p.to_le_bytes());
11287 }
11288 for &(idx, v) in &outliers {
11289 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11290 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11291 }
11292
11293 let mut refw = vec![0f32; rows * cols];
11294 dequant_q1t(&bytes, rows, cols, &mut refw);
11295 let x: Vec<f32> = (0..cols)
11298 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11299 .collect();
11300 let mut expect = vec![0f32; rows];
11301 for r in 0..rows {
11302 let mut a = 0.0f32;
11303 for j in 0..cols {
11304 a += refw[r * cols + j] * x[j];
11305 }
11306 expect[r] = a;
11307 }
11308 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11309 let mut got = vec![0f32; rows];
11310 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11311 for r in 0..rows {
11312 assert!(
11313 (got[r] - expect[r]).abs() < tol(expect[r]),
11314 "row {r}: {} vs {}",
11315 got[r],
11316 expect[r]
11317 );
11318 }
11319 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
11321 let mut gm = vec![0f32; 2 * rows];
11322 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11323 for r in 0..rows {
11324 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11325 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11326 }
11327 let xb: Vec<f32> = (0..cols)
11331 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11332 .collect();
11333 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11334 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11335 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11336 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11337 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11338 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11339 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11340 }
11341
11342 #[test]
11345 fn q1t_matvec2_odd_gpr_matches_singles() {
11346 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11347 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
11349 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11350 for r in 0..rows {
11351 for g in 0..gpr {
11352 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11353 let mut c = [0u8; 7];
11354 for k in 0..GROUP_SIZE {
11355 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11356 }
11357 bytes.extend_from_slice(&c);
11358 }
11359 }
11360 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11361 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11362 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11363 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11364 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11365 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11366 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11367 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11368 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11369 }
11370
11371 #[test]
11375 #[ignore]
11376 fn q1t_matvec2_speed() {
11377 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11378 use std::time::Instant;
11379 let (rows, cols) = (8192usize, 4096usize);
11380 let gpr = cols / GROUP_SIZE;
11381 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11382 for r in 0..rows {
11383 for g in 0..gpr {
11384 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11385 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11386 let mut c = [0u8; 7];
11387 for k in 0..GROUP_SIZE {
11388 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11389 }
11390 bytes.extend_from_slice(&c);
11391 }
11392 }
11393 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11394 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11395 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11396 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11397 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11399 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11400 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11401 for _ in 0..8 {
11402 let t0 = Instant::now();
11403 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11404 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11405 let t1 = Instant::now();
11406 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11407 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11408 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11409 }
11410 assert_eq!(p1, s1);
11411 assert_eq!(p2, s2);
11412 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11413 }
11414
11415 #[test]
11419 #[ignore]
11420 fn q1t_matvec_speed() {
11421 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11422 use std::time::Instant;
11423 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
11425 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11426 for r in 0..rows {
11427 for g in 0..gpr {
11428 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11429 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11430 let mut c = [0u8; 7];
11431 for k in 0..GROUP_SIZE {
11432 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11433 }
11434 bytes.extend_from_slice(&c);
11435 }
11436 }
11437 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
11439 let mut idx = 0usize;
11440 while idx < n {
11441 row_ptr[idx / cols + 1] += 1;
11442 idx += stride;
11443 }
11444 for r in 0..rows {
11445 row_ptr[r + 1] += row_ptr[r];
11446 }
11447 for &p in &row_ptr {
11448 bytes.extend_from_slice(&p.to_le_bytes());
11449 }
11450 let mut idx = 0usize;
11451 while idx < n {
11452 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11453 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11454 idx += stride;
11455 }
11456 let x: Vec<f32> = (0..cols)
11459 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11460 .collect();
11461 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11462
11463 let slow = |out: &mut [f32]| {
11465 let mut buf = vec![0f32; cols];
11466 for r in 0..rows {
11467 for g in 0..gpr {
11468 let off = (r * gpr + g) * Q1T_TILE;
11469 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11470 let codes = &bytes[off + 2..off + Q1T_TILE];
11471 for k in 0..GROUP_SIZE {
11472 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11473 1 => s,
11474 2 => -s,
11475 _ => 0.0,
11476 };
11477 }
11478 }
11479 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11480 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11481 }
11482 };
11483 let iters = 5;
11484 let mut a = vec![0f32; rows];
11485 slow(&mut a); let t = Instant::now();
11487 for _ in 0..iters {
11488 slow(&mut a);
11489 }
11490 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11491
11492 let mut b = vec![0f32; rows];
11493 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
11495 for _ in 0..iters {
11496 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11497 }
11498 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11499
11500 for r in 0..rows {
11501 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11502 }
11503 println!(
11504 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
11505 slow_ms / fast_ms
11506 );
11507 }
11508}
11509
11510#[cfg(test)]
11511mod gemm_bench {
11512 #[test]
11523 #[ignore]
11524 fn q4tp_matmat_throughput() {
11525 let b: usize = std::env::var("CMF_BENCH_B")
11529 .ok()
11530 .and_then(|v| v.parse().ok())
11531 .unwrap_or(296);
11532 let (rows, cols) = (9216usize, 2304usize);
11533 let (_, _, _) = (rows, cols, b);
11534 let total =
11535 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
11536 .unwrap();
11537 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11542 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11543 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11544 let step = cortiq_core::quant::f32_to_f16(0.1);
11545 for r in 0..rows {
11546 let o = params_off + r * 4;
11547 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11548 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11549 }
11550 let _ = codes_off;
11551 let xs: Vec<f32> = (0..b * cols)
11552 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11553 .collect();
11554 let mut out = vec![0f32; b * rows];
11555 let pool = crate::pool::Pool::from_env();
11556 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11562 let reps: usize = std::env::var("CMF_BENCH_REPS")
11563 .ok()
11564 .and_then(|v| v.parse().ok())
11565 .unwrap_or(10);
11566 let mut best = [f64::MAX; 2];
11567 let mut sums = [0f32; 2];
11568 for _ in 0..reps {
11569 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11570 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11571 let t = std::time::Instant::now();
11572 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11573 best[k] = best[k].min(t.elapsed().as_secs_f64());
11574 sums[k] = out.iter().take(64).sum::<f32>();
11575 }
11576 }
11577 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11578 for (k, name) in ["previous", "tuned "].iter().enumerate() {
11579 println!(
11580 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11581 best[k] * 1e3,
11582 flops / best[k] / 1e9,
11583 sums[k]
11584 );
11585 }
11586 assert!(
11587 (sums[0] - sums[1]).abs() < 1e-2,
11588 "the tuned kernel changed the result: {} vs {}",
11589 sums[0],
11590 sums[1]
11591 );
11592 }
11593
11594 #[test]
11600 fn q4tp_matmat_blocked_matches_scalar() {
11601 use std::sync::atomic::Ordering::Relaxed;
11602 for &(rows, cols, b) in &[
11610 (64usize, 128usize, 7usize),
11611 (33, 96, 4),
11612 (16, 256, 9),
11613 (192, 2304, 37),
11614 ] {
11615 let total = cortiq_core::quant::expected_nbytes(
11616 cortiq_core::TensorDtype::Q4TiledP,
11617 &[rows, cols],
11618 )
11619 .unwrap();
11620 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11621 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11622 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11623 let step = cortiq_core::quant::f32_to_f16(0.1);
11624 for r in 0..rows {
11625 let o = params_off + r * 4;
11626 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11627 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11628 }
11629 let xs: Vec<f32> = (0..b * cols)
11630 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11631 .collect();
11632 let mut got = vec![0f32; b * rows];
11633 let mut want = vec![0f32; b * rows];
11634 let gpr = cols / 32;
11635 let view = super::Q4tpView::new(&bytes, rows, cols);
11636 let pool = crate::pool::Pool::from_env();
11637 super::Q4TP_ALT.store(2, Relaxed);
11638 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11639 super::Q4TP_ALT.store(1, Relaxed);
11640 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11641 super::Q4TP_ALT.store(0, Relaxed);
11642 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11649 let (mut worst, mut at) = (0f32, 0usize);
11650 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11651 if (g - w).abs() > worst {
11652 worst = (g - w).abs();
11653 at = i;
11654 }
11655 }
11656 assert!(
11657 worst <= 1e-4 * scale,
11658 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11659 (scale {scale:.3e}) at cell {at}: {} vs {}",
11660 got[at],
11661 want[at]
11662 );
11663
11664 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11671 for bi in 0..b {
11672 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11673 for r in 0..rows {
11674 let mut sc = vec![0f32; gpr];
11675 view.scales_into(r, gpr, &mut sc);
11676 let mut exact = 0f64;
11677 for j in 0..cols {
11678 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11679 exact += w as f64 * sq as f64 * act.xq[j] as f64;
11680 }
11681 exact *= act.sx as f64;
11682 for &(j, xv) in &act.outliers {
11683 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11684 exact += w as f64 * sq as f64 * xv as f64;
11685 }
11686 let i = bi * rows + r;
11687 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11688 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11689 }
11690 }
11691 println!(
11692 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11693 per-column {e_scalar:.3e}"
11694 );
11695 assert!(
11700 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11701 "{rows}x{cols} b={b}: error against f64 too large — blocked \
11702 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11703 );
11704 }
11705 }
11706
11707 #[test]
11711 #[ignore]
11712 fn q4t_matmat_throughput() {
11713 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11714 let total =
11715 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
11716 .unwrap();
11717 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11720 let sc = cortiq_core::quant::f32_to_f16(0.02);
11721 for t in bytes.chunks_mut(super::Q4_TILE) {
11722 t[..2].copy_from_slice(&sc.to_le_bytes());
11723 }
11724 let xs: Vec<f32> = (0..b * cols)
11725 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11726 .collect();
11727 let mut out = vec![0f32; b * rows];
11728 let pool = crate::pool::Pool::from_env();
11729 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11730 let reps: usize = std::env::var("CMF_BENCH_REPS")
11731 .ok()
11732 .and_then(|v| v.parse().ok())
11733 .unwrap_or(10);
11734 let mut best = f64::MAX;
11735 for _ in 0..reps {
11736 let t = std::time::Instant::now();
11737 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11738 best = best.min(t.elapsed().as_secs_f64());
11739 }
11740 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11741 println!(
11742 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11743 best * 1e3,
11744 flops / best / 1e9,
11745 out.iter().take(64).sum::<f32>()
11746 );
11747 }
11748}