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 {
1877 model,
1878 idx,
1879 dtype,
1880 row_scale,
1881 col_field,
1882 ..
1883 } = self
1884 else {
1885 return false;
1886 };
1887 match *dtype {
1888 TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
1889 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1893 if *dtype == TensorDtype::Q8_2f
1896 && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
1897 && crate::gpu::q8_matmat_2f(
1898 model, *idx, row_scale, col_field, xs, b, rows, cols, out,
1899 )
1900 {
1901 return true;
1902 }
1903 let flat: Vec<f32> = (0..b)
1904 .flat_map(|bi| {
1905 prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
1906 })
1907 .collect();
1908 crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
1909 }
1910 _ => false,
1911 }
1912 }
1913
1914 pub fn matvec_many<const N: usize>(
1921 ts: [&QTensor; N],
1922 x: &[f32],
1923 mut outs: [&mut [f32]; N],
1924 pool: Option<&Pool>,
1925 ) {
1926 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1927 let uniform_q8 = ts.iter().all(|t| {
1928 matches!(
1929 t,
1930 Self::Mapped {
1931 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1932 ..
1933 }
1934 )
1935 });
1936 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1937 let uniform_q4 = ts.iter().all(|t| {
1938 matches!(
1939 t,
1940 Self::Mapped {
1941 dtype: TensorDtype::Q4Block,
1942 ..
1943 }
1944 )
1945 });
1946 let uniform_vbit = ts.iter().all(|t| {
1947 matches!(
1948 t,
1949 Self::Mapped {
1950 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1951 ..
1952 }
1953 )
1954 });
1955 let uniform_q1 = ts.iter().all(|t| {
1956 matches!(
1957 t,
1958 Self::Mapped {
1959 dtype: TensorDtype::Q1,
1960 ..
1961 }
1962 )
1963 });
1964 let uniform_q1t = ts.iter().all(|t| {
1965 matches!(
1966 t,
1967 Self::Mapped {
1968 dtype: TensorDtype::Q1T,
1969 ..
1970 }
1971 )
1972 });
1973 let uniform_q4tp = ts.iter().all(|t| {
1978 matches!(
1979 t,
1980 Self::Mapped {
1981 dtype: TensorDtype::Q4TiledP,
1982 ..
1983 }
1984 )
1985 }) && ts
1986 .iter()
1987 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1988 let Some(pool) = pool else {
1989 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1990 t.matvec(x, o, None);
1991 }
1992 return;
1993 };
1994 if total_rows < 256
1995 || !(uniform_q8
1996 || uniform_f32
1997 || uniform_q4
1998 || uniform_vbit
1999 || uniform_q1
2000 || uniform_q1t
2001 || uniform_q4tp)
2002 {
2003 for (t, o) in ts.iter().zip(outs.iter_mut()) {
2004 t.matvec(x, o, Some(pool));
2005 }
2006 return;
2007 }
2008
2009 if uniform_q4tp {
2010 let cols = ts[0].cols();
2016 let gpr = cols / GROUP_SIZE;
2017 let views: Vec<Q4tpView> = ts
2018 .iter()
2019 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2020 .collect();
2021 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2022 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2023 let locate = |flat: usize| -> (usize, usize) {
2025 let mut acc = 0;
2026 for (i, &r) in rows_of.iter().enumerate() {
2027 if flat < acc + r {
2028 return (i, flat - acc);
2029 }
2030 acc += r;
2031 }
2032 (rows_of.len() - 1, 0)
2033 };
2034 let (views, outs_addr) = (&views, &outs_addr);
2035 if a8w8_enabled() {
2036 let act = split_act(x);
2037 let act = &act;
2038 let run = |start: usize, end: usize| {
2039 let mut sc = vec![0f32; gpr];
2040 for flat in start..end {
2041 let (t, r) = locate(flat);
2042 let v = &views[t];
2043 v.scales_into(r, gpr, &mut sc);
2044 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2045 for &(j, xv) in &act.outliers {
2046 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2047 acc += w * s * xv;
2048 }
2049 unsafe { *outs_addr[t].at(r) = acc };
2051 }
2052 };
2053 pool.run_rows(total_rows, &run);
2054 } else {
2055 let run = |start: usize, end: usize| {
2056 let mut sc = vec![0f32; gpr];
2057 for flat in start..end {
2058 let (t, r) = locate(flat);
2059 let v = &views[t];
2060 v.scales_into(r, gpr, &mut sc);
2061 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2063 }
2064 };
2065 pool.run_rows(total_rows, &run);
2066 }
2067 return;
2068 }
2069
2070 if uniform_q1 {
2071 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2074 if a8w8_enabled() {
2075 let act = split_act(x);
2076 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2077 let (act, gsum) = (&act, &gsum);
2078 let closures: [_; N] = std::array::from_fn(|i| {
2079 let (bytes, gpr, out) =
2080 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2081 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2082 });
2083 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2084 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2085 pool.run_many(&parts);
2086 } else {
2087 let closures: [_; N] = std::array::from_fn(|i| {
2088 let (bytes, gpr, out) =
2089 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2090 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2091 });
2092 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2093 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2094 pool.run_many(&parts);
2095 }
2096 return;
2097 }
2098
2099 if uniform_q1t {
2100 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2104 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2105 if a8w8_enabled() {
2106 let act = split_act(x);
2107 let act = &act;
2108 let x_ref = x;
2109 let closures: [_; N] = std::array::from_fn(|i| {
2110 let bytes = ts[i].quant_bytes();
2111 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2112 let gpr = cols / GROUP_SIZE;
2113 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2114 let out = outs_addr[i];
2115 move |s: usize, e: usize| {
2116 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2117 }
2118 });
2119 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2120 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2121 pool.run_many(&parts);
2122 } else {
2123 let x_ref = x;
2124 let closures: [_; N] = std::array::from_fn(|i| {
2125 let bytes = ts[i].quant_bytes();
2126 let (rows, cols) = (ts[i].rows(), ts[i].cols());
2127 let gpr = cols / GROUP_SIZE;
2128 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2129 let out = outs_addr[i];
2130 move |s: usize, e: usize| {
2131 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2132 }
2133 });
2134 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2135 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2136 pool.run_many(&parts);
2137 }
2138 return;
2139 }
2140
2141 if uniform_q4 || uniform_vbit {
2142 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2143 if a8w8_enabled() {
2145 let act = split_act(x);
2146 let act = &act;
2147 if uniform_q4 {
2148 let closures: [_; N] = std::array::from_fn(|i| {
2149 let (packed, scales) =
2150 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2151 let (gpr, cols, out) =
2152 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2153 move |s: usize, e: usize| {
2154 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2155 }
2156 });
2157 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2158 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2159 pool.run_many(&parts);
2160 } else {
2161 let closures: [_; N] = std::array::from_fn(|i| {
2162 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2163 unreachable!()
2164 };
2165 let (bytes, rows, cols, out) = (
2166 ts[i].quant_bytes(),
2167 ts[i].rows(),
2168 ts[i].cols(),
2169 outs_addr[i],
2170 );
2171 move |s: usize, e: usize| {
2172 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2173 }
2174 });
2175 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2176 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2177 pool.run_many(&parts);
2178 }
2179 return;
2180 }
2181 if uniform_q4 {
2182 let closures: [_; N] = std::array::from_fn(|i| {
2183 let (packed, scales) =
2184 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2185 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2186 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2187 });
2188 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2189 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2190 pool.run_many(&parts);
2191 } else {
2192 let closures: [_; N] = std::array::from_fn(|i| {
2193 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2194 unreachable!()
2195 };
2196 let (bytes, rows, cols, out) = (
2197 ts[i].quant_bytes(),
2198 ts[i].rows(),
2199 ts[i].cols(),
2200 outs_addr[i],
2201 );
2202 move |s: usize, e: usize| {
2203 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2204 }
2205 });
2206 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2207 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2208 pool.run_many(&parts);
2209 }
2210 return;
2211 }
2212
2213 if uniform_f32 {
2214 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2215 let closures: [_; N] = std::array::from_fn(|i| {
2216 let Self::F32 { data, cols, .. } = ts[i] else {
2217 unreachable!()
2218 };
2219 let out = outs_addr[i];
2220 move |start: usize, end: usize| {
2221 for o in start..end {
2222 let row = &data[o * cols..(o + 1) * cols];
2223 let mut sum = 0.0f32;
2224 for j in 0..*cols {
2225 sum += row[j] * x[j];
2226 }
2227 unsafe { *out.at(o) = sum };
2229 }
2230 }
2231 });
2232 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2233 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2234 pool.run_many(&parts);
2235 return;
2236 }
2237
2238 struct Ctx<'a> {
2241 bytes: &'a [u8],
2242 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2243 rep: &'a [u8],
2244 row_scale: &'a [f32],
2245 cols: usize,
2246 xs: std::borrow::Cow<'a, [f32]>,
2247 }
2248 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2249 let Self::Mapped {
2250 dtype,
2251 cols,
2252 row_scale,
2253 col_field,
2254 repack,
2255 ..
2256 } = ts[i]
2257 else {
2258 unreachable!()
2259 };
2260 Ctx {
2261 bytes: ts[i].quant_bytes(),
2262 rep: repack,
2263 row_scale,
2264 cols: *cols,
2265 xs: prescale(x, col_field, *dtype),
2266 }
2267 });
2268 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2269 #[cfg(target_arch = "aarch64")]
2270 if sdot_enabled() {
2271 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2272 let closures: [_; N] = std::array::from_fn(|i| {
2273 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2274 move |start: usize, end: usize| {
2275 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2276 }
2277 });
2278 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2279 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2280 pool.run_many(&parts);
2281 return;
2282 }
2283 #[cfg(target_arch = "x86_64")]
2284 if avx2_a8w8_enabled() {
2285 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2286 let closures: [_; N] = std::array::from_fn(|i| {
2287 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2288 move |start: usize, end: usize| {
2289 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2290 }
2291 });
2292 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2293 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2294 pool.run_many(&parts);
2295 return;
2296 }
2297 let closures: [_; N] = std::array::from_fn(|i| {
2298 let (c, out) = (&ctxs[i], outs_addr[i]);
2299 move |start: usize, end: usize| {
2300 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2301 }
2302 });
2303 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2304 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2305 pool.run_many(&parts);
2306 }
2307}
2308
2309impl QTensor {
2310 #[allow(clippy::needless_range_loop)]
2315 pub fn matvec2_many<const N: usize>(
2316 ts: [&QTensor; N],
2317 x1: &[f32],
2318 x2: &[f32],
2319 mut o1s: [&mut [f32]; N],
2320 mut o2s: [&mut [f32]; N],
2321 pool: Option<&Pool>,
2322 ) {
2323 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2324 let uniform_q8 = ts.iter().all(|t| {
2325 matches!(
2326 t,
2327 Self::Mapped {
2328 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2329 ..
2330 }
2331 )
2332 });
2333 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2334 let uniform_q4 = ts.iter().all(|t| {
2335 matches!(
2336 t,
2337 Self::Mapped {
2338 dtype: TensorDtype::Q4Block,
2339 ..
2340 }
2341 )
2342 });
2343 let uniform_vbit = ts.iter().all(|t| {
2344 matches!(
2345 t,
2346 Self::Mapped {
2347 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2348 ..
2349 }
2350 )
2351 });
2352 let fusable = pool.is_some()
2353 && total_rows >= 256
2354 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2355 if !fusable {
2356 for i in 0..N {
2357 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2358 }
2359 return;
2360 }
2361 let pool = pool.unwrap();
2362
2363 if uniform_q4 || uniform_vbit {
2364 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2365 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2366 if a8w8_enabled() {
2368 let a1 = split_act(x1);
2369 let a2 = split_act(x2);
2370 let (a1, a2) = (&a1, &a2);
2371 if uniform_q4 {
2372 let closures: [_; N] = std::array::from_fn(|i| {
2373 let (packed, scales) =
2374 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2375 let (gpr, cols, o1, o2) =
2376 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2377 move |s: usize, e: usize| {
2378 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2379 }
2380 });
2381 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2382 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2383 pool.run_many(&parts);
2384 } else {
2385 let closures: [_; N] = std::array::from_fn(|i| {
2386 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2387 unreachable!()
2388 };
2389 let (bytes, rows, cols, o1, o2) = (
2390 ts[i].quant_bytes(),
2391 ts[i].rows(),
2392 ts[i].cols(),
2393 p1[i],
2394 p2[i],
2395 );
2396 move |s: usize, e: usize| {
2397 vbit_range2_a8w8(
2398 bytes,
2399 vbit_offsets,
2400 x1,
2401 x2,
2402 a1,
2403 a2,
2404 rows,
2405 cols,
2406 o1,
2407 o2,
2408 s,
2409 e,
2410 )
2411 }
2412 });
2413 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2414 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2415 pool.run_many(&parts);
2416 }
2417 return;
2418 }
2419 if uniform_q4 {
2420 let closures: [_; N] = std::array::from_fn(|i| {
2421 let (packed, scales) =
2422 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2423 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2424 move |s: usize, e: usize| {
2425 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2426 }
2427 });
2428 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2429 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2430 pool.run_many(&parts);
2431 } else {
2432 let closures: [_; N] = std::array::from_fn(|i| {
2433 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2434 unreachable!()
2435 };
2436 let (bytes, rows, cols, o1, o2) = (
2437 ts[i].quant_bytes(),
2438 ts[i].rows(),
2439 ts[i].cols(),
2440 p1[i],
2441 p2[i],
2442 );
2443 move |s: usize, e: usize| {
2444 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2445 }
2446 });
2447 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2448 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2449 pool.run_many(&parts);
2450 }
2451 return;
2452 }
2453
2454 if uniform_f32 {
2455 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2456 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2457 let closures: [_; N] = std::array::from_fn(|i| {
2458 let Self::F32 { data, cols, .. } = ts[i] else {
2459 unreachable!()
2460 };
2461 let (o1, o2) = (p1[i], p2[i]);
2462 move |start: usize, end: usize| {
2463 for o in start..end {
2464 let row = &data[o * cols..(o + 1) * cols];
2465 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2466 for j in 0..*cols {
2467 s1 += row[j] * x1[j];
2468 s2 += row[j] * x2[j];
2469 }
2470 unsafe {
2472 *o1.at(o) = s1;
2473 *o2.at(o) = s2;
2474 }
2475 }
2476 }
2477 });
2478 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2479 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2480 pool.run_many(&parts);
2481 return;
2482 }
2483
2484 struct Ctx<'a> {
2485 bytes: &'a [u8],
2486 row_scale: &'a [f32],
2487 cols: usize,
2488 xs1: std::borrow::Cow<'a, [f32]>,
2489 xs2: std::borrow::Cow<'a, [f32]>,
2490 }
2491 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2492 let Self::Mapped {
2493 dtype,
2494 cols,
2495 row_scale,
2496 col_field,
2497 ..
2498 } = ts[i]
2499 else {
2500 unreachable!()
2501 };
2502 Ctx {
2503 bytes: ts[i].quant_bytes(),
2504 row_scale,
2505 cols: *cols,
2506 xs1: prescale(x1, col_field, *dtype),
2507 xs2: prescale(x2, col_field, *dtype),
2508 }
2509 });
2510 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2511 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2512 #[cfg(target_arch = "aarch64")]
2513 if sdot_enabled() {
2514 let acts: [(SplitAct, SplitAct); N] =
2515 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2516 let closures: [_; N] = std::array::from_fn(|i| {
2517 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2518 move |start: usize, end: usize| {
2519 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2520 }
2521 });
2522 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2523 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2524 pool.run_many(&parts);
2525 return;
2526 }
2527 #[cfg(target_arch = "x86_64")]
2528 if avx2_a8w8_enabled() {
2529 let acts: [(SplitAct, SplitAct); N] =
2530 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2531 let closures: [_; N] = std::array::from_fn(|i| {
2532 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2533 move |start: usize, end: usize| {
2534 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2535 }
2536 });
2537 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2538 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2539 pool.run_many(&parts);
2540 return;
2541 }
2542 let closures: [_; N] = std::array::from_fn(|i| {
2543 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2544 move |start: usize, end: usize| {
2545 q8_range2_f32(
2546 c.bytes,
2547 c.row_scale,
2548 &c.xs1,
2549 &c.xs2,
2550 c.cols,
2551 o1,
2552 o2,
2553 start,
2554 end,
2555 )
2556 }
2557 });
2558 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2559 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2560 pool.run_many(&parts);
2561 }
2562
2563 pub fn matvec_silu_mul(
2568 gate: &QTensor,
2569 up: &QTensor,
2570 x: &[f32],
2571 out: &mut [f32],
2572 pool: Option<&Pool>,
2573 ) -> bool {
2574 let inter = gate.rows();
2575 debug_assert_eq!(up.rows(), inter);
2576 debug_assert_eq!(out.len(), inter);
2577 debug_assert_eq!(gate.cols(), up.cols());
2578 if !a8w8_enabled() {
2579 return false;
2580 }
2581 let act = split_act(x);
2582 let act = &act;
2583 let x_ref = x;
2584 let out_addr = SendMut(out.as_mut_ptr());
2585
2586 match (gate, up) {
2587 (
2589 Self::Mapped {
2590 dtype: TensorDtype::Q4Block,
2591 ..
2592 },
2593 Self::Mapped {
2594 dtype: TensorDtype::Q4Block,
2595 ..
2596 },
2597 ) => {
2598 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2599 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2600 let gpr = gate.cols() / GROUP_SIZE;
2601 let cols = gate.cols();
2602 let run = move |start: usize, end: usize| {
2603 for r in start..end {
2604 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2605 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2606 for &(j, xv) in &act.outliers {
2607 let flat = r * cols + j;
2608 let gb = gp[flat / 2];
2609 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2610 let gsc = f16_to_f32(u16::from_le_bytes([
2611 gs[(flat / GROUP_SIZE) * 2],
2612 gs[(flat / GROUP_SIZE) * 2 + 1],
2613 ]));
2614 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2615 let ub = up_p[flat / 2];
2616 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2617 let usc = f16_to_f32(u16::from_le_bytes([
2618 up_s[(flat / GROUP_SIZE) * 2],
2619 up_s[(flat / GROUP_SIZE) * 2 + 1],
2620 ]));
2621 uv += ((un as i32 - 8) as f32) * usc * xv;
2622 }
2623 let silu_g = gv / (1.0 + (-gv).exp());
2624 unsafe { *out_addr.at(r) = silu_g * uv };
2626 }
2627 };
2628 dispatch_rows(pool, inter, &run);
2629 true
2630 }
2631 (
2635 Self::Mapped {
2636 dtype: TensorDtype::Q4Tiled,
2637 ..
2638 },
2639 Self::Mapped {
2640 dtype: TensorDtype::Q4Tiled,
2641 ..
2642 },
2643 ) => {
2644 let g_bytes = gate.quant_bytes();
2645 let u_bytes = up.quant_bytes();
2646 let gpr = gate.cols() / GROUP_SIZE;
2647 let run = move |start: usize, end: usize| {
2648 for r in start..end {
2649 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2650 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2651 for &(j, xv) in &act.outliers {
2652 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2653 gv += w * s * xv;
2654 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2655 uv += w * s * xv;
2656 }
2657 let silu_g = gv / (1.0 + (-gv).exp());
2658 unsafe { *out_addr.at(r) = silu_g * uv };
2660 }
2661 };
2662 dispatch_rows(pool, inter, &run);
2663 true
2664 }
2665 (
2668 Self::Mapped {
2669 dtype: TensorDtype::Q4TiledP,
2670 ..
2671 },
2672 Self::Mapped {
2673 dtype: TensorDtype::Q4TiledP,
2674 ..
2675 },
2676 ) => {
2677 let cols = gate.cols();
2678 let gpr = cols / GROUP_SIZE;
2679 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2680 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2681 let run = |start: usize, end: usize| {
2682 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2683 for r in start..end {
2684 gv_view.scales_into(r, gpr, &mut gsc);
2685 uv_view.scales_into(r, gpr, &mut usc);
2686 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2687 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2688 for &(j, xv) in &act.outliers {
2689 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2690 gv += w * s * xv;
2691 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2692 uv += w * s * xv;
2693 }
2694 let silu_g = gv / (1.0 + (-gv).exp());
2695 unsafe { *out_addr.at(r) = silu_g * uv };
2697 }
2698 };
2699 dispatch_rows(pool, inter, &run);
2700 true
2701 }
2702 (
2708 Self::Mapped {
2709 dtype: TensorDtype::Q1,
2710 ..
2711 },
2712 Self::Mapped {
2713 dtype: TensorDtype::Q1,
2714 ..
2715 },
2716 ) => {
2717 let g_bytes = gate.quant_bytes();
2718 let u_bytes = up.quant_bytes();
2719 let gpr = gate.cols() / GROUP_SIZE;
2720 let gsum = q1_group_sums(&act.xq, gpr);
2721 let gsum = &gsum;
2722 let run = move |start: usize, end: usize| {
2723 for r in start..end {
2724 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2725 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2726 for &(j, xv) in &act.outliers {
2727 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2728 gv += w * s * xv;
2729 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2730 uv += w * s * xv;
2731 }
2732 let silu_g = gv / (1.0 + (-gv).exp());
2733 unsafe { *out_addr.at(r) = silu_g * uv };
2735 }
2736 };
2737 dispatch_rows(pool, inter, &run);
2738 true
2739 }
2740 (
2744 Self::Mapped {
2745 dtype: TensorDtype::Q2TiledP,
2746 ..
2747 },
2748 Self::Mapped {
2749 dtype: TensorDtype::Q2TiledP,
2750 ..
2751 },
2752 ) => {
2753 let cols = gate.cols();
2754 let gpr = cols / GROUP_SIZE;
2755 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2756 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2757 let gsum = q1_group_sums(&act.xq, gpr);
2758 let gsum = &gsum;
2759 let run = move |start: usize, end: usize| {
2760 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2761 for r in start..end {
2762 gv_view.scales_into(r, gpr, &mut gsc);
2763 uv_view.scales_into(r, gpr, &mut usc);
2764 let mut gv =
2765 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2766 let mut uv =
2767 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2768 for &(j, xv) in &act.outliers {
2769 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2770 gv += w * s * xv;
2771 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2772 uv += w * s * xv;
2773 }
2774 let silu_g = gv / (1.0 + (-gv).exp());
2775 unsafe { *out_addr.at(r) = silu_g * uv };
2777 }
2778 };
2779 dispatch_rows(pool, inter, &run);
2780 true
2781 }
2782 (
2787 Self::Mapped {
2788 dtype: TensorDtype::Q8Row,
2789 row_scale: g_rs,
2790 ..
2791 },
2792 Self::Mapped {
2793 dtype: TensorDtype::Q8Row,
2794 row_scale: u_rs,
2795 ..
2796 },
2797 ) => {
2798 let g_bytes = gate.quant_bytes();
2799 let u_bytes = up.quant_bytes();
2800 let cols = gate.cols();
2801 let run = move |start: usize, end: usize| {
2802 for r in start..end {
2803 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2804 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2805 let silu_g = gv / (1.0 + (-gv).exp());
2806 unsafe { *out_addr.at(r) = silu_g * uv };
2808 }
2809 };
2810 dispatch_rows(pool, inter, &run);
2811 true
2812 }
2813 (
2815 Self::Mapped {
2816 dtype: TensorDtype::Q1T,
2817 ..
2818 },
2819 Self::Mapped {
2820 dtype: TensorDtype::Q1T,
2821 ..
2822 },
2823 ) => {
2824 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2825 let g_bytes = gate.quant_bytes();
2826 let u_bytes = up.quant_bytes();
2827 let gpr = gate.cols() / GROUP_SIZE;
2828 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2829 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2830 let run = move |start: usize, end: usize| {
2831 for r in start..end {
2832 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2833 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2834 for &(j, xv) in &act.outliers {
2835 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2836 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2837 }
2838 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2839 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2840 let silu_g = gv / (1.0 + (-gv).exp());
2841 unsafe { *out_addr.at(r) = silu_g * uv };
2843 }
2844 };
2845 dispatch_rows(pool, inter, &run);
2846 true
2847 }
2848 _ => false,
2849 }
2850 }
2851
2852 pub fn moe_gate_up_many(
2866 pairs: &[(&QTensor, &QTensor)],
2867 x: &[f32],
2868 outs: &mut [Vec<f32>],
2869 pool: Option<&Pool>,
2870 ) -> bool {
2871 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2872 return false;
2873 }
2874 let inter = pairs[0].0.rows();
2875 let cols = pairs[0].0.cols();
2876 if cols % GROUP_SIZE != 0 {
2877 return false;
2878 }
2879 let gpr = cols / GROUP_SIZE;
2880 let q2 = matches!(
2883 pairs[0].0,
2884 Self::Mapped {
2885 dtype: TensorDtype::Q2TiledP,
2886 ..
2887 }
2888 );
2889 let want = if q2 {
2890 TensorDtype::Q2TiledP
2891 } else {
2892 TensorDtype::Q4TiledP
2893 };
2894 let mut views = Vec::with_capacity(pairs.len() * 2);
2895 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2896 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2897 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2898 if !both
2899 || g.rows() != inter
2900 || u.rows() != inter
2901 || g.cols() != cols
2902 || u.cols() != cols
2903 || o.len() != inter
2904 {
2905 return false;
2906 }
2907 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2908 views.push(mk(g.quant_bytes(), inter, cols));
2909 views.push(mk(u.quant_bytes(), inter, cols));
2910 }
2911 let act = split_act(x);
2912 let gsum = if q2 {
2913 q1_group_sums(&act.xq, gpr)
2914 } else {
2915 Vec::new()
2916 };
2917 let (act, gsum) = (&act, &gsum);
2918 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2919 let (views, ptrs) = (&views, &ptrs);
2920 let run = |start: usize, end: usize| {
2921 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2922 for flat in start..end {
2923 let (e, r) = (flat / inter, flat % inter);
2924 let gv_view = &views[e * 2];
2925 let uv_view = &views[e * 2 + 1];
2926 gv_view.scales_into(r, gpr, &mut gsc);
2927 uv_view.scales_into(r, gpr, &mut usc);
2928 let (mut gv, mut uv) = if q2 {
2929 (
2930 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2931 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2932 )
2933 } else {
2934 (
2935 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2936 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2937 )
2938 };
2939 for &(j, xv) in &act.outliers {
2940 let (og, ou) = if q2 {
2941 (
2942 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2943 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2944 )
2945 } else {
2946 (
2947 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2948 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2949 )
2950 };
2951 gv += og.0 * og.1 * xv;
2952 uv += ou.0 * ou.1 * xv;
2953 }
2954 let silu_g = gv / (1.0 + (-gv).exp());
2955 unsafe { *ptrs[e].at(r) = silu_g * uv };
2957 }
2958 };
2959 dispatch_rows(pool, pairs.len() * inter, &run);
2960 true
2961 }
2962
2963 pub fn moe_down_many(
2972 downs: &[&QTensor],
2973 gs: &[Vec<f32>],
2974 weights: &[f32],
2975 out: &mut [f32],
2976 pool: Option<&Pool>,
2977 ) -> bool {
2978 if downs.is_empty()
2979 || downs.len() != gs.len()
2980 || downs.len() != weights.len()
2981 || !a8w8_enabled()
2982 {
2983 return false;
2984 }
2985 let rows = out.len();
2986 let cols = downs[0].cols();
2987 if cols % GROUP_SIZE != 0 {
2988 return false;
2989 }
2990 let gpr = cols / GROUP_SIZE;
2991 let mut views = Vec::with_capacity(downs.len());
2992 for (d, g) in downs.iter().zip(gs.iter()) {
2993 if !matches!(
2994 d,
2995 Self::Mapped {
2996 dtype: TensorDtype::Q4TiledP,
2997 ..
2998 }
2999 ) || d.rows() != rows
3000 || d.cols() != cols
3001 || g.len() != cols
3002 {
3003 return false;
3004 }
3005 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3006 }
3007 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3009 let out_addr = SendMut(out.as_mut_ptr());
3017 let (views, acts, weights) = (&views, &acts, &weights);
3018 let run = |start: usize, end: usize| {
3019 let mut sc = vec![0f32; gpr];
3020 for r in start..end {
3021 let mut acc = 0f32;
3022 for (e, v) in views.iter().enumerate() {
3023 v.scales_into(r, gpr, &mut sc);
3024 let a = &acts[e];
3025 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3026 for &(j, xv) in &a.outliers {
3027 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3028 d += w * s * xv;
3029 }
3030 acc += weights[e] * d;
3031 }
3032 unsafe { *out_addr.at(r) = acc };
3034 }
3035 };
3036 dispatch_rows(pool, rows, &run);
3037 true
3038 }
3039}
3040
3041#[cfg(target_os = "macos")]
3046mod accel_blas {
3047 #[link(name = "Accelerate", kind = "framework")]
3048 unsafe extern "C" {
3049 pub fn cblas_sgemm(
3050 order: i32,
3051 trans_a: i32,
3052 trans_b: i32,
3053 m: i32,
3054 n: i32,
3055 k: i32,
3056 alpha: f32,
3057 a: *const f32,
3058 lda: i32,
3059 b: *const f32,
3060 ldb: i32,
3061 beta: f32,
3062 c: *mut f32,
3063 ldc: i32,
3064 );
3065 }
3066}
3067
3068#[cfg(target_os = "macos")]
3069pub(crate) fn accel_gemm_enabled() -> bool {
3070 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3071 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3072}
3073
3074#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3077pub(crate) fn accel_gemm_enabled() -> bool {
3078 true
3079}
3080
3081#[cfg(target_arch = "aarch64")]
3088#[allow(clippy::too_many_arguments)]
3089pub(crate) fn neon_gemm_rm(
3090 m: usize,
3091 n: usize,
3092 k: usize,
3093 alpha: f32,
3094 a: &[f32],
3095 lda: usize,
3096 b_mat: &[f32],
3097 ldb: usize,
3098 b_rows_are_n: bool,
3099 c: &mut [f32],
3100 ldc: usize,
3101) {
3102 debug_assert!(a.len() >= (m - 1) * lda + k);
3103 debug_assert!(c.len() >= (m - 1) * ldc + n);
3104 unsafe {
3106 use core::arch::aarch64::*;
3107 let mut i = 0usize;
3108 while i < m {
3109 let mi = (m - i).min(4);
3110 let mut j = 0usize;
3111 while j < n {
3112 let nj = (n - j).min(8);
3113 if mi == 4 && nj == 8 {
3114 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3115 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3116 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3117 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3118 for p in 0..k {
3119 let (b0, b1) = if b_rows_are_n {
3120 let base = b_mat.as_ptr().add(j * ldb + p);
3123 let g = |o: usize| *base.add(o * ldb);
3124 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3125 } else {
3126 let base = b_mat.as_ptr().add(p * ldb + j);
3127 (
3128 [*base, *base.add(1), *base.add(2), *base.add(3)],
3129 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3130 )
3131 };
3132 let bv0 = vld1q_f32(b0.as_ptr());
3133 let bv1 = vld1q_f32(b1.as_ptr());
3134 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3135 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3136 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3137 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3138 c0a = vfmaq_f32(c0a, a0, bv0);
3139 c0b = vfmaq_f32(c0b, a0, bv1);
3140 c1a = vfmaq_f32(c1a, a1, bv0);
3141 c1b = vfmaq_f32(c1b, a1, bv1);
3142 c2a = vfmaq_f32(c2a, a2, bv0);
3143 c2b = vfmaq_f32(c2b, a2, bv1);
3144 c3a = vfmaq_f32(c3a, a3, bv0);
3145 c3b = vfmaq_f32(c3b, a3, bv1);
3146 }
3147 let al = vdupq_n_f32(alpha);
3148 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3149 .iter()
3150 .enumerate()
3151 {
3152 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3153 vst1q_f32(dst, vmulq_f32(*ca, al));
3154 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3155 }
3156 } else {
3157 for r in 0..mi {
3158 for q in 0..nj {
3159 let mut acc = 0f32;
3160 for p in 0..k {
3161 let bv = if b_rows_are_n {
3162 b_mat[(j + q) * ldb + p]
3163 } else {
3164 b_mat[p * ldb + j + q]
3165 };
3166 acc += a[(i + r) * lda + p] * bv;
3167 }
3168 c[(i + r) * ldc + j + q] = acc * alpha;
3169 }
3170 }
3171 }
3172 j += nj;
3173 }
3174 i += mi;
3175 }
3176 }
3177}
3178
3179#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3181#[allow(clippy::too_many_arguments)]
3182pub(crate) fn sgemm_rm(
3183 m: usize,
3184 n: usize,
3185 k: usize,
3186 alpha: f32,
3187 a: &[f32],
3188 lda: usize,
3189 b_mat: &[f32],
3190 ldb: usize,
3191 b_rows_are_n: bool,
3192 c: &mut [f32],
3193 ldc: usize,
3194) {
3195 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3196}
3197
3198#[allow(clippy::too_many_arguments)]
3202pub fn sgemm_public(
3203 m: usize,
3204 n: usize,
3205 k: usize,
3206 alpha: f32,
3207 a: &[f32],
3208 lda: usize,
3209 b_mat: &[f32],
3210 ldb: usize,
3211 b_rows_are_n: bool,
3212 c: &mut [f32],
3213 ldc: usize,
3214) {
3215 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3216 {
3217 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3218 }
3219 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3224 {
3225 for i in 0..m {
3226 for j in 0..n {
3227 let mut acc = 0f32;
3228 for p in 0..k {
3229 let bv = if b_rows_are_n {
3230 b_mat[j * ldb + p]
3231 } else {
3232 b_mat[p * ldb + j]
3233 };
3234 acc += a[i * lda + p] * bv;
3235 }
3236 c[i * ldc + j] = alpha * acc;
3237 }
3238 }
3239 }
3240}
3241
3242#[cfg(target_os = "macos")]
3245#[allow(clippy::too_many_arguments)]
3246pub(crate) fn sgemm_rm(
3247 m: usize,
3248 n: usize,
3249 k: usize,
3250 alpha: f32,
3251 a: &[f32],
3252 lda: usize,
3253 b_mat: &[f32],
3254 ldb: usize,
3255 b_rows_are_n: bool,
3256 c: &mut [f32],
3257 ldc: usize,
3258) {
3259 debug_assert!(a.len() >= (m - 1) * lda + k);
3260 debug_assert!(c.len() >= (m - 1) * ldc + n);
3261 #[cfg(target_arch = "aarch64")]
3266 if std::env::var("CMF_FORCE_NEON_GEMM")
3267 .map(|v| v == "1")
3268 .unwrap_or(false)
3269 {
3270 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3271 }
3272 unsafe {
3273 accel_blas::cblas_sgemm(
3274 101, 111, if b_rows_are_n { 112 } else { 111 },
3277 m as i32,
3278 n as i32,
3279 k as i32,
3280 alpha,
3281 a.as_ptr(),
3282 lda as i32,
3283 b_mat.as_ptr(),
3284 ldb as i32,
3285 0.0,
3286 c.as_mut_ptr(),
3287 ldc as i32,
3288 );
3289 }
3290}
3291
3292#[cfg(target_os = "macos")]
3299fn qmatmat_accel(
3300 q: &[u8],
3301 row_scale: &[f32],
3302 pre: &[std::borrow::Cow<'_, [f32]>],
3303 rows: usize,
3304 cols: usize,
3305 out: &mut [f32],
3306 pool: Option<&Pool>,
3307) {
3308 const TR: usize = 2048;
3313 let b = pre.len();
3314 thread_local! {
3315 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3316 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3317 }
3318 XPANEL.with(|xp| {
3319 WTILE.with(|wt| {
3320 let mut xpanel = xp.borrow_mut();
3321 xpanel.clear();
3322 for x in pre {
3323 xpanel.extend_from_slice(x);
3324 }
3325 let mut wtile = wt.borrow_mut();
3326 wtile.resize(TR * cols, 0.0);
3327 let mut r0 = 0usize;
3328 while r0 < rows {
3329 let tr = TR.min(rows - r0);
3330 let wt_addr = SendMut(wtile.as_mut_ptr());
3332 let run = |start: usize, end: usize| {
3333 for r in start..end {
3334 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3335 let s = row_scale[r0 + r];
3336 let dst =
3338 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3339 for (d, &v) in dst.iter_mut().zip(row) {
3340 *d = (v as i8) as f32 * s;
3341 }
3342 }
3343 };
3344 dispatch_rows(pool, tr, &run);
3345 unsafe {
3347 accel_blas::cblas_sgemm(
3348 101, 111, 112, b as i32,
3352 tr as i32,
3353 cols as i32,
3354 1.0,
3355 xpanel.as_ptr(),
3356 cols as i32,
3357 wtile.as_ptr(),
3358 cols as i32,
3359 0.0,
3360 out.as_mut_ptr().add(r0),
3361 rows as i32,
3362 );
3363 }
3364 r0 += tr;
3365 }
3366 })
3367 });
3368}
3369
3370fn qmatmat(
3371 q: &[u8],
3372 row_scale: &[f32],
3373 pre: &[std::borrow::Cow<'_, [f32]>],
3374 rows: usize,
3375 cols: usize,
3376 out: &mut [f32],
3377 pool: Option<&Pool>,
3378) {
3379 let b = pre.len();
3380 debug_assert_eq!(out.len(), b * rows);
3381 #[cfg(target_os = "macos")]
3386 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3387 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3388 return;
3389 }
3390 #[cfg(target_arch = "aarch64")]
3391 if sdot_enabled() {
3392 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3393 let out_addr = SendMut(out.as_mut_ptr());
3394 let blocked_ok = blocked_enabled();
3397 let use_i8mm = i8mm_enabled();
3398 if blocked_ok {
3399 let run = |start: usize, end: usize| {
3400 let mut o = start;
3401 while o < end {
3402 if o + 2 <= end {
3403 let r0 = &q[o * cols..(o + 1) * cols];
3404 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3405 let mut bi = 0usize;
3406 while bi + 4 <= acts.len() {
3407 let xs = [
3408 acts[bi].xq.as_slice(),
3409 acts[bi + 1].xq.as_slice(),
3410 acts[bi + 2].xq.as_slice(),
3411 acts[bi + 3].xq.as_slice(),
3412 ];
3413 let d = if use_i8mm {
3414 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3415 } else {
3416 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3417 };
3418 for (r, row) in [r0, r1].into_iter().enumerate() {
3419 for k in 0..4 {
3420 let act = &acts[bi + k];
3421 let mut v = d[r][k] as f32 * act.sx;
3422 for &(j, xv) in &act.outliers {
3423 v += (row[j] as i8) as f32 * xv;
3424 }
3425 unsafe {
3426 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3427 };
3428 }
3429 }
3430 bi += 4;
3431 }
3432 while bi < acts.len() {
3433 for (r, row) in [r0, r1].into_iter().enumerate() {
3434 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3435 unsafe { *out_addr.at(bi * rows + o + r) = v };
3436 }
3437 bi += 1;
3438 }
3439 o += 2;
3440 } else {
3441 let row = &q[o * cols..(o + 1) * cols];
3442 for (bi, act) in acts.iter().enumerate() {
3443 let v = row_dot_sdot(row, act) * row_scale[o];
3444 unsafe { *out_addr.at(bi * rows + o) = v };
3445 }
3446 o += 1;
3447 }
3448 }
3449 };
3450 dispatch_rows(pool, rows, &run);
3451 return;
3452 }
3453 let run = |start: usize, end: usize| {
3454 for o in start..end {
3455 let row = &q[o * cols..(o + 1) * cols];
3456 for (bi, act) in acts.iter().enumerate() {
3457 let v = row_dot_sdot(row, act) * row_scale[o];
3458 unsafe { *out_addr.at(bi * rows + o) = v };
3459 }
3460 }
3461 };
3462 dispatch_rows(pool, rows, &run);
3463 return;
3464 }
3465 #[cfg(target_arch = "x86_64")]
3470 if avx2_a8w8_enabled() {
3471 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3472 let out_addr = SendMut(out.as_mut_ptr());
3473 let blocked_ok = blocked_enabled();
3476 if !avx512vnni_enabled() && blocked_ok {
3477 let run = |start: usize, end: usize| {
3478 let mut o = start;
3479 while o < end {
3480 if o + 2 <= end {
3481 let r0 = &q[o * cols..(o + 1) * cols];
3482 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3483 let mut bi = 0usize;
3484 while bi + 4 <= acts.len() {
3485 let xs = [
3486 acts[bi].xq.as_slice(),
3487 acts[bi + 1].xq.as_slice(),
3488 acts[bi + 2].xq.as_slice(),
3489 acts[bi + 3].xq.as_slice(),
3490 ];
3491 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3492 for (r, row) in [r0, r1].into_iter().enumerate() {
3493 for k in 0..4 {
3494 let act = &acts[bi + k];
3495 let mut v = d[r][k] as f32 * act.sx;
3496 for &(j, xv) in &act.outliers {
3497 v += (row[j] as i8) as f32 * xv;
3498 }
3499 unsafe {
3500 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3501 };
3502 }
3503 }
3504 bi += 4;
3505 }
3506 while bi < acts.len() {
3507 for (r, row) in [r0, r1].into_iter().enumerate() {
3508 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3509 unsafe { *out_addr.at(bi * rows + o + r) = v };
3510 }
3511 bi += 1;
3512 }
3513 o += 2;
3514 } else {
3515 let row = &q[o * cols..(o + 1) * cols];
3516 for (bi, act) in acts.iter().enumerate() {
3517 let v = row_dot_avx2(row, act) * row_scale[o];
3518 unsafe { *out_addr.at(bi * rows + o) = v };
3519 }
3520 o += 1;
3521 }
3522 }
3523 };
3524 dispatch_rows(pool, rows, &run);
3525 return;
3526 }
3527 let run = |start: usize, end: usize| {
3528 for o in start..end {
3529 let row = &q[o * cols..(o + 1) * cols];
3530 for (bi, act) in acts.iter().enumerate() {
3531 let v = row_dot_avx2(row, act) * row_scale[o];
3532 unsafe { *out_addr.at(bi * rows + o) = v };
3533 }
3534 }
3535 };
3536 dispatch_rows(pool, rows, &run);
3537 return;
3538 }
3539 let out_addr = SendMut(out.as_mut_ptr());
3540 let run = |start: usize, end: usize| {
3541 for o in start..end {
3542 let row = &q[o * cols..(o + 1) * cols];
3543 for (bi, x) in pre.iter().enumerate() {
3544 let mut acc = 0f32;
3545 for j in 0..cols {
3546 acc += (row[j] as i8) as f32 * x[j];
3547 }
3548 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3549 }
3550 }
3551 };
3552 dispatch_rows(pool, rows, &run);
3553}
3554
3555fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3558 match pool {
3559 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3560 _ => run(0, rows),
3561 }
3562}
3563
3564fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3566 let groups = rows * cols / GROUP_SIZE;
3567 bytes.split_at(groups * 16)
3568}
3569
3570#[inline]
3575fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3576 #[cfg(target_arch = "aarch64")]
3577 unsafe {
3578 return vbit_fill4_neon(data, buf);
3579 }
3580 #[cfg(target_arch = "x86_64")]
3581 if avx2_enabled() {
3582 return unsafe { vbit_fill4_avx2(data, buf) };
3583 }
3584 #[allow(unreachable_code)]
3585 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3586 let u = unpack8::<4>(&data[blk * 4..]);
3587 for k in 0..8 {
3588 chunk[k] = (u[k] - 7) as i8 as u8;
3589 }
3590 }
3591}
3592
3593#[cfg(target_arch = "aarch64")]
3594#[target_feature(enable = "neon")]
3595unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3596 unsafe {
3599 use core::arch::aarch64::*;
3600 let n = buf.len();
3601 let mask = vdupq_n_u8(0x0F);
3602 let seven = vdupq_n_s8(7);
3603 let mut g = 0usize;
3604 while g * 32 + 32 <= n {
3605 let b = vld1q_u8(data.as_ptr().add(g * 16));
3606 let hi = vshrq_n_u8::<4>(b);
3607 let lo = vandq_u8(b, mask);
3608 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3609 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3610 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3611 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3612 g += 1;
3613 }
3614 }
3615}
3616
3617#[cfg(target_arch = "x86_64")]
3618#[target_feature(enable = "avx2")]
3619unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3620 unsafe {
3622 use core::arch::x86_64::*;
3623 let n = buf.len();
3624 let mask = _mm_set1_epi8(0x0F);
3625 let seven = _mm256_set1_epi8(7);
3626 let mut g = 0usize;
3627 while g * 32 + 32 <= n {
3628 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3629 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3630 let lo = _mm_and_si128(b, mask);
3631 let z = _mm256_sub_epi8(
3632 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3633 seven,
3634 );
3635 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3636 g += 1;
3637 }
3638 }
3639}
3640
3641#[inline(always)]
3646fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3647 let mut acc = 0u64;
3648 for i in 0..B {
3649 acc = (acc << 8) | data[i] as u64;
3650 }
3651 let mask = (1u64 << B) - 1;
3652 let mut out = [0i32; 8];
3653 for (k, o) in out.iter_mut().enumerate() {
3654 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3655 }
3656 out
3657}
3658
3659#[allow(clippy::too_many_arguments)]
3665fn vbitmatvec(
3666 bytes: &[u8],
3667 offsets: &[usize],
3668 x: &[f32],
3669 rows: usize,
3670 cols: usize,
3671 out: &mut [f32],
3672 pool: Option<&Pool>,
3673) {
3674 debug_assert_eq!(out.len(), rows);
3675 debug_assert_eq!(offsets.len(), rows + 1);
3676
3677 if a8w8_enabled() {
3681 let act = split_act(x);
3682 let out_addr = SendMut(out.as_mut_ptr());
3683 let run = move |start: usize, end: usize| {
3684 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3685 };
3686 dispatch_rows(pool, rows, &run);
3687 return;
3688 }
3689
3690 let out_addr = SendMut(out.as_mut_ptr());
3691 let run = move |start: usize, end: usize| {
3692 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3693 };
3694 dispatch_rows(pool, rows, &run);
3695}
3696
3697#[allow(clippy::too_many_arguments)]
3701fn vbit_range_a8w8(
3702 bytes: &[u8],
3703 offsets: &[usize],
3704 x: &[f32],
3705 act: &SplitAct,
3706 rows: usize,
3707 cols: usize,
3708 out: SendMut,
3709 start: usize,
3710 end: usize,
3711) {
3712 let ng = cols / GROUP_SIZE;
3713 let bits = &bytes[..rows];
3714 let sc_off = rows;
3715 let row_dot = |r: usize| -> f32 {
3716 let b = bits[r] as usize;
3717 let l = (1i32 << (b - 1)) - 1;
3718 let mask = (1u64 << b) - 1;
3719 let data = &bytes[offsets[r]..offsets[r + 1]];
3720 if b == 8 {
3721 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3723 let mut dot = 0f32;
3724 for g in 0..ng {
3725 let so = (r * ng + g) * 2;
3726 let sgf = f16_to_f32(u16::from_le_bytes([
3727 bytes[sc_off + so],
3728 bytes[sc_off + so + 1],
3729 ]));
3730 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3731 let mut gd = 0f32;
3732 for &xv in xg.iter() {
3733 if nbits < 8 {
3734 acc = (acc << 8) | data[idx] as u64;
3735 idx += 1;
3736 nbits += 8;
3737 }
3738 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3739 nbits -= 8;
3740 gd += (u - l) as f32 * xv;
3741 }
3742 dot += gd * sgf;
3743 }
3744 return dot;
3745 }
3746 thread_local! {
3750 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3751 const { std::cell::RefCell::new(Vec::new()) };
3752 }
3753 #[inline(always)]
3754 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3755 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3756 let u = unpack8::<B>(&data[blk * B..]);
3757 for k in 0..8 {
3758 chunk[k] = (u[k] - l) as i8 as u8;
3759 }
3760 }
3761 }
3762 let _ = mask;
3763 VBIT_SCRATCH.with(|scratch| {
3764 let mut buf = scratch.borrow_mut();
3765 buf.resize(cols, 0);
3766 match b {
3767 3 => fill::<3>(data, l, &mut buf),
3768 4 => vbit_fill4(data, &mut buf),
3769 5 => fill::<5>(data, l, &mut buf),
3770 6 => fill::<6>(data, l, &mut buf),
3771 _ => unreachable!(),
3772 }
3773 let mut dot = 0f32;
3774 for g in 0..ng {
3775 let so = (r * ng + g) * 2;
3776 let s = f16_to_f32(u16::from_le_bytes([
3777 bytes[sc_off + so],
3778 bytes[sc_off + so + 1],
3779 ]));
3780 let d = dot_i8_i8(
3781 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3782 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3783 ) as f32
3784 * act.sx;
3785 dot += d * s;
3786 }
3787 for &(j, xv) in &act.outliers {
3788 let so = (r * ng + j / GROUP_SIZE) * 2;
3789 let s = f16_to_f32(u16::from_le_bytes([
3790 bytes[sc_off + so],
3791 bytes[sc_off + so + 1],
3792 ]));
3793 dot += (buf[j] as i8) as f32 * s * xv;
3795 }
3796 dot
3797 })
3798 };
3799 for r in start..end {
3800 unsafe { *out.at(r) = row_dot(r) };
3802 }
3803}
3804
3805#[allow(clippy::too_many_arguments)]
3807fn vbit_range_f32(
3808 bytes: &[u8],
3809 offsets: &[usize],
3810 x: &[f32],
3811 rows: usize,
3812 cols: usize,
3813 out: SendMut,
3814 start: usize,
3815 end: usize,
3816) {
3817 let ng = cols / GROUP_SIZE;
3818 let bits = &bytes[..rows];
3819 let sc_off = rows;
3820 #[inline(always)]
3824 fn dot_row<const B: usize>(
3825 data: &[u8],
3826 bytes: &[u8],
3827 sc_off: usize,
3828 r: usize,
3829 ng: usize,
3830 x: &[f32],
3831 ) -> f32 {
3832 let l = ((1i32 << (B - 1)) - 1) as f32;
3833 let gbytes = GROUP_SIZE * B / 8;
3834 let mut dot = 0f32;
3835 for g in 0..ng {
3836 let so = (r * ng + g) * 2;
3837 let s = f16_to_f32(u16::from_le_bytes([
3838 bytes[sc_off + so],
3839 bytes[sc_off + so + 1],
3840 ]));
3841 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3842 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3843 let mut gd = 0f32;
3844 for blk in 0..GROUP_SIZE / 8 {
3845 let u = unpack8::<B>(&gd0[blk * B..]);
3846 let xb = &xg[blk * 8..blk * 8 + 8];
3847 for k in 0..8 {
3848 gd += (u[k] as f32 - l) * xb[k];
3849 }
3850 }
3851 dot += gd * s;
3852 }
3853 dot
3854 }
3855 for r in start..end {
3856 let data = &bytes[offsets[r]..offsets[r + 1]];
3857 let v = match bits[r] {
3858 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3859 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3860 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3861 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3862 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3863 b => unreachable!("vbit bit-width {b} (validated at load)"),
3864 };
3865 unsafe { *out.at(r) = v };
3867 }
3868}
3869
3870#[allow(clippy::too_many_arguments)]
3875fn vbitmatvec2(
3876 bytes: &[u8],
3877 offsets: &[usize],
3878 x1: &[f32],
3879 x2: &[f32],
3880 rows: usize,
3881 cols: usize,
3882 o1: &mut [f32],
3883 o2: &mut [f32],
3884 pool: Option<&Pool>,
3885) {
3886 debug_assert_eq!(o1.len(), rows);
3887 debug_assert_eq!(o2.len(), rows);
3888
3889 if a8w8_enabled() {
3890 let a1 = split_act(x1);
3891 let a2 = split_act(x2);
3892 let p1 = SendMut(o1.as_mut_ptr());
3893 let p2 = SendMut(o2.as_mut_ptr());
3894 let run = move |start: usize, end: usize| {
3895 vbit_range2_a8w8(
3896 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3897 )
3898 };
3899 dispatch_rows(pool, rows, &run);
3900 return;
3901 }
3902
3903 let p1 = SendMut(o1.as_mut_ptr());
3904 let p2 = SendMut(o2.as_mut_ptr());
3905 let run = move |start: usize, end: usize| {
3906 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3907 };
3908 dispatch_rows(pool, rows, &run);
3909}
3910
3911#[allow(clippy::too_many_arguments)]
3915fn vbit_range2_a8w8(
3916 bytes: &[u8],
3917 offsets: &[usize],
3918 x1: &[f32],
3919 x2: &[f32],
3920 a1: &SplitAct,
3921 a2: &SplitAct,
3922 rows: usize,
3923 cols: usize,
3924 p1: SendMut,
3925 p2: SendMut,
3926 start: usize,
3927 end: usize,
3928) {
3929 let ng = cols / GROUP_SIZE;
3930 let bits = &bytes[..rows];
3931 let sc_off = rows;
3932 let row_dots = |r: usize| -> (f32, f32) {
3933 let b = bits[r] as usize;
3934 let l = (1i32 << (b - 1)) - 1;
3935 let data = &bytes[offsets[r]..offsets[r + 1]];
3936 if b == 8 {
3937 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3940 let (mut d1, mut d2) = (0f32, 0f32);
3941 for g in 0..ng {
3942 let so = (r * ng + g) * 2;
3943 let sgf = f16_to_f32(u16::from_le_bytes([
3944 bytes[sc_off + so],
3945 bytes[sc_off + so + 1],
3946 ]));
3947 let (mut g1, mut g2) = (0f32, 0f32);
3948 for k in 0..GROUP_SIZE {
3949 if nbits < 8 {
3950 acc = (acc << 8) | data[idx] as u64;
3951 idx += 1;
3952 nbits += 8;
3953 }
3954 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3955 nbits -= 8;
3956 let w = (u - l) as f32;
3957 g1 += w * x1[g * GROUP_SIZE + k];
3958 g2 += w * x2[g * GROUP_SIZE + k];
3959 }
3960 d1 += g1 * sgf;
3961 d2 += g2 * sgf;
3962 }
3963 return (d1, d2);
3964 }
3965 thread_local! {
3966 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3967 const { std::cell::RefCell::new(Vec::new()) };
3968 }
3969 #[inline(always)]
3970 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3971 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3972 let u = unpack8::<B>(&data[blk * B..]);
3973 for k in 0..8 {
3974 chunk[k] = (u[k] - l) as i8 as u8;
3975 }
3976 }
3977 }
3978 VBIT_SCRATCH2.with(|scratch| {
3979 let mut buf = scratch.borrow_mut();
3980 buf.resize(cols, 0);
3981 match b {
3982 3 => fill::<3>(data, l, &mut buf),
3983 4 => vbit_fill4(data, &mut buf),
3984 5 => fill::<5>(data, l, &mut buf),
3985 6 => fill::<6>(data, l, &mut buf),
3986 _ => unreachable!(),
3987 }
3988 let (mut d1, mut d2) = (0f32, 0f32);
3989 for g in 0..ng {
3990 let so = (r * ng + g) * 2;
3991 let s = f16_to_f32(u16::from_le_bytes([
3992 bytes[sc_off + so],
3993 bytes[sc_off + so + 1],
3994 ]));
3995 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3996 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3997 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3998 d1 += v1 * s;
3999 d2 += v2 * s;
4000 }
4001 for &(j, xv) in &a1.outliers {
4002 let so = (r * ng + j / GROUP_SIZE) * 2;
4003 let s = f16_to_f32(u16::from_le_bytes([
4004 bytes[sc_off + so],
4005 bytes[sc_off + so + 1],
4006 ]));
4007 d1 += (buf[j] as i8) as f32 * s * xv;
4008 }
4009 for &(j, xv) in &a2.outliers {
4010 let so = (r * ng + j / GROUP_SIZE) * 2;
4011 let s = f16_to_f32(u16::from_le_bytes([
4012 bytes[sc_off + so],
4013 bytes[sc_off + so + 1],
4014 ]));
4015 d2 += (buf[j] as i8) as f32 * s * xv;
4016 }
4017 (d1, d2)
4018 })
4019 };
4020 for r in start..end {
4021 let (v1, v2) = row_dots(r);
4022 unsafe {
4024 *p1.at(r) = v1;
4025 *p2.at(r) = v2;
4026 }
4027 }
4028}
4029
4030#[allow(clippy::too_many_arguments)]
4034fn vbit_range2_f32(
4035 bytes: &[u8],
4036 offsets: &[usize],
4037 x1: &[f32],
4038 x2: &[f32],
4039 rows: usize,
4040 cols: usize,
4041 p1: SendMut,
4042 p2: SendMut,
4043 start: usize,
4044 end: usize,
4045) {
4046 let ng = cols / GROUP_SIZE;
4047 let bits = &bytes[..rows];
4048 let sc_off = rows;
4049 #[inline(always)]
4050 #[allow(clippy::too_many_arguments)]
4051 fn dot_row2<const B: usize>(
4052 data: &[u8],
4053 bytes: &[u8],
4054 sc_off: usize,
4055 r: usize,
4056 ng: usize,
4057 x1: &[f32],
4058 x2: &[f32],
4059 ) -> (f32, f32) {
4060 let l = ((1i32 << (B - 1)) - 1) as f32;
4061 let gbytes = GROUP_SIZE * B / 8;
4062 let (mut d1, mut d2) = (0f32, 0f32);
4063 for g in 0..ng {
4064 let so = (r * ng + g) * 2;
4065 let s = f16_to_f32(u16::from_le_bytes([
4066 bytes[sc_off + so],
4067 bytes[sc_off + so + 1],
4068 ]));
4069 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4070 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4071 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4072 let (mut g1, mut g2) = (0f32, 0f32);
4073 for blk in 0..GROUP_SIZE / 8 {
4074 let u = unpack8::<B>(&gd0[blk * B..]);
4075 for k in 0..8 {
4076 let w = u[k] as f32 - l;
4077 g1 += w * x1g[blk * 8 + k];
4078 g2 += w * x2g[blk * 8 + k];
4079 }
4080 }
4081 d1 += g1 * s;
4082 d2 += g2 * s;
4083 }
4084 (d1, d2)
4085 }
4086 for r in start..end {
4087 let data = &bytes[offsets[r]..offsets[r + 1]];
4088 let (v1, v2) = match bits[r] {
4089 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4090 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4091 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4092 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4093 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4094 b => unreachable!("vbit bit-width {b} (validated at load)"),
4095 };
4096 unsafe {
4098 *p1.at(r) = v1;
4099 *p2.at(r) = v2;
4100 }
4101 }
4102}
4103
4104#[inline]
4111#[allow(unreachable_code)]
4112fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4113 #[cfg(target_arch = "aarch64")]
4114 unsafe {
4115 return dot_q4t_row_sdot(bytes, r, gpr, xq);
4116 }
4117 #[cfg(target_arch = "x86_64")]
4118 unsafe {
4119 if vnni_tiles_enabled() {
4120 return dot_q4t_row_vnni(bytes, r, gpr, xq);
4121 }
4122 return dot_q4t_row_avx2(bytes, r, gpr, xq);
4123 }
4124 let mut acc = 0f32;
4125 for gi in 0..gpr {
4126 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4127 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4128 let mut d = 0i32;
4129 for (k, &b) in tile[2..].iter().enumerate() {
4130 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4131 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4132 }
4133 acc += d as f32 * s;
4134 }
4135 acc
4136}
4137
4138#[cfg(target_arch = "aarch64")]
4139#[target_feature(enable = "neon,dotprod")]
4140unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4141 unsafe {
4144 use core::arch::aarch64::*;
4145 use core::arch::asm;
4146 let lomask = vdupq_n_u8(0x0F);
4147 let eight = vdupq_n_s8(8);
4148 let mut acc = 0f32;
4149 for gi in 0..gpr {
4150 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4151 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4152 let b = vld1q_u8(t.add(2));
4153 let lo = vandq_u8(b, lomask);
4154 let hi = vshrq_n_u8::<4>(b);
4155 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4156 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4157 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4158 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4159 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4160 asm!(
4161 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4162 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4163 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4164 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4165 options(pure, nomem, nostack),
4166 );
4167 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4168 }
4169 acc
4170 }
4171}
4172
4173#[cfg(target_arch = "x86_64")]
4174#[target_feature(enable = "avx2")]
4175unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4176 unsafe {
4178 use core::arch::x86_64::*;
4179 let lomask = _mm_set1_epi8(0x0F);
4180 let eight = _mm256_set1_epi8(8);
4181 let ones = _mm256_set1_epi16(1);
4182 let mut acc = 0f32;
4183 for gi in 0..gpr {
4184 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4185 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4186 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4187 let lo = _mm_and_si128(b, lomask);
4188 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4189 let w = _mm256_sub_epi8(
4190 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4191 eight,
4192 );
4193 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4194 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4195 let d = _mm256_madd_epi16(p16, ones);
4196 let hi128 = _mm256_extracti128_si256::<1>(d);
4197 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4198 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4199 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4200 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4201 }
4202 acc
4203 }
4204}
4205
4206#[cfg(target_arch = "x86_64")]
4210#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4211unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4212 unsafe {
4214 use core::arch::x86_64::*;
4215 let lomask = _mm_set1_epi8(0x0F);
4216 let eight = _mm256_set1_epi8(8);
4217 let mut acc = 0f32;
4218 for gi in 0..gpr {
4219 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4220 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4221 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4222 let lo = _mm_and_si128(b, lomask);
4223 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4224 let w = _mm256_sub_epi8(
4225 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4226 eight,
4227 );
4228 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4229 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4230 acc += d as f32 * s;
4231 }
4232 acc
4233 }
4234}
4235
4236#[cfg(target_arch = "x86_64")]
4241#[target_feature(enable = "avx2,fma")]
4246unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4247 unsafe {
4249 use core::arch::x86_64::*;
4250 let lomask = _mm_set1_epi8(0x0F);
4251 let eight = _mm256_set1_epi8(8);
4252 let ones = _mm256_set1_epi16(1);
4253 let mut f0 = _mm256_setzero_ps();
4266 let mut f1 = _mm256_setzero_ps();
4267 let mut f2 = _mm256_setzero_ps();
4268 let mut f3 = _mm256_setzero_ps();
4269 for gi in 0..gpr {
4270 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4271 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4272 let sv = _mm256_set1_ps(s);
4273 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4274 let lo = _mm_and_si128(bb, lomask);
4275 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4276 let w = _mm256_sub_epi8(
4277 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4278 eight,
4279 );
4280 let aw = _mm256_abs_epi8(w);
4281 let off = gi * GROUP_SIZE;
4282 let dot = |xq: &[i8]| {
4283 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4284 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4285 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4286 };
4287 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4288 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4289 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4290 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4291 }
4292 [
4293 hsum256_ps(f0),
4294 hsum256_ps(f1),
4295 hsum256_ps(f2),
4296 hsum256_ps(f3),
4297 ]
4298 }
4299}
4300
4301#[cfg(target_arch = "x86_64")]
4304#[target_feature(enable = "avx2")]
4305#[inline]
4306unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4307 unsafe {
4309 use core::arch::x86_64::*;
4310 let hi = _mm256_extractf128_ps::<1>(v);
4311 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4312 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4313 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4314 _mm_cvtss_f32(s)
4315 }
4316}
4317
4318#[cfg(target_arch = "x86_64")]
4320#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4321unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4322 unsafe {
4324 use core::arch::x86_64::*;
4325 let lomask = _mm_set1_epi8(0x0F);
4326 let eight = _mm256_set1_epi8(8);
4327 let mut f0 = _mm256_setzero_ps();
4330 let mut f1 = _mm256_setzero_ps();
4331 let mut f2 = _mm256_setzero_ps();
4332 let mut f3 = _mm256_setzero_ps();
4333 for gi in 0..gpr {
4334 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4335 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4336 let sv = _mm256_set1_ps(s);
4337 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4338 let lo = _mm_and_si128(bb, lomask);
4339 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4340 let w = _mm256_sub_epi8(
4341 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4342 eight,
4343 );
4344 let aw = _mm256_abs_epi8(w);
4345 let off = gi * GROUP_SIZE;
4346 let dot = |xq: &[i8]| {
4347 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4348 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4349 _mm256_setzero_si256(),
4350 aw,
4351 _mm256_sign_epi8(x, w),
4352 ))
4353 };
4354 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4355 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4356 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4357 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4358 }
4359 let acc = [
4360 hsum256_ps(f0),
4361 hsum256_ps(f1),
4362 hsum256_ps(f2),
4363 hsum256_ps(f3),
4364 ];
4365 acc
4366 }
4367}
4368
4369#[cfg(target_arch = "aarch64")]
4374#[target_feature(enable = "neon,dotprod")]
4375unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4376 unsafe {
4378 use core::arch::aarch64::*;
4379 use core::arch::asm;
4380 let lomask = vdupq_n_u8(0x0F);
4381 let eight = vdupq_n_s8(8);
4382 let mut acc = [0f32; 4];
4383 for gi in 0..gpr {
4384 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4385 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4386 let b = vld1q_u8(t.add(2));
4387 let lo = vandq_u8(b, lomask);
4388 let hi = vshrq_n_u8::<4>(b);
4389 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4390 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4391 for (k, xq) in xs.iter().enumerate() {
4392 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4393 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4394 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4395 asm!(
4396 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4397 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4398 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4399 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4400 options(pure, nomem, nostack),
4401 );
4402 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4403 }
4404 }
4405 acc
4406 }
4407}
4408
4409#[inline]
4411fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4412 let gi = j / GROUP_SIZE;
4413 let k = j % GROUP_SIZE;
4414 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4415 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4416 let byte = tile[2 + k / 2];
4417 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4418 ((nib as i32 - 8) as f32, s)
4419}
4420
4421#[inline]
4424fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4425 let mut acc = 0f32;
4426 for gi in 0..gpr {
4427 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4428 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4429 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4430 let mut ga = 0f32;
4431 for (k, &b) in tile[2..].iter().enumerate() {
4432 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4433 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4434 }
4435 acc += ga * s;
4436 }
4437 acc
4438}
4439
4440struct Q4tpView<'a> {
4444 nib: &'a [u8],
4445 params: &'a [u8],
4446 codes: &'a [u8],
4447 stride: usize,
4448 zero_rung: bool,
4450}
4451
4452impl<'a> Q4tpView<'a> {
4453 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4454 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4455 Self {
4456 nib: &bytes[..params_off],
4457 params: &bytes[params_off..codes_off],
4458 codes: &bytes[codes_off..],
4459 stride,
4460 zero_rung: false,
4461 }
4462 }
4463
4464 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4466 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4467 Self {
4468 nib: &bytes[..params_off],
4469 params: &bytes[params_off..codes_off],
4470 codes: &bytes[codes_off..],
4471 stride,
4472 zero_rung: true,
4473 }
4474 }
4475
4476 #[inline]
4492 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4493 let tab = if self.zero_rung {
4494 q2tp_ladder(self.params, r)
4495 } else {
4496 q4tp_ladder(self.params, r)
4497 };
4498 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4499 let out = &mut out[..gpr];
4500 let mut chunks = out.chunks_exact_mut(8);
4501 let mut ci = 0usize;
4502 for c in &mut chunks {
4503 let w = u64::from(codes[ci])
4504 | u64::from(codes[ci + 1]) << 8
4505 | u64::from(codes[ci + 2]) << 16
4506 | u64::from(codes[ci + 3]) << 24
4507 | u64::from(codes[ci + 4]) << 32;
4508 for (k, o) in c.iter_mut().enumerate() {
4509 *o = tab[((w >> (5 * k)) & 31) as usize];
4510 }
4511 ci += 5;
4512 }
4513 let tail = &codes[ci..];
4516 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4517 *o = tab[q4tp_code(tail, k)];
4518 }
4519 }
4520}
4521
4522#[inline]
4523fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4524 #[cfg(target_arch = "aarch64")]
4525 unsafe {
4526 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4527 }
4528 #[cfg(target_arch = "x86_64")]
4529 unsafe {
4530 if vnni_tiles_enabled() {
4531 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4532 }
4533 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4534 }
4535 #[allow(unreachable_code)]
4536 {
4537 let mut acc = 0f32;
4538 for gi in 0..gpr {
4539 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4540 let s = scales[gi];
4541 let mut d = 0i32;
4542 for (k, &b) in tile.iter().enumerate() {
4543 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4544 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4545 }
4546 acc += d as f32 * s;
4547 }
4548 acc
4549 }
4550}
4551
4552#[cfg(target_arch = "aarch64")]
4555#[target_feature(enable = "neon,dotprod")]
4556unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4557 unsafe {
4560 use core::arch::aarch64::*;
4561 use core::arch::asm;
4562 let lomask = vdupq_n_u8(0x0F);
4563 let eight = vdupq_n_s8(8);
4564 let mut acc = 0f32;
4565 for gi in 0..gpr {
4566 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4567 let s = *scales.get_unchecked(gi);
4568 let b = vld1q_u8(t);
4569 let lo = vandq_u8(b, lomask);
4570 let hi = vshrq_n_u8::<4>(b);
4571 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4572 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4573 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4574 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4575 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4576 asm!(
4577 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4578 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4579 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4580 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4581 options(pure, nomem, nostack),
4582 );
4583 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4584 }
4585 acc
4586 }
4587}
4588
4589#[cfg(target_arch = "x86_64")]
4590#[target_feature(enable = "avx2")]
4591unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4592 unsafe {
4594 use core::arch::x86_64::*;
4595 let lomask = _mm_set1_epi8(0x0F);
4596 let eight = _mm256_set1_epi8(8);
4597 let ones = _mm256_set1_epi16(1);
4598 let mut acc = 0f32;
4599 for gi in 0..gpr {
4600 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4601 let s = *scales.get_unchecked(gi);
4602 let b = _mm_loadu_si128(t as *const __m128i);
4603 let lo = _mm_and_si128(b, lomask);
4604 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4605 let w = _mm256_sub_epi8(
4606 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4607 eight,
4608 );
4609 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4610 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4611 let d = _mm256_madd_epi16(p16, ones);
4612 let hi128 = _mm256_extracti128_si256::<1>(d);
4613 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4614 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4615 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4616 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4617 }
4618 acc
4619 }
4620}
4621
4622#[cfg(target_arch = "x86_64")]
4625#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4626unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4627 unsafe {
4629 use core::arch::x86_64::*;
4630 let lomask = _mm_set1_epi8(0x0F);
4631 let eight = _mm256_set1_epi8(8);
4632 let mut acc = 0f32;
4633 for gi in 0..gpr {
4634 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4635 let s = *scales.get_unchecked(gi);
4636 let b = _mm_loadu_si128(t as *const __m128i);
4637 let lo = _mm_and_si128(b, lomask);
4638 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4639 let w = _mm256_sub_epi8(
4640 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4641 eight,
4642 );
4643 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4644 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4645 }
4646 acc
4647 }
4648}
4649
4650#[inline]
4653fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4654 let mut acc = 0f32;
4655 for gi in 0..gpr {
4656 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4657 let s = scales[gi];
4658 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4659 let mut ga = 0f32;
4660 for (k, &b) in tile.iter().enumerate() {
4661 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4662 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4663 }
4664 acc += ga * s;
4665 }
4666 acc
4667}
4668
4669#[inline]
4672fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4673 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4674 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4675 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4676 ((n as i32 - 8) as f32, scales[gi])
4677}
4678
4679fn q4tp_matvec(
4681 bytes: &[u8],
4682 x: &[f32],
4683 rows: usize,
4684 cols: usize,
4685 out: &mut [f32],
4686 pool: Option<&Pool>,
4687) {
4688 debug_assert_eq!(out.len(), rows);
4689 let gpr = cols / GROUP_SIZE;
4690 let v = Q4tpView::new(bytes, rows, cols);
4691 let out_addr = SendMut(out.as_mut_ptr());
4692 if a8w8_enabled() {
4693 let act = split_act(x);
4694 let run = |start: usize, end: usize| {
4695 with_krow(gpr, |sc| {
4697 for r in start..end {
4698 v.scales_into(r, gpr, sc);
4699 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4700 for &(j, xv) in &act.outliers {
4701 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4702 acc += w * s * xv;
4703 }
4704 unsafe { *out_addr.at(r) = acc };
4706 }
4707 })
4708 };
4709 dispatch_rows(pool, rows, &run);
4710 return;
4711 }
4712 let run = |start: usize, end: usize| {
4713 with_krow(gpr, |sc| {
4714 for r in start..end {
4715 v.scales_into(r, gpr, sc);
4716 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4718 }
4719 })
4720 };
4721 dispatch_rows(pool, rows, &run);
4722}
4723
4724#[allow(clippy::too_many_arguments)]
4727fn q4tp_matvec2(
4728 bytes: &[u8],
4729 x1: &[f32],
4730 x2: &[f32],
4731 rows: usize,
4732 cols: usize,
4733 o1: &mut [f32],
4734 o2: &mut [f32],
4735 pool: Option<&Pool>,
4736) {
4737 let gpr = cols / GROUP_SIZE;
4738 let v = Q4tpView::new(bytes, rows, cols);
4739 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4740 let run = |start: usize, end: usize| {
4741 let mut sc = vec![0f32; gpr];
4742 for r in start..end {
4743 v.scales_into(r, gpr, &mut sc);
4744 unsafe {
4746 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4747 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4748 }
4749 }
4750 };
4751 dispatch_rows(pool, rows, &run);
4752}
4753
4754#[inline]
4757fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4758 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4759 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4760 let c = (byte >> (2 * (k % 4))) & 3;
4761 (c as f32 - 1.5, scales[gi])
4762}
4763
4764#[cfg(target_arch = "x86_64")]
4765const Q2TP_DECODE_U32: [u32; 256] = {
4766 let mut tab = [0u32; 256];
4767 let mut b = 0usize;
4768 while b < 256 {
4769 tab[b] = ((b as u32) & 3)
4770 | ((((b as u32) >> 2) & 3) << 8)
4771 | ((((b as u32) >> 4) & 3) << 16)
4772 | ((((b as u32) >> 6) & 3) << 24);
4773 b += 1;
4774 }
4775 tab
4776};
4777
4778#[cfg(target_arch = "x86_64")]
4782#[target_feature(enable = "avx2")]
4783unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
4784 use core::arch::x86_64::*;
4785 debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
4786 let codes = _mm256_setr_epi32(
4787 Q2TP_DECODE_U32[ch[0] as usize] as i32,
4788 Q2TP_DECODE_U32[ch[1] as usize] as i32,
4789 Q2TP_DECODE_U32[ch[2] as usize] as i32,
4790 Q2TP_DECODE_U32[ch[3] as usize] as i32,
4791 Q2TP_DECODE_U32[ch[4] as usize] as i32,
4792 Q2TP_DECODE_U32[ch[5] as usize] as i32,
4793 Q2TP_DECODE_U32[ch[6] as usize] as i32,
4794 Q2TP_DECODE_U32[ch[7] as usize] as i32,
4795 );
4796 let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
4797 let pair = _mm256_maddubs_epi16(codes, xv);
4798 let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
4799 let sum128 = _mm_add_epi32(
4800 _mm256_castsi256_si128(quad),
4801 _mm256_extracti128_si256(quad, 1),
4802 );
4803 let sum64 = _mm_hadd_epi32(sum128, sum128);
4804 _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
4805}
4806
4807#[inline]
4814fn dot_q2tp_row_i8(
4815 chunks: &[u8],
4816 r: usize,
4817 gpr: usize,
4818 xq: &[i8],
4819 gsum: &[i32],
4820 scales: &[f32],
4821) -> f32 {
4822 let mut acc = 0f32;
4823 let base = r * gpr * Q2TP_CHUNK;
4824 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
4825 let mut codes = [0i8; GROUP_SIZE];
4826 #[cfg(target_arch = "x86_64")]
4827 let avx2 = std::arch::is_x86_feature_detected!("avx2");
4828 for gi in 0..gpr {
4829 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4830 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4831 #[cfg(target_arch = "aarch64")]
4832 let dot = unsafe {
4838 use core::arch::aarch64::*;
4839 let b = vld1_u8(ch.as_ptr());
4840 let three = vdup_n_u8(3);
4841 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4842 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4843 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4844 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4845 let x4 = vld4_s8(xg.as_ptr());
4846 let mut acc4 = vdupq_n_s32(0);
4847 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4848 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4849 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4850 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4851 vaddvq_s32(acc4)
4852 };
4853 #[cfg(target_arch = "x86_64")]
4854 let dot: i32 = if avx2 {
4855 unsafe { q2tp_code_dot_avx2(ch, xg) }
4858 } else {
4859 ch.iter()
4860 .enumerate()
4861 .map(|(k, &b)| {
4862 ((b & 3) as i32) * xg[k * 4] as i32
4863 + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
4864 + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
4865 + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
4866 })
4867 .sum()
4868 };
4869 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
4870 let dot: i32 = {
4871 for (k, &b) in ch.iter().enumerate() {
4872 codes[k * 4] = (b & 3) as i8;
4873 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4874 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4875 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4876 }
4877 codes
4878 .iter()
4879 .zip(xg)
4880 .map(|(&c, &x)| c as i32 * x as i32)
4881 .sum()
4882 };
4883 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4884 }
4885 acc
4886}
4887
4888fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4892 let mut acc = 0f32;
4893 for gi in 0..gpr {
4894 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4895 let s = scales[gi];
4896 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4897 let mut g = 0f32;
4898 for (k, &b) in ch.iter().enumerate() {
4899 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4900 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4901 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4902 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4903 }
4904 acc += s * g;
4905 }
4906 acc
4907}
4908
4909fn q2tp_matvec(
4910 bytes: &[u8],
4911 x: &[f32],
4912 rows: usize,
4913 cols: usize,
4914 out: &mut [f32],
4915 pool: Option<&Pool>,
4916) {
4917 debug_assert_eq!(out.len(), rows);
4918 let gpr = cols / GROUP_SIZE;
4919 let v = Q4tpView::new_q2(bytes, rows, cols);
4920 let out_addr = SendMut(out.as_mut_ptr());
4921 if a8w8_enabled() {
4926 let act = split_act(x);
4927 let gsum = q1_group_sums(&act.xq, gpr);
4928 let (act, gsum) = (&act, &gsum);
4929 let run = move |start: usize, end: usize| {
4930 with_krow(gpr, |sc| {
4931 for r in start..end {
4932 v.scales_into(r, gpr, sc);
4933 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4934 for &(j, xv) in &act.outliers {
4935 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4936 acc += w * s * xv;
4937 }
4938 unsafe { *out_addr.at(r) = acc };
4940 }
4941 })
4942 };
4943 dispatch_rows(pool, rows, &run);
4944 return;
4945 }
4946 let run = |start: usize, end: usize| {
4947 with_krow(gpr, |sc| {
4948 for r in start..end {
4949 v.scales_into(r, gpr, sc);
4950 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4952 }
4953 })
4954 };
4955 dispatch_rows(pool, rows, &run);
4956}
4957
4958#[allow(clippy::too_many_arguments)]
4960fn q2tp_matvec2(
4961 bytes: &[u8],
4962 x1: &[f32],
4963 x2: &[f32],
4964 rows: usize,
4965 cols: usize,
4966 o1: &mut [f32],
4967 o2: &mut [f32],
4968 pool: Option<&Pool>,
4969) {
4970 let gpr = cols / GROUP_SIZE;
4971 let v = Q4tpView::new_q2(bytes, rows, cols);
4972 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4973 let run = |start: usize, end: usize| {
4974 let mut sc = vec![0f32; gpr];
4975 for r in start..end {
4976 v.scales_into(r, gpr, &mut sc);
4977 unsafe {
4979 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4980 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4981 }
4982 }
4983 };
4984 dispatch_rows(pool, rows, &run);
4985}
4986
4987pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4994 let gpr = cols / GROUP_SIZE;
4999 let v = Q4tpView::new_q2(bytes, rows, cols);
5000 with_krow(gpr, |sc| {
5001 for r in 0..rows {
5002 v.scales_into(r, gpr, sc);
5003 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
5004 }
5005 });
5006}
5007
5008pub fn q2tp_matmat_for_test(
5009 bytes: &[u8],
5010 xs_all: &[f32],
5011 b: usize,
5012 rows: usize,
5013 cols: usize,
5014 out: &mut [f32],
5015) {
5016 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
5017}
5018
5019fn q2tp_matmat(
5020 bytes: &[u8],
5021 xs_all: &[f32],
5022 b: usize,
5023 rows: usize,
5024 cols: usize,
5025 out: &mut [f32],
5026 pool: Option<&Pool>,
5027) {
5028 debug_assert_eq!(out.len(), b * rows);
5029 let gpr = cols / GROUP_SIZE;
5030 let v = Q4tpView::new_q2(bytes, rows, cols);
5031 let out_addr = SendMut(out.as_mut_ptr());
5032 let run = |start: usize, end: usize| {
5033 let mut sc = vec![0f32; gpr];
5034 for r in start..end {
5035 v.scales_into(r, gpr, &mut sc);
5036 for bi in 0..b {
5037 let x = &xs_all[bi * cols..(bi + 1) * cols];
5038 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
5040 }
5041 }
5042 };
5043 dispatch_rows(pool, rows, &run);
5044}
5045
5046#[cfg(target_arch = "aarch64")]
5050#[target_feature(enable = "neon,dotprod")]
5051unsafe fn dot_q4tp_row_1x4_sdot_v1(
5052 nib: &[u8],
5053 r: usize,
5054 gpr: usize,
5055 xs: [&[i8]; 4],
5056 scales: &[f32],
5057) -> [f32; 4] {
5058 unsafe {
5059 use core::arch::aarch64::*;
5060 use core::arch::asm;
5061 let lomask = vdupq_n_u8(0x0F);
5062 let eight = vdupq_n_s8(8);
5063 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
5064 for gi in 0..gpr {
5065 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5066 let s = *scales.get_unchecked(gi);
5067 let bb = vld1q_u8(t);
5068 let lo = vandq_u8(bb, lomask);
5069 let hi = vshrq_n_u8::<4>(bb);
5070 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5071 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5072 let mut d = [0f32; 4];
5073 for (k, dk) in d.iter_mut().enumerate() {
5074 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5075 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5076 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5077 asm!(
5078 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5079 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5080 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5081 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5082 options(pure, nomem, nostack),
5083 );
5084 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5085 }
5086 f0 += d[0];
5087 f1 += d[1];
5088 f2 += d[2];
5089 f3 += d[3];
5090 }
5091 [f0, f1, f2, f3]
5092 }
5093}
5094
5095#[allow(dead_code)]
5102static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
5103
5104#[cfg(target_arch = "x86_64")]
5110fn q4tp_blocked_x86() -> bool {
5111 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5112 1 => false,
5113 2 => avx512vnni_enabled(),
5118 _ => blocked_enabled() && avx512vnni_enabled(),
5122 }
5123}
5124
5125#[cfg(target_arch = "aarch64")]
5127#[allow(dead_code)]
5128fn q4tp_v1() -> bool {
5129 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5130 1 => true,
5131 2 => false,
5132 _ => {
5133 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5134 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
5135 }
5136 }
5137}
5138
5139#[cfg(target_arch = "x86_64")]
5151#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5152unsafe fn dot_q4tp_2x8_avx512(
5153 nib: &[u8],
5154 r0: usize,
5155 gpr: usize,
5156 xs: [&[i8]; 8],
5157 sc0: &[f32],
5158 sc1: &[f32],
5159) -> [[f32; 8]; 2] {
5160 unsafe {
5163 use core::arch::x86_64::*;
5164 let lomask = _mm256_set1_epi8(0x0F);
5165 let eight = _mm256_set1_epi8(8);
5166 let zero = _mm512_setzero_si512();
5167 let mut v0 = [_mm512_setzero_ps(); 8];
5168 let mut v1 = [_mm512_setzero_ps(); 8];
5169 let pairs = gpr / 2;
5170 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
5171 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5172 let bb = _mm256_loadu_si256(t as *const __m256i);
5173 let lo = _mm256_and_si256(bb, lomask);
5174 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5175 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5176 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5177 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5178 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5179 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5180 };
5181 for gp in 0..pairs {
5182 let gi = gp * 2;
5183 let (wa0, neg0) = unpack(r0, gi);
5184 let (wa1, neg1) = unpack(r0 + 1, gi);
5185 let off = gi * GROUP_SIZE;
5186 let sv = |sc: &[f32]| {
5187 _mm512_insertf32x8::<1>(
5188 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5189 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5190 )
5191 };
5192 let s0 = sv(sc0);
5193 let s1 = sv(sc1);
5194 for k in 0..8 {
5195 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5196 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5197 zero,
5198 wa0,
5199 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5200 ));
5201 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5202 zero,
5203 wa1,
5204 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5205 ));
5206 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5207 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5208 }
5209 }
5210 let mut acc = [[0f32; 8]; 2];
5211 for k in 0..8 {
5212 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5213 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5214 }
5215 if gpr % 2 == 1 {
5216 let off = (gpr - 1) * GROUP_SIZE;
5217 for j in off..off + GROUP_SIZE {
5218 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5219 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5220 for k in 0..8 {
5221 let x = *xs[k].get_unchecked(j) as f32;
5222 acc[0][k] += w0 * sa * x;
5223 acc[1][k] += w1 * sb * x;
5224 }
5225 }
5226 }
5227 acc
5228 }
5229}
5230
5231#[cfg(target_arch = "x86_64")]
5236#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5237unsafe fn dot_q4tp_row_1x8_avx512(
5238 nib: &[u8],
5239 r: usize,
5240 gpr: usize,
5241 xs: [&[i8]; 8],
5242 scales: &[f32],
5243) -> [f32; 8] {
5244 unsafe {
5246 use core::arch::x86_64::*;
5247 let lomask = _mm256_set1_epi8(0x0F);
5248 let eight = _mm256_set1_epi8(8);
5249 let zero = _mm512_setzero_si512();
5250 let (mut v0, mut v1, mut v2, mut v3) = (
5251 _mm512_setzero_ps(),
5252 _mm512_setzero_ps(),
5253 _mm512_setzero_ps(),
5254 _mm512_setzero_ps(),
5255 );
5256 let (mut v4, mut v5, mut v6, mut v7) = (
5257 _mm512_setzero_ps(),
5258 _mm512_setzero_ps(),
5259 _mm512_setzero_ps(),
5260 _mm512_setzero_ps(),
5261 );
5262 let pairs = gpr / 2;
5263 for gp in 0..pairs {
5264 let gi = gp * 2;
5265 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5266 let bb = _mm256_loadu_si256(t as *const __m256i);
5267 let lo = _mm256_and_si256(bb, lomask);
5268 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5269 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5274 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5275 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5276 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5277 let wabs = _mm512_abs_epi8(w);
5278 let neg = _mm512_movepi8_mask(w);
5279 let off = gi * GROUP_SIZE;
5280 let sv = _mm512_insertf32x8::<1>(
5281 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5282 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5283 );
5284 let dot = |x: &[i8]| -> __m512 {
5285 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5286 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5287 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5288 };
5289 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5290 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5291 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5292 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5293 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5294 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5295 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5296 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5297 }
5298 let mut acc = [
5299 _mm512_reduce_add_ps(v0),
5300 _mm512_reduce_add_ps(v1),
5301 _mm512_reduce_add_ps(v2),
5302 _mm512_reduce_add_ps(v3),
5303 _mm512_reduce_add_ps(v4),
5304 _mm512_reduce_add_ps(v5),
5305 _mm512_reduce_add_ps(v6),
5306 _mm512_reduce_add_ps(v7),
5307 ];
5308 if gpr % 2 == 1 {
5311 let off = (gpr - 1) * GROUP_SIZE;
5312 for j in off..off + GROUP_SIZE {
5313 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5314 let ws = w * s;
5315 for k in 0..8 {
5316 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5317 }
5318 }
5319 }
5320 acc
5321 }
5322}
5323
5324#[cfg(target_arch = "x86_64")]
5337#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5338unsafe fn dot_q4tp_row_1x4_avx512(
5339 nib: &[u8],
5340 r: usize,
5341 gpr: usize,
5342 xs: [&[i8]; 4],
5343 scales: &[f32],
5344) -> [f32; 4] {
5345 unsafe {
5347 use core::arch::x86_64::*;
5348 let lomask = _mm256_set1_epi8(0x0F);
5349 let eight = _mm256_set1_epi8(8);
5350 let zero = _mm512_setzero_si512();
5351 let (mut v0, mut v1, mut v2, mut v3) = (
5352 _mm512_setzero_ps(),
5353 _mm512_setzero_ps(),
5354 _mm512_setzero_ps(),
5355 _mm512_setzero_ps(),
5356 );
5357 let pairs = gpr / 2;
5358 for gp in 0..pairs {
5359 let gi = gp * 2;
5360 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5361 let bb = _mm256_loadu_si256(t as *const __m256i);
5362 let lo = _mm256_and_si256(bb, lomask);
5363 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5364 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5369 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5370 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5371 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5372 let wabs = _mm512_abs_epi8(w);
5373 let neg = _mm512_movepi8_mask(w);
5374 let off = gi * GROUP_SIZE;
5375 let sv = _mm512_insertf32x8::<1>(
5376 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5377 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5378 );
5379 let dot = |x: &[i8]| -> __m512 {
5380 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5381 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5382 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5383 };
5384 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5385 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5386 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5387 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5388 }
5389 let mut acc = [
5390 _mm512_reduce_add_ps(v0),
5391 _mm512_reduce_add_ps(v1),
5392 _mm512_reduce_add_ps(v2),
5393 _mm512_reduce_add_ps(v3),
5394 ];
5395 if gpr % 2 == 1 {
5398 let off = (gpr - 1) * GROUP_SIZE;
5399 for j in off..off + GROUP_SIZE {
5400 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5401 let ws = w * s;
5402 for k in 0..4 {
5403 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5404 }
5405 }
5406 }
5407 acc
5408 }
5409}
5410
5411#[cfg(target_arch = "aarch64")]
5415#[target_feature(enable = "neon,dotprod")]
5416unsafe fn dot_q4tp_row_1x4_sdot(
5417 nib: &[u8],
5418 r: usize,
5419 gpr: usize,
5420 xs: [&[i8]; 4],
5421 scales: &[f32],
5422) -> [f32; 4] {
5423 unsafe {
5425 use core::arch::aarch64::*;
5426 use core::arch::asm;
5427 let lomask = vdupq_n_u8(0x0F);
5428 let eight = vdupq_n_s8(8);
5429 let (mut v0, mut v1, mut v2, mut v3) = (
5445 vdupq_n_f32(0.0),
5446 vdupq_n_f32(0.0),
5447 vdupq_n_f32(0.0),
5448 vdupq_n_f32(0.0),
5449 );
5450 for gi in 0..gpr {
5451 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5452 let s = *scales.get_unchecked(gi);
5453 let bb = vld1q_u8(t);
5454 let lo = vandq_u8(bb, lomask);
5455 let hi = vshrq_n_u8::<4>(bb);
5456 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5457 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5458 let off = gi * GROUP_SIZE;
5459 let dot4 = |x: &[i8]| -> int32x4_t {
5460 let x0 = vld1q_s8(x.as_ptr().add(off));
5461 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5462 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5463 asm!(
5464 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5465 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5466 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5467 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5468 options(pure, nomem, nostack),
5469 );
5470 vaddq_s32(a0, a1)
5471 };
5472 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5473 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5474 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5475 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5476 }
5477 [
5478 vaddvq_f32(v0),
5479 vaddvq_f32(v1),
5480 vaddvq_f32(v2),
5481 vaddvq_f32(v3),
5482 ]
5483 }
5484}
5485
5486fn q4tp_matmat(
5490 bytes: &[u8],
5491 xs_all: &[f32],
5492 b: usize,
5493 rows: usize,
5494 cols: usize,
5495 out: &mut [f32],
5496 pool: Option<&Pool>,
5497) {
5498 debug_assert_eq!(out.len(), b * rows);
5499 let gpr = cols / GROUP_SIZE;
5500 let v = Q4tpView::new(bytes, rows, cols);
5501
5502 #[cfg(target_os = "macos")]
5504 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5505 dequant_matmat_accel(
5506 &|r, dst| {
5507 let mut sc = [0f32; 32];
5508 let mut scv;
5509 let s: &[f32] = if gpr <= 32 {
5510 v.scales_into(r, gpr, &mut sc);
5511 &sc[..gpr]
5512 } else {
5513 scv = vec![0f32; gpr];
5514 v.scales_into(r, gpr, &mut scv);
5515 &scv
5516 };
5517 for gi in 0..gpr {
5518 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5519 for (k, &bb) in tile.iter().enumerate() {
5520 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5521 dst[gi * GROUP_SIZE + k * 2 + 1] =
5522 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5523 }
5524 }
5525 },
5526 xs_all,
5527 b,
5528 rows,
5529 cols,
5530 out,
5531 pool,
5532 );
5533 return;
5534 }
5535
5536 let out_addr = SendMut(out.as_mut_ptr());
5537 if a8w8_enabled() {
5538 let acts: Vec<SplitAct> = (0..b)
5539 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5540 .collect();
5541 let acts = &acts;
5542 #[cfg(target_arch = "aarch64")]
5543 let blocked_ok = sdot_enabled() && blocked_enabled();
5544 #[cfg(target_arch = "x86_64")]
5551 let blocked_ok = q4tp_blocked_x86();
5552 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5553 let blocked_ok = false;
5554 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5564 .ok()
5565 .and_then(|v| v.parse().ok())
5566 .filter(|v| *v > 0)
5567 .unwrap_or(256);
5568 let run = |start: usize, end: usize| {
5569 for abase in (0..acts.len()).step_by(panel_cols) {
5570 let alen = (acts.len() - abase).min(panel_cols);
5571 let mut sc = vec![0f32; gpr];
5572 #[cfg(target_arch = "x86_64")]
5573 let mut r_lo = start;
5574 #[cfg(target_arch = "x86_64")]
5575 if blocked_ok && alen >= 8 {
5576 let mut sc1 = vec![0f32; gpr];
5577 while r_lo + 2 <= end {
5578 v.scales_into(r_lo, gpr, &mut sc);
5579 v.scales_into(r_lo + 1, gpr, &mut sc1);
5580 let mut bi = 0usize;
5581 while bi + 8 <= alen {
5582 let xs = [
5583 acts[abase + bi].xq.as_slice(),
5584 acts[abase + bi + 1].xq.as_slice(),
5585 acts[abase + bi + 2].xq.as_slice(),
5586 acts[abase + bi + 3].xq.as_slice(),
5587 acts[abase + bi + 4].xq.as_slice(),
5588 acts[abase + bi + 5].xq.as_slice(),
5589 acts[abase + bi + 6].xq.as_slice(),
5590 acts[abase + bi + 7].xq.as_slice(),
5591 ];
5592 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5593 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
5594 for k in 0..8 {
5595 let act = &acts[abase + bi + k];
5596 let mut acc = dr[k] * act.sx;
5597 for &(j, xv) in &act.outliers {
5598 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5599 acc += w * s * xv;
5600 }
5601 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5603 }
5604 }
5605 bi += 8;
5606 }
5607 for row in [r_lo, r_lo + 1] {
5610 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5611 for b2 in bi..alen {
5612 let act = &acts[abase + b2];
5613 let xs4 = [
5614 act.xq.as_slice(),
5615 act.xq.as_slice(),
5616 act.xq.as_slice(),
5617 act.xq.as_slice(),
5618 ];
5619 let d =
5620 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5621 let mut acc = d[0] * act.sx;
5622 for &(j, xv) in &act.outliers {
5623 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5624 acc += w * s * xv;
5625 }
5626 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5628 }
5629 }
5630 r_lo += 2;
5631 }
5632 }
5633 #[cfg(target_arch = "x86_64")]
5634 let row_start = r_lo;
5635 #[cfg(not(target_arch = "x86_64"))]
5636 let row_start = start;
5637 for r in row_start..end {
5638 v.scales_into(r, gpr, &mut sc);
5639 let mut bi = 0usize;
5640 #[cfg(target_arch = "x86_64")]
5641 if blocked_ok {
5642 while bi + 8 <= alen {
5643 let xs = [
5644 acts[abase + bi].xq.as_slice(),
5645 acts[abase + bi + 1].xq.as_slice(),
5646 acts[abase + bi + 2].xq.as_slice(),
5647 acts[abase + bi + 3].xq.as_slice(),
5648 acts[abase + bi + 4].xq.as_slice(),
5649 acts[abase + bi + 5].xq.as_slice(),
5650 acts[abase + bi + 6].xq.as_slice(),
5651 acts[abase + bi + 7].xq.as_slice(),
5652 ];
5653 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5654 for k in 0..8 {
5655 let act = &acts[abase + bi + k];
5656 let mut acc = d[k] * act.sx;
5657 for &(j, xv) in &act.outliers {
5658 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5659 acc += w * s * xv;
5660 }
5661 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5663 }
5664 bi += 8;
5665 }
5666 while bi + 4 <= alen {
5667 let xs = [
5668 acts[abase + bi].xq.as_slice(),
5669 acts[abase + bi + 1].xq.as_slice(),
5670 acts[abase + bi + 2].xq.as_slice(),
5671 acts[abase + bi + 3].xq.as_slice(),
5672 ];
5673 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5674 for k in 0..4 {
5675 let act = &acts[abase + bi + k];
5676 let mut acc = d[k] * act.sx;
5677 for &(j, xv) in &act.outliers {
5678 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5679 acc += w * s * xv;
5680 }
5681 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5683 }
5684 bi += 4;
5685 }
5686 }
5687 #[cfg(target_arch = "aarch64")]
5688 if blocked_ok {
5689 while bi + 4 <= alen {
5690 let xs = [
5691 acts[abase + bi].xq.as_slice(),
5692 acts[abase + bi + 1].xq.as_slice(),
5693 acts[abase + bi + 2].xq.as_slice(),
5694 acts[abase + bi + 3].xq.as_slice(),
5695 ];
5696 let d = unsafe {
5697 if q4tp_v1() {
5698 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5699 } else {
5700 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5701 }
5702 };
5703 for k in 0..4 {
5704 let act = &acts[abase + bi + k];
5705 let mut acc = d[k] * act.sx;
5706 for &(j, xv) in &act.outliers {
5707 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5708 acc += w * s * xv;
5709 }
5710 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5712 }
5713 bi += 4;
5714 }
5715 }
5716 let _ = blocked_ok;
5717 while bi < alen {
5718 let act = &acts[abase + bi];
5719 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5720 for &(j, xv) in &act.outliers {
5721 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5722 acc += w * s * xv;
5723 }
5724 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5726 bi += 1;
5727 }
5728 }
5729 }
5730 };
5731 dispatch_rows(pool, rows, &run);
5732 return;
5733 }
5734
5735 let run = |start: usize, end: usize| {
5736 let mut sc = vec![0f32; gpr];
5737 for r in start..end {
5738 v.scales_into(r, gpr, &mut sc);
5739 for bi in 0..b {
5740 let x = &xs_all[bi * cols..(bi + 1) * cols];
5741 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5743 }
5744 }
5745 };
5746 dispatch_rows(pool, rows, &run);
5747}
5748
5749fn q4t_matvec(
5751 bytes: &[u8],
5752 x: &[f32],
5753 rows: usize,
5754 cols: usize,
5755 out: &mut [f32],
5756 pool: Option<&Pool>,
5757) {
5758 debug_assert_eq!(out.len(), rows);
5759 let gpr = cols / GROUP_SIZE;
5760 let out_addr = SendMut(out.as_mut_ptr());
5761 if a8w8_enabled() {
5762 let act = split_act(x);
5763 let run = move |start: usize, end: usize| {
5764 for r in start..end {
5765 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5766 for &(j, xv) in &act.outliers {
5767 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5768 acc += w * s * xv;
5769 }
5770 unsafe { *out_addr.at(r) = acc };
5772 }
5773 };
5774 dispatch_rows(pool, rows, &run);
5775 return;
5776 }
5777 let run = move |start: usize, end: usize| {
5778 for r in start..end {
5779 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5781 }
5782 };
5783 dispatch_rows(pool, rows, &run);
5784}
5785
5786#[allow(clippy::too_many_arguments)]
5788fn q4t_matvec2(
5789 bytes: &[u8],
5790 x1: &[f32],
5791 x2: &[f32],
5792 rows: usize,
5793 cols: usize,
5794 o1: &mut [f32],
5795 o2: &mut [f32],
5796 pool: Option<&Pool>,
5797) {
5798 let gpr = cols / GROUP_SIZE;
5799 let p1 = SendMut(o1.as_mut_ptr());
5800 let p2 = SendMut(o2.as_mut_ptr());
5801 if a8w8_enabled() {
5802 let a1 = split_act(x1);
5803 let a2 = split_act(x2);
5804 let run = move |start: usize, end: usize| {
5805 for r in start..end {
5806 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5807 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5808 for &(j, xv) in &a1.outliers {
5809 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5810 v1 += w * s * xv;
5811 }
5812 for &(j, xv) in &a2.outliers {
5813 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5814 v2 += w * s * xv;
5815 }
5816 unsafe {
5818 *p1.at(r) = v1;
5819 *p2.at(r) = v2;
5820 }
5821 }
5822 };
5823 dispatch_rows(pool, rows, &run);
5824 return;
5825 }
5826 let run = move |start: usize, end: usize| {
5827 for r in start..end {
5828 unsafe {
5830 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5831 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5832 }
5833 }
5834 };
5835 dispatch_rows(pool, rows, &run);
5836}
5837
5838#[allow(clippy::too_many_arguments)]
5840#[cfg(target_os = "macos")]
5846fn dequant_matmat_accel(
5847 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5848 xs_all: &[f32],
5849 b: usize,
5850 rows: usize,
5851 cols: usize,
5852 out: &mut [f32],
5853 pool: Option<&Pool>,
5854) {
5855 const TR: usize = 2048;
5856 thread_local! {
5857 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5858 }
5859 WTILE.with(|wt| {
5860 let mut wtile = wt.borrow_mut();
5861 wtile.resize(TR * cols, 0.0);
5862 let mut r0 = 0usize;
5863 while r0 < rows {
5864 let tr = TR.min(rows - r0);
5865 let wt_addr = SendMut(wtile.as_mut_ptr());
5866 let run = |start: usize, end: usize| {
5867 for r in start..end {
5868 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5870 dequant_row(r0 + r, dst);
5871 }
5872 };
5873 dispatch_rows(pool, tr, &run);
5874 unsafe {
5875 accel_blas::cblas_sgemm(
5876 101, 111, 112, b as i32,
5880 tr as i32,
5881 cols as i32,
5882 1.0,
5883 xs_all.as_ptr(),
5884 cols as i32,
5885 wtile.as_ptr(),
5886 cols as i32,
5887 0.0,
5888 out.as_mut_ptr().add(r0),
5889 rows as i32,
5890 );
5891 }
5892 r0 += tr;
5893 }
5894 });
5895}
5896
5897fn q4t_matmat(
5898 bytes: &[u8],
5899 xs_all: &[f32],
5900 b: usize,
5901 rows: usize,
5902 cols: usize,
5903 out: &mut [f32],
5904 pool: Option<&Pool>,
5905) {
5906 debug_assert_eq!(out.len(), b * rows);
5907 let gpr = cols / GROUP_SIZE;
5908 #[cfg(target_os = "macos")]
5912 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5913 dequant_matmat_accel(
5914 &|r, dst| {
5915 for gi in 0..gpr {
5916 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5917 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5918 for (k, &bb) in tile[2..].iter().enumerate() {
5919 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5920 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5921 }
5922 }
5923 },
5924 xs_all,
5925 b,
5926 rows,
5927 cols,
5928 out,
5929 pool,
5930 );
5931 return;
5932 }
5933 let out_addr = SendMut(out.as_mut_ptr());
5934 if a8w8_enabled() {
5935 let acts: Vec<SplitAct> = (0..b)
5936 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5937 .collect();
5938 let acts = &acts;
5939 #[cfg(target_arch = "x86_64")]
5940 let blocked_ok = avx2_enabled() && blocked_enabled();
5941 #[cfg(target_arch = "aarch64")]
5942 let blocked_ok = sdot_enabled() && blocked_enabled();
5943 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5944 let blocked_ok = false;
5945 let run = move |start: usize, end: usize| {
5946 for r in start..end {
5947 let mut bi = 0usize;
5948 #[cfg(target_arch = "aarch64")]
5949 if blocked_ok {
5950 while bi + 4 <= acts.len() {
5951 let xs = [
5952 acts[bi].xq.as_slice(),
5953 acts[bi + 1].xq.as_slice(),
5954 acts[bi + 2].xq.as_slice(),
5955 acts[bi + 3].xq.as_slice(),
5956 ];
5957 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5958 for k in 0..4 {
5959 let act = &acts[bi + k];
5960 let mut acc = d[k] * act.sx;
5961 for &(j, xv) in &act.outliers {
5962 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5963 acc += w * sc * xv;
5964 }
5965 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5967 }
5968 bi += 4;
5969 }
5970 }
5971 #[cfg(target_arch = "x86_64")]
5972 if blocked_ok {
5973 while bi + 4 <= acts.len() {
5974 let xs = [
5975 acts[bi].xq.as_slice(),
5976 acts[bi + 1].xq.as_slice(),
5977 acts[bi + 2].xq.as_slice(),
5978 acts[bi + 3].xq.as_slice(),
5979 ];
5980 let d = unsafe {
5981 if vnni_tiles_enabled() {
5982 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5983 } else {
5984 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5985 }
5986 };
5987 for k in 0..4 {
5988 let act = &acts[bi + k];
5989 let mut acc = d[k] * act.sx;
5990 for &(j, xv) in &act.outliers {
5991 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5992 acc += w * sc * xv;
5993 }
5994 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5996 }
5997 bi += 4;
5998 }
5999 }
6000 let _ = blocked_ok;
6001 while bi < acts.len() {
6002 let act = &acts[bi];
6003 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6004 for &(j, xv) in &act.outliers {
6005 let (w, s) = q4t_outlier(bytes, r, gpr, j);
6006 acc += w * s * xv;
6007 }
6008 unsafe { *out_addr.at(bi * rows + r) = acc };
6010 bi += 1;
6011 }
6012 }
6013 };
6014 dispatch_rows(pool, rows, &run);
6015 return;
6016 }
6017 let run = move |start: usize, end: usize| {
6018 for r in start..end {
6019 for bi in 0..b {
6020 let x = &xs_all[bi * cols..(bi + 1) * cols];
6021 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
6023 }
6024 }
6025 };
6026 dispatch_rows(pool, rows, &run);
6027}
6028
6029fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
6038 (0..gpr)
6039 .map(|gi| {
6040 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
6041 .iter()
6042 .map(|&v| v as i32)
6043 .sum()
6044 })
6045 .collect()
6046}
6047
6048#[inline]
6052#[allow(unreachable_code)]
6053#[cfg(target_arch = "x86_64")]
6058#[target_feature(enable = "avx2")]
6059unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6060 unsafe {
6062 use core::arch::x86_64::*;
6063 let expand = _mm256_setr_epi8(
6065 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,
6066 3, 3, 3,
6067 );
6068 let bitsel = _mm256_setr_epi8(
6069 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6070 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6071 );
6072 let ones8 = _mm256_set1_epi8(1);
6073 let ones16 = _mm256_set1_epi16(1);
6074 let mut acc = 0f32;
6075 for gi in 0..gpr {
6076 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6077 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6078 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6079 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6080 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6081 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6082 let sel = _mm256_and_si256(x, mask);
6083 let p16 = _mm256_maddubs_epi16(ones8, sel);
6085 let d32 = _mm256_madd_epi16(p16, ones16);
6086 let hi128 = _mm256_extracti128_si256::<1>(d32);
6087 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6088 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6089 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6090 let msum = _mm_cvtsi128_si32(s32);
6091 let d = 2 * msum - gsum[gi];
6094 acc += d as f32 * s;
6095 }
6096 acc
6097 }
6098}
6099
6100#[cfg(target_arch = "x86_64")]
6103#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6104unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6105 unsafe {
6107 use core::arch::x86_64::*;
6108 let expand = _mm256_setr_epi8(
6109 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,
6110 3, 3, 3,
6111 );
6112 let bitsel = _mm256_setr_epi8(
6113 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6114 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6115 );
6116 let ones8 = _mm256_set1_epi8(1);
6117 let mut acc = 0f32;
6118 for gi in 0..gpr {
6119 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6120 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6121 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6122 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6123 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6124 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6125 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6126 let d = 2 * msum - gsum[gi];
6127 acc += d as f32 * s;
6128 }
6129 acc
6130 }
6131}
6132
6133#[cfg(target_arch = "x86_64")]
6135#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6136unsafe fn dot_q1_row_1x4_vnni(
6137 bytes: &[u8],
6138 r: usize,
6139 gpr: usize,
6140 xs: [&[i8]; 4],
6141 gsums: [&[i32]; 4],
6142) -> [f32; 4] {
6143 unsafe {
6145 use core::arch::x86_64::*;
6146 let expand = _mm256_setr_epi8(
6147 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,
6148 3, 3, 3,
6149 );
6150 let bitsel = _mm256_setr_epi8(
6151 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6152 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6153 );
6154 let ones8 = _mm256_set1_epi8(1);
6155 let mut acc = [0f32; 4];
6156 for gi in 0..gpr {
6157 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6158 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6159 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6160 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6161 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6162 for (k, xq) in xs.iter().enumerate() {
6163 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6164 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6165 let d = 2 * msum - gsums[k][gi];
6166 acc[k] += d as f32 * s;
6167 }
6168 }
6169 acc
6170 }
6171}
6172
6173#[cfg(target_arch = "x86_64")]
6176#[target_feature(enable = "avx2")]
6177unsafe fn dot_q1_row_1x4_avx2(
6178 bytes: &[u8],
6179 r: usize,
6180 gpr: usize,
6181 xs: [&[i8]; 4],
6182 gsums: [&[i32]; 4],
6183) -> [f32; 4] {
6184 unsafe {
6186 use core::arch::x86_64::*;
6187 let expand = _mm256_setr_epi8(
6188 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,
6189 3, 3, 3,
6190 );
6191 let bitsel = _mm256_setr_epi8(
6192 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6193 -128, 1, 2, 4, 8, 16, 32, 64, -128,
6194 );
6195 let ones8 = _mm256_set1_epi8(1);
6196 let ones16 = _mm256_set1_epi16(1);
6197 let mut acc = [0f32; 4];
6198 for gi in 0..gpr {
6199 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6200 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6201 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6202 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6203 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6204 for (k, xq) in xs.iter().enumerate() {
6205 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6206 let sel = _mm256_and_si256(x, mask);
6207 let p16 = _mm256_maddubs_epi16(ones8, sel);
6208 let d32 = _mm256_madd_epi16(p16, ones16);
6209 let hi128 = _mm256_extracti128_si256::<1>(d32);
6210 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6211 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6212 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6213 let msum = _mm_cvtsi128_si32(s32);
6214 let d = 2 * msum - gsums[k][gi];
6215 acc[k] += d as f32 * s;
6216 }
6217 }
6218 acc
6219 }
6220}
6221
6222#[allow(unreachable_code)]
6223fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6224 #[cfg(target_arch = "aarch64")]
6225 unsafe {
6226 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6227 }
6228 #[cfg(target_arch = "x86_64")]
6229 if avx2_enabled() {
6230 unsafe {
6231 if vnni_tiles_enabled() {
6232 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6233 }
6234 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6235 }
6236 }
6237 let _ = gsum;
6238 let mut acc = 0f32;
6239 for gi in 0..gpr {
6240 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6241 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6242 let mut d = 0i32;
6243 for (j, &b) in tile[2..].iter().enumerate() {
6244 for k in 0..8 {
6245 let w = ((b >> k) & 1) as i32 * 2 - 1;
6246 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6247 }
6248 }
6249 acc += d as f32 * s;
6250 }
6251 acc
6252}
6253
6254#[cfg(target_arch = "aarch64")]
6263#[target_feature(enable = "neon,dotprod")]
6264unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6265 unsafe {
6268 use core::arch::aarch64::*;
6269 use core::arch::asm;
6270 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6271 let m = vld1q_u8(MASKS.as_ptr());
6272 macro_rules! tile_dot {
6274 ($t:expr, $x:expr) => {{
6275 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6276 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6277 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6278 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6279 let x0 = vld1q_s8($x);
6280 let x1 = vld1q_s8($x.add(16));
6281 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6282 asm!(
6283 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6284 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6285 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6286 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6287 options(pure, nomem, nostack),
6288 );
6289 vaddq_s32(a0, a1)
6290 }};
6291 }
6292 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6301 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6302 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6303 const IW11: [u8; 16] = [
6304 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6305 ];
6306 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6307 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6308 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6309 let isc = vld1_u8(ISC.as_ptr());
6310 macro_rules! tile_dot_tbl {
6312 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6313 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6314 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6315 let x0 = vld1q_s8($x);
6316 let x1 = vld1q_s8($x.add(16));
6317 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6318 asm!(
6319 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6320 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6321 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6322 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6323 options(pure, nomem, nostack),
6324 );
6325 vaddq_s32(a0, a1)
6326 }};
6327 }
6328 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6329 let row_base = r * gpr * Q1_TILE;
6330 let abs_end = bytes.len();
6331 let xp = xq.as_ptr();
6332 let gp = gsum.as_ptr();
6333 let mut accv = vdupq_n_f32(0.0);
6334 let mut gi = 0;
6335 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6338 let t0 = base.add(gi * Q1_TILE);
6339 let ld_a = vld1q_u8(t0);
6340 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6341 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6342 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6343 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6344 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6345 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6347 let g = vld1q_s32(gp.add(gi));
6348 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6349 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6350 let scf: float32x4_t;
6351 asm!(
6352 "fcvtl {o:v}.4s, {i:v}.4h",
6353 o = out(vreg) scf, i = in(vreg) sc16,
6354 options(pure, nomem, nostack),
6355 );
6356 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6357 gi += 4;
6358 }
6359 let mut acc = vaddvq_f32(accv);
6360 while gi < gpr {
6361 let t = base.add(gi * Q1_TILE);
6362 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6363 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6364 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6365 gi += 1;
6366 }
6367 acc
6368 }
6369}
6370
6371#[cfg(target_arch = "aarch64")]
6376#[target_feature(enable = "neon,dotprod")]
6377unsafe fn dot_q1_row_1x4_sdot(
6378 bytes: &[u8],
6379 r: usize,
6380 gpr: usize,
6381 xs: [&[i8]; 4],
6382 gs: [&[i32]; 4],
6383) -> [f32; 4] {
6384 unsafe {
6386 use core::arch::aarch64::*;
6387 use core::arch::asm;
6388 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6389 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6390 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6391 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6392 const IW11: [u8; 16] = [
6393 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6394 ];
6395 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6396 let m = vld1q_u8(MASKS.as_ptr());
6397 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6398 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6399 let isc = vld1_u8(ISC.as_ptr());
6400 macro_rules! sdot2 {
6401 ($w0:expr, $w1:expr, $x:expr) => {{
6402 let x0 = vld1q_s8($x);
6403 let x1 = vld1q_s8($x.add(16));
6404 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6405 asm!(
6406 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6407 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6408 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6409 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6410 options(pure, nomem, nostack),
6411 );
6412 vaddq_s32(a0, a1)
6413 }};
6414 }
6415 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6416 let row_base = r * gpr * Q1_TILE;
6417 let abs_end = bytes.len();
6418 let mut accv = [vdupq_n_f32(0.0); 4];
6419 let mut gi = 0;
6420 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6421 let t0 = base.add(gi * Q1_TILE);
6422 let ld_a = vld1q_u8(t0);
6423 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6424 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6426 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6427 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6428 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6429 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6430 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6431 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6432 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6433 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6434 let scf: float32x4_t;
6435 asm!(
6436 "fcvtl {o:v}.4s, {i:v}.4h",
6437 o = out(vreg) scf, i = in(vreg) sc16,
6438 options(pure, nomem, nostack),
6439 );
6440 for k in 0..4 {
6441 let xp = xs[k].as_ptr();
6442 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6443 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6444 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6445 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6446 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6447 let g = vld1q_s32(gs[k].as_ptr().add(gi));
6448 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6449 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6450 }
6451 gi += 4;
6452 }
6453 let mut acc = [
6454 vaddvq_f32(accv[0]),
6455 vaddvq_f32(accv[1]),
6456 vaddvq_f32(accv[2]),
6457 vaddvq_f32(accv[3]),
6458 ];
6459 while gi < gpr {
6460 let t = base.add(gi * Q1_TILE);
6461 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6462 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6463 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6464 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6465 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6466 for k in 0..4 {
6467 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6468 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6469 }
6470 gi += 1;
6471 }
6472 acc
6473 }
6474}
6475
6476#[inline]
6478fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6479 let gi = j / GROUP_SIZE;
6480 let k = j % GROUP_SIZE;
6481 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6482 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6483 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6484 ((bit as i32 * 2 - 1) as f32, s)
6485}
6486
6487#[inline]
6489fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6490 let mut acc = 0f32;
6491 for gi in 0..gpr {
6492 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6493 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6494 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6495 let mut ga = 0f32;
6496 for (j, &b) in tile[2..].iter().enumerate() {
6497 for k in 0..8 {
6498 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6499 }
6500 }
6501 acc += ga * s;
6502 }
6503 acc
6504}
6505
6506#[allow(clippy::too_many_arguments)]
6509fn q1_range_a8w8(
6510 bytes: &[u8],
6511 gpr: usize,
6512 act: &SplitAct,
6513 gsum: &[i32],
6514 out: SendMut,
6515 start: usize,
6516 end: usize,
6517) {
6518 for r in start..end {
6519 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6520 for &(j, xv) in &act.outliers {
6521 let (w, s) = q1_outlier(bytes, r, gpr, j);
6522 acc += w * s * xv;
6523 }
6524 unsafe { *out.at(r) = acc };
6526 }
6527}
6528
6529fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6531 for r in start..end {
6532 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6534 }
6535}
6536
6537fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6542 let entries = base_len + (rows + 1) * 4;
6543 (base_len, entries, entries <= bytes.len())
6544}
6545
6546#[inline]
6548fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6549 let o = rp_off + r * 4;
6550 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6551}
6552
6553const SIGN5: [[f32; 5]; 256] = {
6557 let mut lut = [[0.0f32; 5]; 256];
6558 let pow3 = [1u16, 3, 9, 27, 81];
6559 let mut byte = 0usize;
6560 while byte < 256 {
6561 let mut i = 0usize;
6562 while i < 5 {
6563 let code = (byte as u16 / pow3[i]) % 3;
6564 lut[byte][i] = if code == 1 {
6565 1.0
6566 } else if code == 2 {
6567 -1.0
6568 } else {
6569 0.0
6570 };
6571 i += 1;
6572 }
6573 byte += 1;
6574 }
6575 lut
6576};
6577
6578const SIGN5_I8: [[i8; 5]; 256] = {
6580 let mut lut = [[0i8; 5]; 256];
6581 let pow3 = [1u16, 3, 9, 27, 81];
6582 let mut byte = 0usize;
6583 while byte < 256 {
6584 let mut i = 0usize;
6585 while i < 5 {
6586 let code = (byte as u16 / pow3[i]) % 3;
6587 lut[byte][i] = if code == 1 {
6588 1
6589 } else if code == 2 {
6590 -1
6591 } else {
6592 0
6593 };
6594 i += 1;
6595 }
6596 byte += 1;
6597 }
6598 lut
6599};
6600
6601const SIGN5_U64: [u64; 256] = {
6607 let mut lut = [0u64; 256];
6608 let pow3 = [1u16, 3, 9, 27, 81];
6609 let mut byte = 0usize;
6610 while byte < 256 {
6611 let mut v = 0u64;
6612 let mut i = 0usize;
6613 while i < 5 {
6614 let code = (byte as u16 / pow3[i]) % 3;
6615 let s: u8 = if code == 1 {
6616 1
6617 } else if code == 2 {
6618 0xFF
6619 } else {
6620 0
6621 };
6622 v |= (s as u64) << (i * 8);
6623 i += 1;
6624 }
6625 lut[byte] = v;
6626 byte += 1;
6627 }
6628 lut
6629};
6630
6631#[inline]
6636fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6637 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6638 let off = (r * gpr + j / GROUP_SIZE) * TILE;
6639 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6640 let within = j % GROUP_SIZE;
6641 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6642}
6643
6644#[cfg(target_arch = "aarch64")]
6647#[target_feature(enable = "neon,dotprod")]
6648#[inline]
6649unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6650 unsafe {
6652 use core::arch::aarch64::*;
6653 use core::arch::asm;
6654 let w0 = vld1q_s8(w);
6655 let w1 = vld1q_s8(w.add(16));
6656 let x0 = vld1q_s8(x);
6657 let x1 = vld1q_s8(x.add(16));
6658 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6659 asm!(
6660 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6661 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6662 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6663 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6664 options(pure, nomem, nostack),
6665 );
6666 vaddvq_s32(vaddq_s32(a0, a1))
6667 }
6668}
6669
6670#[cfg(target_arch = "x86_64")]
6673#[target_feature(enable = "avx2")]
6674#[inline]
6675unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6676 unsafe {
6678 use core::arch::x86_64::*;
6679 let wv = _mm256_loadu_si256(w as *const __m256i);
6680 let xv = _mm256_loadu_si256(x as *const __m256i);
6681 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6682 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6683 let hi128 = _mm256_extracti128_si256::<1>(d);
6684 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6685 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6686 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6687 _mm_cvtsi128_si32(s32)
6688 }
6689}
6690
6691#[inline]
6696fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6697 debug_assert!(dst.len() >= 40);
6698 unsafe {
6701 let p = dst.as_mut_ptr();
6702 for bi in 0..7 {
6703 core::ptr::write_unaligned(
6704 p.add(bi * 5) as *mut u64,
6705 SIGN5_U64[*codes.add(bi) as usize],
6706 );
6707 }
6708 }
6709}
6710
6711#[inline]
6716fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6717 #[cfg(target_arch = "aarch64")]
6718 unsafe {
6719 return sdot32_i8(w, x);
6720 }
6721 #[cfg(target_arch = "x86_64")]
6722 unsafe {
6723 return i8dot32_avx2(w, x);
6724 }
6725 #[allow(unreachable_code)]
6726 unsafe {
6727 let mut s = 0i32;
6728 for k in 0..GROUP_SIZE {
6729 s += *w.add(k) as i32 * *x.add(k) as i32;
6730 }
6731 s
6732 }
6733}
6734
6735#[inline]
6736unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6737 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6738 (
6739 SIGN5_U64[*codes as usize],
6740 SIGN5_U64[*codes.add(1) as usize],
6741 SIGN5_U64[*codes.add(2) as usize],
6742 SIGN5_U64[*codes.add(3) as usize],
6743 SIGN5_U64[*codes.add(4) as usize],
6744 SIGN5_U64[*codes.add(5) as usize],
6745 SIGN5_U64[*codes.add(6) as usize],
6746 )
6747 };
6748
6749 let u0 = s0 | (s1 << 40);
6750 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6751 let u2 = (s3 >> 8) | (s4 << 32);
6752 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6753
6754 (u0, u1, u2, u3)
6755}
6756
6757#[cfg(target_arch = "aarch64")]
6761#[target_feature(enable = "neon,dotprod")]
6762unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6763 use core::arch::aarch64::*;
6764 use core::arch::asm;
6765 unsafe {
6766 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6767 let mut acc = 0f32;
6768 let bytes_ptr = bytes.as_ptr();
6769 let xq_ptr = xq.as_ptr();
6770 let row_off = r * gpr * TILE;
6771
6772 let gpr2 = gpr & !1;
6773 let mut gi = 0;
6774 while gi < gpr2 {
6775 let off0 = row_off + gi * TILE;
6776 let off1 = off0 + TILE;
6777 let s0 = f16_to_f32(u16::from_le_bytes([
6778 *bytes_ptr.add(off0),
6779 *bytes_ptr.add(off0 + 1),
6780 ]));
6781 let s1 = f16_to_f32(u16::from_le_bytes([
6782 *bytes_ptr.add(off1),
6783 *bytes_ptr.add(off1 + 1),
6784 ]));
6785
6786 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6787 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6788
6789 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6790 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6791 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6792 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6793
6794 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6795 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6796 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6797 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6798
6799 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6800 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6801 asm!(
6802 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6803 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6804 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6805 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6806 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6807 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6808 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6809 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6810 options(pure, nomem, nostack),
6811 );
6812 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6813 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6814 acc += d0 as f32 * s0 + d1 as f32 * s1;
6815 gi += 2;
6816 }
6817
6818 if gi < gpr {
6819 let off = row_off + gi * TILE;
6820 let s = f16_to_f32(u16::from_le_bytes([
6821 *bytes_ptr.add(off),
6822 *bytes_ptr.add(off + 1),
6823 ]));
6824 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6825 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6826 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6827 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6828 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6829 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6830 asm!(
6831 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6832 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6833 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6834 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6835 options(pure, nomem, nostack),
6836 );
6837 let d = vaddvq_s32(vaddq_s32(a0, a1));
6838 acc += d as f32 * s;
6839 }
6840 acc
6841 }
6842}
6843
6844#[cfg(target_arch = "x86_64")]
6846#[target_feature(enable = "avx2")]
6847unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6848 use core::arch::x86_64::*;
6849 unsafe {
6850 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6851 let mut acc = 0f32;
6852 let bytes_ptr = bytes.as_ptr();
6853 let xq_ptr = xq.as_ptr();
6854 let row_off = r * gpr * TILE;
6855
6856 let ones = _mm256_set1_epi16(1);
6857 for gi in 0..gpr {
6858 let off = row_off + gi * TILE;
6859 let s = f16_to_f32(u16::from_le_bytes([
6860 *bytes_ptr.add(off),
6861 *bytes_ptr.add(off + 1),
6862 ]));
6863 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6864 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6865 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6866 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6867 let d256 = _mm256_madd_epi16(p16, ones);
6868 let d128 = _mm_add_epi32(
6869 _mm256_castsi256_si128(d256),
6870 _mm256_extracti128_si256(d256, 1),
6871 );
6872 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6873 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6874 acc += d32 as f32 * s;
6875 }
6876 acc
6877 }
6878}
6879
6880#[cfg(target_arch = "x86_64")]
6882#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6883unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6884 use core::arch::x86_64::*;
6885 unsafe {
6887 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6888 let mut acc = 0f32;
6889 let bytes_ptr = bytes.as_ptr();
6890 let xq_ptr = xq.as_ptr();
6891 let row_off = r * gpr * TILE;
6892 for gi in 0..gpr {
6893 let off = row_off + gi * TILE;
6894 let s = f16_to_f32(u16::from_le_bytes([
6895 *bytes_ptr.add(off),
6896 *bytes_ptr.add(off + 1),
6897 ]));
6898 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6899 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6900 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6901 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6902 acc += d as f32 * s;
6903 }
6904 acc
6905 }
6906}
6907
6908#[inline]
6912fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6913 #[cfg(target_arch = "aarch64")]
6914 unsafe {
6915 return q1t_dot_row_sdot(bytes, r, gpr, xq);
6916 }
6917 #[cfg(target_arch = "x86_64")]
6918 unsafe {
6919 if vnni_tiles_enabled() {
6920 return q1t_dot_row_vnni(bytes, r, gpr, xq);
6921 }
6922 return q1t_dot_row_avx2(bytes, r, gpr, xq);
6923 }
6924 #[allow(unreachable_code)]
6925 {
6926 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6927 let mut acc = 0f32;
6928 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
6930 let off = (r * gpr + gi) * TILE;
6931 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6932 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6933 let mut d = 0i32;
6934 for k in 0..GROUP_SIZE {
6935 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6936 }
6937 acc += d as f32 * s;
6938 }
6939 acc
6940 }
6941}
6942
6943fn q1t_row_outlier_correction(
6950 bytes: &[u8],
6951 r: usize,
6952 rp_off: usize,
6953 entries_off: usize,
6954 has_ov: bool,
6955 x: &[f32],
6956) -> f32 {
6957 if !has_ov {
6958 return 0.0;
6959 }
6960 let (c0, c1) = (
6961 q1t_rowptr(bytes, rp_off, r),
6962 q1t_rowptr(bytes, rp_off, r + 1),
6963 );
6964 let mut corr = 0f32;
6965 for p in c0..c1 {
6966 let e = entries_off + p * 4;
6967 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6968 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6969 corr += val * x[col];
6970 }
6971 corr
6972}
6973
6974fn q1t_dequant_row(
6978 bytes: &[u8],
6979 r: usize,
6980 gpr: usize,
6981 rp_off: usize,
6982 entries_off: usize,
6983 has_ov: bool,
6984 buf: &mut [f32],
6985) {
6986 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6987 for g in 0..gpr {
6988 let off = (r * gpr + g) * TILE;
6989 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6990 let codes = &bytes[off + 2..off + TILE];
6991 let bc = g * GROUP_SIZE;
6992 for bi in 0..6 {
6994 let lut = &SIGN5[codes[bi] as usize];
6995 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6996 for i in 0..5 {
6997 d[i] = lut[i] * s;
6998 }
6999 }
7000 let lut = &SIGN5[codes[6] as usize];
7001 buf[bc + 30] = lut[0] * s;
7002 buf[bc + 31] = lut[1] * s;
7003 }
7004 if !has_ov {
7005 return;
7006 }
7007 let (c0, c1) = (
7008 q1t_rowptr(bytes, rp_off, r),
7009 q1t_rowptr(bytes, rp_off, r + 1),
7010 );
7011 for p in c0..c1 {
7012 let e = entries_off + p * 4;
7013 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7014 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7015 }
7016}
7017
7018fn q1t_add_overlay(
7022 bytes: &[u8],
7023 x: &[f32],
7024 rows: usize,
7025 cols: usize,
7026 out: &mut [f32],
7027 pool: Option<&Pool>,
7028) {
7029 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7030 let gpr = cols / GROUP_SIZE;
7031 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7032 if !has_ov {
7033 return;
7034 }
7035 let out_addr = SendMut(out.as_mut_ptr());
7036 let run = move |start: usize, end: usize| {
7037 for r in start..end {
7038 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7039 unsafe { *out_addr.at(r) += corr };
7041 }
7042 };
7043 dispatch_rows(pool, rows, &run);
7044}
7045
7046#[allow(clippy::too_many_arguments)]
7049fn q1t_range_a8w8(
7050 bytes: &[u8],
7051 gpr: usize,
7052 rp_off: usize,
7053 ent_off: usize,
7054 has_ov: bool,
7055 act: &SplitAct,
7056 x: &[f32],
7057 out: SendMut,
7058 start: usize,
7059 end: usize,
7060) {
7061 for r in start..end {
7062 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7063 for &(j, xv) in &act.outliers {
7064 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7065 }
7066 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7067 unsafe { *out.at(r) = acc };
7069 }
7070}
7071
7072#[allow(clippy::too_many_arguments)]
7075fn q1t_range_f32_batch(
7076 bytes: &[u8],
7077 gpr: usize,
7078 rp_off: usize,
7079 ent_off: usize,
7080 has_ov: bool,
7081 x: &[f32],
7082 out: SendMut,
7083 start: usize,
7084 end: usize,
7085) {
7086 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7087 let mut sg = [0f32; GROUP_SIZE];
7088 for r in start..end {
7089 let mut acc = 0f32;
7090 for g in 0..gpr {
7091 let off = (r * gpr + g) * TILE;
7092 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7093 let codes = &bytes[off + 2..off + TILE];
7094 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7095 for bi in 0..6 {
7096 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7097 }
7098 let lut = &SIGN5[codes[6] as usize];
7099 sg[30] = lut[0];
7100 sg[31] = lut[1];
7101 let mut gsum = 0f32;
7102 for k in 0..GROUP_SIZE {
7103 gsum += sg[k] * xg[k];
7104 }
7105 acc += s * gsum;
7106 }
7107 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7108 unsafe { *out.at(r) = acc };
7110 }
7111}
7112
7113fn q1t_matvec(
7117 bytes: &[u8],
7118 x: &[f32],
7119 rows: usize,
7120 cols: usize,
7121 out: &mut [f32],
7122 pool: Option<&Pool>,
7123) {
7124 debug_assert_eq!(out.len(), rows);
7125 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7126 let gpr = cols / GROUP_SIZE;
7127 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7128 let out_addr = SendMut(out.as_mut_ptr());
7129 if a8w8_enabled() {
7133 let act = split_act(x);
7134 let act = &act;
7135 let run = move |start: usize, end: usize| {
7136 for r in start..end {
7137 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7138 for &(j, xv) in &act.outliers {
7139 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7140 }
7141 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7142 unsafe { *out_addr.at(r) = acc };
7144 }
7145 };
7146 dispatch_rows(pool, rows, &run);
7147 return;
7148 }
7149 let run = move |start: usize, end: usize| {
7150 let mut sg = [0f32; GROUP_SIZE];
7154 for r in start..end {
7155 let mut acc = 0f32;
7156 for g in 0..gpr {
7157 let off = (r * gpr + g) * TILE;
7158 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7159 let codes = &bytes[off + 2..off + TILE];
7160 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7161 for bi in 0..6 {
7162 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7163 }
7164 let lut = &SIGN5[codes[6] as usize];
7165 sg[30] = lut[0];
7166 sg[31] = lut[1];
7167 let mut gsum = 0f32;
7168 for k in 0..GROUP_SIZE {
7169 gsum += sg[k] * xg[k];
7170 }
7171 acc += s * gsum;
7172 }
7173 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7174 unsafe { *out_addr.at(r) = acc };
7175 }
7176 };
7177 dispatch_rows(pool, rows, &run);
7178}
7179
7180#[cfg(target_arch = "aarch64")]
7186#[target_feature(enable = "neon,dotprod")]
7187unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7188 use core::arch::aarch64::*;
7189 use core::arch::asm;
7190 unsafe {
7192 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7193 let bytes_ptr = bytes.as_ptr();
7194 let row_off = r * gpr * TILE;
7195 let xp = [xa.as_ptr(), xb.as_ptr()];
7196 let mut acc = [0f32; 2];
7197 macro_rules! sdot2 {
7198 ($w0:expr, $w1:expr, $x:expr) => {{
7199 let x0 = vld1q_s8($x);
7200 let x1 = vld1q_s8($x.add(16));
7201 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7202 asm!(
7203 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7204 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7205 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7206 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7207 options(pure, nomem, nostack),
7208 );
7209 vaddvq_s32(vaddq_s32(a0, a1))
7210 }};
7211 }
7212 let gpr2 = gpr & !1;
7213 let mut gi = 0;
7214 while gi < gpr2 {
7215 let off0 = row_off + gi * TILE;
7216 let off1 = off0 + TILE;
7217 let s0 = f16_to_f32(u16::from_le_bytes([
7218 *bytes_ptr.add(off0),
7219 *bytes_ptr.add(off0 + 1),
7220 ]));
7221 let s1 = f16_to_f32(u16::from_le_bytes([
7222 *bytes_ptr.add(off1),
7223 *bytes_ptr.add(off1 + 1),
7224 ]));
7225 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7226 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7227 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7228 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7229 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7230 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7231 for k in 0..2 {
7232 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7233 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7234 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7235 }
7236 gi += 2;
7237 }
7238 if gi < gpr {
7239 let off = row_off + gi * TILE;
7240 let s = f16_to_f32(u16::from_le_bytes([
7241 *bytes_ptr.add(off),
7242 *bytes_ptr.add(off + 1),
7243 ]));
7244 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7245 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7246 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7247 for k in 0..2 {
7248 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7249 acc[k] += d as f32 * s;
7250 }
7251 }
7252 acc
7253 }
7254}
7255
7256fn q1t_matvec2(
7262 bytes: &[u8],
7263 x1: &[f32],
7264 x2: &[f32],
7265 rows: usize,
7266 cols: usize,
7267 o1: &mut [f32],
7268 o2: &mut [f32],
7269 pool: Option<&Pool>,
7270) {
7271 debug_assert_eq!(o1.len(), rows);
7272 debug_assert_eq!(o2.len(), rows);
7273 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7274 let gpr = cols / GROUP_SIZE;
7275 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7276 let out1 = SendMut(o1.as_mut_ptr());
7277 let out2 = SendMut(o2.as_mut_ptr());
7278 if a8w8_enabled() {
7279 let a1 = split_act(x1);
7280 let a2 = split_act(x2);
7281 let (a1, a2) = (&a1, &a2);
7282 let run = move |start: usize, end: usize| {
7283 for r in start..end {
7284 #[cfg(target_arch = "aarch64")]
7285 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7288 #[cfg(not(target_arch = "aarch64"))]
7289 let ds = [
7290 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7291 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7292 ];
7293 let mut acc1 = ds[0] * a1.sx;
7294 for &(j, xv) in &a1.outliers {
7295 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7296 }
7297 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7298 let mut acc2 = ds[1] * a2.sx;
7299 for &(j, xv) in &a2.outliers {
7300 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7301 }
7302 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7303 unsafe {
7305 *out1.at(r) = acc1;
7306 *out2.at(r) = acc2;
7307 }
7308 }
7309 };
7310 dispatch_rows(pool, rows, &run);
7311 return;
7312 }
7313 let run = move |start: usize, end: usize| {
7314 let mut sg = [0f32; GROUP_SIZE];
7317 for r in start..end {
7318 let mut acc1 = 0f32;
7319 let mut acc2 = 0f32;
7320 for g in 0..gpr {
7321 let off = (r * gpr + g) * TILE;
7322 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7323 let codes = &bytes[off + 2..off + TILE];
7324 for bi in 0..6 {
7325 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7326 }
7327 let lut = &SIGN5[codes[6] as usize];
7328 sg[30] = lut[0];
7329 sg[31] = lut[1];
7330 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7331 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7332 let mut gsum1 = 0f32;
7333 for k in 0..GROUP_SIZE {
7334 gsum1 += sg[k] * xg1[k];
7335 }
7336 acc1 += s * gsum1;
7337 let mut gsum2 = 0f32;
7338 for k in 0..GROUP_SIZE {
7339 gsum2 += sg[k] * xg2[k];
7340 }
7341 acc2 += s * gsum2;
7342 }
7343 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7344 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7345 unsafe {
7347 *out1.at(r) = acc1;
7348 *out2.at(r) = acc2;
7349 }
7350 }
7351 };
7352 dispatch_rows(pool, rows, &run);
7353}
7354
7355fn q1t_matmat(
7358 bytes: &[u8],
7359 xs: &[f32],
7360 b: usize,
7361 rows: usize,
7362 cols: usize,
7363 out: &mut [f32],
7364 pool: Option<&Pool>,
7365) {
7366 debug_assert_eq!(out.len(), b * rows);
7367 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7368 let gpr = cols / GROUP_SIZE;
7369 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7370 let out_addr = SendMut(out.as_mut_ptr());
7371 if a8w8_enabled() {
7375 let acts: Vec<SplitAct> = (0..b)
7376 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7377 .collect();
7378 let acts = &acts;
7379 let run = move |start: usize, end: usize| {
7380 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
7384 for g in 0..gpr {
7385 let off = (r * gpr + g) * TILE;
7386 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7387 q1t_unpack_group_i8(
7388 bytes.as_ptr().wrapping_add(off + 2),
7389 &mut sg[g * GROUP_SIZE..],
7390 );
7391 }
7392 for bi in 0..b {
7393 let act = &acts[bi];
7394 let mut isum = 0f32;
7395 for g in 0..gpr {
7396 let d = q1t_i8dot32(
7397 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7398 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7399 );
7400 isum += d as f32 * sc[g];
7401 }
7402 let mut acc = isum * act.sx;
7403 for &(j, xv) in &act.outliers {
7404 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7405 }
7406 accs[bi] = acc;
7407 }
7408 if has_ov {
7412 let (c0, c1) = (
7413 q1t_rowptr(bytes, rp_off, r),
7414 q1t_rowptr(bytes, rp_off, r + 1),
7415 );
7416 for p in c0..c1 {
7417 let e = ent_off + p * 4;
7418 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7419 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7420 for bi in 0..b {
7421 accs[bi] += val * xs[bi * cols + col];
7422 }
7423 }
7424 }
7425 for bi in 0..b {
7426 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7427 }
7428 }
7429 };
7430 dispatch_rows(pool, rows, &run);
7431 return;
7432 }
7433 let run = move |start: usize, end: usize| {
7434 let mut buf = vec![0f32; cols];
7435 for r in start..end {
7436 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7437 for bi in 0..b {
7438 let xr = &xs[bi * cols..(bi + 1) * cols];
7439 let mut acc = 0f32;
7440 for j in 0..cols {
7441 acc += buf[j] * xr[j];
7442 }
7443 unsafe { *out_addr.at(bi * rows + r) = acc };
7444 }
7445 }
7446 };
7447 dispatch_rows(pool, rows, &run);
7448}
7449
7450fn q1_matvec(
7451 bytes: &[u8],
7452 x: &[f32],
7453 rows: usize,
7454 cols: usize,
7455 out: &mut [f32],
7456 pool: Option<&Pool>,
7457) {
7458 debug_assert_eq!(out.len(), rows);
7459 let gpr = cols / GROUP_SIZE;
7460 let out_addr = SendMut(out.as_mut_ptr());
7461 if a8w8_enabled() {
7462 let act = split_act(x);
7463 let gsum = q1_group_sums(&act.xq, gpr);
7464 let (act, gsum) = (&act, &gsum);
7465 let run = move |start: usize, end: usize| {
7466 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7467 };
7468 dispatch_rows(pool, rows, &run);
7469 return;
7470 }
7471 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7472 dispatch_rows(pool, rows, &run);
7473}
7474
7475#[allow(clippy::too_many_arguments)]
7477fn q1_matvec2(
7478 bytes: &[u8],
7479 x1: &[f32],
7480 x2: &[f32],
7481 rows: usize,
7482 cols: usize,
7483 o1: &mut [f32],
7484 o2: &mut [f32],
7485 pool: Option<&Pool>,
7486) {
7487 let gpr = cols / GROUP_SIZE;
7488 let p1 = SendMut(o1.as_mut_ptr());
7489 let p2 = SendMut(o2.as_mut_ptr());
7490 if a8w8_enabled() {
7491 let a1 = split_act(x1);
7492 let a2 = split_act(x2);
7493 let g1 = q1_group_sums(&a1.xq, gpr);
7494 let g2 = q1_group_sums(&a2.xq, gpr);
7495 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7496 let run = move |start: usize, end: usize| {
7497 for r in start..end {
7498 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7499 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7500 for &(j, xv) in &a1.outliers {
7501 let (w, s) = q1_outlier(bytes, r, gpr, j);
7502 v1 += w * s * xv;
7503 }
7504 for &(j, xv) in &a2.outliers {
7505 let (w, s) = q1_outlier(bytes, r, gpr, j);
7506 v2 += w * s * xv;
7507 }
7508 unsafe {
7510 *p1.at(r) = v1;
7511 *p2.at(r) = v2;
7512 }
7513 }
7514 };
7515 dispatch_rows(pool, rows, &run);
7516 return;
7517 }
7518 let run = move |start: usize, end: usize| {
7519 for r in start..end {
7520 unsafe {
7522 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7523 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7524 }
7525 }
7526 };
7527 dispatch_rows(pool, rows, &run);
7528}
7529
7530#[allow(clippy::too_many_arguments)]
7532fn q1_matmat(
7533 bytes: &[u8],
7534 xs_all: &[f32],
7535 b: usize,
7536 rows: usize,
7537 cols: usize,
7538 out: &mut [f32],
7539 pool: Option<&Pool>,
7540) {
7541 debug_assert_eq!(out.len(), b * rows);
7542 let gpr = cols / GROUP_SIZE;
7543 let out_addr = SendMut(out.as_mut_ptr());
7544 if a8w8_enabled() {
7545 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7546 .map(|bi| {
7547 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7548 let gsum = q1_group_sums(&act.xq, gpr);
7549 (act, gsum)
7550 })
7551 .collect();
7552 let acts = &acts;
7553 #[cfg(target_arch = "x86_64")]
7554 let blocked_ok = avx2_enabled() && blocked_enabled();
7555 #[cfg(target_arch = "aarch64")]
7556 let blocked_ok = sdot_enabled() && blocked_enabled();
7557 let run = move |start: usize, end: usize| {
7558 for r in start..end {
7559 let mut bi = 0usize;
7560 #[cfg(target_arch = "aarch64")]
7563 if blocked_ok {
7564 while bi + 4 <= acts.len() {
7565 let xs = [
7566 acts[bi].0.xq.as_slice(),
7567 acts[bi + 1].0.xq.as_slice(),
7568 acts[bi + 2].0.xq.as_slice(),
7569 acts[bi + 3].0.xq.as_slice(),
7570 ];
7571 let gs = [
7572 acts[bi].1.as_slice(),
7573 acts[bi + 1].1.as_slice(),
7574 acts[bi + 2].1.as_slice(),
7575 acts[bi + 3].1.as_slice(),
7576 ];
7577 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7578 for k in 0..4 {
7579 let (act, _) = &acts[bi + k];
7580 let mut acc = d[k] * act.sx;
7581 for &(j, xv) in &act.outliers {
7582 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7583 acc += w * sc * xv;
7584 }
7585 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7587 }
7588 bi += 4;
7589 }
7590 }
7591 #[cfg(target_arch = "x86_64")]
7592 if blocked_ok {
7593 while bi + 4 <= acts.len() {
7594 let xs = [
7595 acts[bi].0.xq.as_slice(),
7596 acts[bi + 1].0.xq.as_slice(),
7597 acts[bi + 2].0.xq.as_slice(),
7598 acts[bi + 3].0.xq.as_slice(),
7599 ];
7600 let gs = [
7601 acts[bi].1.as_slice(),
7602 acts[bi + 1].1.as_slice(),
7603 acts[bi + 2].1.as_slice(),
7604 acts[bi + 3].1.as_slice(),
7605 ];
7606 let d = unsafe {
7607 if vnni_tiles_enabled() {
7608 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7609 } else {
7610 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7611 }
7612 };
7613 for k in 0..4 {
7614 let (act, _) = &acts[bi + k];
7615 let mut acc = d[k] * act.sx;
7616 for &(j, xv) in &act.outliers {
7617 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7618 acc += w * sc * xv;
7619 }
7620 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7622 }
7623 bi += 4;
7624 }
7625 }
7626 while bi < acts.len() {
7627 let (act, gsum) = &acts[bi];
7628 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7629 for &(j, xv) in &act.outliers {
7630 let (w, s) = q1_outlier(bytes, r, gpr, j);
7631 acc += w * s * xv;
7632 }
7633 unsafe { *out_addr.at(bi * rows + r) = acc };
7635 bi += 1;
7636 }
7637 }
7638 };
7639 dispatch_rows(pool, rows, &run);
7640 return;
7641 }
7642 let run = move |start: usize, end: usize| {
7643 for r in start..end {
7644 for bi in 0..b {
7645 let x = &xs_all[bi * cols..(bi + 1) * cols];
7646 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7648 }
7649 }
7650 };
7651 dispatch_rows(pool, rows, &run);
7652}
7653
7654fn q4matvec(
7660 bytes: &[u8],
7661 x: &[f32],
7662 rows: usize,
7663 cols: usize,
7664 out: &mut [f32],
7665 pool: Option<&Pool>,
7666) {
7667 debug_assert_eq!(out.len(), rows);
7668 let (packed, scales) = q4_split(bytes, rows, cols);
7669 let gpr = cols / GROUP_SIZE;
7670 let out_addr = SendMut(out.as_mut_ptr());
7671
7672 if a8w8_enabled() {
7673 let act = split_act(x);
7674 let run = move |start: usize, end: usize| {
7675 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7676 };
7677 dispatch_rows(pool, rows, &run);
7678 return;
7679 }
7680
7681 let run =
7682 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7683 dispatch_rows(pool, rows, &run);
7684}
7685
7686#[inline]
7689#[allow(unreachable_code)]
7690#[cfg(target_arch = "x86_64")]
7695#[target_feature(enable = "avx2")]
7696unsafe fn dot_q4b_row_1x4_avx2(
7697 buf: &[u8],
7698 scales: &[u8],
7699 g0: usize,
7700 gpr: usize,
7701 xs: [&[i8]; 4],
7702) -> [f32; 4] {
7703 unsafe {
7705 use core::arch::x86_64::*;
7706 let ones = _mm256_set1_epi16(1);
7707 let mut acc = [0f32; 4];
7708 for gi in 0..gpr {
7709 let s = f16_to_f32(u16::from_le_bytes([
7710 scales[(g0 + gi) * 2],
7711 scales[(g0 + gi) * 2 + 1],
7712 ]));
7713 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7714 let aw = _mm256_abs_epi8(w);
7715 for (k, xq) in xs.iter().enumerate() {
7716 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7717 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7718 let d = _mm256_madd_epi16(p16, ones);
7719 let hi128 = _mm256_extracti128_si256::<1>(d);
7720 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7721 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7722 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7723 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7724 }
7725 }
7726 acc
7727 }
7728}
7729
7730#[cfg(target_arch = "x86_64")]
7732#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7733unsafe fn dot_q4b_row_1x4_vnni(
7734 buf: &[u8],
7735 scales: &[u8],
7736 g0: usize,
7737 gpr: usize,
7738 xs: [&[i8]; 4],
7739) -> [f32; 4] {
7740 unsafe {
7742 use core::arch::x86_64::*;
7743 let mut acc = [0f32; 4];
7744 for gi in 0..gpr {
7745 let s = f16_to_f32(u16::from_le_bytes([
7746 scales[(g0 + gi) * 2],
7747 scales[(g0 + gi) * 2 + 1],
7748 ]));
7749 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7750 let aw = _mm256_abs_epi8(w);
7751 for (k, xq) in xs.iter().enumerate() {
7752 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7753 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7754 acc[k] += d as f32 * s;
7755 }
7756 }
7757 acc
7758 }
7759}
7760
7761#[cfg(target_arch = "x86_64")]
7767#[target_feature(enable = "avx2")]
7768unsafe fn dot_q4b_row_1x4_sx_avx2(
7769 buf: &[u8],
7770 scales: &[u8],
7771 g0: usize,
7772 gpr: usize,
7773 xs: [&[i8]; 4],
7774 sxs: [f32; 4],
7775) -> [f32; 4] {
7776 unsafe {
7778 use core::arch::x86_64::*;
7779 let ones = _mm256_set1_epi16(1);
7780 let mut acc = [0f32; 4];
7781 for gi in 0..gpr {
7782 let s = f16_to_f32(u16::from_le_bytes([
7783 scales[(g0 + gi) * 2],
7784 scales[(g0 + gi) * 2 + 1],
7785 ]));
7786 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7787 let aw = _mm256_abs_epi8(w);
7788 for (k, xq) in xs.iter().enumerate() {
7789 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7790 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7791 let d = _mm256_madd_epi16(p16, ones);
7792 let hi128 = _mm256_extracti128_si256::<1>(d);
7793 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7794 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7795 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7796 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7797 }
7798 }
7799 acc
7800 }
7801}
7802
7803#[cfg(target_arch = "x86_64")]
7806#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7807unsafe fn dot_q4b_row_1x4_sx_vnni(
7808 buf: &[u8],
7809 scales: &[u8],
7810 g0: usize,
7811 gpr: usize,
7812 xs: [&[i8]; 4],
7813 sxs: [f32; 4],
7814) -> [f32; 4] {
7815 unsafe {
7817 use core::arch::x86_64::*;
7818 let mut acc = [0f32; 4];
7819 for gi in 0..gpr {
7820 let s = f16_to_f32(u16::from_le_bytes([
7821 scales[(g0 + gi) * 2],
7822 scales[(g0 + gi) * 2 + 1],
7823 ]));
7824 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7825 let aw = _mm256_abs_epi8(w);
7826 for (k, xq) in xs.iter().enumerate() {
7827 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7828 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7829 acc[k] += (d as f32 * sxs[k]) * s;
7830 }
7831 }
7832 acc
7833 }
7834}
7835
7836#[allow(unreachable_code)]
7837fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7838 #[cfg(target_arch = "aarch64")]
7839 unsafe {
7840 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7841 }
7842 #[cfg(target_arch = "x86_64")]
7843 unsafe {
7844 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7845 }
7846 let mut acc = 0f32;
7847 for gi in 0..gpr {
7848 let g = g0 + gi;
7849 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7850 let mut d = 0i32;
7851 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7852 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7853 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7854 }
7855 acc += d as f32 * s;
7856 }
7857 acc
7858}
7859
7860#[inline]
7862#[allow(unreachable_code)]
7863fn dot_q4_row_i8_2(
7864 packed: &[u8],
7865 scales: &[u8],
7866 g0: usize,
7867 gpr: usize,
7868 xq1: &[i8],
7869 xq2: &[i8],
7870) -> (f32, f32) {
7871 #[cfg(target_arch = "aarch64")]
7872 unsafe {
7873 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7874 }
7875 #[cfg(target_arch = "x86_64")]
7876 unsafe {
7877 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7878 }
7879 (
7880 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7881 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7882 )
7883}
7884
7885#[allow(clippy::too_many_arguments)]
7888fn q4_range_a8w8(
7889 packed: &[u8],
7890 scales: &[u8],
7891 gpr: usize,
7892 cols: usize,
7893 act: &SplitAct,
7894 out: SendMut,
7895 start: usize,
7896 end: usize,
7897) {
7898 for r in start..end {
7899 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7900 for &(j, xv) in &act.outliers {
7902 let flat = r * cols + j;
7903 let byte = packed[flat / 2];
7904 let nib = if flat & 1 == 0 {
7905 byte & 0x0F
7906 } else {
7907 byte >> 4
7908 };
7909 let s = f16_to_f32(u16::from_le_bytes([
7910 scales[(flat / GROUP_SIZE) * 2],
7911 scales[(flat / GROUP_SIZE) * 2 + 1],
7912 ]));
7913 acc += ((nib as i32 - 8) as f32) * s * xv;
7914 }
7915 unsafe { *out.at(r) = acc };
7917 }
7918}
7919
7920#[allow(clippy::too_many_arguments)]
7923fn q4_range2_a8w8(
7924 packed: &[u8],
7925 scales: &[u8],
7926 gpr: usize,
7927 cols: usize,
7928 a1: &SplitAct,
7929 a2: &SplitAct,
7930 p1: SendMut,
7931 p2: SendMut,
7932 start: usize,
7933 end: usize,
7934) {
7935 for r in start..end {
7936 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7937 let mut acc1 = s1 * a1.sx;
7938 let mut acc2 = s2 * a2.sx;
7939 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7941 for &(j, xv) in outliers {
7942 let flat = r * cols + j;
7943 let byte = packed[flat / 2];
7944 let nib = if flat & 1 == 0 {
7945 byte & 0x0F
7946 } else {
7947 byte >> 4
7948 };
7949 let s = f16_to_f32(u16::from_le_bytes([
7950 scales[(flat / GROUP_SIZE) * 2],
7951 scales[(flat / GROUP_SIZE) * 2 + 1],
7952 ]));
7953 *acc += ((nib as i32 - 8) as f32) * s * xv;
7954 }
7955 };
7956 fix(&a1.outliers, &mut acc1);
7957 fix(&a2.outliers, &mut acc2);
7958 unsafe {
7960 *p1.at(r) = acc1;
7961 *p2.at(r) = acc2;
7962 }
7963 }
7964}
7965
7966fn q4_range_f32(
7968 packed: &[u8],
7969 scales: &[u8],
7970 gpr: usize,
7971 x: &[f32],
7972 out: SendMut,
7973 start: usize,
7974 end: usize,
7975) {
7976 for r in start..end {
7977 let mut acc = 0f32;
7978 for gi in 0..gpr {
7979 let g = r * gpr + gi;
7980 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7981 let pk = &packed[g * 16..(g + 1) * 16];
7982 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7983 let mut ga = 0f32;
7984 for (k, &b) in pk.iter().enumerate() {
7985 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7986 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7987 }
7988 acc += ga * s;
7989 }
7990 unsafe { *out.at(r) = acc };
7992 }
7993}
7994
7995#[allow(clippy::too_many_arguments)]
7999fn q4matvec2(
8000 bytes: &[u8],
8001 x1: &[f32],
8002 x2: &[f32],
8003 rows: usize,
8004 cols: usize,
8005 o1: &mut [f32],
8006 o2: &mut [f32],
8007 pool: Option<&Pool>,
8008) {
8009 debug_assert_eq!(o1.len(), rows);
8010 debug_assert_eq!(o2.len(), rows);
8011 let (packed, scales) = q4_split(bytes, rows, cols);
8012 let gpr = cols / GROUP_SIZE;
8013
8014 if a8w8_enabled() {
8015 let a1 = split_act(x1);
8016 let a2 = split_act(x2);
8017 let p1 = SendMut(o1.as_mut_ptr());
8018 let p2 = SendMut(o2.as_mut_ptr());
8019 let run = move |start: usize, end: usize| {
8020 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
8021 };
8022 dispatch_rows(pool, rows, &run);
8023 return;
8024 }
8025
8026 let p1 = SendMut(o1.as_mut_ptr());
8027 let p2 = SendMut(o2.as_mut_ptr());
8028 let run = move |start: usize, end: usize| {
8029 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
8030 };
8031 dispatch_rows(pool, rows, &run);
8032}
8033
8034#[allow(clippy::too_many_arguments)]
8036fn q4_range2_f32(
8037 packed: &[u8],
8038 scales: &[u8],
8039 gpr: usize,
8040 x1: &[f32],
8041 x2: &[f32],
8042 p1: SendMut,
8043 p2: SendMut,
8044 start: usize,
8045 end: usize,
8046) {
8047 for r in start..end {
8048 let (mut acc1, mut acc2) = (0f32, 0f32);
8049 for gi in 0..gpr {
8050 let g = r * gpr + gi;
8051 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8052 let pk = &packed[g * 16..(g + 1) * 16];
8053 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8054 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8055 let (mut g1, mut g2) = (0f32, 0f32);
8056 for (k, &b) in pk.iter().enumerate() {
8057 let wl = (b & 0x0F) as f32 - 8.0;
8058 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
8059 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
8060 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
8061 }
8062 acc1 += g1 * s;
8063 acc2 += g2 * s;
8064 }
8065 unsafe {
8067 *p1.at(r) = acc1;
8068 *p2.at(r) = acc2;
8069 }
8070 }
8071}
8072
8073thread_local! {
8074 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8077 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8078}
8079
8080#[allow(clippy::too_many_arguments)]
8086fn q4matmat(
8087 bytes: &[u8],
8088 xs_all: &[f32],
8089 b: usize,
8090 rows: usize,
8091 cols: usize,
8092 out: &mut [f32],
8093 pool: Option<&Pool>,
8094) {
8095 debug_assert_eq!(xs_all.len(), b * cols);
8096 debug_assert_eq!(out.len(), b * rows);
8097 let (packed, scales) = q4_split(bytes, rows, cols);
8098 let gpr = cols / GROUP_SIZE;
8099 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8100
8101 if a8w8_enabled() {
8102 let acts: Vec<SplitAct> = (0..b)
8103 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8104 .collect();
8105 let acts = &acts;
8106 let out_addr = SendMut(out.as_mut_ptr());
8107 let run = move |start: usize, end: usize| {
8108 ROW_I8.with(|rb| {
8109 let mut buf = rb.borrow_mut();
8110 buf.resize(cols, 0);
8111 for r in start..end {
8112 for gi in 0..gpr {
8116 let g = r * gpr + gi;
8117 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8118 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
8119 buf[gi * GROUP_SIZE + k * 2 + 1] =
8120 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
8121 }
8122 }
8123 let mut bi = 0usize;
8124 #[cfg(target_arch = "x86_64")]
8125 if avx2_enabled() && blocked_enabled() {
8126 while bi + 4 <= acts.len() {
8127 let xs = [
8128 acts[bi].xq.as_slice(),
8129 acts[bi + 1].xq.as_slice(),
8130 acts[bi + 2].xq.as_slice(),
8131 acts[bi + 3].xq.as_slice(),
8132 ];
8133 let d = unsafe {
8134 if vnni_tiles_enabled() {
8135 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
8136 } else {
8137 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
8138 }
8139 };
8140 for k in 0..4 {
8141 let act = &acts[bi + k];
8142 let mut acc = d[k] * act.sx;
8143 for &(j, xv) in &act.outliers {
8144 acc += (buf[j] as i8) as f32
8145 * gscale((r * cols + j) / GROUP_SIZE)
8146 * xv;
8147 }
8148 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8150 }
8151 bi += 4;
8152 }
8153 }
8154 while bi < acts.len() {
8155 let act = &acts[bi];
8156 let mut acc = 0f32;
8157 for gi in 0..gpr {
8158 let d = dot_i8_i8(
8159 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8160 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8161 );
8162 acc += d as f32 * gscale(r * gpr + gi);
8163 }
8164 acc *= act.sx;
8165 for &(j, xv) in &act.outliers {
8167 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
8168 }
8169 unsafe { *out_addr.at(bi * rows + r) = acc };
8171 bi += 1;
8172 }
8173 }
8174 })
8175 };
8176 dispatch_rows(pool, rows, &run);
8177 return;
8178 }
8179
8180 let out_addr = SendMut(out.as_mut_ptr());
8181 let run = move |start: usize, end: usize| {
8182 ROW_F32.with(|rb| {
8183 let mut buf = rb.borrow_mut();
8184 buf.resize(cols, 0.0);
8185 for r in start..end {
8186 for gi in 0..gpr {
8189 let g = r * gpr + gi;
8190 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8191 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8192 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8193 }
8194 }
8195 for bi in 0..b {
8196 let x = &xs_all[bi * cols..(bi + 1) * cols];
8197 let mut acc = 0f32;
8198 for gi in 0..gpr {
8199 let mut ga = 0f32;
8200 for k in 0..GROUP_SIZE / 2 {
8205 let e = gi * GROUP_SIZE + k * 2;
8206 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8207 }
8208 acc += ga * gscale(r * gpr + gi);
8209 }
8210 unsafe { *out_addr.at(bi * rows + r) = acc };
8212 }
8213 }
8214 })
8215 };
8216 dispatch_rows(pool, rows, &run);
8217}
8218
8219#[allow(clippy::too_many_arguments)]
8224fn vbitmatmat(
8225 bytes: &[u8],
8226 offsets: &[usize],
8227 xs_all: &[f32],
8228 b: usize,
8229 rows: usize,
8230 cols: usize,
8231 out: &mut [f32],
8232 pool: Option<&Pool>,
8233) {
8234 debug_assert_eq!(xs_all.len(), b * cols);
8235 debug_assert_eq!(out.len(), b * rows);
8236 debug_assert_eq!(offsets.len(), rows + 1);
8237 let ng = cols / GROUP_SIZE;
8238 let bits = &bytes[..rows];
8239 let sc_off = rows;
8240 let gscale = |r: usize, g: usize| {
8241 let so = (r * ng + g) * 2;
8242 f16_to_f32(u16::from_le_bytes([
8243 bytes[sc_off + so],
8244 bytes[sc_off + so + 1],
8245 ]))
8246 };
8247
8248 let decode_f32 = |r: usize, dst: &mut [f32]| {
8250 let bw = bits[r] as usize;
8251 let l = ((1i32 << (bw - 1)) - 1) as f32;
8252 let data = &bytes[offsets[r]..offsets[r + 1]];
8253 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8254 for d in dst.iter_mut() {
8255 while nbits < bw {
8256 acc = (acc << 8) | data[idx] as u64;
8257 idx += 1;
8258 nbits += 8;
8259 }
8260 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8261 nbits -= bw;
8262 *d = u - l;
8263 }
8264 };
8265
8266 if a8w8_enabled() {
8267 let acts: Vec<SplitAct> = (0..b)
8268 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8269 .collect();
8270 let acts = &acts;
8271 let out_addr = SendMut(out.as_mut_ptr());
8272 let run = move |start: usize, end: usize| {
8273 for r in start..end {
8274 let bw = bits[r] as usize;
8275 if bw == 8 {
8276 ROW_F32.with(|rb| {
8279 let mut buf = rb.borrow_mut();
8280 buf.resize(cols, 0.0);
8281 decode_f32(r, &mut buf);
8282 for bi in 0..b {
8283 let x = &xs_all[bi * cols..(bi + 1) * cols];
8284 let mut dot = 0f32;
8285 for g in 0..ng {
8286 let mut gd = 0f32;
8287 for k in 0..GROUP_SIZE {
8288 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8289 }
8290 dot += gd * gscale(r, g);
8291 }
8292 unsafe { *out_addr.at(bi * rows + r) = dot };
8294 }
8295 });
8296 continue;
8297 }
8298 let l = (1i32 << (bw - 1)) - 1;
8299 let data = &bytes[offsets[r]..offsets[r + 1]];
8300 ROW_I8.with(|rb| {
8301 let mut buf = rb.borrow_mut();
8302 buf.resize(cols, 0);
8303 #[inline(always)]
8304 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8305 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8306 let u = unpack8::<B>(&data[blk * B..]);
8307 for k in 0..8 {
8308 chunk[k] = (u[k] - l) as i8 as u8;
8309 }
8310 }
8311 }
8312 match bw {
8313 3 => fill::<3>(data, l, &mut buf),
8314 4 => vbit_fill4(data, &mut buf),
8315 5 => fill::<5>(data, l, &mut buf),
8316 6 => fill::<6>(data, l, &mut buf),
8317 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8318 }
8319 let mut bi = 0usize;
8320 #[cfg(target_arch = "x86_64")]
8324 if avx2_enabled() && blocked_enabled() {
8325 while bi + 4 <= acts.len() {
8326 let xs = [
8327 acts[bi].xq.as_slice(),
8328 acts[bi + 1].xq.as_slice(),
8329 acts[bi + 2].xq.as_slice(),
8330 acts[bi + 3].xq.as_slice(),
8331 ];
8332 let sxs = [
8333 acts[bi].sx,
8334 acts[bi + 1].sx,
8335 acts[bi + 2].sx,
8336 acts[bi + 3].sx,
8337 ];
8338 let d = unsafe {
8339 if vnni_tiles_enabled() {
8340 dot_q4b_row_1x4_sx_vnni(
8341 &buf,
8342 &bytes[sc_off..],
8343 r * ng,
8344 ng,
8345 xs,
8346 sxs,
8347 )
8348 } else {
8349 dot_q4b_row_1x4_sx_avx2(
8350 &buf,
8351 &bytes[sc_off..],
8352 r * ng,
8353 ng,
8354 xs,
8355 sxs,
8356 )
8357 }
8358 };
8359 for k in 0..4 {
8360 let act = &acts[bi + k];
8361 let mut dot = d[k];
8362 for &(j, xv) in &act.outliers {
8363 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8364 }
8365 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8367 }
8368 bi += 4;
8369 }
8370 }
8371 while bi < acts.len() {
8372 let act = &acts[bi];
8373 let mut dot = 0f32;
8374 for g in 0..ng {
8375 let d = dot_i8_i8(
8376 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8377 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8378 ) as f32
8379 * act.sx;
8380 dot += d * gscale(r, g);
8381 }
8382 for &(j, xv) in &act.outliers {
8383 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8384 }
8385 unsafe { *out_addr.at(bi * rows + r) = dot };
8387 bi += 1;
8388 }
8389 });
8390 }
8391 };
8392 dispatch_rows(pool, rows, &run);
8393 return;
8394 }
8395
8396 let out_addr = SendMut(out.as_mut_ptr());
8397 let run = move |start: usize, end: usize| {
8398 ROW_F32.with(|rb| {
8399 let mut buf = rb.borrow_mut();
8400 buf.resize(cols, 0.0);
8401 for r in start..end {
8402 decode_f32(r, &mut buf);
8403 for bi in 0..b {
8404 let x = &xs_all[bi * cols..(bi + 1) * cols];
8405 let mut dot = 0f32;
8406 for g in 0..ng {
8407 let mut gd = 0f32;
8408 for k in 0..GROUP_SIZE {
8409 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8410 }
8411 dot += gd * gscale(r, g);
8412 }
8413 unsafe { *out_addr.at(bi * rows + r) = dot };
8415 }
8416 }
8417 })
8418 };
8419 dispatch_rows(pool, rows, &run);
8420}
8421
8422pub(crate) fn gpu_batch_job<'a>(
8426 t: &'a QTensor,
8427 x: &[f32],
8428) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8429 match t {
8430 QTensor::Mapped {
8431 model,
8432 idx,
8433 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8434 rows,
8435 cols,
8436 row_scale,
8437 col_field,
8438 ..
8439 } => Some((
8440 model.clone(),
8441 crate::gpu::BatchJob {
8442 idx: *idx,
8443 rows: *rows,
8444 cols: *cols,
8445 row_scale,
8446 xs: prescale(x, col_field, *dt).into_owned(),
8447 layout: crate::gpu::BatchLayout::Q8,
8448 },
8449 )),
8450 QTensor::Mapped {
8452 model,
8453 idx,
8454 dtype: TensorDtype::Q1,
8455 rows,
8456 cols,
8457 ..
8458 } => Some((
8459 model.clone(),
8460 crate::gpu::BatchJob {
8461 idx: *idx,
8462 rows: *rows,
8463 cols: *cols,
8464 row_scale: &[],
8465 xs: x.to_vec(),
8466 layout: crate::gpu::BatchLayout::Q1,
8467 },
8468 )),
8469 QTensor::Mapped {
8474 model,
8475 idx,
8476 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8477 rows,
8478 cols,
8479 ..
8480 } => Some((
8481 model.clone(),
8482 crate::gpu::BatchJob {
8483 idx: *idx,
8484 rows: *rows,
8485 cols: *cols,
8486 row_scale: &[],
8487 xs: x.to_vec(),
8488 layout: if *dt == TensorDtype::Q4Tiled {
8489 crate::gpu::BatchLayout::Q4t
8490 } else {
8491 crate::gpu::BatchLayout::Q4tp
8492 },
8493 },
8494 )),
8495 _ => None,
8496 }
8497}
8498
8499thread_local! {
8500 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8501 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8502}
8503
8504pub(crate) fn prescale<'a>(
8505 x: &'a [f32],
8506 col_field: &[f32],
8507 dtype: TensorDtype,
8508) -> std::borrow::Cow<'a, [f32]> {
8509 if dtype == TensorDtype::Q8_2f {
8510 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8511 } else {
8512 std::borrow::Cow::Borrowed(x)
8513 }
8514}
8515
8516pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8519 x: &[f32],
8520 col_field: &[f32],
8521 dtype: TensorDtype,
8522 buf_id: u8,
8523 f: F,
8524) -> R {
8525 if dtype == TensorDtype::Q8_2f {
8526 if buf_id == 1 {
8527 PRESCALE_BUF1.with(|b| {
8528 let mut buf = b.borrow_mut();
8529 buf.clear();
8530 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8531 f(&buf)
8532 })
8533 } else {
8534 PRESCALE_BUF2.with(|b| {
8535 let mut buf = b.borrow_mut();
8536 buf.clear();
8537 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8538 f(&buf)
8539 })
8540 }
8541 } else {
8542 f(x)
8543 }
8544}
8545
8546#[cfg(target_arch = "x86_64")]
8551pub(crate) fn avx2_enabled() -> bool {
8552 use std::sync::OnceLock;
8553 static ON: OnceLock<bool> = OnceLock::new();
8554 *ON.get_or_init(|| {
8555 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8556 && std::arch::is_x86_feature_detected!("avx2")
8557 && std::arch::is_x86_feature_detected!("fma")
8558 })
8559}
8560
8561#[cfg(target_arch = "x86_64")]
8566fn avx2_a8w8_enabled() -> bool {
8567 use std::sync::OnceLock;
8568 static ON: OnceLock<bool> = OnceLock::new();
8569 *ON.get_or_init(|| {
8570 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8571 })
8572}
8573
8574#[inline]
8578pub(crate) fn a8w8_enabled() -> bool {
8579 #[cfg(target_arch = "aarch64")]
8580 {
8581 sdot_enabled()
8582 }
8583 #[cfg(target_arch = "x86_64")]
8584 {
8585 avx2_a8w8_enabled()
8586 }
8587 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8588 {
8589 false
8590 }
8591}
8592
8593#[inline]
8596#[allow(unreachable_code)]
8597fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8598 #[cfg(target_arch = "aarch64")]
8599 unsafe {
8600 return dot_i8_sdot(w, xq);
8601 }
8602 #[cfg(target_arch = "x86_64")]
8603 unsafe {
8604 if avx512vnni_enabled() {
8605 return dot_i8_i8_vnni(w, xq);
8606 }
8607 return dot_i8_i8_avx2(w, xq);
8608 }
8609 w.iter()
8610 .zip(xq)
8611 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8612 .sum()
8613}
8614
8615#[cfg(target_arch = "x86_64")]
8619fn avx512vnni_enabled() -> bool {
8620 use std::sync::OnceLock;
8621 static ON: OnceLock<bool> = OnceLock::new();
8622 *ON.get_or_init(|| {
8623 std::env::var("CMF_AVX512")
8624 .map(|v| v != "0")
8625 .unwrap_or(true)
8626 && std::arch::is_x86_feature_detected!("avx512f")
8627 && std::arch::is_x86_feature_detected!("avx512bw")
8628 && std::arch::is_x86_feature_detected!("avx512vl")
8629 && std::arch::is_x86_feature_detected!("avx512vnni")
8630 })
8631}
8632
8633#[cfg(target_arch = "x86_64")]
8641fn vnni_tiles_enabled() -> bool {
8642 use std::sync::OnceLock;
8643 static ON: OnceLock<bool> = OnceLock::new();
8644 *ON.get_or_init(|| {
8645 std::env::var("CMF_VNNI_TILES")
8646 .map(|v| v != "0")
8647 .unwrap_or(true)
8648 && avx512vnni_enabled()
8649 })
8650}
8651
8652#[cfg(target_arch = "x86_64")]
8657#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8658#[inline]
8659unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8660 unsafe {
8662 use core::arch::x86_64::*;
8663 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8664 let hi128 = _mm256_extracti128_si256::<1>(d);
8665 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8666 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8667 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8668 _mm_cvtsi128_si32(s32)
8669 }
8670}
8671
8672#[cfg(target_arch = "x86_64")]
8677#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8678unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8679 unsafe {
8681 use core::arch::x86_64::*;
8682 let n = w.len();
8683 let mut j = 0usize;
8684 let mut total: i32;
8685 {
8690 #[inline(always)]
8691 unsafe fn step(
8692 w: *const u8,
8693 x: *const i8,
8694 acc: core::arch::x86_64::__m512i,
8695 ) -> core::arch::x86_64::__m512i {
8696 unsafe {
8697 use core::arch::x86_64::*;
8698 let wv = _mm512_loadu_si512(w as *const _);
8699 let xv = _mm512_loadu_si512(x as *const _);
8700 let aw = _mm512_abs_epi8(wv);
8701 let neg = _mm512_movepi8_mask(wv);
8702 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8703 _mm512_dpbusd_epi32(acc, aw, sx)
8704 }
8705 }
8706 let (mut a0, mut a1, mut a2, mut a3) = (
8707 _mm512_setzero_si512(),
8708 _mm512_setzero_si512(),
8709 _mm512_setzero_si512(),
8710 _mm512_setzero_si512(),
8711 );
8712 while j + 256 <= n {
8713 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8714 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8715 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8716 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8717 j += 256;
8718 }
8719 while j + 64 <= n {
8720 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8721 j += 64;
8722 }
8723 let s01 = _mm512_add_epi32(a0, a1);
8724 let s23 = _mm512_add_epi32(a2, a3);
8725 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8726 }
8727 if j + 32 <= n {
8729 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8730 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8731 let d = _mm256_dpbusd_epi32(
8732 _mm256_setzero_si256(),
8733 _mm256_abs_epi8(wv),
8734 _mm256_sign_epi8(xv, wv),
8735 );
8736 let hi128 = _mm256_extracti128_si256::<1>(d);
8737 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8738 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8739 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8740 total += _mm_cvtsi128_si32(s32);
8741 j += 32;
8742 }
8743 while j < n {
8744 total += (w[j] as i8) as i32 * xq[j] as i32;
8745 j += 1;
8746 }
8747 total
8748 }
8749}
8750
8751#[cfg(target_arch = "x86_64")]
8753#[target_feature(enable = "avx2,fma")]
8754unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8755 unsafe {
8757 use core::arch::x86_64::*;
8758 let n = x.len();
8759 let wp = w.as_ptr();
8760 let xp = x.as_ptr();
8761 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8762 let mut j = 0usize;
8763 while j + 16 <= n {
8764 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8765 let lo = _mm256_cvtepi8_epi32(wb);
8766 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8767 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8768 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8769 j += 16;
8770 }
8771 let acc = _mm256_add_ps(a0, a1);
8772 let hi128 = _mm256_extractf128_ps::<1>(acc);
8773 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8774 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8775 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8776 let mut sum = _mm_cvtss_f32(s32);
8777 while j < n {
8778 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8779 j += 1;
8780 }
8781 sum
8782 }
8783}
8784
8785#[cfg(target_arch = "x86_64")]
8790#[target_feature(enable = "avx2")]
8791unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8792 unsafe {
8794 use core::arch::x86_64::*;
8795 let n = w.len();
8796 let ones = _mm256_set1_epi16(1);
8797 let mut acc = _mm256_setzero_si256();
8798 let mut j = 0usize;
8799 while j + 32 <= n {
8800 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8801 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8802 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8803 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8804 j += 32;
8805 }
8806 let hi128 = _mm256_extracti128_si256::<1>(acc);
8807 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8808 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8809 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8810 let mut s = _mm_cvtsi128_si32(s32);
8811 while j < n {
8812 s += (w[j] as i8) as i32 * xq[j] as i32;
8813 j += 1;
8814 }
8815 s
8816 }
8817}
8818
8819#[cfg(target_arch = "aarch64")]
8823#[target_feature(enable = "neon,i8mm")]
8824unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8825 unsafe {
8827 use core::arch::aarch64::*;
8828 use core::arch::asm;
8829 let n = w0.len();
8830 let w0p = w0.as_ptr() as *const i8;
8831 let w1p = w1.as_ptr() as *const i8;
8832 let mut acc01 = vdupq_n_s32(0);
8835 let mut acc23 = vdupq_n_s32(0);
8836 let mut i = 0usize;
8837 while i + 8 <= n {
8838 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8839 let xb01 = vcombine_s8(
8840 vld1_s8(xs[0].as_ptr().add(i)),
8841 vld1_s8(xs[1].as_ptr().add(i)),
8842 );
8843 let xb23 = vcombine_s8(
8844 vld1_s8(xs[2].as_ptr().add(i)),
8845 vld1_s8(xs[3].as_ptr().add(i)),
8846 );
8847 asm!(
8848 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8849 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8850 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8851 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8852 options(pure, nomem, nostack),
8853 );
8854 i += 8;
8855 }
8856 let mut out = [[0i32; 4]; 2];
8857 let a01: [i32; 4] = core::mem::transmute(acc01);
8858 let a23: [i32; 4] = core::mem::transmute(acc23);
8859 out[0][0] = a01[0];
8860 out[0][1] = a01[1];
8861 out[1][0] = a01[2];
8862 out[1][1] = a01[3];
8863 out[0][2] = a23[0];
8864 out[0][3] = a23[1];
8865 out[1][2] = a23[2];
8866 out[1][3] = a23[3];
8867 if i < n {
8868 for (k, x) in xs.iter().enumerate() {
8869 for j in i..n {
8870 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8871 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8872 }
8873 }
8874 }
8875 out
8876 }
8877}
8878
8879#[cfg(target_arch = "aarch64")]
8883#[target_feature(enable = "neon,dotprod")]
8884unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8885 unsafe {
8887 use core::arch::aarch64::*;
8888 use core::arch::asm;
8889 let n = w0.len();
8890 let w0p = w0.as_ptr() as *const i8;
8891 let w1p = w1.as_ptr() as *const i8;
8892 let mut acc = [[vdupq_n_s32(0); 4]; 2];
8893 let mut i = 0usize;
8894 while i + 16 <= n {
8895 let wv0 = vld1q_s8(w0p.add(i));
8896 let wv1 = vld1q_s8(w1p.add(i));
8897 for (k, x) in xs.iter().enumerate() {
8898 let xv = vld1q_s8(x.as_ptr().add(i));
8899 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8900 asm!(
8901 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8902 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8903 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8904 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8905 options(pure, nomem, nostack),
8906 );
8907 acc[0][k] = a0;
8908 acc[1][k] = a1;
8909 }
8910 i += 16;
8911 }
8912 let mut out = [[0i32; 4]; 2];
8913 for r in 0..2 {
8914 for k in 0..4 {
8915 out[r][k] = vaddvq_s32(acc[r][k]);
8916 }
8917 }
8918 if i < n {
8919 for (k, x) in xs.iter().enumerate() {
8920 for j in i..n {
8921 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8922 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8923 }
8924 }
8925 }
8926 out
8927 }
8928}
8929
8930#[cfg(target_arch = "x86_64")]
8936#[target_feature(enable = "avx2")]
8937unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8938 unsafe {
8940 use core::arch::x86_64::*;
8941 let n = w0.len();
8942 let ones = _mm256_set1_epi16(1);
8943 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8944 let mut j = 0usize;
8945 while j + 32 <= n {
8946 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8947 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8948 let aw0 = _mm256_abs_epi8(wv0);
8949 let aw1 = _mm256_abs_epi8(wv1);
8950 for (k, x) in xs.iter().enumerate() {
8951 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8952 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8953 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8954 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8955 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8956 }
8957 j += 32;
8958 }
8959 let mut out = [[0i32; 4]; 2];
8960 for r in 0..2 {
8961 for k in 0..4 {
8962 let a = acc[r][k];
8963 let hi128 = _mm256_extracti128_si256::<1>(a);
8964 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8965 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8966 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8967 out[r][k] = _mm_cvtsi128_si32(s32);
8968 }
8969 }
8970 if j < n {
8971 for (k, x) in xs.iter().enumerate() {
8972 for i in j..n {
8973 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8974 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8975 }
8976 }
8977 }
8978 out
8979 }
8980}
8981
8982#[cfg(target_arch = "x86_64")]
8987#[inline]
8988fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8989 let dot = if avx512vnni_enabled() && row.len() >= 64 {
8990 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8991 } else {
8992 unsafe { dot_i8_i8_avx2(row, &act.xq) }
8993 };
8994 let mut acc = dot as f32 * act.sx;
8995 for &(j, xv) in &act.outliers {
8996 acc += (row[j] as i8) as f32 * xv;
8997 }
8998 acc
8999}
9000
9001#[cfg(target_arch = "x86_64")]
9005#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9006unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9007 unsafe {
9009 use core::arch::x86_64::*;
9010 let n = w.len();
9011 let flip = _mm512_set1_epi8(-128); #[inline(always)]
9013 unsafe fn step(
9014 w: *const u8,
9015 x: *const i8,
9016 flip: core::arch::x86_64::__m512i,
9017 acc: core::arch::x86_64::__m512i,
9018 ) -> core::arch::x86_64::__m512i {
9019 unsafe {
9020 use core::arch::x86_64::*;
9021 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
9022 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
9023 }
9024 }
9025 let (mut a0, mut a1, mut a2, mut a3) = (
9026 _mm512_setzero_si512(),
9027 _mm512_setzero_si512(),
9028 _mm512_setzero_si512(),
9029 _mm512_setzero_si512(),
9030 );
9031 let mut j = 0usize;
9032 while j + 256 <= n {
9033 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9034 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
9035 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
9036 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
9037 j += 256;
9038 }
9039 while j + 64 <= n {
9040 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9041 j += 64;
9042 }
9043 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
9044 _mm512_add_epi32(a0, a1),
9045 _mm512_add_epi32(a2, a3),
9046 ));
9047 while j < n {
9049 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
9050 j += 1;
9051 }
9052 total
9053 }
9054}
9055
9056#[cfg(target_arch = "x86_64")]
9062#[target_feature(enable = "avx2")]
9063unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9064 unsafe {
9067 use core::arch::x86_64::*;
9068 let lomask = _mm_set1_epi8(0x0F);
9069 let eight = _mm256_set1_epi8(8);
9070 let ones = _mm256_set1_epi16(1);
9071 let mut acc = 0f32;
9072 for gi in 0..gpr {
9073 let g = g0 + gi;
9074 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9075 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9076 let lo = _mm_and_si128(b, lomask);
9077 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9078 let w = _mm256_sub_epi8(
9079 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9080 eight,
9081 );
9082 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9083 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
9084 let d = _mm256_madd_epi16(p16, ones);
9085 let hi128 = _mm256_extracti128_si256::<1>(d);
9086 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9087 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9088 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9089 acc += _mm_cvtsi128_si32(s32) as f32 * s;
9090 }
9091 acc
9092 }
9093}
9094
9095#[cfg(target_arch = "x86_64")]
9098#[target_feature(enable = "avx2")]
9099unsafe fn dot_q4_row_avx2_2(
9100 packed: &[u8],
9101 scales: &[u8],
9102 g0: usize,
9103 gpr: usize,
9104 xq1: &[i8],
9105 xq2: &[i8],
9106) -> (f32, f32) {
9107 unsafe {
9109 use core::arch::x86_64::*;
9110 let lomask = _mm_set1_epi8(0x0F);
9111 let eight = _mm256_set1_epi8(8);
9112 let ones = _mm256_set1_epi16(1);
9113 let (mut acc1, mut acc2) = (0f32, 0f32);
9114 #[inline(always)]
9115 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
9116 unsafe {
9117 use core::arch::x86_64::*;
9118 let hi128 = _mm256_extracti128_si256::<1>(d);
9119 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9120 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9121 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9122 _mm_cvtsi128_si32(s32)
9123 }
9124 }
9125 for gi in 0..gpr {
9126 let g = g0 + gi;
9127 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9128 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9129 let lo = _mm_and_si128(b, lomask);
9130 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9131 let w = _mm256_sub_epi8(
9132 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9133 eight,
9134 );
9135 let aw = _mm256_abs_epi8(w);
9136 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9137 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9138 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
9139 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
9140 acc1 += hsum(d1) as f32 * s;
9141 acc2 += hsum(d2) as f32 * s;
9142 }
9143 (acc1, acc2)
9144 }
9145}
9146
9147#[cfg(target_arch = "x86_64")]
9149fn q8_range_avx2(
9150 q: &[u8],
9151 row_scale: &[f32],
9152 act: &SplitAct,
9153 cols: usize,
9154 out_addr: SendMut,
9155 start: usize,
9156 end: usize,
9157) {
9158 for o in start..end {
9159 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9160 unsafe { *out_addr.at(o) = v };
9162 }
9163}
9164
9165#[cfg(target_arch = "x86_64")]
9167#[allow(clippy::too_many_arguments)]
9168fn q8_range2_avx2(
9169 q: &[u8],
9170 row_scale: &[f32],
9171 a1: &SplitAct,
9172 a2: &SplitAct,
9173 cols: usize,
9174 p1: SendMut,
9175 p2: SendMut,
9176 start: usize,
9177 end: usize,
9178) {
9179 for o in start..end {
9180 let row = &q[o * cols..(o + 1) * cols];
9181 unsafe {
9183 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9184 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9185 }
9186 }
9187}
9188
9189#[cfg(target_arch = "aarch64")]
9200fn i8mm_enabled() -> bool {
9201 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9202 *ON.get_or_init(|| {
9203 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9204 && std::arch::is_aarch64_feature_detected!("i8mm")
9205 })
9206}
9207
9208#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9212fn sdot_enabled() -> bool {
9213 use std::sync::OnceLock;
9214 static ON: OnceLock<bool> = OnceLock::new();
9215 *ON.get_or_init(|| {
9216 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9217 if !want {
9218 return false;
9219 }
9220
9221 #[cfg(target_arch = "aarch64")]
9222 {
9223 if std::arch::is_aarch64_feature_detected!("dotprod") {
9224 return true;
9225 }
9226 #[cfg(target_os = "android")]
9227 {
9228 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9229 if cpuinfo.lines().any(|l| {
9230 (l.starts_with("Features") || l.starts_with("features"))
9231 && l.contains("asimddp")
9232 }) {
9233 return true;
9234 }
9235 }
9236 }
9237 false
9238 }
9239 #[cfg(not(target_arch = "aarch64"))]
9240 {
9241 false
9242 }
9243 })
9244}
9245
9246struct SplitAct {
9251 xq: Vec<i8>,
9252 sx: f32,
9253 outliers: Vec<(usize, f32)>,
9254 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9257 xsum: i32,
9258}
9259
9260thread_local! {
9261 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9264 const { std::cell::RefCell::new(Vec::new()) };
9265}
9266
9267impl Drop for SplitAct {
9268 fn drop(&mut self) {
9269 let buf = std::mem::take(&mut self.xq);
9270 if buf.capacity() > 0 {
9271 XQ_FREE.with(|f| {
9272 let mut f = f.borrow_mut();
9273 if f.len() < 16 {
9274 f.push(buf);
9275 }
9276 });
9277 }
9278 }
9279}
9280
9281thread_local! {
9282 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9289}
9290
9291#[inline]
9294fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9295 KROW.with(|s| {
9296 let mut b = s.borrow_mut();
9297 if b.len() < n {
9298 b.resize(n, 0.0);
9299 }
9300 f(&mut b[..n])
9301 })
9302}
9303
9304fn split_act(x: &[f32]) -> SplitAct {
9305 let n = x.len();
9306 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9307 let thr = 8.0 * rms;
9308 let mut outliers: Vec<(usize, f32)> = Vec::new();
9312 let mut amax = 0f32;
9313 for (j, &v) in x.iter().enumerate() {
9314 let a = v.abs();
9315 if a > thr {
9316 outliers.push((j, v));
9317 } else if a > amax {
9318 amax = a;
9319 }
9320 }
9321 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9322 let inv = 1.0 / sx;
9323 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9324 xq.clear();
9325 xq.reserve(n);
9326 if outliers.is_empty() {
9327 xq.extend(
9328 x.iter()
9329 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9330 );
9331 } else {
9332 xq.extend(x.iter().map(|&v| {
9334 if v.abs() > thr {
9335 0
9336 } else {
9337 (v * inv).round().clamp(-127.0, 127.0) as i8
9338 }
9339 }));
9340 }
9341 let xsum = xq.iter().map(|&v| v as i32).sum();
9342 SplitAct {
9343 xq,
9344 sx,
9345 outliers,
9346 xsum,
9347 }
9348}
9349
9350fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9351 let n = x.len();
9352 let rms = (x
9353 .iter()
9354 .zip(col)
9355 .map(|(&a, &c)| {
9356 let v = a * c;
9357 (v * v) as f64
9358 })
9359 .sum::<f64>()
9360 / n.max(1) as f64)
9361 .sqrt() as f32;
9362 let thr = 8.0 * rms;
9363
9364 let mut outliers = Vec::new();
9365 let mut amax = 0f32;
9366 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9367 let v = a * c;
9368 let s = v.abs();
9369 if s > thr {
9370 outliers.push((j, v));
9371 } else if s > amax {
9372 amax = s;
9373 }
9374 }
9375
9376 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9377 let inv = 1.0 / sx;
9378 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9379 xq.clear();
9380 xq.reserve(n);
9381 if outliers.is_empty() {
9382 xq.extend(
9383 x.iter()
9384 .zip(col)
9385 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9386 );
9387 } else {
9388 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9389 let v = a * c;
9390 if v.abs() > thr {
9391 0
9392 } else {
9393 (v * inv).round().clamp(-127.0, 127.0) as i8
9394 }
9395 }));
9396 }
9397 let xsum = xq.iter().map(|&v| v as i32).sum();
9398 SplitAct {
9399 xq,
9400 sx,
9401 outliers,
9402 xsum,
9403 }
9404}
9405
9406#[cfg(target_arch = "aarch64")]
9409#[target_feature(enable = "neon,dotprod")]
9410unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9411 unsafe {
9413 use core::arch::aarch64::*;
9414 use core::arch::asm;
9415 let wp = w.as_ptr() as *const i8;
9416 let n = w.len();
9417 let (mut a0, mut a1, mut a2, mut a3) = (
9418 vdupq_n_s32(0),
9419 vdupq_n_s32(0),
9420 vdupq_n_s32(0),
9421 vdupq_n_s32(0),
9422 );
9423 let mut i = 0;
9424 while i + 64 <= n {
9425 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9426 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9427 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9428 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9429 asm!(
9430 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9431 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9432 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9433 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9434 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9435 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9436 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9437 options(pure, nomem, nostack),
9438 );
9439 i += 64;
9440 }
9441 while i + 16 <= n {
9442 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9443 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9444 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9445 i += 16;
9446 }
9447 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9448 while i < n {
9449 s += (*wp.add(i)) as i32 * xq[i] as i32;
9450 i += 1;
9451 }
9452 s
9453 }
9454}
9455
9456#[cfg(target_arch = "aarch64")]
9460#[target_feature(enable = "neon,dotprod")]
9461unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9462 unsafe {
9464 use core::arch::aarch64::*;
9465 use core::arch::asm;
9466 let n = xq.len();
9467 let px = xq.as_ptr();
9468 let (p0, p1, p2, p3) = (
9469 w0.as_ptr() as *const i8,
9470 w1.as_ptr() as *const i8,
9471 w2.as_ptr() as *const i8,
9472 w3.as_ptr() as *const i8,
9473 );
9474 let (mut a0, mut a1, mut a2, mut a3) = (
9475 vdupq_n_s32(0),
9476 vdupq_n_s32(0),
9477 vdupq_n_s32(0),
9478 vdupq_n_s32(0),
9479 );
9480 let mut i = 0;
9481 while i + 16 <= n {
9482 let x = vld1q_s8(px.add(i));
9483 let v0 = vld1q_s8(p0.add(i));
9484 let v1 = vld1q_s8(p1.add(i));
9485 let v2 = vld1q_s8(p2.add(i));
9486 let v3 = vld1q_s8(p3.add(i));
9487 asm!(
9488 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9489 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9490 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9491 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9492 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9493 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9494 options(pure, nomem, nostack),
9495 );
9496 i += 16;
9497 }
9498 let mut r = [
9499 vaddvq_s32(a0),
9500 vaddvq_s32(a1),
9501 vaddvq_s32(a2),
9502 vaddvq_s32(a3),
9503 ];
9504 while i < n {
9505 let xi = *px.add(i) as i32;
9506 r[0] += (*p0.add(i)) as i32 * xi;
9507 r[1] += (*p1.add(i)) as i32 * xi;
9508 r[2] += (*p2.add(i)) as i32 * xi;
9509 r[3] += (*p3.add(i)) as i32 * xi;
9510 i += 1;
9511 }
9512 r
9513 }
9514}
9515
9516#[cfg(target_arch = "aarch64")]
9523#[target_feature(enable = "neon,dotprod")]
9524unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9525 unsafe {
9528 use core::arch::aarch64::*;
9529 use core::arch::asm;
9530 let n = xq.len();
9531 let px = xq.as_ptr();
9532 let pg = g.as_ptr() as *const i8;
9533 let (mut a0, mut a1, mut a2, mut a3) = (
9534 vdupq_n_s32(0),
9535 vdupq_n_s32(0),
9536 vdupq_n_s32(0),
9537 vdupq_n_s32(0),
9538 );
9539 let mut i = 0;
9540 while i + 16 <= n {
9541 let x = vld1q_s8(px.add(i));
9542 let base = pg.add(4 * i);
9543 let v0 = vld1q_s8(base);
9544 let v1 = vld1q_s8(base.add(16));
9545 let v2 = vld1q_s8(base.add(32));
9546 let v3 = vld1q_s8(base.add(48));
9547 asm!(
9548 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9549 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9550 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9551 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9552 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9553 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9554 options(pure, nomem, nostack),
9555 );
9556 i += 16;
9557 }
9558 [
9559 vaddvq_s32(a0),
9560 vaddvq_s32(a1),
9561 vaddvq_s32(a2),
9562 vaddvq_s32(a3),
9563 ]
9564 }
9565}
9566
9567#[cfg(target_arch = "aarch64")]
9573fn q8_range_sdot(
9574 q: &[u8],
9575 rep: &[u8],
9576 row_scale: &[f32],
9577 act: &SplitAct,
9578 cols: usize,
9579 out_addr: SendMut,
9580 start: usize,
9581 end: usize,
9582) {
9583 let mut o = start;
9584 if !rep.is_empty() {
9587 while o < end && o % 4 != 0 {
9588 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9589 unsafe { *out_addr.at(o) = v };
9590 o += 1;
9591 }
9592 }
9593 while o + 4 <= end {
9594 let r = if rep.is_empty() {
9595 unsafe {
9596 dot_i8_sdot_4rows(
9597 &q[o * cols..(o + 1) * cols],
9598 &q[(o + 1) * cols..(o + 2) * cols],
9599 &q[(o + 2) * cols..(o + 3) * cols],
9600 &q[(o + 3) * cols..(o + 4) * cols],
9601 &act.xq,
9602 )
9603 }
9604 } else {
9605 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9606 };
9607 for k in 0..4 {
9608 let mut acc = r[k] as f32 * act.sx;
9609 for &(j, xv) in &act.outliers {
9610 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9611 }
9612 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9614 }
9615 o += 4;
9616 }
9617 while o < end {
9618 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9619 unsafe { *out_addr.at(o) = v };
9620 o += 1;
9621 }
9622}
9623
9624#[cfg(target_arch = "aarch64")]
9627#[allow(clippy::too_many_arguments)]
9628fn q8_range2_sdot(
9629 q: &[u8],
9630 row_scale: &[f32],
9631 a1: &SplitAct,
9632 a2: &SplitAct,
9633 cols: usize,
9634 p1: SendMut,
9635 p2: SendMut,
9636 start: usize,
9637 end: usize,
9638) {
9639 for o in start..end {
9640 let row = &q[o * cols..(o + 1) * cols];
9641 unsafe {
9643 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9644 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9645 }
9646 }
9647}
9648
9649#[allow(clippy::too_many_arguments)]
9651fn q8_range2_f32(
9652 q: &[u8],
9653 row_scale: &[f32],
9654 x1: &[f32],
9655 x2: &[f32],
9656 cols: usize,
9657 p1: SendMut,
9658 p2: SendMut,
9659 start: usize,
9660 end: usize,
9661) {
9662 for o in start..end {
9663 let row = &q[o * cols..(o + 1) * cols];
9664 unsafe {
9666 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9667 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9668 }
9669 }
9670}
9671
9672fn q8_range_f32(
9674 q: &[u8],
9675 row_scale: &[f32],
9676 xs: &[f32],
9677 cols: usize,
9678 out_addr: SendMut,
9679 start: usize,
9680 end: usize,
9681) {
9682 for o in start..end {
9683 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9684 unsafe { *out_addr.at(o) = v };
9686 }
9687}
9688
9689#[inline]
9693fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9694 #[cfg(target_arch = "aarch64")]
9695 return row_dot_sdot(row, act);
9696 #[cfg(target_arch = "x86_64")]
9697 return row_dot_avx2(row, act);
9698 #[allow(unreachable_code)]
9699 q8_row_dot_scalar(row, act)
9700}
9701
9702#[allow(dead_code)]
9703fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9704 let mut acc = 0i32;
9705 for (k, &b) in row.iter().enumerate() {
9706 acc += (b as i8) as i32 * act.xq[k] as i32;
9707 }
9708 let mut acc = acc as f32 * act.sx;
9709 for &(j, xv) in &act.outliers {
9710 acc += (row[j] as i8) as f32 * xv;
9711 }
9712 acc
9713}
9714
9715#[cfg(target_arch = "aarch64")]
9718#[inline]
9719fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9720 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9721 for &(j, xv) in &act.outliers {
9722 acc += (row[j] as i8) as f32 * xv;
9723 }
9724 acc
9725}
9726
9727#[cfg(target_arch = "aarch64")]
9735#[target_feature(enable = "neon,dotprod")]
9736unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9737 unsafe {
9740 use core::arch::aarch64::*;
9741 use core::arch::asm;
9742 let lomask = vdupq_n_u8(0x0F);
9743 let eight = vdupq_n_s8(8);
9744 let mut acc = 0f32;
9745 for gi in 0..gpr {
9746 let g = g0 + gi;
9747 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9748 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9749 let lo = vandq_u8(b, lomask);
9750 let hi = vshrq_n_u8::<4>(b);
9751 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9752 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9753 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9754 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9755 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9756 asm!(
9757 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9758 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9759 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9760 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9761 options(pure, nomem, nostack),
9762 );
9763 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9764 }
9765 acc
9766 }
9767}
9768
9769#[cfg(target_arch = "aarch64")]
9774#[target_feature(enable = "neon,dotprod")]
9775unsafe fn dot_q4_row_sdot2(
9776 packed: &[u8],
9777 scales: &[u8],
9778 g0: usize,
9779 gpr: usize,
9780 xq1: &[i8],
9781 xq2: &[i8],
9782) -> (f32, f32) {
9783 unsafe {
9786 use core::arch::aarch64::*;
9787 use core::arch::asm;
9788 let lomask = vdupq_n_u8(0x0F);
9789 let eight = vdupq_n_s8(8);
9790 let (mut acc1, mut acc2) = (0f32, 0f32);
9791 for gi in 0..gpr {
9792 let g = g0 + gi;
9793 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9794 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9795 let lo = vandq_u8(b, lomask);
9796 let hi = vshrq_n_u8::<4>(b);
9797 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9798 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9799 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9800 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9801 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9802 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9803 let (mut a0, mut a1, mut b0, mut b1) = (
9804 vdupq_n_s32(0),
9805 vdupq_n_s32(0),
9806 vdupq_n_s32(0),
9807 vdupq_n_s32(0),
9808 );
9809 asm!(
9810 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9811 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9812 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9813 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9814 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9815 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9816 e0 = in(vreg) e0, e1 = in(vreg) e1,
9817 x10 = in(vreg) x10, x11 = in(vreg) x11,
9818 x20 = in(vreg) x20, x21 = in(vreg) x21,
9819 options(pure, nomem, nostack),
9820 );
9821 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9822 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9823 }
9824 (acc1, acc2)
9825 }
9826}
9827
9828#[inline]
9833pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9834 #[cfg(target_arch = "aarch64")]
9835 unsafe {
9836 return axpy_i8_f32_neon(acc, row, w);
9837 }
9838 #[cfg(target_arch = "x86_64")]
9839 if avx2_enabled() {
9840 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9841 }
9842 #[allow(unreachable_code)]
9843 {
9844 for (a, &b) in acc.iter_mut().zip(row) {
9845 *a += w * b as f32;
9846 }
9847 }
9848}
9849
9850#[cfg(target_arch = "x86_64")]
9852#[target_feature(enable = "avx2,fma")]
9853unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9854 unsafe {
9856 use core::arch::x86_64::*;
9857 let n = acc.len().min(row.len());
9858 let ap = acc.as_mut_ptr();
9859 let rp = row.as_ptr();
9860 let wv = _mm256_set1_ps(w);
9861 let mut j = 0usize;
9862 while j + 16 <= n {
9863 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9864 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9865 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9866 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9867 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9868 _mm256_storeu_ps(ap.add(j), v0);
9869 _mm256_storeu_ps(ap.add(j + 8), v1);
9870 j += 16;
9871 }
9872 while j < n {
9873 *ap.add(j) += w * (*rp.add(j)) as f32;
9874 j += 1;
9875 }
9876 }
9877}
9878
9879#[cfg(target_arch = "aarch64")]
9880#[target_feature(enable = "neon")]
9881unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9882 unsafe {
9884 use core::arch::aarch64::*;
9885 let n = acc.len().min(row.len());
9886 let ap = acc.as_mut_ptr();
9887 let rp = row.as_ptr();
9888 let wv = vdupq_n_f32(w);
9889 let mut j = 0usize;
9890 while j + 16 <= n {
9891 let rb = vld1q_s8(rp.add(j));
9892 let lo = vmovl_s8(vget_low_s8(rb));
9893 let hi = vmovl_s8(vget_high_s8(rb));
9894 for (off, half) in [(0, lo), (8, hi)] {
9895 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9896 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9897 let o = j + off;
9898 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9899 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9900 }
9901 j += 16;
9902 }
9903 while j < n {
9904 *ap.add(j) += w * (*rp.add(j)) as f32;
9905 j += 1;
9906 }
9907 }
9908}
9909
9910#[inline]
9913pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9914 #[cfg(target_arch = "aarch64")]
9915 unsafe {
9916 return dot_i8_f32_neon(w, x);
9917 }
9918 #[cfg(target_arch = "x86_64")]
9919 if avx2_enabled() {
9920 return unsafe { dot_i8_f32_avx2(w, x) };
9921 }
9922 #[allow(unreachable_code)]
9923 {
9924 let mut sum = 0.0f32;
9925 for (j, &b) in w.iter().enumerate() {
9926 sum += (b as i8) as f32 * x[j];
9927 }
9928 sum
9929 }
9930}
9931
9932#[inline]
9936fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9937 #[cfg(target_arch = "aarch64")]
9938 unsafe {
9939 return dot_i8_col_f32_neon(w, x, col);
9940 }
9941 #[allow(unreachable_code)]
9942 {
9943 let mut sum = 0.0f32;
9944 for (j, &b) in w.iter().enumerate() {
9945 sum += (b as i8) as f32 * x[j] * col[j];
9946 }
9947 sum
9948 }
9949}
9950
9951#[cfg(target_arch = "aarch64")]
9952#[target_feature(enable = "neon")]
9953unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9954 unsafe {
9956 use core::arch::aarch64::*;
9957 let n = x.len();
9958 let wp = w.as_ptr() as *const i8;
9959 let xp = x.as_ptr();
9960 let cp = col.as_ptr();
9961 let (mut a0, mut a1, mut a2, mut a3) = (
9962 vdupq_n_f32(0.0),
9963 vdupq_n_f32(0.0),
9964 vdupq_n_f32(0.0),
9965 vdupq_n_f32(0.0),
9966 );
9967 let mut j = 0usize;
9968 while j + 16 <= n {
9969 let wb = vld1q_s8(wp.add(j));
9970 let lo = vmovl_s8(vget_low_s8(wb));
9971 let hi = vmovl_s8(vget_high_s8(wb));
9972 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9973 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9974 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9975 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9976 a0 = vfmaq_f32(
9977 a0,
9978 w0,
9979 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9980 );
9981 a1 = vfmaq_f32(
9982 a1,
9983 w1,
9984 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9985 );
9986 a2 = vfmaq_f32(
9987 a2,
9988 w2,
9989 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9990 );
9991 a3 = vfmaq_f32(
9992 a3,
9993 w3,
9994 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9995 );
9996 j += 16;
9997 }
9998 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9999 while j < n {
10000 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
10001 j += 1;
10002 }
10003 sum
10004 }
10005}
10006
10007#[cfg(target_arch = "aarch64")]
10008#[target_feature(enable = "neon")]
10009unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
10010 unsafe {
10012 use core::arch::aarch64::*;
10013 let n = x.len();
10014 let wp = w.as_ptr() as *const i8;
10015 let xp = x.as_ptr();
10016 let (mut a0, mut a1, mut a2, mut a3) = (
10017 vdupq_n_f32(0.0),
10018 vdupq_n_f32(0.0),
10019 vdupq_n_f32(0.0),
10020 vdupq_n_f32(0.0),
10021 );
10022 let mut j = 0usize;
10023 while j + 16 <= n {
10024 let wb = vld1q_s8(wp.add(j));
10025 let lo = vmovl_s8(vget_low_s8(wb));
10026 let hi = vmovl_s8(vget_high_s8(wb));
10027 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10028 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10029 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10030 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10031 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
10032 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
10033 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
10034 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
10035 j += 16;
10036 }
10037 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10038 while j < n {
10039 sum += (*wp.add(j)) as f32 * *xp.add(j);
10040 j += 1;
10041 }
10042 sum
10043 }
10044}
10045
10046#[allow(clippy::too_many_arguments)]
10047fn qmatvec(
10048 q: &[u8],
10049 rep: &[u8],
10050 row_scale: &[f32],
10051 x: &[f32],
10052 col_field: &[f32],
10053 dtype: TensorDtype,
10054 rows: usize,
10055 cols: usize,
10056 out: &mut [f32],
10057 pool: Option<&Pool>,
10058) {
10059 debug_assert_eq!(out.len(), rows);
10060 #[cfg(not(target_arch = "aarch64"))]
10061 let _ = rep;
10062
10063 #[cfg(target_arch = "aarch64")]
10064 if sdot_enabled() {
10065 let act = if dtype == TensorDtype::Q8_2f {
10066 split_act_q8_2f(x, col_field)
10067 } else {
10068 split_act(x)
10069 };
10070 let out_addr = SendMut(out.as_mut_ptr());
10071 let run_range = |start: usize, end: usize| {
10072 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
10073 };
10074 match pool {
10075 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10076 _ => run_range(0, rows),
10077 }
10078 return;
10079 }
10080 #[cfg(target_arch = "x86_64")]
10083 if avx2_a8w8_enabled() {
10084 let act = if dtype == TensorDtype::Q8_2f {
10085 split_act_q8_2f(x, col_field)
10086 } else {
10087 split_act(x)
10088 };
10089 let out_addr = SendMut(out.as_mut_ptr());
10090 let run_range = |start: usize, end: usize| {
10091 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
10092 };
10093 match pool {
10094 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10095 _ => run_range(0, rows),
10096 }
10097 return;
10098 }
10099
10100 prescale_with(x, col_field, dtype, 1, |xs| {
10101 let out_addr = SendMut(out.as_mut_ptr());
10102 let run_range = move |start: usize, end: usize| {
10103 for o in start..end {
10104 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10105 unsafe { *out_addr.at(o) = v };
10107 }
10108 };
10109 match pool {
10110 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10111 _ => run_range(0, rows),
10112 }
10113 });
10114}
10115
10116#[allow(clippy::too_many_arguments)]
10117fn qmatvec2(
10118 q: &[u8],
10119 row_scale: &[f32],
10120 x1: &[f32],
10121 x2: &[f32],
10122 col_field: &[f32],
10123 dtype: TensorDtype,
10124 rows: usize,
10125 cols: usize,
10126 o1: &mut [f32],
10127 o2: &mut [f32],
10128 pool: Option<&Pool>,
10129) {
10130 #[cfg(target_arch = "aarch64")]
10131 if sdot_enabled() {
10132 let a1s = if dtype == TensorDtype::Q8_2f {
10133 split_act_q8_2f(x1, col_field)
10134 } else {
10135 split_act(x1)
10136 };
10137 let a2s = if dtype == TensorDtype::Q8_2f {
10138 split_act_q8_2f(x2, col_field)
10139 } else {
10140 split_act(x2)
10141 };
10142 let p1 = SendMut(o1.as_mut_ptr());
10143 let p2 = SendMut(o2.as_mut_ptr());
10144 let run_range = |start: usize, end: usize| {
10145 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10146 };
10147 match pool {
10148 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10149 _ => run_range(0, rows),
10150 }
10151 return;
10152 }
10153 #[cfg(target_arch = "x86_64")]
10154 if avx2_a8w8_enabled() {
10155 let a1s = if dtype == TensorDtype::Q8_2f {
10156 split_act_q8_2f(x1, col_field)
10157 } else {
10158 split_act(x1)
10159 };
10160 let a2s = if dtype == TensorDtype::Q8_2f {
10161 split_act_q8_2f(x2, col_field)
10162 } else {
10163 split_act(x2)
10164 };
10165 let p1 = SendMut(o1.as_mut_ptr());
10166 let p2 = SendMut(o2.as_mut_ptr());
10167 let run_range = |start: usize, end: usize| {
10168 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10169 };
10170 match pool {
10171 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10172 _ => run_range(0, rows),
10173 }
10174 return;
10175 }
10176
10177 prescale_with(x1, col_field, dtype, 1, |x1s| {
10178 prescale_with(x2, col_field, dtype, 2, |x2s| {
10179 let p1 = SendMut(o1.as_mut_ptr());
10180 let p2 = SendMut(o2.as_mut_ptr());
10181 let run_range = move |start: usize, end: usize| {
10182 for o in start..end {
10183 let row = &q[o * cols..(o + 1) * cols];
10184 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10185 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10186 unsafe {
10188 *p1.at(o) = s1;
10189 *p2.at(o) = s2;
10190 }
10191 }
10192 };
10193 match pool {
10194 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10195 _ => run_range(0, rows),
10196 }
10197 });
10198 });
10199}
10200
10201#[derive(Clone, Copy)]
10202struct SendMut(*mut f32);
10203unsafe impl Send for SendMut {}
10204unsafe impl Sync for SendMut {}
10205
10206impl SendMut {
10207 #[inline]
10208 fn at(self, i: usize) -> *mut f32 {
10209 unsafe { self.0.add(i) }
10210 }
10211}
10212
10213#[cfg(test)]
10214mod tests {
10215 use super::*;
10216
10217 #[test]
10218 fn q2tp_i8_dot_matches_exact_on_grid() {
10219 let (rows, cols) = (5, 64);
10223 let gpr = cols / GROUP_SIZE;
10224 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10228 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10229 .collect();
10230 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10231 let x: Vec<f32> = (0..cols)
10232 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10233 .collect();
10234 let act = split_act(&x);
10235 assert!(
10236 act.outliers.is_empty(),
10237 "on-grid input must have no outliers"
10238 );
10239 let gsum = q1_group_sums(&act.xq, gpr);
10240 for r in 0..rows {
10241 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10242 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10243 assert!(
10244 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10245 "row {r}: exact {exact} vs i8 {fast}"
10246 );
10247 }
10248 }
10249
10250 #[test]
10251 fn q8_row_dot_fast_matches_scalar() {
10252 let cols = 96;
10255 let row: Vec<u8> = (0..cols)
10256 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10257 .collect();
10258 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10259 let act = split_act(&x);
10260 let fast = q8_row_dot(&row, &act);
10261 let scalar = q8_row_dot_scalar(&row, &act);
10262 assert!(
10263 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10264 "fast {fast} vs scalar {scalar}"
10265 );
10266 }
10267
10268 #[test]
10269 fn f32_matvec_matches_matvec_rows_bitexact() {
10270 let (rows, cols) = (300, 40);
10271 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10272 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10273 let qt = QTensor::from_f32(w.clone(), rows, cols);
10274
10275 let mut a = vec![0.0f32; rows];
10276 matvec_rows(None, &w, &x, &mut a);
10277 let mut b = vec![0.0f32; rows];
10278 qt.matvec(&x, &mut b, None);
10279 assert_eq!(a, b);
10280 }
10281
10282 #[test]
10283 fn sdot_kernel_exact_on_grid() {
10284 eprintln!("sdot_enabled = {}", sdot_enabled());
10289 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
10291 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10292 .collect();
10293 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10294 let x: Vec<f32> = (0..cols)
10295 .map(|i| match i % 3 {
10296 0 => 1.0,
10297 1 => -1.0,
10298 _ => 0.0,
10299 })
10300 .collect();
10301 let mut a = vec![0.0f32; rows];
10302 qmatvec(
10303 &w,
10304 &[],
10305 &scales,
10306 &x,
10307 &[],
10308 TensorDtype::Q8Row,
10309 rows,
10310 cols,
10311 &mut a,
10312 None,
10313 );
10314 for o in 0..rows {
10315 let mut acc = 0.0f32;
10316 for j in 0..cols {
10317 acc += (w[o * cols + j] as i8) as f32 * x[j];
10318 }
10319 let expect = acc * scales[o];
10320 assert!(
10321 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10322 "row {o}: {} vs {expect}",
10323 a[o]
10324 );
10325 }
10326 }
10327
10328 #[test]
10329 fn q1_tbl_fast_path_matches_reference() {
10330 let (rows, cols) = (5, 256);
10335 let gpr = cols / GROUP_SIZE;
10336 let mut bytes = Vec::new();
10337 for t in 0..rows * gpr {
10338 let s = 0.007 + (t % 11) as f32 * 0.004;
10339 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10340 for j in 0..4 {
10341 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10342 }
10343 }
10344 let x: Vec<f32> = (0..cols)
10345 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10346 .collect();
10347 let mut w = vec![0.0f32; rows * cols];
10348 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10349 let mut got = vec![0.0f32; rows];
10350 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10351 for o in 0..rows {
10352 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10353 assert!(
10354 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10355 "row {o}: {} vs {expect}",
10356 got[o]
10357 );
10358 }
10359 let b = 5usize;
10362 let mut xs_all = Vec::new();
10363 for bi in 0..b {
10364 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10365 }
10366 let mut mm = vec![0.0f32; b * rows];
10367 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10368 for bi in 0..b {
10369 let mut single = vec![0.0f32; rows];
10370 q1_matvec(
10371 &bytes,
10372 &xs_all[bi * cols..(bi + 1) * cols],
10373 rows,
10374 cols,
10375 &mut single,
10376 None,
10377 );
10378 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10379 }
10380 }
10381
10382 #[test]
10383 fn q1_kernels_match_exact_reference() {
10384 let (rows, cols) = (7, 96);
10386 let gpr = cols / GROUP_SIZE;
10387 let mut bytes = Vec::new();
10388 for t in 0..rows * gpr {
10389 let s = 0.01 + (t % 13) as f32 * 0.003;
10390 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10391 for j in 0..4 {
10392 bytes.push(((t * 31 + j * 97) % 251) as u8);
10393 }
10394 }
10395 let x: Vec<f32> = (0..cols)
10397 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10398 .collect();
10399 let mut w = vec![0.0f32; rows * cols];
10401 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10402 let mut expect = vec![0.0f32; rows];
10403 for o in 0..rows {
10404 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10405 }
10406 let mut got = vec![0.0f32; rows];
10407 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10408 for o in 0..rows {
10409 assert!(
10410 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10411 "row {o}: {} vs {}",
10412 got[o],
10413 expect[o]
10414 );
10415 }
10416 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10418 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10419 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10420 assert_eq!(a1, got);
10421 let mut xs = x.clone();
10422 xs.extend_from_slice(&x2);
10423 let mut mm = vec![0.0f32; 2 * rows];
10424 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10425 assert_eq!(&mm[..rows], got.as_slice());
10426 assert_eq!(&mm[rows..], a2.as_slice());
10427 }
10428
10429 #[test]
10430 fn repack_is_bit_identical() {
10431 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
10437 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10438 .collect();
10439 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10440 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10441 let rep = q8_repack_layout(&w, rows, cols);
10442 for g in 0..rows / 4 {
10444 for c in 0..cols / 16 {
10445 for lane in 0..4 {
10446 assert_eq!(
10447 &rep[g * 4 * cols + c * 64 + lane * 16
10448 ..g * 4 * cols + c * 64 + lane * 16 + 16],
10449 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10450 );
10451 }
10452 }
10453 }
10454 let mut a = vec![0.0f32; rows];
10455 qmatvec(
10456 &w,
10457 &[],
10458 &scales,
10459 &x,
10460 &[],
10461 TensorDtype::Q8Row,
10462 rows,
10463 cols,
10464 &mut a,
10465 None,
10466 );
10467 let mut b = vec![0.0f32; rows];
10468 qmatvec(
10469 &w,
10470 &rep,
10471 &scales,
10472 &x,
10473 &[],
10474 TensorDtype::Q8Row,
10475 rows,
10476 cols,
10477 &mut b,
10478 None,
10479 );
10480 assert_eq!(a, b, "full-range repack output diverged");
10481
10482 #[cfg(target_arch = "aarch64")]
10483 if sdot_enabled() {
10484 let act = split_act(&x);
10486 let mut c1 = vec![0.0f32; rows];
10487 let mut c2 = vec![0.0f32; rows];
10488 q8_range_sdot(
10489 &w,
10490 &[],
10491 &scales,
10492 &act,
10493 cols,
10494 SendMut(c1.as_mut_ptr()),
10495 3,
10496 rows - 2,
10497 );
10498 q8_range_sdot(
10499 &w,
10500 &rep,
10501 &scales,
10502 &act,
10503 cols,
10504 SendMut(c2.as_mut_ptr()),
10505 3,
10506 rows - 2,
10507 );
10508 assert_eq!(c1, c2, "unaligned-range repack output diverged");
10509 }
10510 }
10511
10512 #[test]
10513 fn sdot_a8w8_noise_is_bounded() {
10514 let (rows, cols) = (16, 512);
10518 let w: Vec<u8> = (0..rows * cols)
10519 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10520 .collect();
10521 let scales = vec![0.01f32; rows];
10522 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10523 let mut a = vec![0.0f32; rows];
10524 qmatvec(
10525 &w,
10526 &[],
10527 &scales,
10528 &x,
10529 &[],
10530 TensorDtype::Q8Row,
10531 rows,
10532 cols,
10533 &mut a,
10534 None,
10535 );
10536 let (mut num, mut den) = (0f64, 0f64);
10537 for o in 0..rows {
10538 let mut acc = 0.0f32;
10539 for j in 0..cols {
10540 acc += (w[o * cols + j] as i8) as f32 * x[j];
10541 }
10542 let expect = acc * scales[o];
10543 num += ((a[o] - expect) as f64).powi(2);
10544 den += (expect as f64).powi(2);
10545 }
10546 let rel = (num / den.max(1e-12)).sqrt();
10547 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10548 }
10549
10550 #[test]
10551 fn i8_dot_neon_matches_scalar() {
10552 let n = 100;
10553 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10554 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10555 let mut scalar = 0.0f32;
10556 for j in 0..n {
10557 scalar += (w[j] as i8) as f32 * x[j];
10558 }
10559 let fast = dot_i8_f32(&w, &x);
10560 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10561 }
10562
10563 #[test]
10565 fn vbitmatvec_matches_full_dequant() {
10566 let (rows, cols) = (6, 64);
10567 let ng = cols / GROUP_SIZE;
10568 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10570 let mut bytes = bits.clone();
10571 for g in 0..rows * ng {
10572 let s = 0.02 + 0.001 * g as f32;
10573 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10574 }
10575 for r in 0..rows {
10576 let b = bits[r] as usize;
10577 let (mut acc, mut nb) = (0u64, 0usize);
10578 let mut rowbytes = Vec::new();
10579 for i in 0..cols {
10580 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10581 acc = (acc << b) | v;
10582 nb += b;
10583 while nb >= 8 {
10584 nb -= 8;
10585 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10586 }
10587 }
10588 if nb > 0 {
10589 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10590 }
10591 bytes.extend_from_slice(&rowbytes);
10592 }
10593 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10594
10595 let mut reference = vec![0f32; rows * cols];
10596 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10597 let mut expect = vec![0f32; rows];
10598 for r in 0..rows {
10599 expect[r] = reference[r * cols..(r + 1) * cols]
10600 .iter()
10601 .zip(&x)
10602 .map(|(w, xv)| w * xv)
10603 .sum();
10604 }
10605 let mut got = vec![0f32; rows];
10606 let offsets = vbit_row_offsets(&bytes, rows, cols);
10607 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10608 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10612 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10613 for r in 0..rows {
10614 assert!(
10615 (got[r] - expect[r]).abs() < tol * scale,
10616 "row {r}: {} vs {}",
10617 got[r],
10618 expect[r]
10619 );
10620 }
10621 }
10622
10623 #[test]
10628 #[cfg(target_arch = "x86_64")]
10629 fn vbit_matmat_blocked_matches_per_row() {
10630 let (rows, cols, b) = (64usize, 128usize, 9usize);
10631 let ng = cols / GROUP_SIZE;
10632 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10633 let mut bytes = bits.clone();
10634 for g in 0..rows * ng {
10635 let sc = 0.02 + 0.0005 * g as f32;
10636 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10637 }
10638 for r in 0..rows {
10639 let bw = bits[r] as usize;
10640 let (mut acc, mut nb) = (0u64, 0usize);
10641 let mut rowbytes = Vec::new();
10642 for i in 0..cols {
10643 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10644 acc = (acc << bw) | v;
10645 nb += bw;
10646 while nb >= 8 {
10647 nb -= 8;
10648 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10649 }
10650 }
10651 if nb > 0 {
10652 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10653 }
10654 bytes.extend_from_slice(&rowbytes);
10655 }
10656 let x: Vec<f32> = (0..b * cols)
10657 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10658 .collect();
10659 let offsets = vbit_row_offsets(&bytes, rows, cols);
10660 let mut y_a = vec![0f32; b * rows];
10661 let mut y_b = vec![0f32; b * rows];
10662 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10663 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10664 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10665 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10666 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10667 let max_d = y_a
10668 .iter()
10669 .zip(&y_b)
10670 .map(|(p, q)| (p - q).abs())
10671 .fold(0.0f32, f32::max);
10672 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10673 }
10674
10675 #[test]
10683 fn q4t_matmat_blocked_matches_per_row() {
10684 let (rows, cols, b) = (16usize, 64usize, 9usize);
10685 let gpr = cols / GROUP_SIZE;
10686 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10687 for r in 0..rows {
10688 for g in 0..gpr {
10689 let t = (r * gpr + g) * Q4_TILE;
10690 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10691 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10692 for k in 0..16 {
10693 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10694 }
10695 }
10696 }
10697 let x: Vec<f32> = (0..b * cols)
10698 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10699 .collect();
10700 let mut y_blk = vec![0f32; b * rows];
10701 let mut y_row = vec![0f32; b * rows];
10702 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10703 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10704 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10705 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10706 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10707 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10708 }
10709
10710 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10717 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10718 let gpr = cols / GROUP_SIZE;
10719 let stride = q4tp_code_stride(gpr);
10720 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10721 let mut b = vec![0u8; codes_off + rows * stride];
10722 for r in 0..rows {
10723 for g in 0..gpr {
10724 let t = (r * gpr + g) * Q4TP_NIB;
10725 for k in 0..16 {
10726 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10727 }
10728 }
10729 let lo = -6.0 - 0.03 * (r % 17) as f32;
10730 let step = 0.01 + 0.004 * (r % 11) as f32;
10731 let p = params_off + r * 4;
10732 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10733 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10734 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10735 for g in 0..gpr {
10736 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10737 }
10738 }
10739 b
10740 }
10741
10742 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10746 let gpr = cols / GROUP_SIZE;
10747 let v = Q4tpView::new(bytes, rows, cols);
10748 let mut out = vec![0u8; rows * gpr * Q4_TILE];
10749 let mut sc = vec![0f32; gpr];
10750 for r in 0..rows {
10751 v.scales_into(r, gpr, &mut sc);
10752 for g in 0..gpr {
10753 let t = (r * gpr + g) * Q4_TILE;
10754 let s = sc[g];
10755 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10756 let src = (r * gpr + g) * Q4TP_NIB;
10757 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10758 }
10759 }
10760 out
10761 }
10762
10763 #[test]
10769 fn q4tp_exact_path_matches_dequant_reference() {
10770 let (rows, cols) = (256usize, 512usize);
10771 let gpr = cols / GROUP_SIZE;
10772 let bytes = synth_q4tp(rows, cols);
10773 let mut w = vec![0f32; rows * cols];
10774 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10775
10776 let x: Vec<f32> = (0..cols)
10777 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10778 .collect();
10779 let v = Q4tpView::new(&bytes, rows, cols);
10780 let mut sc = vec![0f32; gpr];
10781 for r in 0..rows {
10782 v.scales_into(r, gpr, &mut sc);
10783 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10784 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10785 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10789 assert!(
10790 (got - want).abs() <= 1e-5 * mag,
10791 "row {r}: kernel {got} vs dequant {want}"
10792 );
10793 }
10794 }
10795
10796 #[test]
10802 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10803 let (rows, cols) = (256usize, 512usize);
10804 let bytes = synth_q4tp(rows, cols);
10805 let twin = q4tp_as_q4t(&bytes, rows, cols);
10806 let x: Vec<f32> = (0..cols)
10807 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10808 .collect();
10809
10810 let mut got = vec![0f32; rows];
10811 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10812 let mut want = vec![0f32; rows];
10813 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10814
10815 let mut w = vec![0f32; rows * cols];
10818 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10819 for r in 0..rows {
10820 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10821 assert!(
10822 (got[r] - want[r]).abs() <= 1e-3 * mag,
10823 "row {r}: q4tp {} vs q4t {}",
10824 got[r],
10825 want[r]
10826 );
10827 }
10828 }
10829
10830 #[test]
10835 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10836 let (rows, cols, b) = (256usize, 512usize, 5usize);
10837 let bytes = synth_q4tp(rows, cols);
10838 let twin = q4tp_as_q4t(&bytes, rows, cols);
10839 let xs: Vec<f32> = (0..b * cols)
10840 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10841 .collect();
10842
10843 let mut got = vec![0f32; b * rows];
10844 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10845 let mut want = vec![0f32; b * rows];
10846 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10847
10848 let mut w = vec![0f32; rows * cols];
10849 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10850 for t in 0..b {
10851 for r in 0..rows {
10852 let mag: f32 = (0..cols)
10853 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10854 .sum();
10855 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10856 assert!(
10857 (g - wa).abs() <= 1e-3 * mag,
10858 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10859 );
10860 }
10861 }
10862 }
10863
10864 #[test]
10865 fn q4tp_matvec2_matches_the_single_stream_kernel() {
10866 let (rows, cols) = (128usize, 256usize);
10867 let gpr = cols / GROUP_SIZE;
10868 let bytes = synth_q4tp(rows, cols);
10869 let xs: Vec<f32> = (0..2 * cols)
10870 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10871 .collect();
10872
10873 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10874 q4tp_matvec2(
10875 &bytes,
10876 &xs[..cols],
10877 &xs[cols..],
10878 rows,
10879 cols,
10880 &mut o1,
10881 &mut o2,
10882 None,
10883 );
10884
10885 let v = Q4tpView::new(&bytes, rows, cols);
10888 let mut sc = vec![0f32; gpr];
10889 for r in 0..rows {
10890 v.scales_into(r, gpr, &mut sc);
10891 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10892 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10893 }
10894 }
10895
10896 #[test]
10903 fn q4tp_matvec_keeps_pace_with_q4t() {
10904 let (rows, cols) = (4096usize, 3072usize);
10905 let bytes = synth_q4tp(rows, cols);
10906 let twin = q4tp_as_q4t(&bytes, rows, cols);
10907 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10908 let mut o = vec![0f32; rows];
10909 let n = 12;
10910 let mut best = (f64::MAX, f64::MAX);
10911 for _ in 0..3 {
10914 let t0 = std::time::Instant::now();
10915 for _ in 0..n {
10916 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10917 }
10918 best.0 = best.0.min(t0.elapsed().as_secs_f64());
10919 let t0 = std::time::Instant::now();
10920 for _ in 0..n {
10921 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10922 }
10923 best.1 = best.1.min(t0.elapsed().as_secs_f64());
10924 }
10925 let ratio = best.1 / best.0;
10926 println!(
10927 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10928 best.0 * 1e3 / n as f64,
10929 best.1 * 1e3 / n as f64
10930 );
10931 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10932 }
10933
10934 #[cfg(target_os = "macos")]
10935 #[test]
10936 fn q4t_matmat_accel_matches_dequant_reference() {
10937 if !accel_gemm_enabled() {
10938 return; }
10940 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
10942 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10943 for r in 0..rows {
10944 for g in 0..gpr {
10945 let t = (r * gpr + g) * Q4_TILE;
10946 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10947 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10948 for k in 0..16 {
10949 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10950 }
10951 }
10952 }
10953 let x: Vec<f32> = (0..b * cols)
10954 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10955 .collect();
10956 let mut got = vec![0f32; b * rows];
10957 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10958 let mut w = vec![0f32; rows * cols];
10960 for r in 0..rows {
10961 for g in 0..gpr {
10962 let t = (r * gpr + g) * Q4_TILE;
10963 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10964 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10965 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10966 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10967 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10968 }
10969 }
10970 }
10971 for bi in 0..b {
10972 for r in 0..rows {
10973 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10974 let d = (got[bi * rows + r] - want).abs();
10975 assert!(
10976 d <= want.abs().max(1.0) * 1e-4,
10977 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10978 got[bi * rows + r]
10979 );
10980 }
10981 }
10982 }
10983
10984 #[test]
10985 fn q4matvec_matches_full_dequant() {
10986 let (rows, cols) = (8, 64);
10987 let groups = rows * cols / GROUP_SIZE;
10988 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10990 for i in 0..groups * 16 {
10991 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10992 }
10993 for g in 0..groups {
10994 let s = 0.01 + 0.003 * g as f32;
10995 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10996 }
10997 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10998
10999 let mut reference = vec![0.0f32; rows * cols];
11000 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11001 let mut expect = vec![0.0f32; rows];
11002 for r in 0..rows {
11003 expect[r] = reference[r * cols..(r + 1) * cols]
11004 .iter()
11005 .zip(&x)
11006 .map(|(w, xv)| w * xv)
11007 .sum();
11008 }
11009
11010 let mut got = vec![0.0f32; rows];
11011 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11012 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11016 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11017 for r in 0..rows {
11018 assert!(
11019 (got[r] - expect[r]).abs() < tol * scale,
11020 "row {r}: {} vs {}",
11021 got[r],
11022 expect[r]
11023 );
11024 }
11025 }
11026
11027 #[test]
11030 fn vbitmatvec2_equals_two_singles() {
11031 let (rows, cols) = (6, 64);
11032 let ng = cols / GROUP_SIZE;
11033 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11034 let mut bytes = bits.clone();
11035 for g in 0..rows * ng {
11036 let s = 0.02 + 0.001 * g as f32;
11037 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11038 }
11039 for r in 0..rows {
11040 let b = bits[r] as usize;
11041 let (mut acc, mut nb) = (0u64, 0usize);
11042 let mut rowbytes = Vec::new();
11043 for i in 0..cols {
11044 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11045 acc = (acc << b) | v;
11046 nb += b;
11047 while nb >= 8 {
11048 nb -= 8;
11049 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11050 }
11051 }
11052 if nb > 0 {
11053 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11054 }
11055 bytes.extend_from_slice(&rowbytes);
11056 }
11057 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11058 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
11059 let offsets = vbit_row_offsets(&bytes, rows, cols);
11060
11061 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11062 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
11063 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
11064 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11065 vbitmatvec2(
11066 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
11067 );
11068 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
11069 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
11070 }
11071
11072 #[test]
11074 fn q4matvec2_equals_two_singles() {
11075 let (rows, cols) = (8, 128);
11076 let groups = rows * cols / GROUP_SIZE;
11077 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11078 for i in 0..groups * 16 {
11079 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11080 }
11081 for g in 0..groups {
11082 let s = 0.01 + 0.003 * g as f32;
11083 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11084 }
11085 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11088 x1[9] = 250.0;
11089 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11090
11091 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11092 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
11093 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
11094 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11095 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
11096 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
11097 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
11098 }
11099
11100 #[test]
11103 fn matvec_many_equals_separate_matvecs() {
11104 use crate::pool::Pool;
11105 let (r1, r2, cols) = (300, 200, 64);
11106 let mk = |salt: usize, rows: usize| {
11107 QTensor::from_f32(
11108 (0..rows * cols)
11109 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
11110 .collect(),
11111 rows,
11112 cols,
11113 )
11114 };
11115 let (a, b) = (mk(1, r1), mk(5, r2));
11116 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
11117 let pool = Pool::new(3);
11118
11119 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
11120 a.matvec(&x, &mut ea, Some(&pool));
11121 b.matvec(&x, &mut eb, Some(&pool));
11122 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
11123 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11124 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
11125 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
11126 }
11127
11128 #[test]
11131 fn batched_matmat_equals_per_position_matvec() {
11132 let (rows, cols, b) = (8, 64, 5);
11133 let groups = rows * cols / GROUP_SIZE;
11135 let mut q4 = Vec::new();
11136 for i in 0..groups * 16 {
11137 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11138 }
11139 for g in 0..groups {
11140 q4.extend_from_slice(
11141 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11142 );
11143 }
11144 let ng = cols / GROUP_SIZE;
11146 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
11147 let mut vb = bits.clone();
11148 for g in 0..rows * ng {
11149 vb.extend_from_slice(
11150 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
11151 );
11152 }
11153 for r in 0..rows {
11154 let bw = bits[r] as usize;
11155 let (mut acc, mut nb) = (0u64, 0usize);
11156 let mut rowbytes = Vec::new();
11157 for i in 0..cols {
11158 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11159 acc = (acc << bw) | v;
11160 nb += bw;
11161 while nb >= 8 {
11162 nb -= 8;
11163 rowbytes.push(((acc >> nb) & 0xFF) as u8);
11164 }
11165 }
11166 if nb > 0 {
11167 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11168 }
11169 vb.extend_from_slice(&rowbytes);
11170 }
11171 let offsets = vbit_row_offsets(&vb, rows, cols);
11172
11173 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11174
11175 let mut got = vec![0f32; b * rows];
11177 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
11178 for bi in 0..b {
11179 let mut expect = vec![0f32; rows];
11180 q4matvec(
11181 &q4,
11182 &xs[bi * cols..(bi + 1) * cols],
11183 rows,
11184 cols,
11185 &mut expect,
11186 None,
11187 );
11188 assert_eq!(
11189 &got[bi * rows..(bi + 1) * rows],
11190 &expect[..],
11191 "q4 batch pos {bi}"
11192 );
11193 }
11194
11195 let mut got = vec![0f32; b * rows];
11197 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
11198 for bi in 0..b {
11199 let mut expect = vec![0f32; rows];
11200 vbitmatvec(
11201 &vb,
11202 &offsets,
11203 &xs[bi * cols..(bi + 1) * cols],
11204 rows,
11205 cols,
11206 &mut expect,
11207 None,
11208 );
11209 assert_eq!(
11210 &got[bi * rows..(bi + 1) * rows],
11211 &expect[..],
11212 "vbit batch pos {bi}"
11213 );
11214 }
11215 }
11216
11217 #[test]
11221 fn q4_tiled_matches_q4_block_bitexact() {
11222 let (rows, cols, b) = (8usize, 128usize, 3usize);
11223 let groups = rows * cols / GROUP_SIZE;
11224 let mut split = Vec::with_capacity(groups * 18);
11225 for i in 0..groups * 16 {
11226 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11227 }
11228 for g in 0..groups {
11229 split.extend_from_slice(
11230 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11231 );
11232 }
11233 let (packed, scales) = split.split_at(groups * 16);
11235 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
11236 for g in 0..groups {
11237 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
11238 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
11239 }
11240
11241 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11242 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11244
11245 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
11246 q4matvec(&split, &x1, rows, cols, &mut a, None);
11247 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
11248 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11249
11250 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11251 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11252 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11253 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11254 assert_eq!(a1, t1);
11255 assert_eq!(a2, t2);
11256
11257 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11258 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11259 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11260 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11261 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11262 }
11263
11264 #[test]
11271 fn q4matvec_sdot_outlier_exact() {
11272 let (rows, cols) = (4, 128);
11273 let groups = rows * cols / GROUP_SIZE;
11274 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11275 for i in 0..groups * 16 {
11276 bytes.push(((i * 11 + 5) % 256) as u8);
11277 }
11278 for g in 0..groups {
11279 let s = 0.02 + 0.002 * g as f32;
11280 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11281 }
11282 let mut x: Vec<f32> = (0..cols)
11283 .map(|i| match i % 3 {
11284 0 => 1.0,
11285 1 => -1.0,
11286 _ => 0.0,
11287 })
11288 .collect();
11289 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
11292 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11293 let mut expect = vec![0.0f32; rows];
11294 for r in 0..rows {
11295 expect[r] = reference[r * cols..(r + 1) * cols]
11296 .iter()
11297 .zip(&x)
11298 .map(|(w, xv)| w * xv)
11299 .sum();
11300 }
11301 let mut got = vec![0.0f32; rows];
11302 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11303 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11304 for r in 0..rows {
11305 assert!(
11306 (got[r] - expect[r]).abs() < 2e-3 * scale,
11307 "row {r}: {} vs {} (outlier term must be exact)",
11308 got[r],
11309 expect[r]
11310 );
11311 }
11312 }
11313
11314 #[test]
11318 fn q1t_matvec_matches_reference() {
11319 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11320 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
11322 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11323 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11325 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11326 let mut bytes = Vec::new();
11327 for r in 0..rows {
11328 for g in 0..gpr {
11329 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11330 let mut c = [0u8; 7];
11331 for k in 0..GROUP_SIZE {
11332 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11334 0
11335 } else {
11336 ((k + r * 3 + g) % 3) as u8 };
11338 cortiq_core::quant::q1t_pack(&mut c, k, code);
11339 }
11340 bytes.extend_from_slice(&c);
11341 }
11342 }
11343 let mut row_ptr = vec![0u32; rows + 1];
11346 for &(idx, _) in &outliers {
11347 row_ptr[idx as usize / cols + 1] += 1;
11348 }
11349 for r in 0..rows {
11350 row_ptr[r + 1] += row_ptr[r];
11351 }
11352 for &p in &row_ptr {
11353 bytes.extend_from_slice(&p.to_le_bytes());
11354 }
11355 for &(idx, v) in &outliers {
11356 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11357 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11358 }
11359
11360 let mut refw = vec![0f32; rows * cols];
11361 dequant_q1t(&bytes, rows, cols, &mut refw);
11362 let x: Vec<f32> = (0..cols)
11365 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11366 .collect();
11367 let mut expect = vec![0f32; rows];
11368 for r in 0..rows {
11369 let mut a = 0.0f32;
11370 for j in 0..cols {
11371 a += refw[r * cols + j] * x[j];
11372 }
11373 expect[r] = a;
11374 }
11375 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11376 let mut got = vec![0f32; rows];
11377 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11378 for r in 0..rows {
11379 assert!(
11380 (got[r] - expect[r]).abs() < tol(expect[r]),
11381 "row {r}: {} vs {}",
11382 got[r],
11383 expect[r]
11384 );
11385 }
11386 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
11388 let mut gm = vec![0f32; 2 * rows];
11389 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11390 for r in 0..rows {
11391 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11392 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11393 }
11394 let xb: Vec<f32> = (0..cols)
11398 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11399 .collect();
11400 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11401 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11402 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11403 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11404 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11405 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11406 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11407 }
11408
11409 #[test]
11412 fn q1t_matvec2_odd_gpr_matches_singles() {
11413 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11414 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
11416 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11417 for r in 0..rows {
11418 for g in 0..gpr {
11419 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11420 let mut c = [0u8; 7];
11421 for k in 0..GROUP_SIZE {
11422 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11423 }
11424 bytes.extend_from_slice(&c);
11425 }
11426 }
11427 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11428 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11429 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11430 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11431 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11432 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11433 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11434 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11435 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11436 }
11437
11438 #[test]
11442 #[ignore]
11443 fn q1t_matvec2_speed() {
11444 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11445 use std::time::Instant;
11446 let (rows, cols) = (8192usize, 4096usize);
11447 let gpr = cols / GROUP_SIZE;
11448 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11449 for r in 0..rows {
11450 for g in 0..gpr {
11451 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11452 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11453 let mut c = [0u8; 7];
11454 for k in 0..GROUP_SIZE {
11455 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11456 }
11457 bytes.extend_from_slice(&c);
11458 }
11459 }
11460 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11461 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11462 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11463 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11464 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11466 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11467 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11468 for _ in 0..8 {
11469 let t0 = Instant::now();
11470 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11471 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11472 let t1 = Instant::now();
11473 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11474 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11475 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11476 }
11477 assert_eq!(p1, s1);
11478 assert_eq!(p2, s2);
11479 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11480 }
11481
11482 #[test]
11486 #[ignore]
11487 fn q1t_matvec_speed() {
11488 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11489 use std::time::Instant;
11490 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
11492 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11493 for r in 0..rows {
11494 for g in 0..gpr {
11495 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11496 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11497 let mut c = [0u8; 7];
11498 for k in 0..GROUP_SIZE {
11499 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11500 }
11501 bytes.extend_from_slice(&c);
11502 }
11503 }
11504 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
11506 let mut idx = 0usize;
11507 while idx < n {
11508 row_ptr[idx / cols + 1] += 1;
11509 idx += stride;
11510 }
11511 for r in 0..rows {
11512 row_ptr[r + 1] += row_ptr[r];
11513 }
11514 for &p in &row_ptr {
11515 bytes.extend_from_slice(&p.to_le_bytes());
11516 }
11517 let mut idx = 0usize;
11518 while idx < n {
11519 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11520 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11521 idx += stride;
11522 }
11523 let x: Vec<f32> = (0..cols)
11526 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11527 .collect();
11528 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11529
11530 let slow = |out: &mut [f32]| {
11532 let mut buf = vec![0f32; cols];
11533 for r in 0..rows {
11534 for g in 0..gpr {
11535 let off = (r * gpr + g) * Q1T_TILE;
11536 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11537 let codes = &bytes[off + 2..off + Q1T_TILE];
11538 for k in 0..GROUP_SIZE {
11539 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11540 1 => s,
11541 2 => -s,
11542 _ => 0.0,
11543 };
11544 }
11545 }
11546 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11547 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11548 }
11549 };
11550 let iters = 5;
11551 let mut a = vec![0f32; rows];
11552 slow(&mut a); let t = Instant::now();
11554 for _ in 0..iters {
11555 slow(&mut a);
11556 }
11557 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11558
11559 let mut b = vec![0f32; rows];
11560 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
11562 for _ in 0..iters {
11563 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11564 }
11565 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11566
11567 for r in 0..rows {
11568 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11569 }
11570 println!(
11571 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
11572 slow_ms / fast_ms
11573 );
11574 }
11575}
11576
11577#[cfg(test)]
11578mod gemm_bench {
11579 #[test]
11590 #[ignore]
11591 fn q4tp_matmat_throughput() {
11592 let b: usize = std::env::var("CMF_BENCH_B")
11596 .ok()
11597 .and_then(|v| v.parse().ok())
11598 .unwrap_or(296);
11599 let (rows, cols) = (9216usize, 2304usize);
11600 let (_, _, _) = (rows, cols, b);
11601 let total =
11602 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
11603 .unwrap();
11604 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11609 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11610 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11611 let step = cortiq_core::quant::f32_to_f16(0.1);
11612 for r in 0..rows {
11613 let o = params_off + r * 4;
11614 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11615 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11616 }
11617 let _ = codes_off;
11618 let xs: Vec<f32> = (0..b * cols)
11619 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11620 .collect();
11621 let mut out = vec![0f32; b * rows];
11622 let pool = crate::pool::Pool::from_env();
11623 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11629 let reps: usize = std::env::var("CMF_BENCH_REPS")
11630 .ok()
11631 .and_then(|v| v.parse().ok())
11632 .unwrap_or(10);
11633 let mut best = [f64::MAX; 2];
11634 let mut sums = [0f32; 2];
11635 for _ in 0..reps {
11636 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11637 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11638 let t = std::time::Instant::now();
11639 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11640 best[k] = best[k].min(t.elapsed().as_secs_f64());
11641 sums[k] = out.iter().take(64).sum::<f32>();
11642 }
11643 }
11644 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11645 for (k, name) in ["previous", "tuned "].iter().enumerate() {
11646 println!(
11647 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11648 best[k] * 1e3,
11649 flops / best[k] / 1e9,
11650 sums[k]
11651 );
11652 }
11653 assert!(
11654 (sums[0] - sums[1]).abs() < 1e-2,
11655 "the tuned kernel changed the result: {} vs {}",
11656 sums[0],
11657 sums[1]
11658 );
11659 }
11660
11661 #[test]
11667 fn q4tp_matmat_blocked_matches_scalar() {
11668 use std::sync::atomic::Ordering::Relaxed;
11669 for &(rows, cols, b) in &[
11677 (64usize, 128usize, 7usize),
11678 (33, 96, 4),
11679 (16, 256, 9),
11680 (192, 2304, 37),
11681 ] {
11682 let total = cortiq_core::quant::expected_nbytes(
11683 cortiq_core::TensorDtype::Q4TiledP,
11684 &[rows, cols],
11685 )
11686 .unwrap();
11687 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11688 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11689 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11690 let step = cortiq_core::quant::f32_to_f16(0.1);
11691 for r in 0..rows {
11692 let o = params_off + r * 4;
11693 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11694 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11695 }
11696 let xs: Vec<f32> = (0..b * cols)
11697 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11698 .collect();
11699 let mut got = vec![0f32; b * rows];
11700 let mut want = vec![0f32; b * rows];
11701 let gpr = cols / 32;
11702 let view = super::Q4tpView::new(&bytes, rows, cols);
11703 let pool = crate::pool::Pool::from_env();
11704 super::Q4TP_ALT.store(2, Relaxed);
11705 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11706 super::Q4TP_ALT.store(1, Relaxed);
11707 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11708 super::Q4TP_ALT.store(0, Relaxed);
11709 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11716 let (mut worst, mut at) = (0f32, 0usize);
11717 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11718 if (g - w).abs() > worst {
11719 worst = (g - w).abs();
11720 at = i;
11721 }
11722 }
11723 assert!(
11724 worst <= 1e-4 * scale,
11725 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11726 (scale {scale:.3e}) at cell {at}: {} vs {}",
11727 got[at],
11728 want[at]
11729 );
11730
11731 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11738 for bi in 0..b {
11739 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11740 for r in 0..rows {
11741 let mut sc = vec![0f32; gpr];
11742 view.scales_into(r, gpr, &mut sc);
11743 let mut exact = 0f64;
11744 for j in 0..cols {
11745 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11746 exact += w as f64 * sq as f64 * act.xq[j] as f64;
11747 }
11748 exact *= act.sx as f64;
11749 for &(j, xv) in &act.outliers {
11750 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11751 exact += w as f64 * sq as f64 * xv as f64;
11752 }
11753 let i = bi * rows + r;
11754 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11755 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11756 }
11757 }
11758 println!(
11759 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11760 per-column {e_scalar:.3e}"
11761 );
11762 assert!(
11767 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11768 "{rows}x{cols} b={b}: error against f64 too large — blocked \
11769 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11770 );
11771 }
11772 }
11773
11774 #[test]
11778 #[ignore]
11779 fn q4t_matmat_throughput() {
11780 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11781 let total =
11782 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
11783 .unwrap();
11784 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11787 let sc = cortiq_core::quant::f32_to_f16(0.02);
11788 for t in bytes.chunks_mut(super::Q4_TILE) {
11789 t[..2].copy_from_slice(&sc.to_le_bytes());
11790 }
11791 let xs: Vec<f32> = (0..b * cols)
11792 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11793 .collect();
11794 let mut out = vec![0f32; b * rows];
11795 let pool = crate::pool::Pool::from_env();
11796 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11797 let reps: usize = std::env::var("CMF_BENCH_REPS")
11798 .ok()
11799 .and_then(|v| v.parse().ok())
11800 .unwrap_or(10);
11801 let mut best = f64::MAX;
11802 for _ in 0..reps {
11803 let t = std::time::Instant::now();
11804 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11805 best = best.min(t.elapsed().as_secs_f64());
11806 }
11807 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11808 println!(
11809 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11810 best * 1e3,
11811 flops / best / 1e9,
11812 out.iter().take(64).sum::<f32>()
11813 );
11814 }
11815}