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").map(|v| v != "0").unwrap_or(true)
130 })
131 }
132 }
133}
134
135static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
136
137pub fn set_blocked_override(on: Option<bool>) {
144 let v = match on {
145 None => 0,
146 Some(false) => 1,
147 Some(true) => 2,
148 };
149 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
150}
151
152fn gpu_lmhead_enabled() -> bool {
153 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
154 *ON.get_or_init(|| std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true))
155}
156
157fn gpu_split_frac() -> f32 {
158 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
159 *F.get_or_init(|| {
160 std::env::var("CMF_GPU_SPLIT")
161 .ok()
162 .and_then(|v| v.parse::<f32>().ok())
163 .unwrap_or(0.5)
164 .clamp(0.0, 1.0)
165 })
166}
167
168impl QTensor {
169 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
170 debug_assert_eq!(data.len(), rows * cols);
171 Self::F32 { data, rows, cols }
172 }
173
174 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
177 let idx = model
180 .tensor_index(name)
181 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
182 let entry = &model.tensors[idx];
183 if entry.shape.len() != 2 {
184 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
185 }
186 let (rows, cols) = (entry.shape[0], entry.shape[1]);
187 let bytes = model.entry_bytes(entry);
188
189 match entry.dtype {
190 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
191 let n = rows * cols;
192 let scales_off = n;
193 let row_scale: Vec<f32> = (0..rows)
194 .map(|o| {
195 f16_to_f32(u16::from_le_bytes([
196 bytes[scales_off + o * 2],
197 bytes[scales_off + o * 2 + 1],
198 ]))
199 })
200 .collect();
201 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
202 let col_off = n + rows * 2;
203 (0..cols)
204 .map(|i| {
205 f16_to_f32(u16::from_le_bytes([
206 bytes[col_off + i * 2],
207 bytes[col_off + i * 2 + 1],
208 ]))
209 })
210 .collect()
211 } else {
212 Vec::new()
213 };
214 Ok(Self::Mapped {
215 model: model.clone(),
216 idx,
217 dtype: entry.dtype,
218 rows,
219 cols,
220 row_scale,
221 col_field,
222 vbit_offsets: Vec::new(),
223 repack: q8_repack(bytes, rows, cols),
224 })
225 }
226 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
228 model: model.clone(),
229 idx,
230 dtype: entry.dtype,
231 rows,
232 cols,
233 row_scale: Vec::new(),
234 col_field: Vec::new(),
235 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
236 repack: Vec::new(),
237 }),
238 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
242 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
243 let offsets: Vec<usize> = (0..=rows)
244 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
245 .collect();
246 Ok(Self::Mapped {
247 model: model.clone(),
248 idx,
249 dtype: entry.dtype,
250 rows,
251 cols,
252 row_scale: Vec::new(),
253 col_field: Vec::new(),
254 vbit_offsets: offsets,
255 repack: Vec::new(),
256 })
257 }
258 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
264 model: model.clone(),
265 idx,
266 dtype: entry.dtype,
267 rows,
268 cols,
269 row_scale: Vec::new(),
270 col_field: Vec::new(),
271 vbit_offsets: Vec::new(),
272 repack: Vec::new(),
273 }),
274 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
277 model: model.clone(),
278 idx,
279 dtype: entry.dtype,
280 rows,
281 cols,
282 row_scale: Vec::new(),
283 col_field: Vec::new(),
284 vbit_offsets: Vec::new(),
285 repack: Vec::new(),
286 }),
287 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
289 model: model.clone(),
290 idx,
291 dtype: entry.dtype,
292 rows,
293 cols,
294 row_scale: Vec::new(),
295 col_field: Vec::new(),
296 vbit_offsets: Vec::new(),
297 repack: Vec::new(),
298 }),
299 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
300 model: model.clone(),
301 idx,
302 dtype: entry.dtype,
303 rows,
304 cols,
305 row_scale: Vec::new(),
306 col_field: Vec::new(),
307 vbit_offsets: Vec::new(),
308 repack: Vec::new(),
309 }),
310 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
312 model: model.clone(),
313 idx,
314 dtype: entry.dtype,
315 rows,
316 cols,
317 row_scale: Vec::new(),
318 col_field: Vec::new(),
319 vbit_offsets: Vec::new(),
320 repack: Vec::new(),
321 }),
322 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
326 model: model.clone(),
327 idx,
328 dtype: entry.dtype,
329 rows,
330 cols,
331 row_scale: Vec::new(),
332 col_field: Vec::new(),
333 vbit_offsets: Vec::new(),
334 repack: Vec::new(),
335 }),
336 _ => {
338 let mut data = vec![0.0f32; rows * cols];
339 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
340 Ok(Self::from_f32(data, rows, cols))
341 }
342 }
343 }
344
345 pub(crate) fn is_q1(&self) -> bool {
348 matches!(
349 self,
350 Self::Mapped {
351 dtype: TensorDtype::Q1,
352 ..
353 }
354 )
355 }
356
357 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
360 match self {
361 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
362 _ => None,
363 }
364 }
365
366 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
373 match self {
374 #[cfg(target_os = "macos")]
375 Self::Mapped {
376 dtype: TensorDtype::Q1T,
377 ..
378 } if !crate::gpu::metal_q1t_enabled() => None,
379 Self::Mapped {
380 idx,
381 dtype:
382 TensorDtype::Q1
383 | TensorDtype::Q1T
384 | TensorDtype::Q4Block
385 | TensorDtype::Q4Tiled
386 | TensorDtype::Q4TiledP
390 | TensorDtype::Q8Row
391 | TensorDtype::Q8_2f,
392 rows,
393 cols,
394 ..
395 } => Some((*idx, *rows, *cols)),
396 _ => None,
397 }
398 }
399
400 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
406 match self {
407 Self::Mapped {
408 idx,
409 dtype: TensorDtype::Q4Tiled,
410 rows,
411 cols,
412 ..
413 } => Some((*idx, *rows, *cols)),
414 _ => None,
415 }
416 }
417
418 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
422 match self {
423 Self::Mapped {
424 idx,
425 dtype: TensorDtype::Q4TiledP,
426 rows,
427 cols,
428 ..
429 } => Some((*idx, *rows, *cols)),
430 _ => None,
431 }
432 }
433
434 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
439 match self {
440 Self::Mapped {
441 idx,
442 dtype: TensorDtype::Q8Row,
443 rows,
444 cols,
445 row_scale,
446 col_field,
447 ..
448 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
449 _ => None,
450 }
451 }
452
453 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
457 match self {
458 Self::Mapped { dtype, .. } => Some(*dtype),
459 _ => None,
460 }
461 }
462
463 pub fn model_idx(&self) -> Option<usize> {
468 match self {
469 Self::Mapped { idx, .. } => Some(*idx),
470 _ => None,
471 }
472 }
473
474 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
479 match self {
480 Self::Mapped { model, .. } => Some(model.clone()),
481 _ => None,
482 }
483 }
484
485 pub fn rows(&self) -> usize {
486 match self {
487 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
488 }
489 }
490
491 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
494 match self {
495 Self::Mapped {
496 model,
497 idx,
498 dtype: TensorDtype::Q4Tiled,
499 ..
500 } => Some((model, *idx)),
501 _ => None,
502 }
503 }
504
505 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
508 match self {
509 Self::Mapped {
510 model,
511 idx,
512 dtype: TensorDtype::Q4TiledP,
513 ..
514 } => Some((model, *idx)),
515 _ => None,
516 }
517 }
518
519 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
522 match self {
523 Self::Mapped {
524 model,
525 idx,
526 dtype: TensorDtype::Q2TiledP,
527 ..
528 } => Some((model, *idx)),
529 _ => None,
530 }
531 }
532
533 pub fn cols(&self) -> usize {
534 match self {
535 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
536 }
537 }
538
539 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
542 match self {
543 Self::Mapped {
544 model,
545 idx,
546 dtype: TensorDtype::Q1,
547 ..
548 } => Some((model, *idx)),
549 _ => None,
550 }
551 }
552
553 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
563 match self {
564 Self::Mapped {
565 model,
566 idx,
567 dtype: TensorDtype::Q8Row,
568 row_scale,
569 ..
570 } => Some((model, *idx, 0, row_scale.as_slice())),
571 Self::Mapped {
572 model,
573 idx,
574 dtype: TensorDtype::Q1,
575 ..
576 } => Some((model, *idx, 1, &[])),
577 Self::Mapped {
582 model,
583 idx,
584 dtype: TensorDtype::Q4Tiled,
585 ..
586 } => Some((model, *idx, 5, &[])),
587 Self::Mapped {
591 model,
592 idx,
593 dtype: TensorDtype::Q4TiledP,
594 ..
595 } => Some((model, *idx, 6, &[])),
596 Self::Mapped {
597 model,
598 idx,
599 dtype: TensorDtype::Q4Block,
600 ..
601 } => Some((model, *idx, 2, &[])),
602 Self::Mapped {
607 model,
608 idx,
609 dtype: TensorDtype::Q8_2f,
610 ..
611 } => Some((model, *idx, 7, &[])),
612 Self::Mapped {
613 model,
614 idx,
615 dtype: TensorDtype::Q1T,
616 ..
617 } => Some((model, *idx, 3, &[])),
618 _ => None,
619 }
620 }
621
622 pub fn as_f32(&self) -> Option<&[f32]> {
625 match self {
626 Self::F32 { data, .. } => Some(data),
627 Self::Mapped { .. } => None,
628 }
629 }
630
631 fn quant_bytes(&self) -> &[u8] {
632 match self {
633 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
634 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
635 }
636 }
637
638 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
640 let cols = self.cols();
641 debug_assert_eq!(dst.len(), cols);
642 match self {
643 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
644 Self::Mapped {
645 dtype,
646 row_scale,
647 col_field,
648 vbit_offsets,
649 ..
650 } => {
651 if *dtype == TensorDtype::Q4Tiled {
652 let bytes = self.quant_bytes();
653 let gpr = cols / GROUP_SIZE;
654 for gi in 0..gpr {
655 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
656 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
657 for (k, &b) in tile[2..].iter().enumerate() {
658 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
659 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
660 }
661 }
662 return;
663 }
664 if *dtype == TensorDtype::Q4TiledP {
665 let bytes = self.quant_bytes();
666 let gpr = cols / GROUP_SIZE;
667 let v = Q4tpView::new(bytes, self.rows(), cols);
668 let mut sc = vec![0f32; gpr];
669 v.scales_into(r, gpr, &mut sc);
670 for gi in 0..gpr {
671 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
672 let s = sc[gi];
673 for (k, &b) in tile.iter().enumerate() {
674 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
675 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
676 }
677 }
678 return;
679 }
680 if *dtype == TensorDtype::Q2TiledP {
681 let bytes = self.quant_bytes();
682 let gpr = cols / GROUP_SIZE;
683 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
684 let mut sc = vec![0f32; gpr];
685 v.scales_into(r, gpr, &mut sc);
686 for gi in 0..gpr {
687 let ch =
688 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
689 let s = sc[gi];
690 for (k, &b) in ch.iter().enumerate() {
691 for j in 0..4 {
692 dst[gi * GROUP_SIZE + k * 4 + j] =
693 (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
694 }
695 }
696 }
697 return;
698 }
699 if *dtype == TensorDtype::Q4Block {
700 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
701 let gpr = cols / GROUP_SIZE;
702 for gi in 0..gpr {
703 let g = r * gpr + gi;
704 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
705 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
706 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
707 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
708 }
709 }
710 return;
711 }
712 if *dtype == TensorDtype::Q1 {
713 let bytes = self.quant_bytes();
714 let gpr = cols / GROUP_SIZE;
715 for gi in 0..gpr {
716 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
717 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
718 for (j, &b) in tile[2..].iter().enumerate() {
719 for k in 0..8 {
720 dst[gi * GROUP_SIZE + j * 8 + k] =
721 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
722 }
723 }
724 }
725 return;
726 }
727 if *dtype == TensorDtype::Q1T {
728 let bytes = self.quant_bytes();
729 let gpr = cols / GROUP_SIZE;
730 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
731 for gi in 0..gpr {
732 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
733 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
734 bytes[off],
735 bytes[off + 1],
736 ]));
737 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
738 for k in 0..GROUP_SIZE {
739 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
740 {
741 1 => s,
742 2 => -s,
743 _ => 0.0,
744 };
745 }
746 }
747 let rows = self.rows();
749 let entries = base_len + (rows + 1) * 4;
750 if entries <= bytes.len() {
751 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
752 let r0 = u32::from_le_bytes([
753 ptrs[r * 4],
754 ptrs[r * 4 + 1],
755 ptrs[r * 4 + 2],
756 ptrs[r * 4 + 3],
757 ]) as usize;
758 let r1 = u32::from_le_bytes([
759 ptrs[(r + 1) * 4],
760 ptrs[(r + 1) * 4 + 1],
761 ptrs[(r + 1) * 4 + 2],
762 ptrs[(r + 1) * 4 + 3],
763 ]) as usize;
764 let off = entries + r0 * 4;
765 for i in 0..r1 - r0 {
766 let item = &bytes[off + i * 4..off + i * 4 + 4];
767 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
768 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
769 item[2], item[3],
770 ]));
771 if c < cols {
772 dst[c] = v;
773 }
774 }
775 }
776 return;
777 }
778 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
779 let bytes = self.quant_bytes();
780 let rows = self.rows();
781 let ng = cols / GROUP_SIZE;
782 let bits = &bytes[..rows];
783 let sc_off = rows;
784 let off = vbit_offsets[r];
787 let b = bits[r] as usize;
788 let l = ((1usize << (b - 1)) - 1) as f32;
789 let data = &bytes[off..];
790 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
791 for (i, d) in dst.iter_mut().enumerate() {
792 while nbits < b {
793 acc = (acc << 8) | data[idx] as u64;
794 idx += 1;
795 nbits += 8;
796 }
797 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
798 nbits -= b;
799 let so = (r * ng + i / GROUP_SIZE) * 2;
800 let sv = f16_to_f32(u16::from_le_bytes([
801 bytes[sc_off + so],
802 bytes[sc_off + so + 1],
803 ]));
804 *d = (u - l) * sv;
805 }
806 return;
807 }
808 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
809 let s = row_scale[r];
810 match dtype {
811 TensorDtype::Q8Row => {
812 for (d, &b) in dst.iter_mut().zip(q) {
813 *d = (b as i8) as f32 * s;
814 }
815 }
816 TensorDtype::Q8_2f => {
817 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
818 *d = (b as i8) as f32 * s * col_field[i];
819 }
820 }
821 _ => unreachable!(),
822 }
823 }
824 }
825 }
826
827 pub fn sparse_col_ok(&self) -> bool {
832 match self {
833 Self::F32 { .. } => true,
834 Self::Mapped { dtype, .. } => {
835 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
836 }
837 }
838 }
839
840 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
844 let inter = self.cols();
845 let hidden = self.rows();
846 debug_assert_eq!(out.len(), hidden);
847 match self {
848 Self::F32 { data, .. } => {
849 for (k, o) in out.iter_mut().enumerate() {
850 *o += w * data[k * inter + c];
851 }
852 }
853 Self::Mapped {
854 dtype,
855 row_scale,
856 col_field,
857 ..
858 } => {
859 let q = self.quant_bytes();
860 let colf = if *dtype == TensorDtype::Q8_2f {
861 col_field[c]
862 } else {
863 1.0
864 };
865 let wc = w * colf;
866 for (k, o) in out.iter_mut().enumerate() {
867 let b = q[k * inter + c] as i8 as f32;
868 *o += wc * b * row_scale[k];
869 }
870 }
871 }
872 }
873
874 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
878 let cols = self.cols();
879 match self {
880 Self::F32 { data, .. } => {
881 let row = &data[r * cols..(r + 1) * cols];
882 row.iter().zip(x).map(|(w, v)| w * v).sum()
883 }
884 Self::Mapped {
885 dtype,
886 row_scale,
887 col_field,
888 ..
889 } => match dtype {
890 TensorDtype::Q8Row => {
891 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
892 dot_i8_f32(q, x) * row_scale[r]
893 }
894 TensorDtype::Q8_2f => {
895 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
896 dot_i8_col_f32(q, x, col_field) * row_scale[r]
897 }
898 _ => {
899 self.row_f32(r, scratch);
900 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
901 }
902 },
903 }
904 }
905
906 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
909 match self {
910 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
914 Self::Mapped {
915 model,
916 idx,
917 dtype,
918 rows,
919 cols,
920 row_scale,
921 col_field,
922 vbit_offsets,
923 repack,
924 } => {
925 let _ = (model, idx);
926 assert!(
936 out.len() >= *rows && x.len() >= *cols,
937 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
938 out.len(),
939 x.len(),
940 );
941 if *dtype == TensorDtype::Q4Block {
942 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
946 let t0 = std::time::Instant::now();
947 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
948 crate::gpu::ProbeArm::Gpu => {
949 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
950 crate::gpu::probe_record(
951 crate::gpu::OpClass::Matvec,
952 true,
953 t0.elapsed(),
954 );
955 return;
956 }
957 }
958 crate::gpu::ProbeArm::CpuTimed => {
959 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
960 crate::gpu::probe_record(
961 crate::gpu::OpClass::Matvec,
962 false,
963 t0.elapsed(),
964 );
965 return;
966 }
967 crate::gpu::ProbeArm::Cpu => {}
968 }
969 }
970 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
971 return;
972 }
973 if *dtype == TensorDtype::Q4Tiled {
974 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
979 let t0 = std::time::Instant::now();
980 let cls = crate::gpu::matvec_class(*rows, *cols);
981 match crate::gpu::probe_arm(cls) {
982 crate::gpu::ProbeArm::Gpu => {
983 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
984 crate::gpu::probe_record(cls, true, t0.elapsed());
985 return;
986 }
987 }
988 crate::gpu::ProbeArm::CpuTimed => {
989 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
990 crate::gpu::probe_record(cls, false, t0.elapsed());
991 return;
992 }
993 crate::gpu::ProbeArm::Cpu => {}
994 }
995 }
996 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
997 return;
998 }
999 if *dtype == TensorDtype::Q4TiledP {
1000 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1006 let t0 = std::time::Instant::now();
1007 let cls = crate::gpu::matvec_class(*rows, *cols);
1008 match crate::gpu::probe_arm(cls) {
1009 crate::gpu::ProbeArm::Gpu => {
1010 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1011 crate::gpu::probe_record(cls, true, t0.elapsed());
1012 return;
1013 }
1014 }
1015 crate::gpu::ProbeArm::CpuTimed => {
1016 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1017 crate::gpu::probe_record(cls, false, t0.elapsed());
1018 return;
1019 }
1020 crate::gpu::ProbeArm::Cpu => {}
1021 }
1022 }
1023 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1024 return;
1025 }
1026 if *dtype == TensorDtype::Q2TiledP {
1027 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1028 return;
1029 }
1030 if *dtype == TensorDtype::Q1 {
1031 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1036 let t0 = std::time::Instant::now();
1037 let arm = if crate::gpu::q1_force() {
1038 crate::gpu::ProbeArm::Gpu
1039 } else {
1040 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1041 };
1042 match arm {
1043 crate::gpu::ProbeArm::Gpu => {
1044 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1045 crate::gpu::probe_record(
1046 crate::gpu::OpClass::Matvec,
1047 true,
1048 t0.elapsed(),
1049 );
1050 return;
1051 }
1052 }
1053 crate::gpu::ProbeArm::CpuTimed => {
1054 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1055 crate::gpu::probe_record(
1056 crate::gpu::OpClass::Matvec,
1057 false,
1058 t0.elapsed(),
1059 );
1060 return;
1061 }
1062 crate::gpu::ProbeArm::Cpu => {}
1063 }
1064 }
1065 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1066 return;
1067 }
1068 if *dtype == TensorDtype::Q1T {
1069 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1073 let t0 = std::time::Instant::now();
1074 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1075 crate::gpu::ProbeArm::Gpu => {
1076 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1077 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1078 crate::gpu::probe_record(
1079 crate::gpu::OpClass::Matvec,
1080 true,
1081 t0.elapsed(),
1082 );
1083 return;
1084 }
1085 }
1086 crate::gpu::ProbeArm::CpuTimed => {
1087 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1088 crate::gpu::probe_record(
1089 crate::gpu::OpClass::Matvec,
1090 false,
1091 t0.elapsed(),
1092 );
1093 return;
1094 }
1095 crate::gpu::ProbeArm::Cpu => {}
1096 }
1097 }
1098 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1099 return;
1100 }
1101 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1102 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1103 return;
1104 }
1105 let xs = prescale(x, col_field, *dtype);
1106 if *rows >= crate::gpu::min_rows()
1111 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1112 && gpu_lmhead_enabled()
1113 && crate::gpu::enabled_here()
1114 {
1115 let t0 = std::time::Instant::now();
1118 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1119 crate::gpu::ProbeArm::Gpu => {}
1120 crate::gpu::ProbeArm::CpuTimed => {
1121 qmatvec(
1122 self.quant_bytes(),
1123 repack,
1124 row_scale,
1125 x,
1126 col_field,
1127 *dtype,
1128 *rows,
1129 *cols,
1130 out,
1131 pool,
1132 );
1133 crate::gpu::probe_record(
1134 crate::gpu::OpClass::Matvec,
1135 false,
1136 t0.elapsed(),
1137 );
1138 return;
1139 }
1140 crate::gpu::ProbeArm::Cpu => {
1141 qmatvec(
1142 self.quant_bytes(),
1143 repack,
1144 row_scale,
1145 x,
1146 col_field,
1147 *dtype,
1148 *rows,
1149 *cols,
1150 out,
1151 pool,
1152 );
1153 return;
1154 }
1155 }
1156 let frac = gpu_split_frac();
1157 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1158 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1159 let bytes = self.quant_bytes();
1160 let ok = std::thread::scope(|sc| {
1161 let g = sc.spawn(|| {
1162 crate::gpu::q8_matvec_range(
1163 model,
1164 *idx,
1165 cpu_rows,
1166 &row_scale[cpu_rows..],
1167 &xs,
1168 *rows - cpu_rows,
1169 *cols,
1170 out_gpu,
1171 )
1172 });
1173 if cpu_rows > 0 {
1174 let rep_cpu = if repack.is_empty() {
1177 &[][..]
1178 } else {
1179 &repack[..(cpu_rows / 4) * 4 * *cols]
1180 };
1181 qmatvec(
1182 &bytes[..cpu_rows * *cols],
1183 rep_cpu,
1184 &row_scale[..cpu_rows],
1185 x,
1186 col_field,
1187 *dtype,
1188 cpu_rows,
1189 *cols,
1190 out_cpu,
1191 pool,
1192 );
1193 }
1194 g.join().unwrap_or(false)
1195 });
1196 if ok {
1197 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1198 return;
1199 }
1200 qmatvec(
1203 &bytes[cpu_rows * *cols..(*rows) * *cols],
1204 &[],
1205 &row_scale[cpu_rows..],
1206 x,
1207 col_field,
1208 *dtype,
1209 *rows - cpu_rows,
1210 *cols,
1211 out_gpu,
1212 pool,
1213 );
1214 return;
1215 }
1216 qmatvec(
1217 self.quant_bytes(),
1218 repack,
1219 row_scale,
1220 x,
1221 col_field,
1222 *dtype,
1223 *rows,
1224 *cols,
1225 out,
1226 pool,
1227 );
1228 }
1229 }
1230 }
1231
1232 pub fn matvec2(
1234 &self,
1235 x1: &[f32],
1236 x2: &[f32],
1237 o1: &mut [f32],
1238 o2: &mut [f32],
1239 pool: Option<&Pool>,
1240 ) {
1241 match self {
1242 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1243 Self::Mapped {
1244 dtype,
1245 rows,
1246 cols,
1247 row_scale,
1248 col_field,
1249 vbit_offsets,
1250 ..
1251 } => {
1252 if *dtype == TensorDtype::Q4Block {
1253 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1254 return;
1255 }
1256 if *dtype == TensorDtype::Q4Tiled {
1257 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1258 return;
1259 }
1260 if *dtype == TensorDtype::Q4TiledP {
1261 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1262 return;
1263 }
1264 if *dtype == TensorDtype::Q2TiledP {
1265 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1266 return;
1267 }
1268 if *dtype == TensorDtype::Q1 {
1269 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1270 return;
1271 }
1272 if *dtype == TensorDtype::Q1T {
1273 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1279 return;
1280 }
1281 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1282 vbitmatvec2(
1283 self.quant_bytes(),
1284 vbit_offsets,
1285 x1,
1286 x2,
1287 *rows,
1288 *cols,
1289 o1,
1290 o2,
1291 pool,
1292 );
1293 return;
1294 }
1295 qmatvec2(
1296 self.quant_bytes(),
1297 row_scale,
1298 x1,
1299 x2,
1300 col_field,
1301 *dtype,
1302 *rows,
1303 *cols,
1304 o1,
1305 o2,
1306 pool,
1307 );
1308 }
1309 }
1310 }
1311}
1312
1313impl QTensor {
1314 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1319 let cols = self.cols();
1320 let rows = self.rows();
1321 debug_assert_eq!(xs_all.len(), b * cols);
1322 debug_assert_eq!(out.len(), b * rows);
1323 if crate::gptq_capture::capturing() {
1327 if let Self::Mapped { model, idx, .. } = self {
1328 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1329 }
1330 }
1331 match self {
1332 Self::F32 { data, .. } => {
1333 let out_addr = SendMut(out.as_mut_ptr());
1334 let run = |start: usize, end: usize| {
1335 for o in start..end {
1336 let row = &data[o * cols..(o + 1) * cols];
1337 for bi in 0..b {
1338 let x = &xs_all[bi * cols..(bi + 1) * cols];
1339 let mut acc = 0f32;
1340 for j in 0..cols {
1341 acc += row[j] * x[j];
1342 }
1343 unsafe { *out_addr.at(bi * rows + o) = acc };
1344 }
1345 }
1346 };
1347 dispatch_rows(pool, rows, &run);
1348 }
1349 Self::Mapped {
1350 dtype,
1351 row_scale,
1352 col_field,
1353 vbit_offsets,
1354 ..
1355 } => {
1356 if *dtype == TensorDtype::Q4Block {
1357 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1358 return;
1359 }
1360 if *dtype == TensorDtype::Q4TiledP {
1361 if b >= 32
1373 && b * rows * cols >= 128_000_000
1374 && cols % 32 == 0
1375 && !crate::gpu::mm_killed()
1376 && crate::gpu::enabled_here()
1377 {
1378 let class = if b >= 128 {
1379 crate::gpu::OpClass::MatmatWide
1380 } else {
1381 crate::gpu::OpClass::Matmat
1382 };
1383 if let Self::Mapped { model, idx, .. } = self {
1384 let t0 = std::time::Instant::now();
1385 let resident = crate::gpu::weight_is_resident(model, *idx);
1389 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1390 crate::gpu::ProbeArm::Gpu => {
1391 if crate::gpu::q4tp_matmat(
1392 model, *idx, xs_all, b, rows, cols, out,
1393 ) {
1394 let el = t0.elapsed();
1395 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1405 let budget = std::time::Duration::from_secs_f64(
1406 flops / 1.5e12 * 8.0 + 0.020,
1407 );
1408 if el > budget && !crate::gpu::probe_was_cold() {
1409 tracing::warn!(
1410 "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1411 device contended, CPU for the rest of the process"
1412 );
1413 crate::gpu::mm_kill();
1414 }
1415 crate::gpu::probe_record(class, true, el);
1416 return;
1417 }
1418 }
1419 crate::gpu::ProbeArm::CpuTimed => {
1420 q4tp_matmat(
1421 self.quant_bytes(),
1422 xs_all,
1423 b,
1424 rows,
1425 cols,
1426 out,
1427 pool,
1428 );
1429 crate::gpu::probe_record(class, false, t0.elapsed());
1430 return;
1431 }
1432 crate::gpu::ProbeArm::Cpu => {}
1433 }
1434 }
1435 }
1436 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1437 return;
1438 }
1439 if *dtype == TensorDtype::Q2TiledP {
1440 if b >= 32
1446 && b * rows * cols >= 128_000_000
1447 && cols % 32 == 0
1448 && !crate::gpu::mm_killed()
1449 && crate::gpu::enabled_here()
1450 {
1451 let class = if b >= 128 {
1452 crate::gpu::OpClass::MatmatWide
1453 } else {
1454 crate::gpu::OpClass::Matmat
1455 };
1456 if let Self::Mapped { model, idx, .. } = self {
1457 let t0 = std::time::Instant::now();
1458 match crate::gpu::probe_arm(class) {
1459 crate::gpu::ProbeArm::Gpu => {
1460 if crate::gpu::q2tp_matmat(
1461 model, *idx, xs_all, b, rows, cols, out,
1462 ) {
1463 crate::gpu::probe_record(class, true, t0.elapsed());
1464 return;
1465 }
1466 }
1467 crate::gpu::ProbeArm::CpuTimed => {
1468 q2tp_matmat(
1469 self.quant_bytes(),
1470 xs_all,
1471 b,
1472 rows,
1473 cols,
1474 out,
1475 pool,
1476 );
1477 crate::gpu::probe_record(class, false, t0.elapsed());
1478 return;
1479 }
1480 crate::gpu::ProbeArm::Cpu => {}
1481 }
1482 }
1483 }
1484 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1489 return;
1490 }
1491 if *dtype == TensorDtype::Q4Tiled {
1492 if b >= 32
1504 && b * rows * cols >= 128_000_000
1505 && cols % 32 == 0
1506 && !crate::gpu::mm_killed()
1507 && crate::gpu::enabled_here()
1508 {
1509 let class = if b >= 128 {
1510 crate::gpu::OpClass::MatmatWide
1511 } else {
1512 crate::gpu::OpClass::Matmat
1513 };
1514 if let Self::Mapped { model, idx, .. } = self {
1515 let t0 = std::time::Instant::now();
1516 match crate::gpu::probe_arm(class) {
1517 crate::gpu::ProbeArm::Gpu => {
1518 if crate::gpu::q4t_matmat(
1519 model, *idx, xs_all, b, rows, cols, out,
1520 ) {
1521 let el = t0.elapsed();
1522 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1532 let budget = std::time::Duration::from_secs_f64(
1533 flops / 1.5e12 * 8.0 + 0.020,
1534 );
1535 if el > budget && !crate::gpu::probe_was_cold() {
1536 tracing::warn!(
1537 "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1538 device contended, CPU for the rest of the process"
1539 );
1540 crate::gpu::mm_kill();
1541 }
1542 crate::gpu::probe_record(class, true, el);
1543 return;
1544 }
1545 }
1546 crate::gpu::ProbeArm::CpuTimed => {
1547 q4t_matmat(
1548 self.quant_bytes(),
1549 xs_all,
1550 b,
1551 rows,
1552 cols,
1553 out,
1554 pool,
1555 );
1556 crate::gpu::probe_record(class, false, t0.elapsed());
1557 return;
1558 }
1559 crate::gpu::ProbeArm::Cpu => {}
1560 }
1561 }
1562 }
1563 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1564 return;
1565 }
1566 if *dtype == TensorDtype::Q1 {
1567 if b >= 32
1570 && b * rows * cols >= 128_000_000
1571 && cols % 64 == 0
1572 && crate::gpu::enabled_here()
1573 {
1574 if let Self::Mapped { model, idx, .. } = self {
1575 let t0 = std::time::Instant::now();
1576 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1577 crate::gpu::ProbeArm::Gpu => {
1578 if crate::gpu::q1_matmat(
1579 model, *idx, xs_all, b, rows, cols, out,
1580 ) {
1581 crate::gpu::probe_record(
1582 crate::gpu::OpClass::Matmat,
1583 true,
1584 t0.elapsed(),
1585 );
1586 return;
1587 }
1588 }
1589 crate::gpu::ProbeArm::CpuTimed => {
1590 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1591 crate::gpu::probe_record(
1592 crate::gpu::OpClass::Matmat,
1593 false,
1594 t0.elapsed(),
1595 );
1596 return;
1597 }
1598 crate::gpu::ProbeArm::Cpu => {}
1599 }
1600 }
1601 }
1602 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1603 return;
1604 }
1605 if *dtype == TensorDtype::Q1T {
1606 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1609 if let Self::Mapped { model, idx, .. } = self {
1610 let t0 = std::time::Instant::now();
1611 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1612 crate::gpu::ProbeArm::Gpu => {
1613 if crate::gpu::q1t_matmat(
1614 model, *idx, xs_all, b, rows, cols, out,
1615 ) {
1616 crate::gpu::probe_record(
1617 crate::gpu::OpClass::Matmat,
1618 true,
1619 t0.elapsed(),
1620 );
1621 return;
1622 }
1623 }
1624 crate::gpu::ProbeArm::CpuTimed => {
1625 q1t_matmat(
1626 self.quant_bytes(),
1627 xs_all,
1628 b,
1629 rows,
1630 cols,
1631 out,
1632 pool,
1633 );
1634 crate::gpu::probe_record(
1635 crate::gpu::OpClass::Matmat,
1636 false,
1637 t0.elapsed(),
1638 );
1639 return;
1640 }
1641 crate::gpu::ProbeArm::Cpu => {}
1642 }
1643 }
1644 }
1645 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1646 return;
1647 }
1648 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1649 vbitmatmat(
1650 self.quant_bytes(),
1651 vbit_offsets,
1652 xs_all,
1653 b,
1654 rows,
1655 cols,
1656 out,
1657 pool,
1658 );
1659 return;
1660 }
1661 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1662 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1663 .collect();
1664 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1670 if let Self::Mapped { model, idx, .. } = self {
1671 let t0 = std::time::Instant::now();
1672 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1673 crate::gpu::ProbeArm::Gpu
1674 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1675 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1676 {
1677 let q = self.quant_bytes();
1681 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1682 return;
1683 }
1684 crate::gpu::ProbeArm::Gpu => {
1685 let flat: Vec<f32> =
1686 pre.iter().flat_map(|v| v.iter().copied()).collect();
1687 if crate::gpu::q8_matmat(
1688 model, *idx, row_scale, &flat, b, rows, cols, out,
1689 ) {
1690 crate::gpu::probe_record(
1691 crate::gpu::OpClass::Matmat,
1692 true,
1693 t0.elapsed(),
1694 );
1695 return;
1696 }
1697 }
1698 crate::gpu::ProbeArm::CpuTimed => {
1699 let q = self.quant_bytes();
1700 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1701 crate::gpu::probe_record(
1702 crate::gpu::OpClass::Matmat,
1703 false,
1704 t0.elapsed(),
1705 );
1706 return;
1707 }
1708 crate::gpu::ProbeArm::Cpu => {}
1709 }
1710 }
1711 }
1712 let q = self.quant_bytes();
1713 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1714 }
1715 }
1716 }
1717}
1718
1719impl QTensor {
1720 pub fn matvec_many<const N: usize>(
1727 ts: [&QTensor; N],
1728 x: &[f32],
1729 mut outs: [&mut [f32]; N],
1730 pool: Option<&Pool>,
1731 ) {
1732 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1733 let uniform_q8 = ts.iter().all(|t| {
1734 matches!(
1735 t,
1736 Self::Mapped {
1737 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1738 ..
1739 }
1740 )
1741 });
1742 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1743 let uniform_q4 = ts.iter().all(|t| {
1744 matches!(
1745 t,
1746 Self::Mapped {
1747 dtype: TensorDtype::Q4Block,
1748 ..
1749 }
1750 )
1751 });
1752 let uniform_vbit = ts.iter().all(|t| {
1753 matches!(
1754 t,
1755 Self::Mapped {
1756 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1757 ..
1758 }
1759 )
1760 });
1761 let uniform_q1 = ts.iter().all(|t| {
1762 matches!(
1763 t,
1764 Self::Mapped {
1765 dtype: TensorDtype::Q1,
1766 ..
1767 }
1768 )
1769 });
1770 let uniform_q1t = ts.iter().all(|t| {
1771 matches!(
1772 t,
1773 Self::Mapped {
1774 dtype: TensorDtype::Q1T,
1775 ..
1776 }
1777 )
1778 });
1779 let uniform_q4tp = ts.iter().all(|t| {
1784 matches!(
1785 t,
1786 Self::Mapped {
1787 dtype: TensorDtype::Q4TiledP,
1788 ..
1789 }
1790 )
1791 }) && ts.iter().all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1792 let Some(pool) = pool else {
1793 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1794 t.matvec(x, o, None);
1795 }
1796 return;
1797 };
1798 if total_rows < 256
1799 || !(uniform_q8
1800 || uniform_f32
1801 || uniform_q4
1802 || uniform_vbit
1803 || uniform_q1
1804 || uniform_q1t
1805 || uniform_q4tp)
1806 {
1807 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1808 t.matvec(x, o, Some(pool));
1809 }
1810 return;
1811 }
1812
1813 if uniform_q4tp {
1814 let cols = ts[0].cols();
1820 let gpr = cols / GROUP_SIZE;
1821 let views: Vec<Q4tpView> = ts
1822 .iter()
1823 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1824 .collect();
1825 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1826 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1827 let locate = |flat: usize| -> (usize, usize) {
1829 let mut acc = 0;
1830 for (i, &r) in rows_of.iter().enumerate() {
1831 if flat < acc + r {
1832 return (i, flat - acc);
1833 }
1834 acc += r;
1835 }
1836 (rows_of.len() - 1, 0)
1837 };
1838 let (views, outs_addr) = (&views, &outs_addr);
1839 if a8w8_enabled() {
1840 let act = split_act(x);
1841 let act = &act;
1842 let run = |start: usize, end: usize| {
1843 let mut sc = vec![0f32; gpr];
1844 for flat in start..end {
1845 let (t, r) = locate(flat);
1846 let v = &views[t];
1847 v.scales_into(r, gpr, &mut sc);
1848 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1849 for &(j, xv) in &act.outliers {
1850 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1851 acc += w * s * xv;
1852 }
1853 unsafe { *outs_addr[t].at(r) = acc };
1855 }
1856 };
1857 pool.run_rows(total_rows, &run);
1858 } else {
1859 let run = |start: usize, end: usize| {
1860 let mut sc = vec![0f32; gpr];
1861 for flat in start..end {
1862 let (t, r) = locate(flat);
1863 let v = &views[t];
1864 v.scales_into(r, gpr, &mut sc);
1865 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1867 }
1868 };
1869 pool.run_rows(total_rows, &run);
1870 }
1871 return;
1872 }
1873
1874 if uniform_q1 {
1875 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1878 if a8w8_enabled() {
1879 let act = split_act(x);
1880 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1881 let (act, gsum) = (&act, &gsum);
1882 let closures: [_; N] = std::array::from_fn(|i| {
1883 let (bytes, gpr, out) =
1884 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1885 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1886 });
1887 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1888 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1889 pool.run_many(&parts);
1890 } else {
1891 let closures: [_; N] = std::array::from_fn(|i| {
1892 let (bytes, gpr, out) =
1893 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1894 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1895 });
1896 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1897 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1898 pool.run_many(&parts);
1899 }
1900 return;
1901 }
1902
1903 if uniform_q1t {
1904 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1908 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1909 if a8w8_enabled() {
1910 let act = split_act(x);
1911 let act = &act;
1912 let x_ref = x;
1913 let closures: [_; N] = std::array::from_fn(|i| {
1914 let bytes = ts[i].quant_bytes();
1915 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1916 let gpr = cols / GROUP_SIZE;
1917 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1918 let out = outs_addr[i];
1919 move |s: usize, e: usize| {
1920 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1921 }
1922 });
1923 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1924 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1925 pool.run_many(&parts);
1926 } else {
1927 let x_ref = x;
1928 let closures: [_; N] = std::array::from_fn(|i| {
1929 let bytes = ts[i].quant_bytes();
1930 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1931 let gpr = cols / GROUP_SIZE;
1932 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1933 let out = outs_addr[i];
1934 move |s: usize, e: usize| {
1935 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1936 }
1937 });
1938 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1939 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1940 pool.run_many(&parts);
1941 }
1942 return;
1943 }
1944
1945 if uniform_q4 || uniform_vbit {
1946 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1947 if a8w8_enabled() {
1949 let act = split_act(x);
1950 let act = &act;
1951 if uniform_q4 {
1952 let closures: [_; N] = std::array::from_fn(|i| {
1953 let (packed, scales) =
1954 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1955 let (gpr, cols, out) =
1956 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1957 move |s: usize, e: usize| {
1958 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1959 }
1960 });
1961 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1962 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1963 pool.run_many(&parts);
1964 } else {
1965 let closures: [_; N] = std::array::from_fn(|i| {
1966 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1967 unreachable!()
1968 };
1969 let (bytes, rows, cols, out) = (
1970 ts[i].quant_bytes(),
1971 ts[i].rows(),
1972 ts[i].cols(),
1973 outs_addr[i],
1974 );
1975 move |s: usize, e: usize| {
1976 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1977 }
1978 });
1979 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1980 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1981 pool.run_many(&parts);
1982 }
1983 return;
1984 }
1985 if uniform_q4 {
1986 let closures: [_; N] = std::array::from_fn(|i| {
1987 let (packed, scales) =
1988 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1989 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1990 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1991 });
1992 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1993 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1994 pool.run_many(&parts);
1995 } else {
1996 let closures: [_; N] = std::array::from_fn(|i| {
1997 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1998 unreachable!()
1999 };
2000 let (bytes, rows, cols, out) = (
2001 ts[i].quant_bytes(),
2002 ts[i].rows(),
2003 ts[i].cols(),
2004 outs_addr[i],
2005 );
2006 move |s: usize, e: usize| {
2007 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2008 }
2009 });
2010 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2011 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2012 pool.run_many(&parts);
2013 }
2014 return;
2015 }
2016
2017 if uniform_f32 {
2018 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2019 let closures: [_; N] = std::array::from_fn(|i| {
2020 let Self::F32 { data, cols, .. } = ts[i] else {
2021 unreachable!()
2022 };
2023 let out = outs_addr[i];
2024 move |start: usize, end: usize| {
2025 for o in start..end {
2026 let row = &data[o * cols..(o + 1) * cols];
2027 let mut sum = 0.0f32;
2028 for j in 0..*cols {
2029 sum += row[j] * x[j];
2030 }
2031 unsafe { *out.at(o) = sum };
2033 }
2034 }
2035 });
2036 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2037 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2038 pool.run_many(&parts);
2039 return;
2040 }
2041
2042 struct Ctx<'a> {
2045 bytes: &'a [u8],
2046 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2047 rep: &'a [u8],
2048 row_scale: &'a [f32],
2049 cols: usize,
2050 xs: std::borrow::Cow<'a, [f32]>,
2051 }
2052 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2053 let Self::Mapped {
2054 dtype,
2055 cols,
2056 row_scale,
2057 col_field,
2058 repack,
2059 ..
2060 } = ts[i]
2061 else {
2062 unreachable!()
2063 };
2064 Ctx {
2065 bytes: ts[i].quant_bytes(),
2066 rep: repack,
2067 row_scale,
2068 cols: *cols,
2069 xs: prescale(x, col_field, *dtype),
2070 }
2071 });
2072 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2073 #[cfg(target_arch = "aarch64")]
2074 if sdot_enabled() {
2075 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2076 let closures: [_; N] = std::array::from_fn(|i| {
2077 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2078 move |start: usize, end: usize| {
2079 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2080 }
2081 });
2082 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2083 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2084 pool.run_many(&parts);
2085 return;
2086 }
2087 #[cfg(target_arch = "x86_64")]
2088 if avx2_a8w8_enabled() {
2089 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2090 let closures: [_; N] = std::array::from_fn(|i| {
2091 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2092 move |start: usize, end: usize| {
2093 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2094 }
2095 });
2096 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2097 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2098 pool.run_many(&parts);
2099 return;
2100 }
2101 let closures: [_; N] = std::array::from_fn(|i| {
2102 let (c, out) = (&ctxs[i], outs_addr[i]);
2103 move |start: usize, end: usize| {
2104 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2105 }
2106 });
2107 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2108 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2109 pool.run_many(&parts);
2110 }
2111}
2112
2113impl QTensor {
2114 #[allow(clippy::needless_range_loop)]
2119 pub fn matvec2_many<const N: usize>(
2120 ts: [&QTensor; N],
2121 x1: &[f32],
2122 x2: &[f32],
2123 mut o1s: [&mut [f32]; N],
2124 mut o2s: [&mut [f32]; N],
2125 pool: Option<&Pool>,
2126 ) {
2127 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2128 let uniform_q8 = ts.iter().all(|t| {
2129 matches!(
2130 t,
2131 Self::Mapped {
2132 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2133 ..
2134 }
2135 )
2136 });
2137 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2138 let uniform_q4 = ts.iter().all(|t| {
2139 matches!(
2140 t,
2141 Self::Mapped {
2142 dtype: TensorDtype::Q4Block,
2143 ..
2144 }
2145 )
2146 });
2147 let uniform_vbit = ts.iter().all(|t| {
2148 matches!(
2149 t,
2150 Self::Mapped {
2151 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2152 ..
2153 }
2154 )
2155 });
2156 let fusable = pool.is_some()
2157 && total_rows >= 256
2158 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2159 if !fusable {
2160 for i in 0..N {
2161 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2162 }
2163 return;
2164 }
2165 let pool = pool.unwrap();
2166
2167 if uniform_q4 || uniform_vbit {
2168 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2169 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2170 if a8w8_enabled() {
2172 let a1 = split_act(x1);
2173 let a2 = split_act(x2);
2174 let (a1, a2) = (&a1, &a2);
2175 if uniform_q4 {
2176 let closures: [_; N] = std::array::from_fn(|i| {
2177 let (packed, scales) =
2178 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2179 let (gpr, cols, o1, o2) =
2180 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2181 move |s: usize, e: usize| {
2182 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2183 }
2184 });
2185 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2186 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2187 pool.run_many(&parts);
2188 } else {
2189 let closures: [_; N] = std::array::from_fn(|i| {
2190 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2191 unreachable!()
2192 };
2193 let (bytes, rows, cols, o1, o2) = (
2194 ts[i].quant_bytes(),
2195 ts[i].rows(),
2196 ts[i].cols(),
2197 p1[i],
2198 p2[i],
2199 );
2200 move |s: usize, e: usize| {
2201 vbit_range2_a8w8(
2202 bytes,
2203 vbit_offsets,
2204 x1,
2205 x2,
2206 a1,
2207 a2,
2208 rows,
2209 cols,
2210 o1,
2211 o2,
2212 s,
2213 e,
2214 )
2215 }
2216 });
2217 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2218 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2219 pool.run_many(&parts);
2220 }
2221 return;
2222 }
2223 if uniform_q4 {
2224 let closures: [_; N] = std::array::from_fn(|i| {
2225 let (packed, scales) =
2226 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2227 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2228 move |s: usize, e: usize| {
2229 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
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 } else {
2236 let closures: [_; N] = std::array::from_fn(|i| {
2237 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2238 unreachable!()
2239 };
2240 let (bytes, rows, cols, o1, o2) = (
2241 ts[i].quant_bytes(),
2242 ts[i].rows(),
2243 ts[i].cols(),
2244 p1[i],
2245 p2[i],
2246 );
2247 move |s: usize, e: usize| {
2248 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2249 }
2250 });
2251 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2252 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2253 pool.run_many(&parts);
2254 }
2255 return;
2256 }
2257
2258 if uniform_f32 {
2259 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2260 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2261 let closures: [_; N] = std::array::from_fn(|i| {
2262 let Self::F32 { data, cols, .. } = ts[i] else {
2263 unreachable!()
2264 };
2265 let (o1, o2) = (p1[i], p2[i]);
2266 move |start: usize, end: usize| {
2267 for o in start..end {
2268 let row = &data[o * cols..(o + 1) * cols];
2269 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2270 for j in 0..*cols {
2271 s1 += row[j] * x1[j];
2272 s2 += row[j] * x2[j];
2273 }
2274 unsafe {
2276 *o1.at(o) = s1;
2277 *o2.at(o) = s2;
2278 }
2279 }
2280 }
2281 });
2282 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2283 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2284 pool.run_many(&parts);
2285 return;
2286 }
2287
2288 struct Ctx<'a> {
2289 bytes: &'a [u8],
2290 row_scale: &'a [f32],
2291 cols: usize,
2292 xs1: std::borrow::Cow<'a, [f32]>,
2293 xs2: std::borrow::Cow<'a, [f32]>,
2294 }
2295 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2296 let Self::Mapped {
2297 dtype,
2298 cols,
2299 row_scale,
2300 col_field,
2301 ..
2302 } = ts[i]
2303 else {
2304 unreachable!()
2305 };
2306 Ctx {
2307 bytes: ts[i].quant_bytes(),
2308 row_scale,
2309 cols: *cols,
2310 xs1: prescale(x1, col_field, *dtype),
2311 xs2: prescale(x2, col_field, *dtype),
2312 }
2313 });
2314 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2315 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2316 #[cfg(target_arch = "aarch64")]
2317 if sdot_enabled() {
2318 let acts: [(SplitAct, SplitAct); N] =
2319 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2320 let closures: [_; N] = std::array::from_fn(|i| {
2321 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2322 move |start: usize, end: usize| {
2323 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2324 }
2325 });
2326 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2327 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2328 pool.run_many(&parts);
2329 return;
2330 }
2331 #[cfg(target_arch = "x86_64")]
2332 if avx2_a8w8_enabled() {
2333 let acts: [(SplitAct, SplitAct); N] =
2334 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2335 let closures: [_; N] = std::array::from_fn(|i| {
2336 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2337 move |start: usize, end: usize| {
2338 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2339 }
2340 });
2341 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2342 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2343 pool.run_many(&parts);
2344 return;
2345 }
2346 let closures: [_; N] = std::array::from_fn(|i| {
2347 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2348 move |start: usize, end: usize| {
2349 q8_range2_f32(
2350 c.bytes,
2351 c.row_scale,
2352 &c.xs1,
2353 &c.xs2,
2354 c.cols,
2355 o1,
2356 o2,
2357 start,
2358 end,
2359 )
2360 }
2361 });
2362 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2363 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2364 pool.run_many(&parts);
2365 }
2366
2367 pub fn matvec_silu_mul(
2372 gate: &QTensor,
2373 up: &QTensor,
2374 x: &[f32],
2375 out: &mut [f32],
2376 pool: Option<&Pool>,
2377 ) -> bool {
2378 let inter = gate.rows();
2379 debug_assert_eq!(up.rows(), inter);
2380 debug_assert_eq!(out.len(), inter);
2381 debug_assert_eq!(gate.cols(), up.cols());
2382 if !a8w8_enabled() {
2383 return false;
2384 }
2385 let act = split_act(x);
2386 let act = &act;
2387 let x_ref = x;
2388 let out_addr = SendMut(out.as_mut_ptr());
2389
2390 match (gate, up) {
2391 (
2393 Self::Mapped {
2394 dtype: TensorDtype::Q4Block,
2395 ..
2396 },
2397 Self::Mapped {
2398 dtype: TensorDtype::Q4Block,
2399 ..
2400 },
2401 ) => {
2402 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2403 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2404 let gpr = gate.cols() / GROUP_SIZE;
2405 let cols = gate.cols();
2406 let run = move |start: usize, end: usize| {
2407 for r in start..end {
2408 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2409 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2410 for &(j, xv) in &act.outliers {
2411 let flat = r * cols + j;
2412 let gb = gp[flat / 2];
2413 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2414 let gsc = f16_to_f32(u16::from_le_bytes([
2415 gs[(flat / GROUP_SIZE) * 2],
2416 gs[(flat / GROUP_SIZE) * 2 + 1],
2417 ]));
2418 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2419 let ub = up_p[flat / 2];
2420 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2421 let usc = f16_to_f32(u16::from_le_bytes([
2422 up_s[(flat / GROUP_SIZE) * 2],
2423 up_s[(flat / GROUP_SIZE) * 2 + 1],
2424 ]));
2425 uv += ((un as i32 - 8) as f32) * usc * xv;
2426 }
2427 let silu_g = gv / (1.0 + (-gv).exp());
2428 unsafe { *out_addr.at(r) = silu_g * uv };
2430 }
2431 };
2432 dispatch_rows(pool, inter, &run);
2433 true
2434 }
2435 (
2439 Self::Mapped {
2440 dtype: TensorDtype::Q4Tiled,
2441 ..
2442 },
2443 Self::Mapped {
2444 dtype: TensorDtype::Q4Tiled,
2445 ..
2446 },
2447 ) => {
2448 let g_bytes = gate.quant_bytes();
2449 let u_bytes = up.quant_bytes();
2450 let gpr = gate.cols() / GROUP_SIZE;
2451 let run = move |start: usize, end: usize| {
2452 for r in start..end {
2453 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2454 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2455 for &(j, xv) in &act.outliers {
2456 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2457 gv += w * s * xv;
2458 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2459 uv += w * s * xv;
2460 }
2461 let silu_g = gv / (1.0 + (-gv).exp());
2462 unsafe { *out_addr.at(r) = silu_g * uv };
2464 }
2465 };
2466 dispatch_rows(pool, inter, &run);
2467 true
2468 }
2469 (
2472 Self::Mapped {
2473 dtype: TensorDtype::Q4TiledP,
2474 ..
2475 },
2476 Self::Mapped {
2477 dtype: TensorDtype::Q4TiledP,
2478 ..
2479 },
2480 ) => {
2481 let cols = gate.cols();
2482 let gpr = cols / GROUP_SIZE;
2483 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2484 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2485 let run = |start: usize, end: usize| {
2486 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2487 for r in start..end {
2488 gv_view.scales_into(r, gpr, &mut gsc);
2489 uv_view.scales_into(r, gpr, &mut usc);
2490 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2491 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2492 for &(j, xv) in &act.outliers {
2493 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2494 gv += w * s * xv;
2495 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2496 uv += w * s * xv;
2497 }
2498 let silu_g = gv / (1.0 + (-gv).exp());
2499 unsafe { *out_addr.at(r) = silu_g * uv };
2501 }
2502 };
2503 dispatch_rows(pool, inter, &run);
2504 true
2505 }
2506 (
2512 Self::Mapped {
2513 dtype: TensorDtype::Q1,
2514 ..
2515 },
2516 Self::Mapped {
2517 dtype: TensorDtype::Q1,
2518 ..
2519 },
2520 ) => {
2521 let g_bytes = gate.quant_bytes();
2522 let u_bytes = up.quant_bytes();
2523 let gpr = gate.cols() / GROUP_SIZE;
2524 let gsum = q1_group_sums(&act.xq, gpr);
2525 let gsum = &gsum;
2526 let run = move |start: usize, end: usize| {
2527 for r in start..end {
2528 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2529 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2530 for &(j, xv) in &act.outliers {
2531 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2532 gv += w * s * xv;
2533 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2534 uv += w * s * xv;
2535 }
2536 let silu_g = gv / (1.0 + (-gv).exp());
2537 unsafe { *out_addr.at(r) = silu_g * uv };
2539 }
2540 };
2541 dispatch_rows(pool, inter, &run);
2542 true
2543 }
2544 (
2548 Self::Mapped {
2549 dtype: TensorDtype::Q2TiledP,
2550 ..
2551 },
2552 Self::Mapped {
2553 dtype: TensorDtype::Q2TiledP,
2554 ..
2555 },
2556 ) => {
2557 let cols = gate.cols();
2558 let gpr = cols / GROUP_SIZE;
2559 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2560 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2561 let gsum = q1_group_sums(&act.xq, gpr);
2562 let gsum = &gsum;
2563 let run = move |start: usize, end: usize| {
2564 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2565 for r in start..end {
2566 gv_view.scales_into(r, gpr, &mut gsc);
2567 uv_view.scales_into(r, gpr, &mut usc);
2568 let mut gv =
2569 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2570 let mut uv =
2571 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2572 for &(j, xv) in &act.outliers {
2573 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2574 gv += w * s * xv;
2575 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2576 uv += w * s * xv;
2577 }
2578 let silu_g = gv / (1.0 + (-gv).exp());
2579 unsafe { *out_addr.at(r) = silu_g * uv };
2581 }
2582 };
2583 dispatch_rows(pool, inter, &run);
2584 true
2585 }
2586 (
2591 Self::Mapped {
2592 dtype: TensorDtype::Q8Row,
2593 row_scale: g_rs,
2594 ..
2595 },
2596 Self::Mapped {
2597 dtype: TensorDtype::Q8Row,
2598 row_scale: u_rs,
2599 ..
2600 },
2601 ) => {
2602 let g_bytes = gate.quant_bytes();
2603 let u_bytes = up.quant_bytes();
2604 let cols = gate.cols();
2605 let run = move |start: usize, end: usize| {
2606 for r in start..end {
2607 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2608 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2609 let silu_g = gv / (1.0 + (-gv).exp());
2610 unsafe { *out_addr.at(r) = silu_g * uv };
2612 }
2613 };
2614 dispatch_rows(pool, inter, &run);
2615 true
2616 }
2617 (
2619 Self::Mapped {
2620 dtype: TensorDtype::Q1T,
2621 ..
2622 },
2623 Self::Mapped {
2624 dtype: TensorDtype::Q1T,
2625 ..
2626 },
2627 ) => {
2628 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2629 let g_bytes = gate.quant_bytes();
2630 let u_bytes = up.quant_bytes();
2631 let gpr = gate.cols() / GROUP_SIZE;
2632 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2633 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2634 let run = move |start: usize, end: usize| {
2635 for r in start..end {
2636 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2637 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2638 for &(j, xv) in &act.outliers {
2639 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2640 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2641 }
2642 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2643 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2644 let silu_g = gv / (1.0 + (-gv).exp());
2645 unsafe { *out_addr.at(r) = silu_g * uv };
2647 }
2648 };
2649 dispatch_rows(pool, inter, &run);
2650 true
2651 }
2652 _ => false,
2653 }
2654 }
2655
2656 pub fn moe_gate_up_many(
2670 pairs: &[(&QTensor, &QTensor)],
2671 x: &[f32],
2672 outs: &mut [Vec<f32>],
2673 pool: Option<&Pool>,
2674 ) -> bool {
2675 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2676 return false;
2677 }
2678 let inter = pairs[0].0.rows();
2679 let cols = pairs[0].0.cols();
2680 if cols % GROUP_SIZE != 0 {
2681 return false;
2682 }
2683 let gpr = cols / GROUP_SIZE;
2684 let q2 = matches!(
2687 pairs[0].0,
2688 Self::Mapped {
2689 dtype: TensorDtype::Q2TiledP,
2690 ..
2691 }
2692 );
2693 let want = if q2 {
2694 TensorDtype::Q2TiledP
2695 } else {
2696 TensorDtype::Q4TiledP
2697 };
2698 let mut views = Vec::with_capacity(pairs.len() * 2);
2699 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2700 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2701 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2702 if !both
2703 || g.rows() != inter
2704 || u.rows() != inter
2705 || g.cols() != cols
2706 || u.cols() != cols
2707 || o.len() != inter
2708 {
2709 return false;
2710 }
2711 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2712 views.push(mk(g.quant_bytes(), inter, cols));
2713 views.push(mk(u.quant_bytes(), inter, cols));
2714 }
2715 let act = split_act(x);
2716 let gsum = if q2 {
2717 q1_group_sums(&act.xq, gpr)
2718 } else {
2719 Vec::new()
2720 };
2721 let (act, gsum) = (&act, &gsum);
2722 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2723 let (views, ptrs) = (&views, &ptrs);
2724 let run = |start: usize, end: usize| {
2725 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2726 for flat in start..end {
2727 let (e, r) = (flat / inter, flat % inter);
2728 let gv_view = &views[e * 2];
2729 let uv_view = &views[e * 2 + 1];
2730 gv_view.scales_into(r, gpr, &mut gsc);
2731 uv_view.scales_into(r, gpr, &mut usc);
2732 let (mut gv, mut uv) = if q2 {
2733 (
2734 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2735 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2736 )
2737 } else {
2738 (
2739 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2740 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2741 )
2742 };
2743 for &(j, xv) in &act.outliers {
2744 let (og, ou) = if q2 {
2745 (
2746 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2747 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2748 )
2749 } else {
2750 (
2751 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2752 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2753 )
2754 };
2755 gv += og.0 * og.1 * xv;
2756 uv += ou.0 * ou.1 * xv;
2757 }
2758 let silu_g = gv / (1.0 + (-gv).exp());
2759 unsafe { *ptrs[e].at(r) = silu_g * uv };
2761 }
2762 };
2763 dispatch_rows(pool, pairs.len() * inter, &run);
2764 true
2765 }
2766
2767 pub fn moe_down_many(
2776 downs: &[&QTensor],
2777 gs: &[Vec<f32>],
2778 weights: &[f32],
2779 out: &mut [f32],
2780 pool: Option<&Pool>,
2781 ) -> bool {
2782 if downs.is_empty()
2783 || downs.len() != gs.len()
2784 || downs.len() != weights.len()
2785 || !a8w8_enabled()
2786 {
2787 return false;
2788 }
2789 let rows = out.len();
2790 let cols = downs[0].cols();
2791 if cols % GROUP_SIZE != 0 {
2792 return false;
2793 }
2794 let gpr = cols / GROUP_SIZE;
2795 let mut views = Vec::with_capacity(downs.len());
2796 for (d, g) in downs.iter().zip(gs.iter()) {
2797 if !matches!(
2798 d,
2799 Self::Mapped {
2800 dtype: TensorDtype::Q4TiledP,
2801 ..
2802 }
2803 ) || d.rows() != rows
2804 || d.cols() != cols
2805 || g.len() != cols
2806 {
2807 return false;
2808 }
2809 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2810 }
2811 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2813 let out_addr = SendMut(out.as_mut_ptr());
2821 let (views, acts, weights) = (&views, &acts, &weights);
2822 let run = |start: usize, end: usize| {
2823 let mut sc = vec![0f32; gpr];
2824 for r in start..end {
2825 let mut acc = 0f32;
2826 for (e, v) in views.iter().enumerate() {
2827 v.scales_into(r, gpr, &mut sc);
2828 let a = &acts[e];
2829 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2830 for &(j, xv) in &a.outliers {
2831 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2832 d += w * s * xv;
2833 }
2834 acc += weights[e] * d;
2835 }
2836 unsafe { *out_addr.at(r) = acc };
2838 }
2839 };
2840 dispatch_rows(pool, rows, &run);
2841 true
2842 }
2843}
2844
2845#[cfg(target_os = "macos")]
2850mod accel_blas {
2851 #[link(name = "Accelerate", kind = "framework")]
2852 unsafe extern "C" {
2853 pub fn cblas_sgemm(
2854 order: i32,
2855 trans_a: i32,
2856 trans_b: i32,
2857 m: i32,
2858 n: i32,
2859 k: i32,
2860 alpha: f32,
2861 a: *const f32,
2862 lda: i32,
2863 b: *const f32,
2864 ldb: i32,
2865 beta: f32,
2866 c: *mut f32,
2867 ldc: i32,
2868 );
2869 }
2870}
2871
2872#[cfg(target_os = "macos")]
2873pub(crate) fn accel_gemm_enabled() -> bool {
2874 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2875 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2876}
2877
2878#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2881pub(crate) fn accel_gemm_enabled() -> bool {
2882 true
2883}
2884
2885#[cfg(target_arch = "aarch64")]
2892#[allow(clippy::too_many_arguments)]
2893pub(crate) fn neon_gemm_rm(
2894 m: usize,
2895 n: usize,
2896 k: usize,
2897 alpha: f32,
2898 a: &[f32],
2899 lda: usize,
2900 b_mat: &[f32],
2901 ldb: usize,
2902 b_rows_are_n: bool,
2903 c: &mut [f32],
2904 ldc: usize,
2905) {
2906 debug_assert!(a.len() >= (m - 1) * lda + k);
2907 debug_assert!(c.len() >= (m - 1) * ldc + n);
2908 unsafe {
2910 use core::arch::aarch64::*;
2911 let mut i = 0usize;
2912 while i < m {
2913 let mi = (m - i).min(4);
2914 let mut j = 0usize;
2915 while j < n {
2916 let nj = (n - j).min(8);
2917 if mi == 4 && nj == 8 {
2918 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2919 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2920 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2921 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2922 for p in 0..k {
2923 let (b0, b1) = if b_rows_are_n {
2924 let base = b_mat.as_ptr().add(j * ldb + p);
2927 let g = |o: usize| *base.add(o * ldb);
2928 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2929 } else {
2930 let base = b_mat.as_ptr().add(p * ldb + j);
2931 (
2932 [*base, *base.add(1), *base.add(2), *base.add(3)],
2933 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2934 )
2935 };
2936 let bv0 = vld1q_f32(b0.as_ptr());
2937 let bv1 = vld1q_f32(b1.as_ptr());
2938 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2939 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2940 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2941 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2942 c0a = vfmaq_f32(c0a, a0, bv0);
2943 c0b = vfmaq_f32(c0b, a0, bv1);
2944 c1a = vfmaq_f32(c1a, a1, bv0);
2945 c1b = vfmaq_f32(c1b, a1, bv1);
2946 c2a = vfmaq_f32(c2a, a2, bv0);
2947 c2b = vfmaq_f32(c2b, a2, bv1);
2948 c3a = vfmaq_f32(c3a, a3, bv0);
2949 c3b = vfmaq_f32(c3b, a3, bv1);
2950 }
2951 let al = vdupq_n_f32(alpha);
2952 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2953 .iter()
2954 .enumerate()
2955 {
2956 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2957 vst1q_f32(dst, vmulq_f32(*ca, al));
2958 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2959 }
2960 } else {
2961 for r in 0..mi {
2962 for q in 0..nj {
2963 let mut acc = 0f32;
2964 for p in 0..k {
2965 let bv = if b_rows_are_n {
2966 b_mat[(j + q) * ldb + p]
2967 } else {
2968 b_mat[p * ldb + j + q]
2969 };
2970 acc += a[(i + r) * lda + p] * bv;
2971 }
2972 c[(i + r) * ldc + j + q] = acc * alpha;
2973 }
2974 }
2975 }
2976 j += nj;
2977 }
2978 i += mi;
2979 }
2980 }
2981}
2982
2983#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2985#[allow(clippy::too_many_arguments)]
2986pub(crate) fn sgemm_rm(
2987 m: usize,
2988 n: usize,
2989 k: usize,
2990 alpha: f32,
2991 a: &[f32],
2992 lda: usize,
2993 b_mat: &[f32],
2994 ldb: usize,
2995 b_rows_are_n: bool,
2996 c: &mut [f32],
2997 ldc: usize,
2998) {
2999 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3000}
3001
3002#[allow(clippy::too_many_arguments)]
3006pub fn sgemm_public(
3007 m: usize,
3008 n: usize,
3009 k: usize,
3010 alpha: f32,
3011 a: &[f32],
3012 lda: usize,
3013 b_mat: &[f32],
3014 ldb: usize,
3015 b_rows_are_n: bool,
3016 c: &mut [f32],
3017 ldc: usize,
3018) {
3019 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3020 {
3021 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3022 }
3023 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3028 {
3029 for i in 0..m {
3030 for j in 0..n {
3031 let mut acc = 0f32;
3032 for p in 0..k {
3033 let bv = if b_rows_are_n {
3034 b_mat[j * ldb + p]
3035 } else {
3036 b_mat[p * ldb + j]
3037 };
3038 acc += a[i * lda + p] * bv;
3039 }
3040 c[i * ldc + j] = alpha * acc;
3041 }
3042 }
3043 }
3044}
3045
3046#[cfg(target_os = "macos")]
3049#[allow(clippy::too_many_arguments)]
3050pub(crate) fn sgemm_rm(
3051 m: usize,
3052 n: usize,
3053 k: usize,
3054 alpha: f32,
3055 a: &[f32],
3056 lda: usize,
3057 b_mat: &[f32],
3058 ldb: usize,
3059 b_rows_are_n: bool,
3060 c: &mut [f32],
3061 ldc: usize,
3062) {
3063 debug_assert!(a.len() >= (m - 1) * lda + k);
3064 debug_assert!(c.len() >= (m - 1) * ldc + n);
3065 #[cfg(target_arch = "aarch64")]
3070 if std::env::var("CMF_FORCE_NEON_GEMM")
3071 .map(|v| v == "1")
3072 .unwrap_or(false)
3073 {
3074 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3075 }
3076 unsafe {
3077 accel_blas::cblas_sgemm(
3078 101, 111, if b_rows_are_n { 112 } else { 111 },
3081 m as i32,
3082 n as i32,
3083 k as i32,
3084 alpha,
3085 a.as_ptr(),
3086 lda as i32,
3087 b_mat.as_ptr(),
3088 ldb as i32,
3089 0.0,
3090 c.as_mut_ptr(),
3091 ldc as i32,
3092 );
3093 }
3094}
3095
3096#[cfg(target_os = "macos")]
3103fn qmatmat_accel(
3104 q: &[u8],
3105 row_scale: &[f32],
3106 pre: &[std::borrow::Cow<'_, [f32]>],
3107 rows: usize,
3108 cols: usize,
3109 out: &mut [f32],
3110 pool: Option<&Pool>,
3111) {
3112 const TR: usize = 2048;
3117 let b = pre.len();
3118 thread_local! {
3119 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3120 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3121 }
3122 XPANEL.with(|xp| {
3123 WTILE.with(|wt| {
3124 let mut xpanel = xp.borrow_mut();
3125 xpanel.clear();
3126 for x in pre {
3127 xpanel.extend_from_slice(x);
3128 }
3129 let mut wtile = wt.borrow_mut();
3130 wtile.resize(TR * cols, 0.0);
3131 let mut r0 = 0usize;
3132 while r0 < rows {
3133 let tr = TR.min(rows - r0);
3134 let wt_addr = SendMut(wtile.as_mut_ptr());
3136 let run = |start: usize, end: usize| {
3137 for r in start..end {
3138 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3139 let s = row_scale[r0 + r];
3140 let dst =
3142 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3143 for (d, &v) in dst.iter_mut().zip(row) {
3144 *d = (v as i8) as f32 * s;
3145 }
3146 }
3147 };
3148 dispatch_rows(pool, tr, &run);
3149 unsafe {
3151 accel_blas::cblas_sgemm(
3152 101, 111, 112, b as i32,
3156 tr as i32,
3157 cols as i32,
3158 1.0,
3159 xpanel.as_ptr(),
3160 cols as i32,
3161 wtile.as_ptr(),
3162 cols as i32,
3163 0.0,
3164 out.as_mut_ptr().add(r0),
3165 rows as i32,
3166 );
3167 }
3168 r0 += tr;
3169 }
3170 })
3171 });
3172}
3173
3174fn qmatmat(
3175 q: &[u8],
3176 row_scale: &[f32],
3177 pre: &[std::borrow::Cow<'_, [f32]>],
3178 rows: usize,
3179 cols: usize,
3180 out: &mut [f32],
3181 pool: Option<&Pool>,
3182) {
3183 let b = pre.len();
3184 debug_assert_eq!(out.len(), b * rows);
3185 #[cfg(target_os = "macos")]
3190 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3191 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3192 return;
3193 }
3194 #[cfg(target_arch = "aarch64")]
3195 if sdot_enabled() {
3196 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3197 let out_addr = SendMut(out.as_mut_ptr());
3198 let blocked_ok = blocked_enabled();
3201 let use_i8mm = i8mm_enabled();
3202 if blocked_ok {
3203 let run = |start: usize, end: usize| {
3204 let mut o = start;
3205 while o < end {
3206 if o + 2 <= end {
3207 let r0 = &q[o * cols..(o + 1) * cols];
3208 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3209 let mut bi = 0usize;
3210 while bi + 4 <= acts.len() {
3211 let xs = [
3212 acts[bi].xq.as_slice(),
3213 acts[bi + 1].xq.as_slice(),
3214 acts[bi + 2].xq.as_slice(),
3215 acts[bi + 3].xq.as_slice(),
3216 ];
3217 let d = if use_i8mm {
3218 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3219 } else {
3220 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3221 };
3222 for (r, row) in [r0, r1].into_iter().enumerate() {
3223 for k in 0..4 {
3224 let act = &acts[bi + k];
3225 let mut v = d[r][k] as f32 * act.sx;
3226 for &(j, xv) in &act.outliers {
3227 v += (row[j] as i8) as f32 * xv;
3228 }
3229 unsafe {
3230 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3231 };
3232 }
3233 }
3234 bi += 4;
3235 }
3236 while bi < acts.len() {
3237 for (r, row) in [r0, r1].into_iter().enumerate() {
3238 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3239 unsafe { *out_addr.at(bi * rows + o + r) = v };
3240 }
3241 bi += 1;
3242 }
3243 o += 2;
3244 } else {
3245 let row = &q[o * cols..(o + 1) * cols];
3246 for (bi, act) in acts.iter().enumerate() {
3247 let v = row_dot_sdot(row, act) * row_scale[o];
3248 unsafe { *out_addr.at(bi * rows + o) = v };
3249 }
3250 o += 1;
3251 }
3252 }
3253 };
3254 dispatch_rows(pool, rows, &run);
3255 return;
3256 }
3257 let run = |start: usize, end: usize| {
3258 for o in start..end {
3259 let row = &q[o * cols..(o + 1) * cols];
3260 for (bi, act) in acts.iter().enumerate() {
3261 let v = row_dot_sdot(row, act) * row_scale[o];
3262 unsafe { *out_addr.at(bi * rows + o) = v };
3263 }
3264 }
3265 };
3266 dispatch_rows(pool, rows, &run);
3267 return;
3268 }
3269 #[cfg(target_arch = "x86_64")]
3274 if avx2_a8w8_enabled() {
3275 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3276 let out_addr = SendMut(out.as_mut_ptr());
3277 let blocked_ok = blocked_enabled();
3280 if !avx512vnni_enabled() && blocked_ok {
3281 let run = |start: usize, end: usize| {
3282 let mut o = start;
3283 while o < end {
3284 if o + 2 <= end {
3285 let r0 = &q[o * cols..(o + 1) * cols];
3286 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3287 let mut bi = 0usize;
3288 while bi + 4 <= acts.len() {
3289 let xs = [
3290 acts[bi].xq.as_slice(),
3291 acts[bi + 1].xq.as_slice(),
3292 acts[bi + 2].xq.as_slice(),
3293 acts[bi + 3].xq.as_slice(),
3294 ];
3295 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3296 for (r, row) in [r0, r1].into_iter().enumerate() {
3297 for k in 0..4 {
3298 let act = &acts[bi + k];
3299 let mut v = d[r][k] as f32 * act.sx;
3300 for &(j, xv) in &act.outliers {
3301 v += (row[j] as i8) as f32 * xv;
3302 }
3303 unsafe {
3304 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3305 };
3306 }
3307 }
3308 bi += 4;
3309 }
3310 while bi < acts.len() {
3311 for (r, row) in [r0, r1].into_iter().enumerate() {
3312 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3313 unsafe { *out_addr.at(bi * rows + o + r) = v };
3314 }
3315 bi += 1;
3316 }
3317 o += 2;
3318 } else {
3319 let row = &q[o * cols..(o + 1) * cols];
3320 for (bi, act) in acts.iter().enumerate() {
3321 let v = row_dot_avx2(row, act) * row_scale[o];
3322 unsafe { *out_addr.at(bi * rows + o) = v };
3323 }
3324 o += 1;
3325 }
3326 }
3327 };
3328 dispatch_rows(pool, rows, &run);
3329 return;
3330 }
3331 let run = |start: usize, end: usize| {
3332 for o in start..end {
3333 let row = &q[o * cols..(o + 1) * cols];
3334 for (bi, act) in acts.iter().enumerate() {
3335 let v = row_dot_avx2(row, act) * row_scale[o];
3336 unsafe { *out_addr.at(bi * rows + o) = v };
3337 }
3338 }
3339 };
3340 dispatch_rows(pool, rows, &run);
3341 return;
3342 }
3343 let out_addr = SendMut(out.as_mut_ptr());
3344 let run = |start: usize, end: usize| {
3345 for o in start..end {
3346 let row = &q[o * cols..(o + 1) * cols];
3347 for (bi, x) in pre.iter().enumerate() {
3348 let mut acc = 0f32;
3349 for j in 0..cols {
3350 acc += (row[j] as i8) as f32 * x[j];
3351 }
3352 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3353 }
3354 }
3355 };
3356 dispatch_rows(pool, rows, &run);
3357}
3358
3359fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3362 match pool {
3363 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3364 _ => run(0, rows),
3365 }
3366}
3367
3368fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3370 let groups = rows * cols / GROUP_SIZE;
3371 bytes.split_at(groups * 16)
3372}
3373
3374#[inline]
3379fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3380 #[cfg(target_arch = "aarch64")]
3381 unsafe {
3382 return vbit_fill4_neon(data, buf);
3383 }
3384 #[cfg(target_arch = "x86_64")]
3385 if avx2_enabled() {
3386 return unsafe { vbit_fill4_avx2(data, buf) };
3387 }
3388 #[allow(unreachable_code)]
3389 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3390 let u = unpack8::<4>(&data[blk * 4..]);
3391 for k in 0..8 {
3392 chunk[k] = (u[k] - 7) as i8 as u8;
3393 }
3394 }
3395}
3396
3397#[cfg(target_arch = "aarch64")]
3398#[target_feature(enable = "neon")]
3399unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3400 unsafe {
3403 use core::arch::aarch64::*;
3404 let n = buf.len();
3405 let mask = vdupq_n_u8(0x0F);
3406 let seven = vdupq_n_s8(7);
3407 let mut g = 0usize;
3408 while g * 32 + 32 <= n {
3409 let b = vld1q_u8(data.as_ptr().add(g * 16));
3410 let hi = vshrq_n_u8::<4>(b);
3411 let lo = vandq_u8(b, mask);
3412 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3413 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3414 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3415 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3416 g += 1;
3417 }
3418 }
3419}
3420
3421#[cfg(target_arch = "x86_64")]
3422#[target_feature(enable = "avx2")]
3423unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3424 unsafe {
3426 use core::arch::x86_64::*;
3427 let n = buf.len();
3428 let mask = _mm_set1_epi8(0x0F);
3429 let seven = _mm256_set1_epi8(7);
3430 let mut g = 0usize;
3431 while g * 32 + 32 <= n {
3432 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3433 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3434 let lo = _mm_and_si128(b, mask);
3435 let z = _mm256_sub_epi8(
3436 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3437 seven,
3438 );
3439 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3440 g += 1;
3441 }
3442 }
3443}
3444
3445#[inline(always)]
3450fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3451 let mut acc = 0u64;
3452 for i in 0..B {
3453 acc = (acc << 8) | data[i] as u64;
3454 }
3455 let mask = (1u64 << B) - 1;
3456 let mut out = [0i32; 8];
3457 for (k, o) in out.iter_mut().enumerate() {
3458 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3459 }
3460 out
3461}
3462
3463#[allow(clippy::too_many_arguments)]
3469fn vbitmatvec(
3470 bytes: &[u8],
3471 offsets: &[usize],
3472 x: &[f32],
3473 rows: usize,
3474 cols: usize,
3475 out: &mut [f32],
3476 pool: Option<&Pool>,
3477) {
3478 debug_assert_eq!(out.len(), rows);
3479 debug_assert_eq!(offsets.len(), rows + 1);
3480
3481 if a8w8_enabled() {
3485 let act = split_act(x);
3486 let out_addr = SendMut(out.as_mut_ptr());
3487 let run = move |start: usize, end: usize| {
3488 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3489 };
3490 dispatch_rows(pool, rows, &run);
3491 return;
3492 }
3493
3494 let out_addr = SendMut(out.as_mut_ptr());
3495 let run = move |start: usize, end: usize| {
3496 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3497 };
3498 dispatch_rows(pool, rows, &run);
3499}
3500
3501#[allow(clippy::too_many_arguments)]
3505fn vbit_range_a8w8(
3506 bytes: &[u8],
3507 offsets: &[usize],
3508 x: &[f32],
3509 act: &SplitAct,
3510 rows: usize,
3511 cols: usize,
3512 out: SendMut,
3513 start: usize,
3514 end: usize,
3515) {
3516 let ng = cols / GROUP_SIZE;
3517 let bits = &bytes[..rows];
3518 let sc_off = rows;
3519 let row_dot = |r: usize| -> f32 {
3520 let b = bits[r] as usize;
3521 let l = (1i32 << (b - 1)) - 1;
3522 let mask = (1u64 << b) - 1;
3523 let data = &bytes[offsets[r]..offsets[r + 1]];
3524 if b == 8 {
3525 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3527 let mut dot = 0f32;
3528 for g in 0..ng {
3529 let so = (r * ng + g) * 2;
3530 let sgf = f16_to_f32(u16::from_le_bytes([
3531 bytes[sc_off + so],
3532 bytes[sc_off + so + 1],
3533 ]));
3534 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3535 let mut gd = 0f32;
3536 for &xv in xg.iter() {
3537 if nbits < 8 {
3538 acc = (acc << 8) | data[idx] as u64;
3539 idx += 1;
3540 nbits += 8;
3541 }
3542 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3543 nbits -= 8;
3544 gd += (u - l) as f32 * xv;
3545 }
3546 dot += gd * sgf;
3547 }
3548 return dot;
3549 }
3550 thread_local! {
3554 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3555 const { std::cell::RefCell::new(Vec::new()) };
3556 }
3557 #[inline(always)]
3558 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3559 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3560 let u = unpack8::<B>(&data[blk * B..]);
3561 for k in 0..8 {
3562 chunk[k] = (u[k] - l) as i8 as u8;
3563 }
3564 }
3565 }
3566 let _ = mask;
3567 VBIT_SCRATCH.with(|scratch| {
3568 let mut buf = scratch.borrow_mut();
3569 buf.resize(cols, 0);
3570 match b {
3571 3 => fill::<3>(data, l, &mut buf),
3572 4 => vbit_fill4(data, &mut buf),
3573 5 => fill::<5>(data, l, &mut buf),
3574 6 => fill::<6>(data, l, &mut buf),
3575 _ => unreachable!(),
3576 }
3577 let mut dot = 0f32;
3578 for g in 0..ng {
3579 let so = (r * ng + g) * 2;
3580 let s = f16_to_f32(u16::from_le_bytes([
3581 bytes[sc_off + so],
3582 bytes[sc_off + so + 1],
3583 ]));
3584 let d = dot_i8_i8(
3585 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3586 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3587 ) as f32
3588 * act.sx;
3589 dot += d * s;
3590 }
3591 for &(j, xv) in &act.outliers {
3592 let so = (r * ng + j / GROUP_SIZE) * 2;
3593 let s = f16_to_f32(u16::from_le_bytes([
3594 bytes[sc_off + so],
3595 bytes[sc_off + so + 1],
3596 ]));
3597 dot += (buf[j] as i8) as f32 * s * xv;
3599 }
3600 dot
3601 })
3602 };
3603 for r in start..end {
3604 unsafe { *out.at(r) = row_dot(r) };
3606 }
3607}
3608
3609#[allow(clippy::too_many_arguments)]
3611fn vbit_range_f32(
3612 bytes: &[u8],
3613 offsets: &[usize],
3614 x: &[f32],
3615 rows: usize,
3616 cols: usize,
3617 out: SendMut,
3618 start: usize,
3619 end: usize,
3620) {
3621 let ng = cols / GROUP_SIZE;
3622 let bits = &bytes[..rows];
3623 let sc_off = rows;
3624 #[inline(always)]
3628 fn dot_row<const B: usize>(
3629 data: &[u8],
3630 bytes: &[u8],
3631 sc_off: usize,
3632 r: usize,
3633 ng: usize,
3634 x: &[f32],
3635 ) -> f32 {
3636 let l = ((1i32 << (B - 1)) - 1) as f32;
3637 let gbytes = GROUP_SIZE * B / 8;
3638 let mut dot = 0f32;
3639 for g in 0..ng {
3640 let so = (r * ng + g) * 2;
3641 let s = f16_to_f32(u16::from_le_bytes([
3642 bytes[sc_off + so],
3643 bytes[sc_off + so + 1],
3644 ]));
3645 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3646 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3647 let mut gd = 0f32;
3648 for blk in 0..GROUP_SIZE / 8 {
3649 let u = unpack8::<B>(&gd0[blk * B..]);
3650 let xb = &xg[blk * 8..blk * 8 + 8];
3651 for k in 0..8 {
3652 gd += (u[k] as f32 - l) * xb[k];
3653 }
3654 }
3655 dot += gd * s;
3656 }
3657 dot
3658 }
3659 for r in start..end {
3660 let data = &bytes[offsets[r]..offsets[r + 1]];
3661 let v = match bits[r] {
3662 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3663 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3664 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3665 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3666 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3667 b => unreachable!("vbit bit-width {b} (validated at load)"),
3668 };
3669 unsafe { *out.at(r) = v };
3671 }
3672}
3673
3674#[allow(clippy::too_many_arguments)]
3679fn vbitmatvec2(
3680 bytes: &[u8],
3681 offsets: &[usize],
3682 x1: &[f32],
3683 x2: &[f32],
3684 rows: usize,
3685 cols: usize,
3686 o1: &mut [f32],
3687 o2: &mut [f32],
3688 pool: Option<&Pool>,
3689) {
3690 debug_assert_eq!(o1.len(), rows);
3691 debug_assert_eq!(o2.len(), rows);
3692
3693 if a8w8_enabled() {
3694 let a1 = split_act(x1);
3695 let a2 = split_act(x2);
3696 let p1 = SendMut(o1.as_mut_ptr());
3697 let p2 = SendMut(o2.as_mut_ptr());
3698 let run = move |start: usize, end: usize| {
3699 vbit_range2_a8w8(
3700 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3701 )
3702 };
3703 dispatch_rows(pool, rows, &run);
3704 return;
3705 }
3706
3707 let p1 = SendMut(o1.as_mut_ptr());
3708 let p2 = SendMut(o2.as_mut_ptr());
3709 let run = move |start: usize, end: usize| {
3710 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3711 };
3712 dispatch_rows(pool, rows, &run);
3713}
3714
3715#[allow(clippy::too_many_arguments)]
3719fn vbit_range2_a8w8(
3720 bytes: &[u8],
3721 offsets: &[usize],
3722 x1: &[f32],
3723 x2: &[f32],
3724 a1: &SplitAct,
3725 a2: &SplitAct,
3726 rows: usize,
3727 cols: usize,
3728 p1: SendMut,
3729 p2: SendMut,
3730 start: usize,
3731 end: usize,
3732) {
3733 let ng = cols / GROUP_SIZE;
3734 let bits = &bytes[..rows];
3735 let sc_off = rows;
3736 let row_dots = |r: usize| -> (f32, f32) {
3737 let b = bits[r] as usize;
3738 let l = (1i32 << (b - 1)) - 1;
3739 let data = &bytes[offsets[r]..offsets[r + 1]];
3740 if b == 8 {
3741 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3744 let (mut d1, mut d2) = (0f32, 0f32);
3745 for g in 0..ng {
3746 let so = (r * ng + g) * 2;
3747 let sgf = f16_to_f32(u16::from_le_bytes([
3748 bytes[sc_off + so],
3749 bytes[sc_off + so + 1],
3750 ]));
3751 let (mut g1, mut g2) = (0f32, 0f32);
3752 for k in 0..GROUP_SIZE {
3753 if nbits < 8 {
3754 acc = (acc << 8) | data[idx] as u64;
3755 idx += 1;
3756 nbits += 8;
3757 }
3758 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3759 nbits -= 8;
3760 let w = (u - l) as f32;
3761 g1 += w * x1[g * GROUP_SIZE + k];
3762 g2 += w * x2[g * GROUP_SIZE + k];
3763 }
3764 d1 += g1 * sgf;
3765 d2 += g2 * sgf;
3766 }
3767 return (d1, d2);
3768 }
3769 thread_local! {
3770 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3771 const { std::cell::RefCell::new(Vec::new()) };
3772 }
3773 #[inline(always)]
3774 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3775 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3776 let u = unpack8::<B>(&data[blk * B..]);
3777 for k in 0..8 {
3778 chunk[k] = (u[k] - l) as i8 as u8;
3779 }
3780 }
3781 }
3782 VBIT_SCRATCH2.with(|scratch| {
3783 let mut buf = scratch.borrow_mut();
3784 buf.resize(cols, 0);
3785 match b {
3786 3 => fill::<3>(data, l, &mut buf),
3787 4 => vbit_fill4(data, &mut buf),
3788 5 => fill::<5>(data, l, &mut buf),
3789 6 => fill::<6>(data, l, &mut buf),
3790 _ => unreachable!(),
3791 }
3792 let (mut d1, mut d2) = (0f32, 0f32);
3793 for g in 0..ng {
3794 let so = (r * ng + g) * 2;
3795 let s = f16_to_f32(u16::from_le_bytes([
3796 bytes[sc_off + so],
3797 bytes[sc_off + so + 1],
3798 ]));
3799 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3800 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3801 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3802 d1 += v1 * s;
3803 d2 += v2 * s;
3804 }
3805 for &(j, xv) in &a1.outliers {
3806 let so = (r * ng + j / GROUP_SIZE) * 2;
3807 let s = f16_to_f32(u16::from_le_bytes([
3808 bytes[sc_off + so],
3809 bytes[sc_off + so + 1],
3810 ]));
3811 d1 += (buf[j] as i8) as f32 * s * xv;
3812 }
3813 for &(j, xv) in &a2.outliers {
3814 let so = (r * ng + j / GROUP_SIZE) * 2;
3815 let s = f16_to_f32(u16::from_le_bytes([
3816 bytes[sc_off + so],
3817 bytes[sc_off + so + 1],
3818 ]));
3819 d2 += (buf[j] as i8) as f32 * s * xv;
3820 }
3821 (d1, d2)
3822 })
3823 };
3824 for r in start..end {
3825 let (v1, v2) = row_dots(r);
3826 unsafe {
3828 *p1.at(r) = v1;
3829 *p2.at(r) = v2;
3830 }
3831 }
3832}
3833
3834#[allow(clippy::too_many_arguments)]
3838fn vbit_range2_f32(
3839 bytes: &[u8],
3840 offsets: &[usize],
3841 x1: &[f32],
3842 x2: &[f32],
3843 rows: usize,
3844 cols: usize,
3845 p1: SendMut,
3846 p2: SendMut,
3847 start: usize,
3848 end: usize,
3849) {
3850 let ng = cols / GROUP_SIZE;
3851 let bits = &bytes[..rows];
3852 let sc_off = rows;
3853 #[inline(always)]
3854 #[allow(clippy::too_many_arguments)]
3855 fn dot_row2<const B: usize>(
3856 data: &[u8],
3857 bytes: &[u8],
3858 sc_off: usize,
3859 r: usize,
3860 ng: usize,
3861 x1: &[f32],
3862 x2: &[f32],
3863 ) -> (f32, f32) {
3864 let l = ((1i32 << (B - 1)) - 1) as f32;
3865 let gbytes = GROUP_SIZE * B / 8;
3866 let (mut d1, mut d2) = (0f32, 0f32);
3867 for g in 0..ng {
3868 let so = (r * ng + g) * 2;
3869 let s = f16_to_f32(u16::from_le_bytes([
3870 bytes[sc_off + so],
3871 bytes[sc_off + so + 1],
3872 ]));
3873 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3874 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3875 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3876 let (mut g1, mut g2) = (0f32, 0f32);
3877 for blk in 0..GROUP_SIZE / 8 {
3878 let u = unpack8::<B>(&gd0[blk * B..]);
3879 for k in 0..8 {
3880 let w = u[k] as f32 - l;
3881 g1 += w * x1g[blk * 8 + k];
3882 g2 += w * x2g[blk * 8 + k];
3883 }
3884 }
3885 d1 += g1 * s;
3886 d2 += g2 * s;
3887 }
3888 (d1, d2)
3889 }
3890 for r in start..end {
3891 let data = &bytes[offsets[r]..offsets[r + 1]];
3892 let (v1, v2) = match bits[r] {
3893 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3894 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3895 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3896 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3897 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3898 b => unreachable!("vbit bit-width {b} (validated at load)"),
3899 };
3900 unsafe {
3902 *p1.at(r) = v1;
3903 *p2.at(r) = v2;
3904 }
3905 }
3906}
3907
3908#[inline]
3915#[allow(unreachable_code)]
3916fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3917 #[cfg(target_arch = "aarch64")]
3918 unsafe {
3919 return dot_q4t_row_sdot(bytes, r, gpr, xq);
3920 }
3921 #[cfg(target_arch = "x86_64")]
3922 unsafe {
3923 if vnni_tiles_enabled() {
3924 return dot_q4t_row_vnni(bytes, r, gpr, xq);
3925 }
3926 return dot_q4t_row_avx2(bytes, r, gpr, xq);
3927 }
3928 let mut acc = 0f32;
3929 for gi in 0..gpr {
3930 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3931 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3932 let mut d = 0i32;
3933 for (k, &b) in tile[2..].iter().enumerate() {
3934 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3935 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3936 }
3937 acc += d as f32 * s;
3938 }
3939 acc
3940}
3941
3942#[cfg(target_arch = "aarch64")]
3943#[target_feature(enable = "neon,dotprod")]
3944unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3945 unsafe {
3948 use core::arch::aarch64::*;
3949 use core::arch::asm;
3950 let lomask = vdupq_n_u8(0x0F);
3951 let eight = vdupq_n_s8(8);
3952 let mut acc = 0f32;
3953 for gi in 0..gpr {
3954 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3955 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3956 let b = vld1q_u8(t.add(2));
3957 let lo = vandq_u8(b, lomask);
3958 let hi = vshrq_n_u8::<4>(b);
3959 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3960 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3961 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3962 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3963 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3964 asm!(
3965 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3966 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3967 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3968 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3969 options(pure, nomem, nostack),
3970 );
3971 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3972 }
3973 acc
3974 }
3975}
3976
3977#[cfg(target_arch = "x86_64")]
3978#[target_feature(enable = "avx2")]
3979unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3980 unsafe {
3982 use core::arch::x86_64::*;
3983 let lomask = _mm_set1_epi8(0x0F);
3984 let eight = _mm256_set1_epi8(8);
3985 let ones = _mm256_set1_epi16(1);
3986 let mut acc = 0f32;
3987 for gi in 0..gpr {
3988 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3989 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3990 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3991 let lo = _mm_and_si128(b, lomask);
3992 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3993 let w = _mm256_sub_epi8(
3994 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3995 eight,
3996 );
3997 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3998 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3999 let d = _mm256_madd_epi16(p16, ones);
4000 let hi128 = _mm256_extracti128_si256::<1>(d);
4001 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4002 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4003 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4004 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4005 }
4006 acc
4007 }
4008}
4009
4010#[cfg(target_arch = "x86_64")]
4014#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4015unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4016 unsafe {
4018 use core::arch::x86_64::*;
4019 let lomask = _mm_set1_epi8(0x0F);
4020 let eight = _mm256_set1_epi8(8);
4021 let mut acc = 0f32;
4022 for gi in 0..gpr {
4023 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4024 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4025 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4026 let lo = _mm_and_si128(b, lomask);
4027 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4028 let w = _mm256_sub_epi8(
4029 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4030 eight,
4031 );
4032 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4033 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4034 acc += d as f32 * s;
4035 }
4036 acc
4037 }
4038}
4039
4040#[cfg(target_arch = "x86_64")]
4045#[target_feature(enable = "avx2,fma")]
4050unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4051 unsafe {
4053 use core::arch::x86_64::*;
4054 let lomask = _mm_set1_epi8(0x0F);
4055 let eight = _mm256_set1_epi8(8);
4056 let ones = _mm256_set1_epi16(1);
4057 let mut f0 = _mm256_setzero_ps();
4070 let mut f1 = _mm256_setzero_ps();
4071 let mut f2 = _mm256_setzero_ps();
4072 let mut f3 = _mm256_setzero_ps();
4073 for gi in 0..gpr {
4074 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4075 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4076 let sv = _mm256_set1_ps(s);
4077 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4078 let lo = _mm_and_si128(bb, lomask);
4079 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4080 let w = _mm256_sub_epi8(
4081 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4082 eight,
4083 );
4084 let aw = _mm256_abs_epi8(w);
4085 let off = gi * GROUP_SIZE;
4086 let dot = |xq: &[i8]| {
4087 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4088 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4089 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4090 };
4091 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4092 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4093 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4094 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4095 }
4096 [
4097 hsum256_ps(f0),
4098 hsum256_ps(f1),
4099 hsum256_ps(f2),
4100 hsum256_ps(f3),
4101 ]
4102 }
4103}
4104
4105#[cfg(target_arch = "x86_64")]
4108#[target_feature(enable = "avx2")]
4109#[inline]
4110unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4111 unsafe {
4113 use core::arch::x86_64::*;
4114 let hi = _mm256_extractf128_ps::<1>(v);
4115 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4116 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4117 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4118 _mm_cvtss_f32(s)
4119 }
4120}
4121
4122#[cfg(target_arch = "x86_64")]
4124#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4125unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4126 unsafe {
4128 use core::arch::x86_64::*;
4129 let lomask = _mm_set1_epi8(0x0F);
4130 let eight = _mm256_set1_epi8(8);
4131 let mut f0 = _mm256_setzero_ps();
4134 let mut f1 = _mm256_setzero_ps();
4135 let mut f2 = _mm256_setzero_ps();
4136 let mut f3 = _mm256_setzero_ps();
4137 for gi in 0..gpr {
4138 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4139 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4140 let sv = _mm256_set1_ps(s);
4141 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4142 let lo = _mm_and_si128(bb, lomask);
4143 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4144 let w = _mm256_sub_epi8(
4145 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4146 eight,
4147 );
4148 let aw = _mm256_abs_epi8(w);
4149 let off = gi * GROUP_SIZE;
4150 let dot = |xq: &[i8]| {
4151 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4152 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4153 _mm256_setzero_si256(),
4154 aw,
4155 _mm256_sign_epi8(x, w),
4156 ))
4157 };
4158 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4159 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4160 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4161 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4162 }
4163 let acc = [
4164 hsum256_ps(f0),
4165 hsum256_ps(f1),
4166 hsum256_ps(f2),
4167 hsum256_ps(f3),
4168 ];
4169 acc
4170 }
4171}
4172
4173#[cfg(target_arch = "aarch64")]
4178#[target_feature(enable = "neon,dotprod")]
4179unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4180 unsafe {
4182 use core::arch::aarch64::*;
4183 use core::arch::asm;
4184 let lomask = vdupq_n_u8(0x0F);
4185 let eight = vdupq_n_s8(8);
4186 let mut acc = [0f32; 4];
4187 for gi in 0..gpr {
4188 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4189 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4190 let b = vld1q_u8(t.add(2));
4191 let lo = vandq_u8(b, lomask);
4192 let hi = vshrq_n_u8::<4>(b);
4193 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4194 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4195 for (k, xq) in xs.iter().enumerate() {
4196 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4197 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4198 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4199 asm!(
4200 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4201 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4202 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4203 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4204 options(pure, nomem, nostack),
4205 );
4206 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4207 }
4208 }
4209 acc
4210 }
4211}
4212
4213#[inline]
4215fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4216 let gi = j / GROUP_SIZE;
4217 let k = j % GROUP_SIZE;
4218 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4219 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4220 let byte = tile[2 + k / 2];
4221 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4222 ((nib as i32 - 8) as f32, s)
4223}
4224
4225#[inline]
4228fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4229 let mut acc = 0f32;
4230 for gi in 0..gpr {
4231 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4232 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4233 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4234 let mut ga = 0f32;
4235 for (k, &b) in tile[2..].iter().enumerate() {
4236 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4237 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4238 }
4239 acc += ga * s;
4240 }
4241 acc
4242}
4243
4244struct Q4tpView<'a> {
4248 nib: &'a [u8],
4249 params: &'a [u8],
4250 codes: &'a [u8],
4251 stride: usize,
4252 zero_rung: bool,
4254}
4255
4256impl<'a> Q4tpView<'a> {
4257 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4258 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4259 Self {
4260 nib: &bytes[..params_off],
4261 params: &bytes[params_off..codes_off],
4262 codes: &bytes[codes_off..],
4263 stride,
4264 zero_rung: false,
4265 }
4266 }
4267
4268 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4270 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4271 Self {
4272 nib: &bytes[..params_off],
4273 params: &bytes[params_off..codes_off],
4274 codes: &bytes[codes_off..],
4275 stride,
4276 zero_rung: true,
4277 }
4278 }
4279
4280 #[inline]
4296 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4297 let tab = if self.zero_rung {
4298 q2tp_ladder(self.params, r)
4299 } else {
4300 q4tp_ladder(self.params, r)
4301 };
4302 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4303 let out = &mut out[..gpr];
4304 let mut chunks = out.chunks_exact_mut(8);
4305 let mut ci = 0usize;
4306 for c in &mut chunks {
4307 let w = u64::from(codes[ci])
4308 | u64::from(codes[ci + 1]) << 8
4309 | u64::from(codes[ci + 2]) << 16
4310 | u64::from(codes[ci + 3]) << 24
4311 | u64::from(codes[ci + 4]) << 32;
4312 for (k, o) in c.iter_mut().enumerate() {
4313 *o = tab[((w >> (5 * k)) & 31) as usize];
4314 }
4315 ci += 5;
4316 }
4317 let tail = &codes[ci..];
4320 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4321 *o = tab[q4tp_code(tail, k)];
4322 }
4323 }
4324}
4325
4326#[inline]
4327fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4328 #[cfg(target_arch = "aarch64")]
4329 unsafe {
4330 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4331 }
4332 #[cfg(target_arch = "x86_64")]
4333 unsafe {
4334 if vnni_tiles_enabled() {
4335 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4336 }
4337 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4338 }
4339 #[allow(unreachable_code)]
4340 {
4341 let mut acc = 0f32;
4342 for gi in 0..gpr {
4343 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4344 let s = scales[gi];
4345 let mut d = 0i32;
4346 for (k, &b) in tile.iter().enumerate() {
4347 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4348 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4349 }
4350 acc += d as f32 * s;
4351 }
4352 acc
4353 }
4354}
4355
4356#[cfg(target_arch = "aarch64")]
4359#[target_feature(enable = "neon,dotprod")]
4360unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4361 unsafe {
4364 use core::arch::aarch64::*;
4365 use core::arch::asm;
4366 let lomask = vdupq_n_u8(0x0F);
4367 let eight = vdupq_n_s8(8);
4368 let mut acc = 0f32;
4369 for gi in 0..gpr {
4370 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4371 let s = *scales.get_unchecked(gi);
4372 let b = vld1q_u8(t);
4373 let lo = vandq_u8(b, lomask);
4374 let hi = vshrq_n_u8::<4>(b);
4375 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4376 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4377 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4378 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4379 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4380 asm!(
4381 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4382 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4383 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4384 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4385 options(pure, nomem, nostack),
4386 );
4387 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4388 }
4389 acc
4390 }
4391}
4392
4393#[cfg(target_arch = "x86_64")]
4394#[target_feature(enable = "avx2")]
4395unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4396 unsafe {
4398 use core::arch::x86_64::*;
4399 let lomask = _mm_set1_epi8(0x0F);
4400 let eight = _mm256_set1_epi8(8);
4401 let ones = _mm256_set1_epi16(1);
4402 let mut acc = 0f32;
4403 for gi in 0..gpr {
4404 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4405 let s = *scales.get_unchecked(gi);
4406 let b = _mm_loadu_si128(t as *const __m128i);
4407 let lo = _mm_and_si128(b, lomask);
4408 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4409 let w = _mm256_sub_epi8(
4410 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4411 eight,
4412 );
4413 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4414 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4415 let d = _mm256_madd_epi16(p16, ones);
4416 let hi128 = _mm256_extracti128_si256::<1>(d);
4417 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4418 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4419 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4420 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4421 }
4422 acc
4423 }
4424}
4425
4426
4427#[cfg(target_arch = "x86_64")]
4430#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4431unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4432 unsafe {
4434 use core::arch::x86_64::*;
4435 let lomask = _mm_set1_epi8(0x0F);
4436 let eight = _mm256_set1_epi8(8);
4437 let mut acc = 0f32;
4438 for gi in 0..gpr {
4439 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4440 let s = *scales.get_unchecked(gi);
4441 let b = _mm_loadu_si128(t as *const __m128i);
4442 let lo = _mm_and_si128(b, lomask);
4443 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4444 let w = _mm256_sub_epi8(
4445 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4446 eight,
4447 );
4448 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4449 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4450 }
4451 acc
4452 }
4453}
4454
4455#[inline]
4458fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4459 let mut acc = 0f32;
4460 for gi in 0..gpr {
4461 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4462 let s = scales[gi];
4463 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4464 let mut ga = 0f32;
4465 for (k, &b) in tile.iter().enumerate() {
4466 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4467 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4468 }
4469 acc += ga * s;
4470 }
4471 acc
4472}
4473
4474#[inline]
4477fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4478 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4479 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4480 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4481 ((n as i32 - 8) as f32, scales[gi])
4482}
4483
4484fn q4tp_matvec(
4486 bytes: &[u8],
4487 x: &[f32],
4488 rows: usize,
4489 cols: usize,
4490 out: &mut [f32],
4491 pool: Option<&Pool>,
4492) {
4493 debug_assert_eq!(out.len(), rows);
4494 let gpr = cols / GROUP_SIZE;
4495 let v = Q4tpView::new(bytes, rows, cols);
4496 let out_addr = SendMut(out.as_mut_ptr());
4497 if a8w8_enabled() {
4498 let act = split_act(x);
4499 let run = |start: usize, end: usize| {
4500 with_krow(gpr, |sc| {
4502 for r in start..end {
4503 v.scales_into(r, gpr, sc);
4504 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4505 for &(j, xv) in &act.outliers {
4506 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4507 acc += w * s * xv;
4508 }
4509 unsafe { *out_addr.at(r) = acc };
4511 }
4512 })
4513 };
4514 dispatch_rows(pool, rows, &run);
4515 return;
4516 }
4517 let run = |start: usize, end: usize| {
4518 with_krow(gpr, |sc| {
4519 for r in start..end {
4520 v.scales_into(r, gpr, sc);
4521 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4523 }
4524 })
4525 };
4526 dispatch_rows(pool, rows, &run);
4527}
4528
4529#[allow(clippy::too_many_arguments)]
4532fn q4tp_matvec2(
4533 bytes: &[u8],
4534 x1: &[f32],
4535 x2: &[f32],
4536 rows: usize,
4537 cols: usize,
4538 o1: &mut [f32],
4539 o2: &mut [f32],
4540 pool: Option<&Pool>,
4541) {
4542 let gpr = cols / GROUP_SIZE;
4543 let v = Q4tpView::new(bytes, rows, cols);
4544 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4545 let run = |start: usize, end: usize| {
4546 let mut sc = vec![0f32; gpr];
4547 for r in start..end {
4548 v.scales_into(r, gpr, &mut sc);
4549 unsafe {
4551 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4552 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4553 }
4554 }
4555 };
4556 dispatch_rows(pool, rows, &run);
4557}
4558
4559#[inline]
4562fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4563 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4564 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4565 let c = (byte >> (2 * (k % 4))) & 3;
4566 (c as f32 - 1.5, scales[gi])
4567}
4568
4569#[inline]
4576fn dot_q2tp_row_i8(
4577 chunks: &[u8],
4578 r: usize,
4579 gpr: usize,
4580 xq: &[i8],
4581 gsum: &[i32],
4582 scales: &[f32],
4583) -> f32 {
4584 let mut acc = 0f32;
4585 let base = r * gpr * Q2TP_CHUNK;
4586 #[cfg(not(target_arch = "aarch64"))]
4587 let mut codes = [0i8; GROUP_SIZE];
4588 for gi in 0..gpr {
4589 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4590 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4591 #[cfg(target_arch = "aarch64")]
4592 let dot = unsafe {
4598 use core::arch::aarch64::*;
4599 let b = vld1_u8(ch.as_ptr());
4600 let three = vdup_n_u8(3);
4601 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4602 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4603 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4604 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4605 let x4 = vld4_s8(xg.as_ptr());
4606 let mut acc4 = vdupq_n_s32(0);
4607 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4608 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4609 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4610 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4611 vaddvq_s32(acc4)
4612 };
4613 #[cfg(not(target_arch = "aarch64"))]
4614 let dot: i32 = {
4615 for (k, &b) in ch.iter().enumerate() {
4616 codes[k * 4] = (b & 3) as i8;
4617 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4618 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4619 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4620 }
4621 codes
4622 .iter()
4623 .zip(xg)
4624 .map(|(&c, &x)| c as i32 * x as i32)
4625 .sum()
4626 };
4627 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4628 }
4629 acc
4630}
4631
4632fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4636 let mut acc = 0f32;
4637 for gi in 0..gpr {
4638 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4639 let s = scales[gi];
4640 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4641 let mut g = 0f32;
4642 for (k, &b) in ch.iter().enumerate() {
4643 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4644 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4645 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4646 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4647 }
4648 acc += s * g;
4649 }
4650 acc
4651}
4652
4653fn q2tp_matvec(
4654 bytes: &[u8],
4655 x: &[f32],
4656 rows: usize,
4657 cols: usize,
4658 out: &mut [f32],
4659 pool: Option<&Pool>,
4660) {
4661 debug_assert_eq!(out.len(), rows);
4662 let gpr = cols / GROUP_SIZE;
4663 let v = Q4tpView::new_q2(bytes, rows, cols);
4664 let out_addr = SendMut(out.as_mut_ptr());
4665 if a8w8_enabled() {
4670 let act = split_act(x);
4671 let gsum = q1_group_sums(&act.xq, gpr);
4672 let (act, gsum) = (&act, &gsum);
4673 let run = move |start: usize, end: usize| {
4674 with_krow(gpr, |sc| {
4675 for r in start..end {
4676 v.scales_into(r, gpr, sc);
4677 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4678 for &(j, xv) in &act.outliers {
4679 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4680 acc += w * s * xv;
4681 }
4682 unsafe { *out_addr.at(r) = acc };
4684 }
4685 })
4686 };
4687 dispatch_rows(pool, rows, &run);
4688 return;
4689 }
4690 let run = |start: usize, end: usize| {
4691 with_krow(gpr, |sc| {
4692 for r in start..end {
4693 v.scales_into(r, gpr, sc);
4694 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4696 }
4697 })
4698 };
4699 dispatch_rows(pool, rows, &run);
4700}
4701
4702#[allow(clippy::too_many_arguments)]
4704fn q2tp_matvec2(
4705 bytes: &[u8],
4706 x1: &[f32],
4707 x2: &[f32],
4708 rows: usize,
4709 cols: usize,
4710 o1: &mut [f32],
4711 o2: &mut [f32],
4712 pool: Option<&Pool>,
4713) {
4714 let gpr = cols / GROUP_SIZE;
4715 let v = Q4tpView::new_q2(bytes, rows, cols);
4716 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4717 let run = |start: usize, end: usize| {
4718 let mut sc = vec![0f32; gpr];
4719 for r in start..end {
4720 v.scales_into(r, gpr, &mut sc);
4721 unsafe {
4723 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4724 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4725 }
4726 }
4727 };
4728 dispatch_rows(pool, rows, &run);
4729}
4730
4731pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4738 let gpr = cols / GROUP_SIZE;
4743 let v = Q4tpView::new_q2(bytes, rows, cols);
4744 with_krow(gpr, |sc| {
4745 for r in 0..rows {
4746 v.scales_into(r, gpr, sc);
4747 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4748 }
4749 });
4750}
4751
4752pub fn q2tp_matmat_for_test(
4753 bytes: &[u8],
4754 xs_all: &[f32],
4755 b: usize,
4756 rows: usize,
4757 cols: usize,
4758 out: &mut [f32],
4759) {
4760 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4761}
4762
4763fn q2tp_matmat(
4764 bytes: &[u8],
4765 xs_all: &[f32],
4766 b: usize,
4767 rows: usize,
4768 cols: usize,
4769 out: &mut [f32],
4770 pool: Option<&Pool>,
4771) {
4772 debug_assert_eq!(out.len(), b * rows);
4773 let gpr = cols / GROUP_SIZE;
4774 let v = Q4tpView::new_q2(bytes, rows, cols);
4775 let out_addr = SendMut(out.as_mut_ptr());
4776 let run = |start: usize, end: usize| {
4777 let mut sc = vec![0f32; gpr];
4778 for r in start..end {
4779 v.scales_into(r, gpr, &mut sc);
4780 for bi in 0..b {
4781 let x = &xs_all[bi * cols..(bi + 1) * cols];
4782 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4784 }
4785 }
4786 };
4787 dispatch_rows(pool, rows, &run);
4788}
4789
4790#[cfg(target_arch = "aarch64")]
4794#[target_feature(enable = "neon,dotprod")]
4795unsafe fn dot_q4tp_row_1x4_sdot_v1(
4796 nib: &[u8],
4797 r: usize,
4798 gpr: usize,
4799 xs: [&[i8]; 4],
4800 scales: &[f32],
4801) -> [f32; 4] {
4802 unsafe {
4803 use core::arch::aarch64::*;
4804 use core::arch::asm;
4805 let lomask = vdupq_n_u8(0x0F);
4806 let eight = vdupq_n_s8(8);
4807 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4808 for gi in 0..gpr {
4809 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4810 let s = *scales.get_unchecked(gi);
4811 let bb = vld1q_u8(t);
4812 let lo = vandq_u8(bb, lomask);
4813 let hi = vshrq_n_u8::<4>(bb);
4814 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4815 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4816 let mut d = [0f32; 4];
4817 for (k, dk) in d.iter_mut().enumerate() {
4818 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4819 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4820 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4821 asm!(
4822 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4823 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4824 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4825 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4826 options(pure, nomem, nostack),
4827 );
4828 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4829 }
4830 f0 += d[0];
4831 f1 += d[1];
4832 f2 += d[2];
4833 f3 += d[3];
4834 }
4835 [f0, f1, f2, f3]
4836 }
4837}
4838
4839#[allow(dead_code)]
4846static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4847
4848#[cfg(target_arch = "x86_64")]
4854fn q4tp_blocked_x86() -> bool {
4855 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4856 1 => false,
4857 2 => avx512vnni_enabled(),
4862 _ => blocked_enabled() && avx512vnni_enabled(),
4866 }
4867}
4868
4869#[cfg(target_arch = "aarch64")]
4871#[allow(dead_code)]
4872fn q4tp_v1() -> bool {
4873 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4874 1 => true,
4875 2 => false,
4876 _ => {
4877 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4878 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4879 }
4880 }
4881}
4882
4883#[cfg(target_arch = "x86_64")]
4895#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4896unsafe fn dot_q4tp_2x8_avx512(
4897 nib: &[u8],
4898 r0: usize,
4899 gpr: usize,
4900 xs: [&[i8]; 8],
4901 sc0: &[f32],
4902 sc1: &[f32],
4903) -> [[f32; 8]; 2] {
4904 unsafe {
4907 use core::arch::x86_64::*;
4908 let lomask = _mm256_set1_epi8(0x0F);
4909 let eight = _mm256_set1_epi8(8);
4910 let zero = _mm512_setzero_si512();
4911 let mut v0 = [_mm512_setzero_ps(); 8];
4912 let mut v1 = [_mm512_setzero_ps(); 8];
4913 let pairs = gpr / 2;
4914 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4915 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4916 let bb = _mm256_loadu_si256(t as *const __m256i);
4917 let lo = _mm256_and_si256(bb, lomask);
4918 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4919 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4920 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4921 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4922 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4923 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
4924 };
4925 for gp in 0..pairs {
4926 let gi = gp * 2;
4927 let (wa0, neg0) = unpack(r0, gi);
4928 let (wa1, neg1) = unpack(r0 + 1, gi);
4929 let off = gi * GROUP_SIZE;
4930 let sv = |sc: &[f32]| {
4931 _mm512_insertf32x8::<1>(
4932 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
4933 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
4934 )
4935 };
4936 let s0 = sv(sc0);
4937 let s1 = sv(sc1);
4938 for k in 0..8 {
4939 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
4940 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4941 zero,
4942 wa0,
4943 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
4944 ));
4945 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4946 zero,
4947 wa1,
4948 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
4949 ));
4950 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
4951 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
4952 }
4953 }
4954 let mut acc = [[0f32; 8]; 2];
4955 for k in 0..8 {
4956 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
4957 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
4958 }
4959 if gpr % 2 == 1 {
4960 let off = (gpr - 1) * GROUP_SIZE;
4961 for j in off..off + GROUP_SIZE {
4962 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
4963 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
4964 for k in 0..8 {
4965 let x = *xs[k].get_unchecked(j) as f32;
4966 acc[0][k] += w0 * sa * x;
4967 acc[1][k] += w1 * sb * x;
4968 }
4969 }
4970 }
4971 acc
4972 }
4973}
4974
4975#[cfg(target_arch = "x86_64")]
4980#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4981unsafe fn dot_q4tp_row_1x8_avx512(
4982 nib: &[u8],
4983 r: usize,
4984 gpr: usize,
4985 xs: [&[i8]; 8],
4986 scales: &[f32],
4987) -> [f32; 8] {
4988 unsafe {
4990 use core::arch::x86_64::*;
4991 let lomask = _mm256_set1_epi8(0x0F);
4992 let eight = _mm256_set1_epi8(8);
4993 let zero = _mm512_setzero_si512();
4994 let (mut v0, mut v1, mut v2, mut v3) = (
4995 _mm512_setzero_ps(),
4996 _mm512_setzero_ps(),
4997 _mm512_setzero_ps(),
4998 _mm512_setzero_ps(),
4999 );
5000 let (mut v4, mut v5, mut v6, mut v7) = (
5001 _mm512_setzero_ps(),
5002 _mm512_setzero_ps(),
5003 _mm512_setzero_ps(),
5004 _mm512_setzero_ps(),
5005 );
5006 let pairs = gpr / 2;
5007 for gp in 0..pairs {
5008 let gi = gp * 2;
5009 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5010 let bb = _mm256_loadu_si256(t as *const __m256i);
5011 let lo = _mm256_and_si256(bb, lomask);
5012 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5013 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5018 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5019 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5020 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5021 let wabs = _mm512_abs_epi8(w);
5022 let neg = _mm512_movepi8_mask(w);
5023 let off = gi * GROUP_SIZE;
5024 let sv = _mm512_insertf32x8::<1>(
5025 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5026 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5027 );
5028 let dot = |x: &[i8]| -> __m512 {
5029 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5030 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5031 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5032 };
5033 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5034 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5035 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5036 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5037 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5038 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5039 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5040 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5041 }
5042 let mut acc = [
5043 _mm512_reduce_add_ps(v0),
5044 _mm512_reduce_add_ps(v1),
5045 _mm512_reduce_add_ps(v2),
5046 _mm512_reduce_add_ps(v3),
5047 _mm512_reduce_add_ps(v4),
5048 _mm512_reduce_add_ps(v5),
5049 _mm512_reduce_add_ps(v6),
5050 _mm512_reduce_add_ps(v7),
5051 ];
5052 if gpr % 2 == 1 {
5055 let off = (gpr - 1) * GROUP_SIZE;
5056 for j in off..off + GROUP_SIZE {
5057 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5058 let ws = w * s;
5059 for k in 0..8 {
5060 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5061 }
5062 }
5063 }
5064 acc
5065 }
5066}
5067
5068#[cfg(target_arch = "x86_64")]
5081#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5082unsafe fn dot_q4tp_row_1x4_avx512(
5083 nib: &[u8],
5084 r: usize,
5085 gpr: usize,
5086 xs: [&[i8]; 4],
5087 scales: &[f32],
5088) -> [f32; 4] {
5089 unsafe {
5091 use core::arch::x86_64::*;
5092 let lomask = _mm256_set1_epi8(0x0F);
5093 let eight = _mm256_set1_epi8(8);
5094 let zero = _mm512_setzero_si512();
5095 let (mut v0, mut v1, mut v2, mut v3) = (
5096 _mm512_setzero_ps(),
5097 _mm512_setzero_ps(),
5098 _mm512_setzero_ps(),
5099 _mm512_setzero_ps(),
5100 );
5101 let pairs = gpr / 2;
5102 for gp in 0..pairs {
5103 let gi = gp * 2;
5104 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5105 let bb = _mm256_loadu_si256(t as *const __m256i);
5106 let lo = _mm256_and_si256(bb, lomask);
5107 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5108 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5113 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5114 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5115 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5116 let wabs = _mm512_abs_epi8(w);
5117 let neg = _mm512_movepi8_mask(w);
5118 let off = gi * GROUP_SIZE;
5119 let sv = _mm512_insertf32x8::<1>(
5120 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5121 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5122 );
5123 let dot = |x: &[i8]| -> __m512 {
5124 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5125 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5126 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5127 };
5128 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5129 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5130 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5131 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5132 }
5133 let mut acc = [
5134 _mm512_reduce_add_ps(v0),
5135 _mm512_reduce_add_ps(v1),
5136 _mm512_reduce_add_ps(v2),
5137 _mm512_reduce_add_ps(v3),
5138 ];
5139 if gpr % 2 == 1 {
5142 let off = (gpr - 1) * GROUP_SIZE;
5143 for j in off..off + GROUP_SIZE {
5144 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5145 let ws = w * s;
5146 for k in 0..4 {
5147 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5148 }
5149 }
5150 }
5151 acc
5152 }
5153}
5154
5155#[cfg(target_arch = "aarch64")]
5159#[target_feature(enable = "neon,dotprod")]
5160unsafe fn dot_q4tp_row_1x4_sdot(
5161 nib: &[u8],
5162 r: usize,
5163 gpr: usize,
5164 xs: [&[i8]; 4],
5165 scales: &[f32],
5166) -> [f32; 4] {
5167 unsafe {
5169 use core::arch::aarch64::*;
5170 use core::arch::asm;
5171 let lomask = vdupq_n_u8(0x0F);
5172 let eight = vdupq_n_s8(8);
5173 let (mut v0, mut v1, mut v2, mut v3) = (
5189 vdupq_n_f32(0.0),
5190 vdupq_n_f32(0.0),
5191 vdupq_n_f32(0.0),
5192 vdupq_n_f32(0.0),
5193 );
5194 for gi in 0..gpr {
5195 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5196 let s = *scales.get_unchecked(gi);
5197 let bb = vld1q_u8(t);
5198 let lo = vandq_u8(bb, lomask);
5199 let hi = vshrq_n_u8::<4>(bb);
5200 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5201 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5202 let off = gi * GROUP_SIZE;
5203 let dot4 = |x: &[i8]| -> int32x4_t {
5204 let x0 = vld1q_s8(x.as_ptr().add(off));
5205 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5206 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5207 asm!(
5208 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5209 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5210 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5211 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5212 options(pure, nomem, nostack),
5213 );
5214 vaddq_s32(a0, a1)
5215 };
5216 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5217 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5218 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5219 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5220 }
5221 [
5222 vaddvq_f32(v0),
5223 vaddvq_f32(v1),
5224 vaddvq_f32(v2),
5225 vaddvq_f32(v3),
5226 ]
5227 }
5228}
5229
5230fn q4tp_matmat(
5234 bytes: &[u8],
5235 xs_all: &[f32],
5236 b: usize,
5237 rows: usize,
5238 cols: usize,
5239 out: &mut [f32],
5240 pool: Option<&Pool>,
5241) {
5242 debug_assert_eq!(out.len(), b * rows);
5243 let gpr = cols / GROUP_SIZE;
5244 let v = Q4tpView::new(bytes, rows, cols);
5245
5246 #[cfg(target_os = "macos")]
5248 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5249 dequant_matmat_accel(
5250 &|r, dst| {
5251 let mut sc = [0f32; 32];
5252 let mut scv;
5253 let s: &[f32] = if gpr <= 32 {
5254 v.scales_into(r, gpr, &mut sc);
5255 &sc[..gpr]
5256 } else {
5257 scv = vec![0f32; gpr];
5258 v.scales_into(r, gpr, &mut scv);
5259 &scv
5260 };
5261 for gi in 0..gpr {
5262 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5263 for (k, &bb) in tile.iter().enumerate() {
5264 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5265 dst[gi * GROUP_SIZE + k * 2 + 1] =
5266 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5267 }
5268 }
5269 },
5270 xs_all,
5271 b,
5272 rows,
5273 cols,
5274 out,
5275 pool,
5276 );
5277 return;
5278 }
5279
5280 let out_addr = SendMut(out.as_mut_ptr());
5281 if a8w8_enabled() {
5282 let acts: Vec<SplitAct> = (0..b)
5283 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5284 .collect();
5285 let acts = &acts;
5286 #[cfg(target_arch = "aarch64")]
5287 let blocked_ok = sdot_enabled() && blocked_enabled();
5288 #[cfg(target_arch = "x86_64")]
5295 let blocked_ok = q4tp_blocked_x86();
5296 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5297 let blocked_ok = false;
5298 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5308 .ok()
5309 .and_then(|v| v.parse().ok())
5310 .filter(|v| *v > 0)
5311 .unwrap_or(256);
5312 let run = |start: usize, end: usize| {
5313 for abase in (0..acts.len()).step_by(panel_cols) {
5314 let alen = (acts.len() - abase).min(panel_cols);
5315 let mut sc = vec![0f32; gpr];
5316 #[cfg(target_arch = "x86_64")]
5317 let mut r_lo = start;
5318 #[cfg(target_arch = "x86_64")]
5319 if blocked_ok && alen >= 8 {
5320 let mut sc1 = vec![0f32; gpr];
5321 while r_lo + 2 <= end {
5322 v.scales_into(r_lo, gpr, &mut sc);
5323 v.scales_into(r_lo + 1, gpr, &mut sc1);
5324 let mut bi = 0usize;
5325 while bi + 8 <= alen {
5326 let xs = [
5327 acts[abase + bi].xq.as_slice(),
5328 acts[abase + bi + 1].xq.as_slice(),
5329 acts[abase + bi + 2].xq.as_slice(),
5330 acts[abase + bi + 3].xq.as_slice(),
5331 acts[abase + bi + 4].xq.as_slice(),
5332 acts[abase + bi + 5].xq.as_slice(),
5333 acts[abase + bi + 6].xq.as_slice(),
5334 acts[abase + bi + 7].xq.as_slice(),
5335 ];
5336 let d =
5337 unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5338 for (row, dr, scr) in
5339 [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)]
5340 {
5341 for k in 0..8 {
5342 let act = &acts[abase + bi + k];
5343 let mut acc = dr[k] * act.sx;
5344 for &(j, xv) in &act.outliers {
5345 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5346 acc += w * s * xv;
5347 }
5348 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5350 }
5351 }
5352 bi += 8;
5353 }
5354 for row in [r_lo, r_lo + 1] {
5357 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5358 for b2 in bi..alen {
5359 let act = &acts[abase + b2];
5360 let xs4 = [
5361 act.xq.as_slice(),
5362 act.xq.as_slice(),
5363 act.xq.as_slice(),
5364 act.xq.as_slice(),
5365 ];
5366 let d =
5367 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5368 let mut acc = d[0] * act.sx;
5369 for &(j, xv) in &act.outliers {
5370 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5371 acc += w * s * xv;
5372 }
5373 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5375 }
5376 }
5377 r_lo += 2;
5378 }
5379 }
5380 #[cfg(target_arch = "x86_64")]
5381 let row_start = r_lo;
5382 #[cfg(not(target_arch = "x86_64"))]
5383 let row_start = start;
5384 for r in row_start..end {
5385 v.scales_into(r, gpr, &mut sc);
5386 let mut bi = 0usize;
5387 #[cfg(target_arch = "x86_64")]
5388 if blocked_ok {
5389 while bi + 8 <= alen {
5390 let xs = [
5391 acts[abase + bi].xq.as_slice(),
5392 acts[abase + bi + 1].xq.as_slice(),
5393 acts[abase + bi + 2].xq.as_slice(),
5394 acts[abase + bi + 3].xq.as_slice(),
5395 acts[abase + bi + 4].xq.as_slice(),
5396 acts[abase + bi + 5].xq.as_slice(),
5397 acts[abase + bi + 6].xq.as_slice(),
5398 acts[abase + bi + 7].xq.as_slice(),
5399 ];
5400 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5401 for k in 0..8 {
5402 let act = &acts[abase + bi + k];
5403 let mut acc = d[k] * act.sx;
5404 for &(j, xv) in &act.outliers {
5405 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5406 acc += w * s * xv;
5407 }
5408 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5410 }
5411 bi += 8;
5412 }
5413 while bi + 4 <= alen {
5414 let xs = [
5415 acts[abase + bi].xq.as_slice(),
5416 acts[abase + bi + 1].xq.as_slice(),
5417 acts[abase + bi + 2].xq.as_slice(),
5418 acts[abase + bi + 3].xq.as_slice(),
5419 ];
5420 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5421 for k in 0..4 {
5422 let act = &acts[abase + bi + k];
5423 let mut acc = d[k] * act.sx;
5424 for &(j, xv) in &act.outliers {
5425 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5426 acc += w * s * xv;
5427 }
5428 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5430 }
5431 bi += 4;
5432 }
5433 }
5434 #[cfg(target_arch = "aarch64")]
5435 if blocked_ok {
5436 while bi + 4 <= alen {
5437 let xs = [
5438 acts[abase + bi].xq.as_slice(),
5439 acts[abase + bi + 1].xq.as_slice(),
5440 acts[abase + bi + 2].xq.as_slice(),
5441 acts[abase + bi + 3].xq.as_slice(),
5442 ];
5443 let d = unsafe {
5444 if q4tp_v1() {
5445 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5446 } else {
5447 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5448 }
5449 };
5450 for k in 0..4 {
5451 let act = &acts[abase + bi + k];
5452 let mut acc = d[k] * act.sx;
5453 for &(j, xv) in &act.outliers {
5454 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5455 acc += w * s * xv;
5456 }
5457 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5459 }
5460 bi += 4;
5461 }
5462 }
5463 let _ = blocked_ok;
5464 while bi < alen {
5465 let act = &acts[abase + bi];
5466 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5467 for &(j, xv) in &act.outliers {
5468 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5469 acc += w * s * xv;
5470 }
5471 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5473 bi += 1;
5474 }
5475 }
5476
5477 }
5478 };
5479 dispatch_rows(pool, rows, &run);
5480 return;
5481 }
5482
5483 let run = |start: usize, end: usize| {
5484 let mut sc = vec![0f32; gpr];
5485 for r in start..end {
5486 v.scales_into(r, gpr, &mut sc);
5487 for bi in 0..b {
5488 let x = &xs_all[bi * cols..(bi + 1) * cols];
5489 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5491 }
5492 }
5493 };
5494 dispatch_rows(pool, rows, &run);
5495}
5496
5497fn q4t_matvec(
5499 bytes: &[u8],
5500 x: &[f32],
5501 rows: usize,
5502 cols: usize,
5503 out: &mut [f32],
5504 pool: Option<&Pool>,
5505) {
5506 debug_assert_eq!(out.len(), rows);
5507 let gpr = cols / GROUP_SIZE;
5508 let out_addr = SendMut(out.as_mut_ptr());
5509 if a8w8_enabled() {
5510 let act = split_act(x);
5511 let run = move |start: usize, end: usize| {
5512 for r in start..end {
5513 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5514 for &(j, xv) in &act.outliers {
5515 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5516 acc += w * s * xv;
5517 }
5518 unsafe { *out_addr.at(r) = acc };
5520 }
5521 };
5522 dispatch_rows(pool, rows, &run);
5523 return;
5524 }
5525 let run = move |start: usize, end: usize| {
5526 for r in start..end {
5527 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5529 }
5530 };
5531 dispatch_rows(pool, rows, &run);
5532}
5533
5534#[allow(clippy::too_many_arguments)]
5536fn q4t_matvec2(
5537 bytes: &[u8],
5538 x1: &[f32],
5539 x2: &[f32],
5540 rows: usize,
5541 cols: usize,
5542 o1: &mut [f32],
5543 o2: &mut [f32],
5544 pool: Option<&Pool>,
5545) {
5546 let gpr = cols / GROUP_SIZE;
5547 let p1 = SendMut(o1.as_mut_ptr());
5548 let p2 = SendMut(o2.as_mut_ptr());
5549 if a8w8_enabled() {
5550 let a1 = split_act(x1);
5551 let a2 = split_act(x2);
5552 let run = move |start: usize, end: usize| {
5553 for r in start..end {
5554 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5555 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5556 for &(j, xv) in &a1.outliers {
5557 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5558 v1 += w * s * xv;
5559 }
5560 for &(j, xv) in &a2.outliers {
5561 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5562 v2 += w * s * xv;
5563 }
5564 unsafe {
5566 *p1.at(r) = v1;
5567 *p2.at(r) = v2;
5568 }
5569 }
5570 };
5571 dispatch_rows(pool, rows, &run);
5572 return;
5573 }
5574 let run = move |start: usize, end: usize| {
5575 for r in start..end {
5576 unsafe {
5578 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5579 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5580 }
5581 }
5582 };
5583 dispatch_rows(pool, rows, &run);
5584}
5585
5586#[allow(clippy::too_many_arguments)]
5588#[cfg(target_os = "macos")]
5594fn dequant_matmat_accel(
5595 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5596 xs_all: &[f32],
5597 b: usize,
5598 rows: usize,
5599 cols: usize,
5600 out: &mut [f32],
5601 pool: Option<&Pool>,
5602) {
5603 const TR: usize = 2048;
5604 thread_local! {
5605 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5606 }
5607 WTILE.with(|wt| {
5608 let mut wtile = wt.borrow_mut();
5609 wtile.resize(TR * cols, 0.0);
5610 let mut r0 = 0usize;
5611 while r0 < rows {
5612 let tr = TR.min(rows - r0);
5613 let wt_addr = SendMut(wtile.as_mut_ptr());
5614 let run = |start: usize, end: usize| {
5615 for r in start..end {
5616 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5618 dequant_row(r0 + r, dst);
5619 }
5620 };
5621 dispatch_rows(pool, tr, &run);
5622 unsafe {
5623 accel_blas::cblas_sgemm(
5624 101, 111, 112, b as i32,
5628 tr as i32,
5629 cols as i32,
5630 1.0,
5631 xs_all.as_ptr(),
5632 cols as i32,
5633 wtile.as_ptr(),
5634 cols as i32,
5635 0.0,
5636 out.as_mut_ptr().add(r0),
5637 rows as i32,
5638 );
5639 }
5640 r0 += tr;
5641 }
5642 });
5643}
5644
5645fn q4t_matmat(
5646 bytes: &[u8],
5647 xs_all: &[f32],
5648 b: usize,
5649 rows: usize,
5650 cols: usize,
5651 out: &mut [f32],
5652 pool: Option<&Pool>,
5653) {
5654 debug_assert_eq!(out.len(), b * rows);
5655 let gpr = cols / GROUP_SIZE;
5656 #[cfg(target_os = "macos")]
5660 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5661 dequant_matmat_accel(
5662 &|r, dst| {
5663 for gi in 0..gpr {
5664 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5665 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5666 for (k, &bb) in tile[2..].iter().enumerate() {
5667 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5668 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5669 }
5670 }
5671 },
5672 xs_all,
5673 b,
5674 rows,
5675 cols,
5676 out,
5677 pool,
5678 );
5679 return;
5680 }
5681 let out_addr = SendMut(out.as_mut_ptr());
5682 if a8w8_enabled() {
5683 let acts: Vec<SplitAct> = (0..b)
5684 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5685 .collect();
5686 let acts = &acts;
5687 #[cfg(target_arch = "x86_64")]
5688 let blocked_ok = avx2_enabled()
5689 && blocked_enabled();
5690 #[cfg(target_arch = "aarch64")]
5691 let blocked_ok = sdot_enabled()
5692 && blocked_enabled();
5693 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5694 let blocked_ok = false;
5695 let run = move |start: usize, end: usize| {
5696 for r in start..end {
5697 let mut bi = 0usize;
5698 #[cfg(target_arch = "aarch64")]
5699 if blocked_ok {
5700 while bi + 4 <= acts.len() {
5701 let xs = [
5702 acts[bi].xq.as_slice(),
5703 acts[bi + 1].xq.as_slice(),
5704 acts[bi + 2].xq.as_slice(),
5705 acts[bi + 3].xq.as_slice(),
5706 ];
5707 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5708 for k in 0..4 {
5709 let act = &acts[bi + k];
5710 let mut acc = d[k] * act.sx;
5711 for &(j, xv) in &act.outliers {
5712 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5713 acc += w * sc * xv;
5714 }
5715 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5717 }
5718 bi += 4;
5719 }
5720 }
5721 #[cfg(target_arch = "x86_64")]
5722 if blocked_ok {
5723 while bi + 4 <= acts.len() {
5724 let xs = [
5725 acts[bi].xq.as_slice(),
5726 acts[bi + 1].xq.as_slice(),
5727 acts[bi + 2].xq.as_slice(),
5728 acts[bi + 3].xq.as_slice(),
5729 ];
5730 let d = unsafe {
5731 if vnni_tiles_enabled() {
5732 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5733 } else {
5734 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5735 }
5736 };
5737 for k in 0..4 {
5738 let act = &acts[bi + k];
5739 let mut acc = d[k] * act.sx;
5740 for &(j, xv) in &act.outliers {
5741 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5742 acc += w * sc * xv;
5743 }
5744 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5746 }
5747 bi += 4;
5748 }
5749 }
5750 let _ = blocked_ok;
5751 while bi < acts.len() {
5752 let act = &acts[bi];
5753 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5754 for &(j, xv) in &act.outliers {
5755 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5756 acc += w * s * xv;
5757 }
5758 unsafe { *out_addr.at(bi * rows + r) = acc };
5760 bi += 1;
5761 }
5762 }
5763 };
5764 dispatch_rows(pool, rows, &run);
5765 return;
5766 }
5767 let run = move |start: usize, end: usize| {
5768 for r in start..end {
5769 for bi in 0..b {
5770 let x = &xs_all[bi * cols..(bi + 1) * cols];
5771 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5773 }
5774 }
5775 };
5776 dispatch_rows(pool, rows, &run);
5777}
5778
5779fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5788 (0..gpr)
5789 .map(|gi| {
5790 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5791 .iter()
5792 .map(|&v| v as i32)
5793 .sum()
5794 })
5795 .collect()
5796}
5797
5798#[inline]
5802#[allow(unreachable_code)]
5803#[cfg(target_arch = "x86_64")]
5808#[target_feature(enable = "avx2")]
5809unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5810 unsafe {
5812 use core::arch::x86_64::*;
5813 let expand = _mm256_setr_epi8(
5815 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,
5816 3, 3, 3,
5817 );
5818 let bitsel = _mm256_setr_epi8(
5819 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5820 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5821 );
5822 let ones8 = _mm256_set1_epi8(1);
5823 let ones16 = _mm256_set1_epi16(1);
5824 let mut acc = 0f32;
5825 for gi in 0..gpr {
5826 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5827 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5828 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5829 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5830 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5831 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5832 let sel = _mm256_and_si256(x, mask);
5833 let p16 = _mm256_maddubs_epi16(ones8, sel);
5835 let d32 = _mm256_madd_epi16(p16, ones16);
5836 let hi128 = _mm256_extracti128_si256::<1>(d32);
5837 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5838 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5839 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5840 let msum = _mm_cvtsi128_si32(s32);
5841 let d = 2 * msum - gsum[gi];
5844 acc += d as f32 * s;
5845 }
5846 acc
5847 }
5848}
5849
5850#[cfg(target_arch = "x86_64")]
5853#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5854unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5855 unsafe {
5857 use core::arch::x86_64::*;
5858 let expand = _mm256_setr_epi8(
5859 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,
5860 3, 3, 3,
5861 );
5862 let bitsel = _mm256_setr_epi8(
5863 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5864 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5865 );
5866 let ones8 = _mm256_set1_epi8(1);
5867 let mut acc = 0f32;
5868 for gi in 0..gpr {
5869 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5870 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5871 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5872 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5873 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5874 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5875 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5876 let d = 2 * msum - gsum[gi];
5877 acc += d as f32 * s;
5878 }
5879 acc
5880 }
5881}
5882
5883#[cfg(target_arch = "x86_64")]
5885#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5886unsafe fn dot_q1_row_1x4_vnni(
5887 bytes: &[u8],
5888 r: usize,
5889 gpr: usize,
5890 xs: [&[i8]; 4],
5891 gsums: [&[i32]; 4],
5892) -> [f32; 4] {
5893 unsafe {
5895 use core::arch::x86_64::*;
5896 let expand = _mm256_setr_epi8(
5897 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,
5898 3, 3, 3,
5899 );
5900 let bitsel = _mm256_setr_epi8(
5901 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5902 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5903 );
5904 let ones8 = _mm256_set1_epi8(1);
5905 let mut acc = [0f32; 4];
5906 for gi in 0..gpr {
5907 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5908 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5909 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5910 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5911 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5912 for (k, xq) in xs.iter().enumerate() {
5913 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5914 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5915 let d = 2 * msum - gsums[k][gi];
5916 acc[k] += d as f32 * s;
5917 }
5918 }
5919 acc
5920 }
5921}
5922
5923#[cfg(target_arch = "x86_64")]
5926#[target_feature(enable = "avx2")]
5927unsafe fn dot_q1_row_1x4_avx2(
5928 bytes: &[u8],
5929 r: usize,
5930 gpr: usize,
5931 xs: [&[i8]; 4],
5932 gsums: [&[i32]; 4],
5933) -> [f32; 4] {
5934 unsafe {
5936 use core::arch::x86_64::*;
5937 let expand = _mm256_setr_epi8(
5938 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,
5939 3, 3, 3,
5940 );
5941 let bitsel = _mm256_setr_epi8(
5942 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5943 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5944 );
5945 let ones8 = _mm256_set1_epi8(1);
5946 let ones16 = _mm256_set1_epi16(1);
5947 let mut acc = [0f32; 4];
5948 for gi in 0..gpr {
5949 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5950 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5951 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5952 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5953 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5954 for (k, xq) in xs.iter().enumerate() {
5955 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5956 let sel = _mm256_and_si256(x, mask);
5957 let p16 = _mm256_maddubs_epi16(ones8, sel);
5958 let d32 = _mm256_madd_epi16(p16, ones16);
5959 let hi128 = _mm256_extracti128_si256::<1>(d32);
5960 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5961 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5962 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5963 let msum = _mm_cvtsi128_si32(s32);
5964 let d = 2 * msum - gsums[k][gi];
5965 acc[k] += d as f32 * s;
5966 }
5967 }
5968 acc
5969 }
5970}
5971
5972#[allow(unreachable_code)]
5973fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5974 #[cfg(target_arch = "aarch64")]
5975 unsafe {
5976 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
5977 }
5978 #[cfg(target_arch = "x86_64")]
5979 if avx2_enabled() {
5980 unsafe {
5981 if vnni_tiles_enabled() {
5982 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
5983 }
5984 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
5985 }
5986 }
5987 let _ = gsum;
5988 let mut acc = 0f32;
5989 for gi in 0..gpr {
5990 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5991 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5992 let mut d = 0i32;
5993 for (j, &b) in tile[2..].iter().enumerate() {
5994 for k in 0..8 {
5995 let w = ((b >> k) & 1) as i32 * 2 - 1;
5996 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
5997 }
5998 }
5999 acc += d as f32 * s;
6000 }
6001 acc
6002}
6003
6004#[cfg(target_arch = "aarch64")]
6013#[target_feature(enable = "neon,dotprod")]
6014unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6015 unsafe {
6018 use core::arch::aarch64::*;
6019 use core::arch::asm;
6020 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6021 let m = vld1q_u8(MASKS.as_ptr());
6022 macro_rules! tile_dot {
6024 ($t:expr, $x:expr) => {{
6025 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6026 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6027 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6028 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6029 let x0 = vld1q_s8($x);
6030 let x1 = vld1q_s8($x.add(16));
6031 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6032 asm!(
6033 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6034 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6035 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6036 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6037 options(pure, nomem, nostack),
6038 );
6039 vaddq_s32(a0, a1)
6040 }};
6041 }
6042 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6051 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6052 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6053 const IW11: [u8; 16] = [
6054 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6055 ];
6056 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6057 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6058 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6059 let isc = vld1_u8(ISC.as_ptr());
6060 macro_rules! tile_dot_tbl {
6062 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6063 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6064 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6065 let x0 = vld1q_s8($x);
6066 let x1 = vld1q_s8($x.add(16));
6067 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6068 asm!(
6069 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6070 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6071 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6072 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6073 options(pure, nomem, nostack),
6074 );
6075 vaddq_s32(a0, a1)
6076 }};
6077 }
6078 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6079 let row_base = r * gpr * Q1_TILE;
6080 let abs_end = bytes.len();
6081 let xp = xq.as_ptr();
6082 let gp = gsum.as_ptr();
6083 let mut accv = vdupq_n_f32(0.0);
6084 let mut gi = 0;
6085 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6088 let t0 = base.add(gi * Q1_TILE);
6089 let ld_a = vld1q_u8(t0);
6090 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6091 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6092 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6093 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6094 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6095 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6097 let g = vld1q_s32(gp.add(gi));
6098 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6099 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6100 let scf: float32x4_t;
6101 asm!(
6102 "fcvtl {o:v}.4s, {i:v}.4h",
6103 o = out(vreg) scf, i = in(vreg) sc16,
6104 options(pure, nomem, nostack),
6105 );
6106 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6107 gi += 4;
6108 }
6109 let mut acc = vaddvq_f32(accv);
6110 while gi < gpr {
6111 let t = base.add(gi * Q1_TILE);
6112 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6113 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6114 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6115 gi += 1;
6116 }
6117 acc
6118 }
6119}
6120
6121#[cfg(target_arch = "aarch64")]
6126#[target_feature(enable = "neon,dotprod")]
6127unsafe fn dot_q1_row_1x4_sdot(
6128 bytes: &[u8],
6129 r: usize,
6130 gpr: usize,
6131 xs: [&[i8]; 4],
6132 gs: [&[i32]; 4],
6133) -> [f32; 4] {
6134 unsafe {
6136 use core::arch::aarch64::*;
6137 use core::arch::asm;
6138 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6139 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6140 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6141 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6142 const IW11: [u8; 16] = [
6143 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6144 ];
6145 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6146 let m = vld1q_u8(MASKS.as_ptr());
6147 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6148 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6149 let isc = vld1_u8(ISC.as_ptr());
6150 macro_rules! sdot2 {
6151 ($w0:expr, $w1:expr, $x:expr) => {{
6152 let x0 = vld1q_s8($x);
6153 let x1 = vld1q_s8($x.add(16));
6154 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6155 asm!(
6156 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6157 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6158 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6159 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6160 options(pure, nomem, nostack),
6161 );
6162 vaddq_s32(a0, a1)
6163 }};
6164 }
6165 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6166 let row_base = r * gpr * Q1_TILE;
6167 let abs_end = bytes.len();
6168 let mut accv = [vdupq_n_f32(0.0); 4];
6169 let mut gi = 0;
6170 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6171 let t0 = base.add(gi * Q1_TILE);
6172 let ld_a = vld1q_u8(t0);
6173 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6174 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6176 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6177 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6178 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6179 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6180 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6181 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6182 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6183 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6184 let scf: float32x4_t;
6185 asm!(
6186 "fcvtl {o:v}.4s, {i:v}.4h",
6187 o = out(vreg) scf, i = in(vreg) sc16,
6188 options(pure, nomem, nostack),
6189 );
6190 for k in 0..4 {
6191 let xp = xs[k].as_ptr();
6192 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6193 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6194 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6195 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6196 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6197 let g = vld1q_s32(gs[k].as_ptr().add(gi));
6198 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6199 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6200 }
6201 gi += 4;
6202 }
6203 let mut acc = [
6204 vaddvq_f32(accv[0]),
6205 vaddvq_f32(accv[1]),
6206 vaddvq_f32(accv[2]),
6207 vaddvq_f32(accv[3]),
6208 ];
6209 while gi < gpr {
6210 let t = base.add(gi * Q1_TILE);
6211 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6212 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6213 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6214 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6215 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6216 for k in 0..4 {
6217 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6218 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6219 }
6220 gi += 1;
6221 }
6222 acc
6223 }
6224}
6225
6226#[inline]
6228fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6229 let gi = j / GROUP_SIZE;
6230 let k = j % GROUP_SIZE;
6231 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6232 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6233 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6234 ((bit as i32 * 2 - 1) as f32, s)
6235}
6236
6237#[inline]
6239fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6240 let mut acc = 0f32;
6241 for gi in 0..gpr {
6242 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6243 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6244 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6245 let mut ga = 0f32;
6246 for (j, &b) in tile[2..].iter().enumerate() {
6247 for k in 0..8 {
6248 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6249 }
6250 }
6251 acc += ga * s;
6252 }
6253 acc
6254}
6255
6256#[allow(clippy::too_many_arguments)]
6259fn q1_range_a8w8(
6260 bytes: &[u8],
6261 gpr: usize,
6262 act: &SplitAct,
6263 gsum: &[i32],
6264 out: SendMut,
6265 start: usize,
6266 end: usize,
6267) {
6268 for r in start..end {
6269 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6270 for &(j, xv) in &act.outliers {
6271 let (w, s) = q1_outlier(bytes, r, gpr, j);
6272 acc += w * s * xv;
6273 }
6274 unsafe { *out.at(r) = acc };
6276 }
6277}
6278
6279fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6281 for r in start..end {
6282 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6284 }
6285}
6286
6287fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6292 let entries = base_len + (rows + 1) * 4;
6293 (base_len, entries, entries <= bytes.len())
6294}
6295
6296#[inline]
6298fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6299 let o = rp_off + r * 4;
6300 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6301}
6302
6303const SIGN5: [[f32; 5]; 256] = {
6307 let mut lut = [[0.0f32; 5]; 256];
6308 let pow3 = [1u16, 3, 9, 27, 81];
6309 let mut byte = 0usize;
6310 while byte < 256 {
6311 let mut i = 0usize;
6312 while i < 5 {
6313 let code = (byte as u16 / pow3[i]) % 3;
6314 lut[byte][i] = if code == 1 {
6315 1.0
6316 } else if code == 2 {
6317 -1.0
6318 } else {
6319 0.0
6320 };
6321 i += 1;
6322 }
6323 byte += 1;
6324 }
6325 lut
6326};
6327
6328const SIGN5_I8: [[i8; 5]; 256] = {
6330 let mut lut = [[0i8; 5]; 256];
6331 let pow3 = [1u16, 3, 9, 27, 81];
6332 let mut byte = 0usize;
6333 while byte < 256 {
6334 let mut i = 0usize;
6335 while i < 5 {
6336 let code = (byte as u16 / pow3[i]) % 3;
6337 lut[byte][i] = if code == 1 {
6338 1
6339 } else if code == 2 {
6340 -1
6341 } else {
6342 0
6343 };
6344 i += 1;
6345 }
6346 byte += 1;
6347 }
6348 lut
6349};
6350
6351const SIGN5_U64: [u64; 256] = {
6357 let mut lut = [0u64; 256];
6358 let pow3 = [1u16, 3, 9, 27, 81];
6359 let mut byte = 0usize;
6360 while byte < 256 {
6361 let mut v = 0u64;
6362 let mut i = 0usize;
6363 while i < 5 {
6364 let code = (byte as u16 / pow3[i]) % 3;
6365 let s: u8 = if code == 1 {
6366 1
6367 } else if code == 2 {
6368 0xFF
6369 } else {
6370 0
6371 };
6372 v |= (s as u64) << (i * 8);
6373 i += 1;
6374 }
6375 lut[byte] = v;
6376 byte += 1;
6377 }
6378 lut
6379};
6380
6381#[inline]
6386fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6387 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6388 let off = (r * gpr + j / GROUP_SIZE) * TILE;
6389 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6390 let within = j % GROUP_SIZE;
6391 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6392}
6393
6394#[cfg(target_arch = "aarch64")]
6397#[target_feature(enable = "neon,dotprod")]
6398#[inline]
6399unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6400 unsafe {
6402 use core::arch::aarch64::*;
6403 use core::arch::asm;
6404 let w0 = vld1q_s8(w);
6405 let w1 = vld1q_s8(w.add(16));
6406 let x0 = vld1q_s8(x);
6407 let x1 = vld1q_s8(x.add(16));
6408 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6409 asm!(
6410 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6411 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6412 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6413 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6414 options(pure, nomem, nostack),
6415 );
6416 vaddvq_s32(vaddq_s32(a0, a1))
6417 }
6418}
6419
6420#[cfg(target_arch = "x86_64")]
6423#[target_feature(enable = "avx2")]
6424#[inline]
6425unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6426 unsafe {
6428 use core::arch::x86_64::*;
6429 let wv = _mm256_loadu_si256(w as *const __m256i);
6430 let xv = _mm256_loadu_si256(x as *const __m256i);
6431 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6432 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6433 let hi128 = _mm256_extracti128_si256::<1>(d);
6434 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6435 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6436 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6437 _mm_cvtsi128_si32(s32)
6438 }
6439}
6440
6441#[inline]
6446fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6447 debug_assert!(dst.len() >= 40);
6448 unsafe {
6451 let p = dst.as_mut_ptr();
6452 for bi in 0..7 {
6453 core::ptr::write_unaligned(
6454 p.add(bi * 5) as *mut u64,
6455 SIGN5_U64[*codes.add(bi) as usize],
6456 );
6457 }
6458 }
6459}
6460
6461#[inline]
6466fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6467 #[cfg(target_arch = "aarch64")]
6468 unsafe {
6469 return sdot32_i8(w, x);
6470 }
6471 #[cfg(target_arch = "x86_64")]
6472 unsafe {
6473 return i8dot32_avx2(w, x);
6474 }
6475 #[allow(unreachable_code)]
6476 unsafe {
6477 let mut s = 0i32;
6478 for k in 0..GROUP_SIZE {
6479 s += *w.add(k) as i32 * *x.add(k) as i32;
6480 }
6481 s
6482 }
6483}
6484
6485#[inline]
6486unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6487 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6488 (
6489 SIGN5_U64[*codes as usize],
6490 SIGN5_U64[*codes.add(1) as usize],
6491 SIGN5_U64[*codes.add(2) as usize],
6492 SIGN5_U64[*codes.add(3) as usize],
6493 SIGN5_U64[*codes.add(4) as usize],
6494 SIGN5_U64[*codes.add(5) as usize],
6495 SIGN5_U64[*codes.add(6) as usize],
6496 )
6497 };
6498
6499 let u0 = s0 | (s1 << 40);
6500 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6501 let u2 = (s3 >> 8) | (s4 << 32);
6502 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6503
6504 (u0, u1, u2, u3)
6505}
6506
6507#[cfg(target_arch = "aarch64")]
6511#[target_feature(enable = "neon,dotprod")]
6512unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6513 use core::arch::aarch64::*;
6514 use core::arch::asm;
6515 unsafe {
6516 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6517 let mut acc = 0f32;
6518 let bytes_ptr = bytes.as_ptr();
6519 let xq_ptr = xq.as_ptr();
6520 let row_off = r * gpr * TILE;
6521
6522 let gpr2 = gpr & !1;
6523 let mut gi = 0;
6524 while gi < gpr2 {
6525 let off0 = row_off + gi * TILE;
6526 let off1 = off0 + TILE;
6527 let s0 = f16_to_f32(u16::from_le_bytes([
6528 *bytes_ptr.add(off0),
6529 *bytes_ptr.add(off0 + 1),
6530 ]));
6531 let s1 = f16_to_f32(u16::from_le_bytes([
6532 *bytes_ptr.add(off1),
6533 *bytes_ptr.add(off1 + 1),
6534 ]));
6535
6536 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6537 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6538
6539 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6540 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6541 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6542 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6543
6544 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6545 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6546 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6547 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6548
6549 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6550 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6551 asm!(
6552 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6553 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6554 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6555 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6556 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6557 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6558 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6559 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6560 options(pure, nomem, nostack),
6561 );
6562 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6563 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6564 acc += d0 as f32 * s0 + d1 as f32 * s1;
6565 gi += 2;
6566 }
6567
6568 if gi < gpr {
6569 let off = row_off + gi * TILE;
6570 let s = f16_to_f32(u16::from_le_bytes([
6571 *bytes_ptr.add(off),
6572 *bytes_ptr.add(off + 1),
6573 ]));
6574 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6575 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6576 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6577 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6578 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6579 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6580 asm!(
6581 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6582 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6583 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6584 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6585 options(pure, nomem, nostack),
6586 );
6587 let d = vaddvq_s32(vaddq_s32(a0, a1));
6588 acc += d as f32 * s;
6589 }
6590 acc
6591 }
6592}
6593
6594#[cfg(target_arch = "x86_64")]
6596#[target_feature(enable = "avx2")]
6597unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6598 use core::arch::x86_64::*;
6599 unsafe {
6600 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6601 let mut acc = 0f32;
6602 let bytes_ptr = bytes.as_ptr();
6603 let xq_ptr = xq.as_ptr();
6604 let row_off = r * gpr * TILE;
6605
6606 let ones = _mm256_set1_epi16(1);
6607 for gi in 0..gpr {
6608 let off = row_off + gi * TILE;
6609 let s = f16_to_f32(u16::from_le_bytes([
6610 *bytes_ptr.add(off),
6611 *bytes_ptr.add(off + 1),
6612 ]));
6613 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6614 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6615 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6616 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6617 let d256 = _mm256_madd_epi16(p16, ones);
6618 let d128 = _mm_add_epi32(
6619 _mm256_castsi256_si128(d256),
6620 _mm256_extracti128_si256(d256, 1),
6621 );
6622 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6623 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6624 acc += d32 as f32 * s;
6625 }
6626 acc
6627 }
6628}
6629
6630#[cfg(target_arch = "x86_64")]
6632#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6633unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6634 use core::arch::x86_64::*;
6635 unsafe {
6637 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6638 let mut acc = 0f32;
6639 let bytes_ptr = bytes.as_ptr();
6640 let xq_ptr = xq.as_ptr();
6641 let row_off = r * gpr * TILE;
6642 for gi in 0..gpr {
6643 let off = row_off + gi * TILE;
6644 let s = f16_to_f32(u16::from_le_bytes([
6645 *bytes_ptr.add(off),
6646 *bytes_ptr.add(off + 1),
6647 ]));
6648 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6649 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6650 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6651 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6652 acc += d as f32 * s;
6653 }
6654 acc
6655 }
6656}
6657
6658#[inline]
6662fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6663 #[cfg(target_arch = "aarch64")]
6664 unsafe {
6665 return q1t_dot_row_sdot(bytes, r, gpr, xq);
6666 }
6667 #[cfg(target_arch = "x86_64")]
6668 unsafe {
6669 if vnni_tiles_enabled() {
6670 return q1t_dot_row_vnni(bytes, r, gpr, xq);
6671 }
6672 return q1t_dot_row_avx2(bytes, r, gpr, xq);
6673 }
6674 #[allow(unreachable_code)]
6675 {
6676 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6677 let mut acc = 0f32;
6678 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
6680 let off = (r * gpr + gi) * TILE;
6681 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6682 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6683 let mut d = 0i32;
6684 for k in 0..GROUP_SIZE {
6685 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6686 }
6687 acc += d as f32 * s;
6688 }
6689 acc
6690 }
6691}
6692
6693fn q1t_row_outlier_correction(
6700 bytes: &[u8],
6701 r: usize,
6702 rp_off: usize,
6703 entries_off: usize,
6704 has_ov: bool,
6705 x: &[f32],
6706) -> f32 {
6707 if !has_ov {
6708 return 0.0;
6709 }
6710 let (c0, c1) = (
6711 q1t_rowptr(bytes, rp_off, r),
6712 q1t_rowptr(bytes, rp_off, r + 1),
6713 );
6714 let mut corr = 0f32;
6715 for p in c0..c1 {
6716 let e = entries_off + p * 4;
6717 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6718 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6719 corr += val * x[col];
6720 }
6721 corr
6722}
6723
6724fn q1t_dequant_row(
6728 bytes: &[u8],
6729 r: usize,
6730 gpr: usize,
6731 rp_off: usize,
6732 entries_off: usize,
6733 has_ov: bool,
6734 buf: &mut [f32],
6735) {
6736 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6737 for g in 0..gpr {
6738 let off = (r * gpr + g) * TILE;
6739 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6740 let codes = &bytes[off + 2..off + TILE];
6741 let bc = g * GROUP_SIZE;
6742 for bi in 0..6 {
6744 let lut = &SIGN5[codes[bi] as usize];
6745 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6746 for i in 0..5 {
6747 d[i] = lut[i] * s;
6748 }
6749 }
6750 let lut = &SIGN5[codes[6] as usize];
6751 buf[bc + 30] = lut[0] * s;
6752 buf[bc + 31] = lut[1] * s;
6753 }
6754 if !has_ov {
6755 return;
6756 }
6757 let (c0, c1) = (
6758 q1t_rowptr(bytes, rp_off, r),
6759 q1t_rowptr(bytes, rp_off, r + 1),
6760 );
6761 for p in c0..c1 {
6762 let e = entries_off + p * 4;
6763 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6764 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6765 }
6766}
6767
6768fn q1t_add_overlay(
6772 bytes: &[u8],
6773 x: &[f32],
6774 rows: usize,
6775 cols: usize,
6776 out: &mut [f32],
6777 pool: Option<&Pool>,
6778) {
6779 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6780 let gpr = cols / GROUP_SIZE;
6781 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6782 if !has_ov {
6783 return;
6784 }
6785 let out_addr = SendMut(out.as_mut_ptr());
6786 let run = move |start: usize, end: usize| {
6787 for r in start..end {
6788 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6789 unsafe { *out_addr.at(r) += corr };
6791 }
6792 };
6793 dispatch_rows(pool, rows, &run);
6794}
6795
6796#[allow(clippy::too_many_arguments)]
6799fn q1t_range_a8w8(
6800 bytes: &[u8],
6801 gpr: usize,
6802 rp_off: usize,
6803 ent_off: usize,
6804 has_ov: bool,
6805 act: &SplitAct,
6806 x: &[f32],
6807 out: SendMut,
6808 start: usize,
6809 end: usize,
6810) {
6811 for r in start..end {
6812 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6813 for &(j, xv) in &act.outliers {
6814 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6815 }
6816 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6817 unsafe { *out.at(r) = acc };
6819 }
6820}
6821
6822#[allow(clippy::too_many_arguments)]
6825fn q1t_range_f32_batch(
6826 bytes: &[u8],
6827 gpr: usize,
6828 rp_off: usize,
6829 ent_off: usize,
6830 has_ov: bool,
6831 x: &[f32],
6832 out: SendMut,
6833 start: usize,
6834 end: usize,
6835) {
6836 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6837 let mut sg = [0f32; GROUP_SIZE];
6838 for r in start..end {
6839 let mut acc = 0f32;
6840 for g in 0..gpr {
6841 let off = (r * gpr + g) * TILE;
6842 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6843 let codes = &bytes[off + 2..off + TILE];
6844 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6845 for bi in 0..6 {
6846 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6847 }
6848 let lut = &SIGN5[codes[6] as usize];
6849 sg[30] = lut[0];
6850 sg[31] = lut[1];
6851 let mut gsum = 0f32;
6852 for k in 0..GROUP_SIZE {
6853 gsum += sg[k] * xg[k];
6854 }
6855 acc += s * gsum;
6856 }
6857 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6858 unsafe { *out.at(r) = acc };
6860 }
6861}
6862
6863fn q1t_matvec(
6867 bytes: &[u8],
6868 x: &[f32],
6869 rows: usize,
6870 cols: usize,
6871 out: &mut [f32],
6872 pool: Option<&Pool>,
6873) {
6874 debug_assert_eq!(out.len(), rows);
6875 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6876 let gpr = cols / GROUP_SIZE;
6877 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6878 let out_addr = SendMut(out.as_mut_ptr());
6879 if a8w8_enabled() {
6883 let act = split_act(x);
6884 let act = &act;
6885 let run = move |start: usize, end: usize| {
6886 for r in start..end {
6887 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6888 for &(j, xv) in &act.outliers {
6889 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6890 }
6891 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6892 unsafe { *out_addr.at(r) = acc };
6894 }
6895 };
6896 dispatch_rows(pool, rows, &run);
6897 return;
6898 }
6899 let run = move |start: usize, end: usize| {
6900 let mut sg = [0f32; GROUP_SIZE];
6904 for r in start..end {
6905 let mut acc = 0f32;
6906 for g in 0..gpr {
6907 let off = (r * gpr + g) * TILE;
6908 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6909 let codes = &bytes[off + 2..off + TILE];
6910 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6911 for bi in 0..6 {
6912 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6913 }
6914 let lut = &SIGN5[codes[6] as usize];
6915 sg[30] = lut[0];
6916 sg[31] = lut[1];
6917 let mut gsum = 0f32;
6918 for k in 0..GROUP_SIZE {
6919 gsum += sg[k] * xg[k];
6920 }
6921 acc += s * gsum;
6922 }
6923 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6924 unsafe { *out_addr.at(r) = acc };
6925 }
6926 };
6927 dispatch_rows(pool, rows, &run);
6928}
6929
6930#[cfg(target_arch = "aarch64")]
6936#[target_feature(enable = "neon,dotprod")]
6937unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6938 use core::arch::aarch64::*;
6939 use core::arch::asm;
6940 unsafe {
6942 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6943 let bytes_ptr = bytes.as_ptr();
6944 let row_off = r * gpr * TILE;
6945 let xp = [xa.as_ptr(), xb.as_ptr()];
6946 let mut acc = [0f32; 2];
6947 macro_rules! sdot2 {
6948 ($w0:expr, $w1:expr, $x:expr) => {{
6949 let x0 = vld1q_s8($x);
6950 let x1 = vld1q_s8($x.add(16));
6951 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6952 asm!(
6953 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6954 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6955 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6956 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6957 options(pure, nomem, nostack),
6958 );
6959 vaddvq_s32(vaddq_s32(a0, a1))
6960 }};
6961 }
6962 let gpr2 = gpr & !1;
6963 let mut gi = 0;
6964 while gi < gpr2 {
6965 let off0 = row_off + gi * TILE;
6966 let off1 = off0 + TILE;
6967 let s0 = f16_to_f32(u16::from_le_bytes([
6968 *bytes_ptr.add(off0),
6969 *bytes_ptr.add(off0 + 1),
6970 ]));
6971 let s1 = f16_to_f32(u16::from_le_bytes([
6972 *bytes_ptr.add(off1),
6973 *bytes_ptr.add(off1 + 1),
6974 ]));
6975 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6976 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6977 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6978 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6979 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6980 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6981 for k in 0..2 {
6982 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
6983 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
6984 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
6985 }
6986 gi += 2;
6987 }
6988 if gi < gpr {
6989 let off = row_off + gi * TILE;
6990 let s = f16_to_f32(u16::from_le_bytes([
6991 *bytes_ptr.add(off),
6992 *bytes_ptr.add(off + 1),
6993 ]));
6994 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6995 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6996 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6997 for k in 0..2 {
6998 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
6999 acc[k] += d as f32 * s;
7000 }
7001 }
7002 acc
7003 }
7004}
7005
7006fn q1t_matvec2(
7012 bytes: &[u8],
7013 x1: &[f32],
7014 x2: &[f32],
7015 rows: usize,
7016 cols: usize,
7017 o1: &mut [f32],
7018 o2: &mut [f32],
7019 pool: Option<&Pool>,
7020) {
7021 debug_assert_eq!(o1.len(), rows);
7022 debug_assert_eq!(o2.len(), rows);
7023 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7024 let gpr = cols / GROUP_SIZE;
7025 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7026 let out1 = SendMut(o1.as_mut_ptr());
7027 let out2 = SendMut(o2.as_mut_ptr());
7028 if a8w8_enabled() {
7029 let a1 = split_act(x1);
7030 let a2 = split_act(x2);
7031 let (a1, a2) = (&a1, &a2);
7032 let run = move |start: usize, end: usize| {
7033 for r in start..end {
7034 #[cfg(target_arch = "aarch64")]
7035 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7038 #[cfg(not(target_arch = "aarch64"))]
7039 let ds = [
7040 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7041 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7042 ];
7043 let mut acc1 = ds[0] * a1.sx;
7044 for &(j, xv) in &a1.outliers {
7045 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7046 }
7047 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7048 let mut acc2 = ds[1] * a2.sx;
7049 for &(j, xv) in &a2.outliers {
7050 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7051 }
7052 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7053 unsafe {
7055 *out1.at(r) = acc1;
7056 *out2.at(r) = acc2;
7057 }
7058 }
7059 };
7060 dispatch_rows(pool, rows, &run);
7061 return;
7062 }
7063 let run = move |start: usize, end: usize| {
7064 let mut sg = [0f32; GROUP_SIZE];
7067 for r in start..end {
7068 let mut acc1 = 0f32;
7069 let mut acc2 = 0f32;
7070 for g in 0..gpr {
7071 let off = (r * gpr + g) * TILE;
7072 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7073 let codes = &bytes[off + 2..off + TILE];
7074 for bi in 0..6 {
7075 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7076 }
7077 let lut = &SIGN5[codes[6] as usize];
7078 sg[30] = lut[0];
7079 sg[31] = lut[1];
7080 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7081 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7082 let mut gsum1 = 0f32;
7083 for k in 0..GROUP_SIZE {
7084 gsum1 += sg[k] * xg1[k];
7085 }
7086 acc1 += s * gsum1;
7087 let mut gsum2 = 0f32;
7088 for k in 0..GROUP_SIZE {
7089 gsum2 += sg[k] * xg2[k];
7090 }
7091 acc2 += s * gsum2;
7092 }
7093 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7094 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7095 unsafe {
7097 *out1.at(r) = acc1;
7098 *out2.at(r) = acc2;
7099 }
7100 }
7101 };
7102 dispatch_rows(pool, rows, &run);
7103}
7104
7105fn q1t_matmat(
7108 bytes: &[u8],
7109 xs: &[f32],
7110 b: usize,
7111 rows: usize,
7112 cols: usize,
7113 out: &mut [f32],
7114 pool: Option<&Pool>,
7115) {
7116 debug_assert_eq!(out.len(), b * rows);
7117 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7118 let gpr = cols / GROUP_SIZE;
7119 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7120 let out_addr = SendMut(out.as_mut_ptr());
7121 if a8w8_enabled() {
7125 let acts: Vec<SplitAct> = (0..b)
7126 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7127 .collect();
7128 let acts = &acts;
7129 let run = move |start: usize, end: usize| {
7130 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
7134 for g in 0..gpr {
7135 let off = (r * gpr + g) * TILE;
7136 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7137 q1t_unpack_group_i8(
7138 bytes.as_ptr().wrapping_add(off + 2),
7139 &mut sg[g * GROUP_SIZE..],
7140 );
7141 }
7142 for bi in 0..b {
7143 let act = &acts[bi];
7144 let mut isum = 0f32;
7145 for g in 0..gpr {
7146 let d = q1t_i8dot32(
7147 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7148 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7149 );
7150 isum += d as f32 * sc[g];
7151 }
7152 let mut acc = isum * act.sx;
7153 for &(j, xv) in &act.outliers {
7154 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7155 }
7156 accs[bi] = acc;
7157 }
7158 if has_ov {
7162 let (c0, c1) = (
7163 q1t_rowptr(bytes, rp_off, r),
7164 q1t_rowptr(bytes, rp_off, r + 1),
7165 );
7166 for p in c0..c1 {
7167 let e = ent_off + p * 4;
7168 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7169 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7170 for bi in 0..b {
7171 accs[bi] += val * xs[bi * cols + col];
7172 }
7173 }
7174 }
7175 for bi in 0..b {
7176 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7177 }
7178 }
7179 };
7180 dispatch_rows(pool, rows, &run);
7181 return;
7182 }
7183 let run = move |start: usize, end: usize| {
7184 let mut buf = vec![0f32; cols];
7185 for r in start..end {
7186 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7187 for bi in 0..b {
7188 let xr = &xs[bi * cols..(bi + 1) * cols];
7189 let mut acc = 0f32;
7190 for j in 0..cols {
7191 acc += buf[j] * xr[j];
7192 }
7193 unsafe { *out_addr.at(bi * rows + r) = acc };
7194 }
7195 }
7196 };
7197 dispatch_rows(pool, rows, &run);
7198}
7199
7200fn q1_matvec(
7201 bytes: &[u8],
7202 x: &[f32],
7203 rows: usize,
7204 cols: usize,
7205 out: &mut [f32],
7206 pool: Option<&Pool>,
7207) {
7208 debug_assert_eq!(out.len(), rows);
7209 let gpr = cols / GROUP_SIZE;
7210 let out_addr = SendMut(out.as_mut_ptr());
7211 if a8w8_enabled() {
7212 let act = split_act(x);
7213 let gsum = q1_group_sums(&act.xq, gpr);
7214 let (act, gsum) = (&act, &gsum);
7215 let run = move |start: usize, end: usize| {
7216 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7217 };
7218 dispatch_rows(pool, rows, &run);
7219 return;
7220 }
7221 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7222 dispatch_rows(pool, rows, &run);
7223}
7224
7225#[allow(clippy::too_many_arguments)]
7227fn q1_matvec2(
7228 bytes: &[u8],
7229 x1: &[f32],
7230 x2: &[f32],
7231 rows: usize,
7232 cols: usize,
7233 o1: &mut [f32],
7234 o2: &mut [f32],
7235 pool: Option<&Pool>,
7236) {
7237 let gpr = cols / GROUP_SIZE;
7238 let p1 = SendMut(o1.as_mut_ptr());
7239 let p2 = SendMut(o2.as_mut_ptr());
7240 if a8w8_enabled() {
7241 let a1 = split_act(x1);
7242 let a2 = split_act(x2);
7243 let g1 = q1_group_sums(&a1.xq, gpr);
7244 let g2 = q1_group_sums(&a2.xq, gpr);
7245 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7246 let run = move |start: usize, end: usize| {
7247 for r in start..end {
7248 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7249 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7250 for &(j, xv) in &a1.outliers {
7251 let (w, s) = q1_outlier(bytes, r, gpr, j);
7252 v1 += w * s * xv;
7253 }
7254 for &(j, xv) in &a2.outliers {
7255 let (w, s) = q1_outlier(bytes, r, gpr, j);
7256 v2 += w * s * xv;
7257 }
7258 unsafe {
7260 *p1.at(r) = v1;
7261 *p2.at(r) = v2;
7262 }
7263 }
7264 };
7265 dispatch_rows(pool, rows, &run);
7266 return;
7267 }
7268 let run = move |start: usize, end: usize| {
7269 for r in start..end {
7270 unsafe {
7272 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7273 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7274 }
7275 }
7276 };
7277 dispatch_rows(pool, rows, &run);
7278}
7279
7280#[allow(clippy::too_many_arguments)]
7282fn q1_matmat(
7283 bytes: &[u8],
7284 xs_all: &[f32],
7285 b: usize,
7286 rows: usize,
7287 cols: usize,
7288 out: &mut [f32],
7289 pool: Option<&Pool>,
7290) {
7291 debug_assert_eq!(out.len(), b * rows);
7292 let gpr = cols / GROUP_SIZE;
7293 let out_addr = SendMut(out.as_mut_ptr());
7294 if a8w8_enabled() {
7295 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7296 .map(|bi| {
7297 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7298 let gsum = q1_group_sums(&act.xq, gpr);
7299 (act, gsum)
7300 })
7301 .collect();
7302 let acts = &acts;
7303 #[cfg(target_arch = "x86_64")]
7304 let blocked_ok = avx2_enabled()
7305 && blocked_enabled();
7306 #[cfg(target_arch = "aarch64")]
7307 let blocked_ok = sdot_enabled()
7308 && blocked_enabled();
7309 let run = move |start: usize, end: usize| {
7310 for r in start..end {
7311 let mut bi = 0usize;
7312 #[cfg(target_arch = "aarch64")]
7315 if blocked_ok {
7316 while bi + 4 <= acts.len() {
7317 let xs = [
7318 acts[bi].0.xq.as_slice(),
7319 acts[bi + 1].0.xq.as_slice(),
7320 acts[bi + 2].0.xq.as_slice(),
7321 acts[bi + 3].0.xq.as_slice(),
7322 ];
7323 let gs = [
7324 acts[bi].1.as_slice(),
7325 acts[bi + 1].1.as_slice(),
7326 acts[bi + 2].1.as_slice(),
7327 acts[bi + 3].1.as_slice(),
7328 ];
7329 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7330 for k in 0..4 {
7331 let (act, _) = &acts[bi + k];
7332 let mut acc = d[k] * act.sx;
7333 for &(j, xv) in &act.outliers {
7334 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7335 acc += w * sc * xv;
7336 }
7337 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7339 }
7340 bi += 4;
7341 }
7342 }
7343 #[cfg(target_arch = "x86_64")]
7344 if blocked_ok {
7345 while bi + 4 <= acts.len() {
7346 let xs = [
7347 acts[bi].0.xq.as_slice(),
7348 acts[bi + 1].0.xq.as_slice(),
7349 acts[bi + 2].0.xq.as_slice(),
7350 acts[bi + 3].0.xq.as_slice(),
7351 ];
7352 let gs = [
7353 acts[bi].1.as_slice(),
7354 acts[bi + 1].1.as_slice(),
7355 acts[bi + 2].1.as_slice(),
7356 acts[bi + 3].1.as_slice(),
7357 ];
7358 let d = unsafe {
7359 if vnni_tiles_enabled() {
7360 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7361 } else {
7362 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7363 }
7364 };
7365 for k in 0..4 {
7366 let (act, _) = &acts[bi + k];
7367 let mut acc = d[k] * act.sx;
7368 for &(j, xv) in &act.outliers {
7369 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7370 acc += w * sc * xv;
7371 }
7372 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7374 }
7375 bi += 4;
7376 }
7377 }
7378 while bi < acts.len() {
7379 let (act, gsum) = &acts[bi];
7380 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7381 for &(j, xv) in &act.outliers {
7382 let (w, s) = q1_outlier(bytes, r, gpr, j);
7383 acc += w * s * xv;
7384 }
7385 unsafe { *out_addr.at(bi * rows + r) = acc };
7387 bi += 1;
7388 }
7389 }
7390 };
7391 dispatch_rows(pool, rows, &run);
7392 return;
7393 }
7394 let run = move |start: usize, end: usize| {
7395 for r in start..end {
7396 for bi in 0..b {
7397 let x = &xs_all[bi * cols..(bi + 1) * cols];
7398 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7400 }
7401 }
7402 };
7403 dispatch_rows(pool, rows, &run);
7404}
7405
7406fn q4matvec(
7412 bytes: &[u8],
7413 x: &[f32],
7414 rows: usize,
7415 cols: usize,
7416 out: &mut [f32],
7417 pool: Option<&Pool>,
7418) {
7419 debug_assert_eq!(out.len(), rows);
7420 let (packed, scales) = q4_split(bytes, rows, cols);
7421 let gpr = cols / GROUP_SIZE;
7422 let out_addr = SendMut(out.as_mut_ptr());
7423
7424 if a8w8_enabled() {
7425 let act = split_act(x);
7426 let run = move |start: usize, end: usize| {
7427 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7428 };
7429 dispatch_rows(pool, rows, &run);
7430 return;
7431 }
7432
7433 let run =
7434 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7435 dispatch_rows(pool, rows, &run);
7436}
7437
7438#[inline]
7441#[allow(unreachable_code)]
7442#[cfg(target_arch = "x86_64")]
7447#[target_feature(enable = "avx2")]
7448unsafe fn dot_q4b_row_1x4_avx2(
7449 buf: &[u8],
7450 scales: &[u8],
7451 g0: usize,
7452 gpr: usize,
7453 xs: [&[i8]; 4],
7454) -> [f32; 4] {
7455 unsafe {
7457 use core::arch::x86_64::*;
7458 let ones = _mm256_set1_epi16(1);
7459 let mut acc = [0f32; 4];
7460 for gi in 0..gpr {
7461 let s = f16_to_f32(u16::from_le_bytes([
7462 scales[(g0 + gi) * 2],
7463 scales[(g0 + gi) * 2 + 1],
7464 ]));
7465 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7466 let aw = _mm256_abs_epi8(w);
7467 for (k, xq) in xs.iter().enumerate() {
7468 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7469 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7470 let d = _mm256_madd_epi16(p16, ones);
7471 let hi128 = _mm256_extracti128_si256::<1>(d);
7472 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7473 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7474 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7475 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7476 }
7477 }
7478 acc
7479 }
7480}
7481
7482#[cfg(target_arch = "x86_64")]
7484#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7485unsafe fn dot_q4b_row_1x4_vnni(
7486 buf: &[u8],
7487 scales: &[u8],
7488 g0: usize,
7489 gpr: usize,
7490 xs: [&[i8]; 4],
7491) -> [f32; 4] {
7492 unsafe {
7494 use core::arch::x86_64::*;
7495 let mut acc = [0f32; 4];
7496 for gi in 0..gpr {
7497 let s = f16_to_f32(u16::from_le_bytes([
7498 scales[(g0 + gi) * 2],
7499 scales[(g0 + gi) * 2 + 1],
7500 ]));
7501 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7502 let aw = _mm256_abs_epi8(w);
7503 for (k, xq) in xs.iter().enumerate() {
7504 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7505 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7506 acc[k] += d as f32 * s;
7507 }
7508 }
7509 acc
7510 }
7511}
7512
7513#[cfg(target_arch = "x86_64")]
7519#[target_feature(enable = "avx2")]
7520unsafe fn dot_q4b_row_1x4_sx_avx2(
7521 buf: &[u8],
7522 scales: &[u8],
7523 g0: usize,
7524 gpr: usize,
7525 xs: [&[i8]; 4],
7526 sxs: [f32; 4],
7527) -> [f32; 4] {
7528 unsafe {
7530 use core::arch::x86_64::*;
7531 let ones = _mm256_set1_epi16(1);
7532 let mut acc = [0f32; 4];
7533 for gi in 0..gpr {
7534 let s = f16_to_f32(u16::from_le_bytes([
7535 scales[(g0 + gi) * 2],
7536 scales[(g0 + gi) * 2 + 1],
7537 ]));
7538 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7539 let aw = _mm256_abs_epi8(w);
7540 for (k, xq) in xs.iter().enumerate() {
7541 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7542 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7543 let d = _mm256_madd_epi16(p16, ones);
7544 let hi128 = _mm256_extracti128_si256::<1>(d);
7545 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7546 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7547 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7548 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7549 }
7550 }
7551 acc
7552 }
7553}
7554
7555#[cfg(target_arch = "x86_64")]
7558#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7559unsafe fn dot_q4b_row_1x4_sx_vnni(
7560 buf: &[u8],
7561 scales: &[u8],
7562 g0: usize,
7563 gpr: usize,
7564 xs: [&[i8]; 4],
7565 sxs: [f32; 4],
7566) -> [f32; 4] {
7567 unsafe {
7569 use core::arch::x86_64::*;
7570 let mut acc = [0f32; 4];
7571 for gi in 0..gpr {
7572 let s = f16_to_f32(u16::from_le_bytes([
7573 scales[(g0 + gi) * 2],
7574 scales[(g0 + gi) * 2 + 1],
7575 ]));
7576 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7577 let aw = _mm256_abs_epi8(w);
7578 for (k, xq) in xs.iter().enumerate() {
7579 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7580 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7581 acc[k] += (d as f32 * sxs[k]) * s;
7582 }
7583 }
7584 acc
7585 }
7586}
7587
7588#[allow(unreachable_code)]
7589fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7590 #[cfg(target_arch = "aarch64")]
7591 unsafe {
7592 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7593 }
7594 #[cfg(target_arch = "x86_64")]
7595 unsafe {
7596 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7597 }
7598 let mut acc = 0f32;
7599 for gi in 0..gpr {
7600 let g = g0 + gi;
7601 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7602 let mut d = 0i32;
7603 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7604 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7605 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7606 }
7607 acc += d as f32 * s;
7608 }
7609 acc
7610}
7611
7612#[inline]
7614#[allow(unreachable_code)]
7615fn dot_q4_row_i8_2(
7616 packed: &[u8],
7617 scales: &[u8],
7618 g0: usize,
7619 gpr: usize,
7620 xq1: &[i8],
7621 xq2: &[i8],
7622) -> (f32, f32) {
7623 #[cfg(target_arch = "aarch64")]
7624 unsafe {
7625 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7626 }
7627 #[cfg(target_arch = "x86_64")]
7628 unsafe {
7629 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7630 }
7631 (
7632 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7633 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7634 )
7635}
7636
7637#[allow(clippy::too_many_arguments)]
7640fn q4_range_a8w8(
7641 packed: &[u8],
7642 scales: &[u8],
7643 gpr: usize,
7644 cols: usize,
7645 act: &SplitAct,
7646 out: SendMut,
7647 start: usize,
7648 end: usize,
7649) {
7650 for r in start..end {
7651 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7652 for &(j, xv) in &act.outliers {
7654 let flat = r * cols + j;
7655 let byte = packed[flat / 2];
7656 let nib = if flat & 1 == 0 {
7657 byte & 0x0F
7658 } else {
7659 byte >> 4
7660 };
7661 let s = f16_to_f32(u16::from_le_bytes([
7662 scales[(flat / GROUP_SIZE) * 2],
7663 scales[(flat / GROUP_SIZE) * 2 + 1],
7664 ]));
7665 acc += ((nib as i32 - 8) as f32) * s * xv;
7666 }
7667 unsafe { *out.at(r) = acc };
7669 }
7670}
7671
7672#[allow(clippy::too_many_arguments)]
7675fn q4_range2_a8w8(
7676 packed: &[u8],
7677 scales: &[u8],
7678 gpr: usize,
7679 cols: usize,
7680 a1: &SplitAct,
7681 a2: &SplitAct,
7682 p1: SendMut,
7683 p2: SendMut,
7684 start: usize,
7685 end: usize,
7686) {
7687 for r in start..end {
7688 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7689 let mut acc1 = s1 * a1.sx;
7690 let mut acc2 = s2 * a2.sx;
7691 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7693 for &(j, xv) in outliers {
7694 let flat = r * cols + j;
7695 let byte = packed[flat / 2];
7696 let nib = if flat & 1 == 0 {
7697 byte & 0x0F
7698 } else {
7699 byte >> 4
7700 };
7701 let s = f16_to_f32(u16::from_le_bytes([
7702 scales[(flat / GROUP_SIZE) * 2],
7703 scales[(flat / GROUP_SIZE) * 2 + 1],
7704 ]));
7705 *acc += ((nib as i32 - 8) as f32) * s * xv;
7706 }
7707 };
7708 fix(&a1.outliers, &mut acc1);
7709 fix(&a2.outliers, &mut acc2);
7710 unsafe {
7712 *p1.at(r) = acc1;
7713 *p2.at(r) = acc2;
7714 }
7715 }
7716}
7717
7718fn q4_range_f32(
7720 packed: &[u8],
7721 scales: &[u8],
7722 gpr: usize,
7723 x: &[f32],
7724 out: SendMut,
7725 start: usize,
7726 end: usize,
7727) {
7728 for r in start..end {
7729 let mut acc = 0f32;
7730 for gi in 0..gpr {
7731 let g = r * gpr + gi;
7732 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7733 let pk = &packed[g * 16..(g + 1) * 16];
7734 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7735 let mut ga = 0f32;
7736 for (k, &b) in pk.iter().enumerate() {
7737 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7738 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7739 }
7740 acc += ga * s;
7741 }
7742 unsafe { *out.at(r) = acc };
7744 }
7745}
7746
7747#[allow(clippy::too_many_arguments)]
7751fn q4matvec2(
7752 bytes: &[u8],
7753 x1: &[f32],
7754 x2: &[f32],
7755 rows: usize,
7756 cols: usize,
7757 o1: &mut [f32],
7758 o2: &mut [f32],
7759 pool: Option<&Pool>,
7760) {
7761 debug_assert_eq!(o1.len(), rows);
7762 debug_assert_eq!(o2.len(), rows);
7763 let (packed, scales) = q4_split(bytes, rows, cols);
7764 let gpr = cols / GROUP_SIZE;
7765
7766 if a8w8_enabled() {
7767 let a1 = split_act(x1);
7768 let a2 = split_act(x2);
7769 let p1 = SendMut(o1.as_mut_ptr());
7770 let p2 = SendMut(o2.as_mut_ptr());
7771 let run = move |start: usize, end: usize| {
7772 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7773 };
7774 dispatch_rows(pool, rows, &run);
7775 return;
7776 }
7777
7778 let p1 = SendMut(o1.as_mut_ptr());
7779 let p2 = SendMut(o2.as_mut_ptr());
7780 let run = move |start: usize, end: usize| {
7781 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7782 };
7783 dispatch_rows(pool, rows, &run);
7784}
7785
7786#[allow(clippy::too_many_arguments)]
7788fn q4_range2_f32(
7789 packed: &[u8],
7790 scales: &[u8],
7791 gpr: usize,
7792 x1: &[f32],
7793 x2: &[f32],
7794 p1: SendMut,
7795 p2: SendMut,
7796 start: usize,
7797 end: usize,
7798) {
7799 for r in start..end {
7800 let (mut acc1, mut acc2) = (0f32, 0f32);
7801 for gi in 0..gpr {
7802 let g = r * gpr + gi;
7803 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7804 let pk = &packed[g * 16..(g + 1) * 16];
7805 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7806 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7807 let (mut g1, mut g2) = (0f32, 0f32);
7808 for (k, &b) in pk.iter().enumerate() {
7809 let wl = (b & 0x0F) as f32 - 8.0;
7810 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7811 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7812 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7813 }
7814 acc1 += g1 * s;
7815 acc2 += g2 * s;
7816 }
7817 unsafe {
7819 *p1.at(r) = acc1;
7820 *p2.at(r) = acc2;
7821 }
7822 }
7823}
7824
7825thread_local! {
7826 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7829 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7830}
7831
7832#[allow(clippy::too_many_arguments)]
7838fn q4matmat(
7839 bytes: &[u8],
7840 xs_all: &[f32],
7841 b: usize,
7842 rows: usize,
7843 cols: usize,
7844 out: &mut [f32],
7845 pool: Option<&Pool>,
7846) {
7847 debug_assert_eq!(xs_all.len(), b * cols);
7848 debug_assert_eq!(out.len(), b * rows);
7849 let (packed, scales) = q4_split(bytes, rows, cols);
7850 let gpr = cols / GROUP_SIZE;
7851 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7852
7853 if a8w8_enabled() {
7854 let acts: Vec<SplitAct> = (0..b)
7855 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7856 .collect();
7857 let acts = &acts;
7858 let out_addr = SendMut(out.as_mut_ptr());
7859 let run = move |start: usize, end: usize| {
7860 ROW_I8.with(|rb| {
7861 let mut buf = rb.borrow_mut();
7862 buf.resize(cols, 0);
7863 for r in start..end {
7864 for gi in 0..gpr {
7868 let g = r * gpr + gi;
7869 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7870 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7871 buf[gi * GROUP_SIZE + k * 2 + 1] =
7872 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7873 }
7874 }
7875 let mut bi = 0usize;
7876 #[cfg(target_arch = "x86_64")]
7877 if avx2_enabled()
7878 && blocked_enabled()
7879 {
7880 while bi + 4 <= acts.len() {
7881 let xs = [
7882 acts[bi].xq.as_slice(),
7883 acts[bi + 1].xq.as_slice(),
7884 acts[bi + 2].xq.as_slice(),
7885 acts[bi + 3].xq.as_slice(),
7886 ];
7887 let d = unsafe {
7888 if vnni_tiles_enabled() {
7889 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7890 } else {
7891 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7892 }
7893 };
7894 for k in 0..4 {
7895 let act = &acts[bi + k];
7896 let mut acc = d[k] * act.sx;
7897 for &(j, xv) in &act.outliers {
7898 acc += (buf[j] as i8) as f32
7899 * gscale((r * cols + j) / GROUP_SIZE)
7900 * xv;
7901 }
7902 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7904 }
7905 bi += 4;
7906 }
7907 }
7908 while bi < acts.len() {
7909 let act = &acts[bi];
7910 let mut acc = 0f32;
7911 for gi in 0..gpr {
7912 let d = dot_i8_i8(
7913 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7914 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7915 );
7916 acc += d as f32 * gscale(r * gpr + gi);
7917 }
7918 acc *= act.sx;
7919 for &(j, xv) in &act.outliers {
7921 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7922 }
7923 unsafe { *out_addr.at(bi * rows + r) = acc };
7925 bi += 1;
7926 }
7927 }
7928 })
7929 };
7930 dispatch_rows(pool, rows, &run);
7931 return;
7932 }
7933
7934 let out_addr = SendMut(out.as_mut_ptr());
7935 let run = move |start: usize, end: usize| {
7936 ROW_F32.with(|rb| {
7937 let mut buf = rb.borrow_mut();
7938 buf.resize(cols, 0.0);
7939 for r in start..end {
7940 for gi in 0..gpr {
7943 let g = r * gpr + gi;
7944 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7945 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7946 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7947 }
7948 }
7949 for bi in 0..b {
7950 let x = &xs_all[bi * cols..(bi + 1) * cols];
7951 let mut acc = 0f32;
7952 for gi in 0..gpr {
7953 let mut ga = 0f32;
7954 for k in 0..GROUP_SIZE / 2 {
7959 let e = gi * GROUP_SIZE + k * 2;
7960 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
7961 }
7962 acc += ga * gscale(r * gpr + gi);
7963 }
7964 unsafe { *out_addr.at(bi * rows + r) = acc };
7966 }
7967 }
7968 })
7969 };
7970 dispatch_rows(pool, rows, &run);
7971}
7972
7973#[allow(clippy::too_many_arguments)]
7978fn vbitmatmat(
7979 bytes: &[u8],
7980 offsets: &[usize],
7981 xs_all: &[f32],
7982 b: usize,
7983 rows: usize,
7984 cols: usize,
7985 out: &mut [f32],
7986 pool: Option<&Pool>,
7987) {
7988 debug_assert_eq!(xs_all.len(), b * cols);
7989 debug_assert_eq!(out.len(), b * rows);
7990 debug_assert_eq!(offsets.len(), rows + 1);
7991 let ng = cols / GROUP_SIZE;
7992 let bits = &bytes[..rows];
7993 let sc_off = rows;
7994 let gscale = |r: usize, g: usize| {
7995 let so = (r * ng + g) * 2;
7996 f16_to_f32(u16::from_le_bytes([
7997 bytes[sc_off + so],
7998 bytes[sc_off + so + 1],
7999 ]))
8000 };
8001
8002 let decode_f32 = |r: usize, dst: &mut [f32]| {
8004 let bw = bits[r] as usize;
8005 let l = ((1i32 << (bw - 1)) - 1) as f32;
8006 let data = &bytes[offsets[r]..offsets[r + 1]];
8007 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8008 for d in dst.iter_mut() {
8009 while nbits < bw {
8010 acc = (acc << 8) | data[idx] as u64;
8011 idx += 1;
8012 nbits += 8;
8013 }
8014 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8015 nbits -= bw;
8016 *d = u - l;
8017 }
8018 };
8019
8020 if a8w8_enabled() {
8021 let acts: Vec<SplitAct> = (0..b)
8022 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8023 .collect();
8024 let acts = &acts;
8025 let out_addr = SendMut(out.as_mut_ptr());
8026 let run = move |start: usize, end: usize| {
8027 for r in start..end {
8028 let bw = bits[r] as usize;
8029 if bw == 8 {
8030 ROW_F32.with(|rb| {
8033 let mut buf = rb.borrow_mut();
8034 buf.resize(cols, 0.0);
8035 decode_f32(r, &mut buf);
8036 for bi in 0..b {
8037 let x = &xs_all[bi * cols..(bi + 1) * cols];
8038 let mut dot = 0f32;
8039 for g in 0..ng {
8040 let mut gd = 0f32;
8041 for k in 0..GROUP_SIZE {
8042 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8043 }
8044 dot += gd * gscale(r, g);
8045 }
8046 unsafe { *out_addr.at(bi * rows + r) = dot };
8048 }
8049 });
8050 continue;
8051 }
8052 let l = (1i32 << (bw - 1)) - 1;
8053 let data = &bytes[offsets[r]..offsets[r + 1]];
8054 ROW_I8.with(|rb| {
8055 let mut buf = rb.borrow_mut();
8056 buf.resize(cols, 0);
8057 #[inline(always)]
8058 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8059 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8060 let u = unpack8::<B>(&data[blk * B..]);
8061 for k in 0..8 {
8062 chunk[k] = (u[k] - l) as i8 as u8;
8063 }
8064 }
8065 }
8066 match bw {
8067 3 => fill::<3>(data, l, &mut buf),
8068 4 => vbit_fill4(data, &mut buf),
8069 5 => fill::<5>(data, l, &mut buf),
8070 6 => fill::<6>(data, l, &mut buf),
8071 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8072 }
8073 let mut bi = 0usize;
8074 #[cfg(target_arch = "x86_64")]
8078 if avx2_enabled()
8079 && blocked_enabled()
8080 {
8081 while bi + 4 <= acts.len() {
8082 let xs = [
8083 acts[bi].xq.as_slice(),
8084 acts[bi + 1].xq.as_slice(),
8085 acts[bi + 2].xq.as_slice(),
8086 acts[bi + 3].xq.as_slice(),
8087 ];
8088 let sxs = [
8089 acts[bi].sx,
8090 acts[bi + 1].sx,
8091 acts[bi + 2].sx,
8092 acts[bi + 3].sx,
8093 ];
8094 let d = unsafe {
8095 if vnni_tiles_enabled() {
8096 dot_q4b_row_1x4_sx_vnni(
8097 &buf,
8098 &bytes[sc_off..],
8099 r * ng,
8100 ng,
8101 xs,
8102 sxs,
8103 )
8104 } else {
8105 dot_q4b_row_1x4_sx_avx2(
8106 &buf,
8107 &bytes[sc_off..],
8108 r * ng,
8109 ng,
8110 xs,
8111 sxs,
8112 )
8113 }
8114 };
8115 for k in 0..4 {
8116 let act = &acts[bi + k];
8117 let mut dot = d[k];
8118 for &(j, xv) in &act.outliers {
8119 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8120 }
8121 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8123 }
8124 bi += 4;
8125 }
8126 }
8127 while bi < acts.len() {
8128 let act = &acts[bi];
8129 let mut dot = 0f32;
8130 for g in 0..ng {
8131 let d = dot_i8_i8(
8132 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8133 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8134 ) as f32
8135 * act.sx;
8136 dot += d * gscale(r, g);
8137 }
8138 for &(j, xv) in &act.outliers {
8139 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8140 }
8141 unsafe { *out_addr.at(bi * rows + r) = dot };
8143 bi += 1;
8144 }
8145 });
8146 }
8147 };
8148 dispatch_rows(pool, rows, &run);
8149 return;
8150 }
8151
8152 let out_addr = SendMut(out.as_mut_ptr());
8153 let run = move |start: usize, end: usize| {
8154 ROW_F32.with(|rb| {
8155 let mut buf = rb.borrow_mut();
8156 buf.resize(cols, 0.0);
8157 for r in start..end {
8158 decode_f32(r, &mut buf);
8159 for bi in 0..b {
8160 let x = &xs_all[bi * cols..(bi + 1) * cols];
8161 let mut dot = 0f32;
8162 for g in 0..ng {
8163 let mut gd = 0f32;
8164 for k in 0..GROUP_SIZE {
8165 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8166 }
8167 dot += gd * gscale(r, g);
8168 }
8169 unsafe { *out_addr.at(bi * rows + r) = dot };
8171 }
8172 }
8173 })
8174 };
8175 dispatch_rows(pool, rows, &run);
8176}
8177
8178pub(crate) fn gpu_batch_job<'a>(
8182 t: &'a QTensor,
8183 x: &[f32],
8184) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8185 match t {
8186 QTensor::Mapped {
8187 model,
8188 idx,
8189 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8190 rows,
8191 cols,
8192 row_scale,
8193 col_field,
8194 ..
8195 } => Some((
8196 model.clone(),
8197 crate::gpu::BatchJob {
8198 idx: *idx,
8199 rows: *rows,
8200 cols: *cols,
8201 row_scale,
8202 xs: prescale(x, col_field, *dt).into_owned(),
8203 layout: crate::gpu::BatchLayout::Q8,
8204 },
8205 )),
8206 QTensor::Mapped {
8208 model,
8209 idx,
8210 dtype: TensorDtype::Q1,
8211 rows,
8212 cols,
8213 ..
8214 } => Some((
8215 model.clone(),
8216 crate::gpu::BatchJob {
8217 idx: *idx,
8218 rows: *rows,
8219 cols: *cols,
8220 row_scale: &[],
8221 xs: x.to_vec(),
8222 layout: crate::gpu::BatchLayout::Q1,
8223 },
8224 )),
8225 QTensor::Mapped {
8230 model,
8231 idx,
8232 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8233 rows,
8234 cols,
8235 ..
8236 } => Some((
8237 model.clone(),
8238 crate::gpu::BatchJob {
8239 idx: *idx,
8240 rows: *rows,
8241 cols: *cols,
8242 row_scale: &[],
8243 xs: x.to_vec(),
8244 layout: if *dt == TensorDtype::Q4Tiled {
8245 crate::gpu::BatchLayout::Q4t
8246 } else {
8247 crate::gpu::BatchLayout::Q4tp
8248 },
8249 },
8250 )),
8251 _ => None,
8252 }
8253}
8254
8255thread_local! {
8256 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8257 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8258}
8259
8260pub(crate) fn prescale<'a>(
8261 x: &'a [f32],
8262 col_field: &[f32],
8263 dtype: TensorDtype,
8264) -> std::borrow::Cow<'a, [f32]> {
8265 if dtype == TensorDtype::Q8_2f {
8266 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8267 } else {
8268 std::borrow::Cow::Borrowed(x)
8269 }
8270}
8271
8272pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8275 x: &[f32],
8276 col_field: &[f32],
8277 dtype: TensorDtype,
8278 buf_id: u8,
8279 f: F,
8280) -> R {
8281 if dtype == TensorDtype::Q8_2f {
8282 if buf_id == 1 {
8283 PRESCALE_BUF1.with(|b| {
8284 let mut buf = b.borrow_mut();
8285 buf.clear();
8286 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8287 f(&buf)
8288 })
8289 } else {
8290 PRESCALE_BUF2.with(|b| {
8291 let mut buf = b.borrow_mut();
8292 buf.clear();
8293 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8294 f(&buf)
8295 })
8296 }
8297 } else {
8298 f(x)
8299 }
8300}
8301
8302#[cfg(target_arch = "x86_64")]
8307pub(crate) fn avx2_enabled() -> bool {
8308 use std::sync::OnceLock;
8309 static ON: OnceLock<bool> = OnceLock::new();
8310 *ON.get_or_init(|| {
8311 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8312 && std::arch::is_x86_feature_detected!("avx2")
8313 && std::arch::is_x86_feature_detected!("fma")
8314 })
8315}
8316
8317#[cfg(target_arch = "x86_64")]
8322fn avx2_a8w8_enabled() -> bool {
8323 use std::sync::OnceLock;
8324 static ON: OnceLock<bool> = OnceLock::new();
8325 *ON.get_or_init(|| {
8326 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8327 })
8328}
8329
8330#[inline]
8334pub(crate) fn a8w8_enabled() -> bool {
8335 #[cfg(target_arch = "aarch64")]
8336 {
8337 sdot_enabled()
8338 }
8339 #[cfg(target_arch = "x86_64")]
8340 {
8341 avx2_a8w8_enabled()
8342 }
8343 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8344 {
8345 false
8346 }
8347}
8348
8349#[inline]
8352#[allow(unreachable_code)]
8353fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8354 #[cfg(target_arch = "aarch64")]
8355 unsafe {
8356 return dot_i8_sdot(w, xq);
8357 }
8358 #[cfg(target_arch = "x86_64")]
8359 unsafe {
8360 if avx512vnni_enabled() {
8361 return dot_i8_i8_vnni(w, xq);
8362 }
8363 return dot_i8_i8_avx2(w, xq);
8364 }
8365 w.iter()
8366 .zip(xq)
8367 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8368 .sum()
8369}
8370
8371#[cfg(target_arch = "x86_64")]
8375fn avx512vnni_enabled() -> bool {
8376 use std::sync::OnceLock;
8377 static ON: OnceLock<bool> = OnceLock::new();
8378 *ON.get_or_init(|| {
8379 std::env::var("CMF_AVX512")
8380 .map(|v| v != "0")
8381 .unwrap_or(true)
8382 && std::arch::is_x86_feature_detected!("avx512f")
8383 && std::arch::is_x86_feature_detected!("avx512bw")
8384 && std::arch::is_x86_feature_detected!("avx512vl")
8385 && std::arch::is_x86_feature_detected!("avx512vnni")
8386 })
8387}
8388
8389#[cfg(target_arch = "x86_64")]
8397fn vnni_tiles_enabled() -> bool {
8398 use std::sync::OnceLock;
8399 static ON: OnceLock<bool> = OnceLock::new();
8400 *ON.get_or_init(|| {
8401 std::env::var("CMF_VNNI_TILES")
8402 .map(|v| v != "0")
8403 .unwrap_or(true)
8404 && avx512vnni_enabled()
8405 })
8406}
8407
8408#[cfg(target_arch = "x86_64")]
8413#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8414#[inline]
8415unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8416 unsafe {
8418 use core::arch::x86_64::*;
8419 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8420 let hi128 = _mm256_extracti128_si256::<1>(d);
8421 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8422 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8423 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8424 _mm_cvtsi128_si32(s32)
8425 }
8426}
8427
8428#[cfg(target_arch = "x86_64")]
8433#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8434unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8435 unsafe {
8437 use core::arch::x86_64::*;
8438 let n = w.len();
8439 let mut j = 0usize;
8440 let mut total: i32;
8441 {
8446 #[inline(always)]
8447 unsafe fn step(
8448 w: *const u8,
8449 x: *const i8,
8450 acc: core::arch::x86_64::__m512i,
8451 ) -> core::arch::x86_64::__m512i {
8452 unsafe {
8453 use core::arch::x86_64::*;
8454 let wv = _mm512_loadu_si512(w as *const _);
8455 let xv = _mm512_loadu_si512(x as *const _);
8456 let aw = _mm512_abs_epi8(wv);
8457 let neg = _mm512_movepi8_mask(wv);
8458 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8459 _mm512_dpbusd_epi32(acc, aw, sx)
8460 }
8461 }
8462 let (mut a0, mut a1, mut a2, mut a3) = (
8463 _mm512_setzero_si512(),
8464 _mm512_setzero_si512(),
8465 _mm512_setzero_si512(),
8466 _mm512_setzero_si512(),
8467 );
8468 while j + 256 <= n {
8469 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8470 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8471 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8472 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8473 j += 256;
8474 }
8475 while j + 64 <= n {
8476 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8477 j += 64;
8478 }
8479 let s01 = _mm512_add_epi32(a0, a1);
8480 let s23 = _mm512_add_epi32(a2, a3);
8481 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8482 }
8483 if j + 32 <= n {
8485 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8486 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8487 let d = _mm256_dpbusd_epi32(
8488 _mm256_setzero_si256(),
8489 _mm256_abs_epi8(wv),
8490 _mm256_sign_epi8(xv, wv),
8491 );
8492 let hi128 = _mm256_extracti128_si256::<1>(d);
8493 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8494 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8495 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8496 total += _mm_cvtsi128_si32(s32);
8497 j += 32;
8498 }
8499 while j < n {
8500 total += (w[j] as i8) as i32 * xq[j] as i32;
8501 j += 1;
8502 }
8503 total
8504 }
8505}
8506
8507#[cfg(target_arch = "x86_64")]
8509#[target_feature(enable = "avx2,fma")]
8510unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8511 unsafe {
8513 use core::arch::x86_64::*;
8514 let n = x.len();
8515 let wp = w.as_ptr();
8516 let xp = x.as_ptr();
8517 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8518 let mut j = 0usize;
8519 while j + 16 <= n {
8520 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8521 let lo = _mm256_cvtepi8_epi32(wb);
8522 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8523 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8524 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8525 j += 16;
8526 }
8527 let acc = _mm256_add_ps(a0, a1);
8528 let hi128 = _mm256_extractf128_ps::<1>(acc);
8529 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8530 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8531 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8532 let mut sum = _mm_cvtss_f32(s32);
8533 while j < n {
8534 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8535 j += 1;
8536 }
8537 sum
8538 }
8539}
8540
8541#[cfg(target_arch = "x86_64")]
8546#[target_feature(enable = "avx2")]
8547unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8548 unsafe {
8550 use core::arch::x86_64::*;
8551 let n = w.len();
8552 let ones = _mm256_set1_epi16(1);
8553 let mut acc = _mm256_setzero_si256();
8554 let mut j = 0usize;
8555 while j + 32 <= n {
8556 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8557 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8558 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8559 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8560 j += 32;
8561 }
8562 let hi128 = _mm256_extracti128_si256::<1>(acc);
8563 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8564 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8565 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8566 let mut s = _mm_cvtsi128_si32(s32);
8567 while j < n {
8568 s += (w[j] as i8) as i32 * xq[j] as i32;
8569 j += 1;
8570 }
8571 s
8572 }
8573}
8574
8575#[cfg(target_arch = "aarch64")]
8579#[target_feature(enable = "neon,i8mm")]
8580unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8581 unsafe {
8583 use core::arch::aarch64::*;
8584 use core::arch::asm;
8585 let n = w0.len();
8586 let w0p = w0.as_ptr() as *const i8;
8587 let w1p = w1.as_ptr() as *const i8;
8588 let mut acc01 = vdupq_n_s32(0);
8591 let mut acc23 = vdupq_n_s32(0);
8592 let mut i = 0usize;
8593 while i + 8 <= n {
8594 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8595 let xb01 = vcombine_s8(
8596 vld1_s8(xs[0].as_ptr().add(i)),
8597 vld1_s8(xs[1].as_ptr().add(i)),
8598 );
8599 let xb23 = vcombine_s8(
8600 vld1_s8(xs[2].as_ptr().add(i)),
8601 vld1_s8(xs[3].as_ptr().add(i)),
8602 );
8603 asm!(
8604 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8605 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8606 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8607 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8608 options(pure, nomem, nostack),
8609 );
8610 i += 8;
8611 }
8612 let mut out = [[0i32; 4]; 2];
8613 let a01: [i32; 4] = core::mem::transmute(acc01);
8614 let a23: [i32; 4] = core::mem::transmute(acc23);
8615 out[0][0] = a01[0];
8616 out[0][1] = a01[1];
8617 out[1][0] = a01[2];
8618 out[1][1] = a01[3];
8619 out[0][2] = a23[0];
8620 out[0][3] = a23[1];
8621 out[1][2] = a23[2];
8622 out[1][3] = a23[3];
8623 if i < n {
8624 for (k, x) in xs.iter().enumerate() {
8625 for j in i..n {
8626 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8627 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8628 }
8629 }
8630 }
8631 out
8632 }
8633}
8634
8635#[cfg(target_arch = "aarch64")]
8639#[target_feature(enable = "neon,dotprod")]
8640unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8641 unsafe {
8643 use core::arch::aarch64::*;
8644 use core::arch::asm;
8645 let n = w0.len();
8646 let w0p = w0.as_ptr() as *const i8;
8647 let w1p = w1.as_ptr() as *const i8;
8648 let mut acc = [[vdupq_n_s32(0); 4]; 2];
8649 let mut i = 0usize;
8650 while i + 16 <= n {
8651 let wv0 = vld1q_s8(w0p.add(i));
8652 let wv1 = vld1q_s8(w1p.add(i));
8653 for (k, x) in xs.iter().enumerate() {
8654 let xv = vld1q_s8(x.as_ptr().add(i));
8655 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8656 asm!(
8657 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8658 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8659 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8660 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8661 options(pure, nomem, nostack),
8662 );
8663 acc[0][k] = a0;
8664 acc[1][k] = a1;
8665 }
8666 i += 16;
8667 }
8668 let mut out = [[0i32; 4]; 2];
8669 for r in 0..2 {
8670 for k in 0..4 {
8671 out[r][k] = vaddvq_s32(acc[r][k]);
8672 }
8673 }
8674 if i < n {
8675 for (k, x) in xs.iter().enumerate() {
8676 for j in i..n {
8677 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8678 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8679 }
8680 }
8681 }
8682 out
8683 }
8684}
8685
8686#[cfg(target_arch = "x86_64")]
8692#[target_feature(enable = "avx2")]
8693unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8694 unsafe {
8696 use core::arch::x86_64::*;
8697 let n = w0.len();
8698 let ones = _mm256_set1_epi16(1);
8699 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8700 let mut j = 0usize;
8701 while j + 32 <= n {
8702 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8703 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8704 let aw0 = _mm256_abs_epi8(wv0);
8705 let aw1 = _mm256_abs_epi8(wv1);
8706 for (k, x) in xs.iter().enumerate() {
8707 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8708 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8709 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8710 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8711 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8712 }
8713 j += 32;
8714 }
8715 let mut out = [[0i32; 4]; 2];
8716 for r in 0..2 {
8717 for k in 0..4 {
8718 let a = acc[r][k];
8719 let hi128 = _mm256_extracti128_si256::<1>(a);
8720 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8721 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8722 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8723 out[r][k] = _mm_cvtsi128_si32(s32);
8724 }
8725 }
8726 if j < n {
8727 for (k, x) in xs.iter().enumerate() {
8728 for i in j..n {
8729 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8730 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8731 }
8732 }
8733 }
8734 out
8735 }
8736}
8737
8738#[cfg(target_arch = "x86_64")]
8743#[inline]
8744fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8745 let dot = if avx512vnni_enabled() && row.len() >= 64 {
8746 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8747 } else {
8748 unsafe { dot_i8_i8_avx2(row, &act.xq) }
8749 };
8750 let mut acc = dot as f32 * act.sx;
8751 for &(j, xv) in &act.outliers {
8752 acc += (row[j] as i8) as f32 * xv;
8753 }
8754 acc
8755}
8756
8757#[cfg(target_arch = "x86_64")]
8761#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8762unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8763 unsafe {
8765 use core::arch::x86_64::*;
8766 let n = w.len();
8767 let flip = _mm512_set1_epi8(-128); #[inline(always)]
8769 unsafe fn step(
8770 w: *const u8,
8771 x: *const i8,
8772 flip: core::arch::x86_64::__m512i,
8773 acc: core::arch::x86_64::__m512i,
8774 ) -> core::arch::x86_64::__m512i {
8775 unsafe {
8776 use core::arch::x86_64::*;
8777 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8778 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8779 }
8780 }
8781 let (mut a0, mut a1, mut a2, mut a3) = (
8782 _mm512_setzero_si512(),
8783 _mm512_setzero_si512(),
8784 _mm512_setzero_si512(),
8785 _mm512_setzero_si512(),
8786 );
8787 let mut j = 0usize;
8788 while j + 256 <= n {
8789 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8790 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8791 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8792 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8793 j += 256;
8794 }
8795 while j + 64 <= n {
8796 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8797 j += 64;
8798 }
8799 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8800 _mm512_add_epi32(a0, a1),
8801 _mm512_add_epi32(a2, a3),
8802 ));
8803 while j < n {
8805 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8806 j += 1;
8807 }
8808 total
8809 }
8810}
8811
8812#[cfg(target_arch = "x86_64")]
8818#[target_feature(enable = "avx2")]
8819unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8820 unsafe {
8823 use core::arch::x86_64::*;
8824 let lomask = _mm_set1_epi8(0x0F);
8825 let eight = _mm256_set1_epi8(8);
8826 let ones = _mm256_set1_epi16(1);
8827 let mut acc = 0f32;
8828 for gi in 0..gpr {
8829 let g = g0 + gi;
8830 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8831 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8832 let lo = _mm_and_si128(b, lomask);
8833 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8834 let w = _mm256_sub_epi8(
8835 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8836 eight,
8837 );
8838 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8839 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8840 let d = _mm256_madd_epi16(p16, ones);
8841 let hi128 = _mm256_extracti128_si256::<1>(d);
8842 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8843 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8844 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8845 acc += _mm_cvtsi128_si32(s32) as f32 * s;
8846 }
8847 acc
8848 }
8849}
8850
8851#[cfg(target_arch = "x86_64")]
8854#[target_feature(enable = "avx2")]
8855unsafe fn dot_q4_row_avx2_2(
8856 packed: &[u8],
8857 scales: &[u8],
8858 g0: usize,
8859 gpr: usize,
8860 xq1: &[i8],
8861 xq2: &[i8],
8862) -> (f32, f32) {
8863 unsafe {
8865 use core::arch::x86_64::*;
8866 let lomask = _mm_set1_epi8(0x0F);
8867 let eight = _mm256_set1_epi8(8);
8868 let ones = _mm256_set1_epi16(1);
8869 let (mut acc1, mut acc2) = (0f32, 0f32);
8870 #[inline(always)]
8871 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8872 unsafe {
8873 use core::arch::x86_64::*;
8874 let hi128 = _mm256_extracti128_si256::<1>(d);
8875 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8876 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8877 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8878 _mm_cvtsi128_si32(s32)
8879 }
8880 }
8881 for gi in 0..gpr {
8882 let g = g0 + gi;
8883 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8884 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8885 let lo = _mm_and_si128(b, lomask);
8886 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8887 let w = _mm256_sub_epi8(
8888 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8889 eight,
8890 );
8891 let aw = _mm256_abs_epi8(w);
8892 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8893 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8894 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8895 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8896 acc1 += hsum(d1) as f32 * s;
8897 acc2 += hsum(d2) as f32 * s;
8898 }
8899 (acc1, acc2)
8900 }
8901}
8902
8903#[cfg(target_arch = "x86_64")]
8905fn q8_range_avx2(
8906 q: &[u8],
8907 row_scale: &[f32],
8908 act: &SplitAct,
8909 cols: usize,
8910 out_addr: SendMut,
8911 start: usize,
8912 end: usize,
8913) {
8914 for o in start..end {
8915 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8916 unsafe { *out_addr.at(o) = v };
8918 }
8919}
8920
8921#[cfg(target_arch = "x86_64")]
8923#[allow(clippy::too_many_arguments)]
8924fn q8_range2_avx2(
8925 q: &[u8],
8926 row_scale: &[f32],
8927 a1: &SplitAct,
8928 a2: &SplitAct,
8929 cols: usize,
8930 p1: SendMut,
8931 p2: SendMut,
8932 start: usize,
8933 end: usize,
8934) {
8935 for o in start..end {
8936 let row = &q[o * cols..(o + 1) * cols];
8937 unsafe {
8939 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8940 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8941 }
8942 }
8943}
8944
8945#[cfg(target_arch = "aarch64")]
8956fn i8mm_enabled() -> bool {
8957 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8958 *ON.get_or_init(|| {
8959 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8960 && std::arch::is_aarch64_feature_detected!("i8mm")
8961 })
8962}
8963
8964#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
8968fn sdot_enabled() -> bool {
8969 use std::sync::OnceLock;
8970 static ON: OnceLock<bool> = OnceLock::new();
8971 *ON.get_or_init(|| {
8972 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
8973 if !want {
8974 return false;
8975 }
8976
8977 #[cfg(target_arch = "aarch64")]
8978 {
8979 if std::arch::is_aarch64_feature_detected!("dotprod") {
8980 return true;
8981 }
8982 #[cfg(target_os = "android")]
8983 {
8984 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
8985 if cpuinfo.lines().any(|l| {
8986 (l.starts_with("Features") || l.starts_with("features"))
8987 && l.contains("asimddp")
8988 }) {
8989 return true;
8990 }
8991 }
8992 }
8993 false
8994 }
8995 #[cfg(not(target_arch = "aarch64"))]
8996 {
8997 false
8998 }
8999 })
9000}
9001
9002struct SplitAct {
9007 xq: Vec<i8>,
9008 sx: f32,
9009 outliers: Vec<(usize, f32)>,
9010 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9013 xsum: i32,
9014}
9015
9016thread_local! {
9017 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9020 const { std::cell::RefCell::new(Vec::new()) };
9021}
9022
9023impl Drop for SplitAct {
9024 fn drop(&mut self) {
9025 let buf = std::mem::take(&mut self.xq);
9026 if buf.capacity() > 0 {
9027 XQ_FREE.with(|f| {
9028 let mut f = f.borrow_mut();
9029 if f.len() < 16 {
9030 f.push(buf);
9031 }
9032 });
9033 }
9034 }
9035}
9036
9037thread_local! {
9038 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9045}
9046
9047#[inline]
9050fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9051 KROW.with(|s| {
9052 let mut b = s.borrow_mut();
9053 if b.len() < n {
9054 b.resize(n, 0.0);
9055 }
9056 f(&mut b[..n])
9057 })
9058}
9059
9060fn split_act(x: &[f32]) -> SplitAct {
9061 let n = x.len();
9062 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9063 let thr = 8.0 * rms;
9064 let mut outliers: Vec<(usize, f32)> = Vec::new();
9068 let mut amax = 0f32;
9069 for (j, &v) in x.iter().enumerate() {
9070 let a = v.abs();
9071 if a > thr {
9072 outliers.push((j, v));
9073 } else if a > amax {
9074 amax = a;
9075 }
9076 }
9077 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9078 let inv = 1.0 / sx;
9079 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9080 xq.clear();
9081 xq.reserve(n);
9082 if outliers.is_empty() {
9083 xq.extend(
9084 x.iter()
9085 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9086 );
9087 } else {
9088 xq.extend(x.iter().map(|&v| {
9090 if v.abs() > thr {
9091 0
9092 } else {
9093 (v * inv).round().clamp(-127.0, 127.0) as i8
9094 }
9095 }));
9096 }
9097 let xsum = xq.iter().map(|&v| v as i32).sum();
9098 SplitAct {
9099 xq,
9100 sx,
9101 outliers,
9102 xsum,
9103 }
9104}
9105
9106fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9107 let n = x.len();
9108 let rms = (x
9109 .iter()
9110 .zip(col)
9111 .map(|(&a, &c)| {
9112 let v = a * c;
9113 (v * v) as f64
9114 })
9115 .sum::<f64>()
9116 / n.max(1) as f64)
9117 .sqrt() as f32;
9118 let thr = 8.0 * rms;
9119
9120 let mut outliers = Vec::new();
9121 let mut amax = 0f32;
9122 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9123 let v = a * c;
9124 let s = v.abs();
9125 if s > thr {
9126 outliers.push((j, v));
9127 } else if s > amax {
9128 amax = s;
9129 }
9130 }
9131
9132 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9133 let inv = 1.0 / sx;
9134 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9135 xq.clear();
9136 xq.reserve(n);
9137 if outliers.is_empty() {
9138 xq.extend(
9139 x.iter()
9140 .zip(col)
9141 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9142 );
9143 } else {
9144 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9145 let v = a * c;
9146 if v.abs() > thr {
9147 0
9148 } else {
9149 (v * inv).round().clamp(-127.0, 127.0) as i8
9150 }
9151 }));
9152 }
9153 let xsum = xq.iter().map(|&v| v as i32).sum();
9154 SplitAct {
9155 xq,
9156 sx,
9157 outliers,
9158 xsum,
9159 }
9160}
9161
9162#[cfg(target_arch = "aarch64")]
9165#[target_feature(enable = "neon,dotprod")]
9166unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9167 unsafe {
9169 use core::arch::aarch64::*;
9170 use core::arch::asm;
9171 let wp = w.as_ptr() as *const i8;
9172 let n = w.len();
9173 let (mut a0, mut a1, mut a2, mut a3) = (
9174 vdupq_n_s32(0),
9175 vdupq_n_s32(0),
9176 vdupq_n_s32(0),
9177 vdupq_n_s32(0),
9178 );
9179 let mut i = 0;
9180 while i + 64 <= n {
9181 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9182 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9183 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9184 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9185 asm!(
9186 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9187 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9188 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9189 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9190 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9191 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9192 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9193 options(pure, nomem, nostack),
9194 );
9195 i += 64;
9196 }
9197 while i + 16 <= n {
9198 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9199 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9200 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9201 i += 16;
9202 }
9203 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9204 while i < n {
9205 s += (*wp.add(i)) as i32 * xq[i] as i32;
9206 i += 1;
9207 }
9208 s
9209 }
9210}
9211
9212#[cfg(target_arch = "aarch64")]
9216#[target_feature(enable = "neon,dotprod")]
9217unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9218 unsafe {
9220 use core::arch::aarch64::*;
9221 use core::arch::asm;
9222 let n = xq.len();
9223 let px = xq.as_ptr();
9224 let (p0, p1, p2, p3) = (
9225 w0.as_ptr() as *const i8,
9226 w1.as_ptr() as *const i8,
9227 w2.as_ptr() as *const i8,
9228 w3.as_ptr() as *const i8,
9229 );
9230 let (mut a0, mut a1, mut a2, mut a3) = (
9231 vdupq_n_s32(0),
9232 vdupq_n_s32(0),
9233 vdupq_n_s32(0),
9234 vdupq_n_s32(0),
9235 );
9236 let mut i = 0;
9237 while i + 16 <= n {
9238 let x = vld1q_s8(px.add(i));
9239 let v0 = vld1q_s8(p0.add(i));
9240 let v1 = vld1q_s8(p1.add(i));
9241 let v2 = vld1q_s8(p2.add(i));
9242 let v3 = vld1q_s8(p3.add(i));
9243 asm!(
9244 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9245 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9246 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9247 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9248 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9249 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9250 options(pure, nomem, nostack),
9251 );
9252 i += 16;
9253 }
9254 let mut r = [
9255 vaddvq_s32(a0),
9256 vaddvq_s32(a1),
9257 vaddvq_s32(a2),
9258 vaddvq_s32(a3),
9259 ];
9260 while i < n {
9261 let xi = *px.add(i) as i32;
9262 r[0] += (*p0.add(i)) as i32 * xi;
9263 r[1] += (*p1.add(i)) as i32 * xi;
9264 r[2] += (*p2.add(i)) as i32 * xi;
9265 r[3] += (*p3.add(i)) as i32 * xi;
9266 i += 1;
9267 }
9268 r
9269 }
9270}
9271
9272#[cfg(target_arch = "aarch64")]
9279#[target_feature(enable = "neon,dotprod")]
9280unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9281 unsafe {
9284 use core::arch::aarch64::*;
9285 use core::arch::asm;
9286 let n = xq.len();
9287 let px = xq.as_ptr();
9288 let pg = g.as_ptr() as *const i8;
9289 let (mut a0, mut a1, mut a2, mut a3) = (
9290 vdupq_n_s32(0),
9291 vdupq_n_s32(0),
9292 vdupq_n_s32(0),
9293 vdupq_n_s32(0),
9294 );
9295 let mut i = 0;
9296 while i + 16 <= n {
9297 let x = vld1q_s8(px.add(i));
9298 let base = pg.add(4 * i);
9299 let v0 = vld1q_s8(base);
9300 let v1 = vld1q_s8(base.add(16));
9301 let v2 = vld1q_s8(base.add(32));
9302 let v3 = vld1q_s8(base.add(48));
9303 asm!(
9304 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9305 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9306 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9307 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9308 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9309 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9310 options(pure, nomem, nostack),
9311 );
9312 i += 16;
9313 }
9314 [
9315 vaddvq_s32(a0),
9316 vaddvq_s32(a1),
9317 vaddvq_s32(a2),
9318 vaddvq_s32(a3),
9319 ]
9320 }
9321}
9322
9323#[cfg(target_arch = "aarch64")]
9329fn q8_range_sdot(
9330 q: &[u8],
9331 rep: &[u8],
9332 row_scale: &[f32],
9333 act: &SplitAct,
9334 cols: usize,
9335 out_addr: SendMut,
9336 start: usize,
9337 end: usize,
9338) {
9339 let mut o = start;
9340 if !rep.is_empty() {
9343 while o < end && o % 4 != 0 {
9344 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9345 unsafe { *out_addr.at(o) = v };
9346 o += 1;
9347 }
9348 }
9349 while o + 4 <= end {
9350 let r = if rep.is_empty() {
9351 unsafe {
9352 dot_i8_sdot_4rows(
9353 &q[o * cols..(o + 1) * cols],
9354 &q[(o + 1) * cols..(o + 2) * cols],
9355 &q[(o + 2) * cols..(o + 3) * cols],
9356 &q[(o + 3) * cols..(o + 4) * cols],
9357 &act.xq,
9358 )
9359 }
9360 } else {
9361 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9362 };
9363 for k in 0..4 {
9364 let mut acc = r[k] as f32 * act.sx;
9365 for &(j, xv) in &act.outliers {
9366 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9367 }
9368 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9370 }
9371 o += 4;
9372 }
9373 while o < end {
9374 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9375 unsafe { *out_addr.at(o) = v };
9376 o += 1;
9377 }
9378}
9379
9380#[cfg(target_arch = "aarch64")]
9383#[allow(clippy::too_many_arguments)]
9384fn q8_range2_sdot(
9385 q: &[u8],
9386 row_scale: &[f32],
9387 a1: &SplitAct,
9388 a2: &SplitAct,
9389 cols: usize,
9390 p1: SendMut,
9391 p2: SendMut,
9392 start: usize,
9393 end: usize,
9394) {
9395 for o in start..end {
9396 let row = &q[o * cols..(o + 1) * cols];
9397 unsafe {
9399 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9400 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9401 }
9402 }
9403}
9404
9405#[allow(clippy::too_many_arguments)]
9407fn q8_range2_f32(
9408 q: &[u8],
9409 row_scale: &[f32],
9410 x1: &[f32],
9411 x2: &[f32],
9412 cols: usize,
9413 p1: SendMut,
9414 p2: SendMut,
9415 start: usize,
9416 end: usize,
9417) {
9418 for o in start..end {
9419 let row = &q[o * cols..(o + 1) * cols];
9420 unsafe {
9422 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9423 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9424 }
9425 }
9426}
9427
9428fn q8_range_f32(
9430 q: &[u8],
9431 row_scale: &[f32],
9432 xs: &[f32],
9433 cols: usize,
9434 out_addr: SendMut,
9435 start: usize,
9436 end: usize,
9437) {
9438 for o in start..end {
9439 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9440 unsafe { *out_addr.at(o) = v };
9442 }
9443}
9444
9445#[inline]
9449fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9450 #[cfg(target_arch = "aarch64")]
9451 return row_dot_sdot(row, act);
9452 #[cfg(target_arch = "x86_64")]
9453 return row_dot_avx2(row, act);
9454 #[allow(unreachable_code)]
9455 q8_row_dot_scalar(row, act)
9456}
9457
9458#[allow(dead_code)]
9459fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9460 let mut acc = 0i32;
9461 for (k, &b) in row.iter().enumerate() {
9462 acc += (b as i8) as i32 * act.xq[k] as i32;
9463 }
9464 let mut acc = acc as f32 * act.sx;
9465 for &(j, xv) in &act.outliers {
9466 acc += (row[j] as i8) as f32 * xv;
9467 }
9468 acc
9469}
9470
9471#[cfg(target_arch = "aarch64")]
9474#[inline]
9475fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9476 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9477 for &(j, xv) in &act.outliers {
9478 acc += (row[j] as i8) as f32 * xv;
9479 }
9480 acc
9481}
9482
9483#[cfg(target_arch = "aarch64")]
9491#[target_feature(enable = "neon,dotprod")]
9492unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9493 unsafe {
9496 use core::arch::aarch64::*;
9497 use core::arch::asm;
9498 let lomask = vdupq_n_u8(0x0F);
9499 let eight = vdupq_n_s8(8);
9500 let mut acc = 0f32;
9501 for gi in 0..gpr {
9502 let g = g0 + gi;
9503 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9504 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9505 let lo = vandq_u8(b, lomask);
9506 let hi = vshrq_n_u8::<4>(b);
9507 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9508 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9509 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9510 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9511 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9512 asm!(
9513 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9514 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9515 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9516 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9517 options(pure, nomem, nostack),
9518 );
9519 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9520 }
9521 acc
9522 }
9523}
9524
9525#[cfg(target_arch = "aarch64")]
9530#[target_feature(enable = "neon,dotprod")]
9531unsafe fn dot_q4_row_sdot2(
9532 packed: &[u8],
9533 scales: &[u8],
9534 g0: usize,
9535 gpr: usize,
9536 xq1: &[i8],
9537 xq2: &[i8],
9538) -> (f32, f32) {
9539 unsafe {
9542 use core::arch::aarch64::*;
9543 use core::arch::asm;
9544 let lomask = vdupq_n_u8(0x0F);
9545 let eight = vdupq_n_s8(8);
9546 let (mut acc1, mut acc2) = (0f32, 0f32);
9547 for gi in 0..gpr {
9548 let g = g0 + gi;
9549 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9550 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9551 let lo = vandq_u8(b, lomask);
9552 let hi = vshrq_n_u8::<4>(b);
9553 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9554 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9555 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9556 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9557 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9558 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9559 let (mut a0, mut a1, mut b0, mut b1) = (
9560 vdupq_n_s32(0),
9561 vdupq_n_s32(0),
9562 vdupq_n_s32(0),
9563 vdupq_n_s32(0),
9564 );
9565 asm!(
9566 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9567 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9568 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9569 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9570 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9571 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9572 e0 = in(vreg) e0, e1 = in(vreg) e1,
9573 x10 = in(vreg) x10, x11 = in(vreg) x11,
9574 x20 = in(vreg) x20, x21 = in(vreg) x21,
9575 options(pure, nomem, nostack),
9576 );
9577 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9578 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9579 }
9580 (acc1, acc2)
9581 }
9582}
9583
9584#[inline]
9589pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9590 #[cfg(target_arch = "aarch64")]
9591 unsafe {
9592 return axpy_i8_f32_neon(acc, row, w);
9593 }
9594 #[cfg(target_arch = "x86_64")]
9595 if avx2_enabled() {
9596 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9597 }
9598 #[allow(unreachable_code)]
9599 {
9600 for (a, &b) in acc.iter_mut().zip(row) {
9601 *a += w * b as f32;
9602 }
9603 }
9604}
9605
9606#[cfg(target_arch = "x86_64")]
9608#[target_feature(enable = "avx2,fma")]
9609unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9610 unsafe {
9612 use core::arch::x86_64::*;
9613 let n = acc.len().min(row.len());
9614 let ap = acc.as_mut_ptr();
9615 let rp = row.as_ptr();
9616 let wv = _mm256_set1_ps(w);
9617 let mut j = 0usize;
9618 while j + 16 <= n {
9619 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9620 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9621 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9622 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9623 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9624 _mm256_storeu_ps(ap.add(j), v0);
9625 _mm256_storeu_ps(ap.add(j + 8), v1);
9626 j += 16;
9627 }
9628 while j < n {
9629 *ap.add(j) += w * (*rp.add(j)) as f32;
9630 j += 1;
9631 }
9632 }
9633}
9634
9635#[cfg(target_arch = "aarch64")]
9636#[target_feature(enable = "neon")]
9637unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9638 unsafe {
9640 use core::arch::aarch64::*;
9641 let n = acc.len().min(row.len());
9642 let ap = acc.as_mut_ptr();
9643 let rp = row.as_ptr();
9644 let wv = vdupq_n_f32(w);
9645 let mut j = 0usize;
9646 while j + 16 <= n {
9647 let rb = vld1q_s8(rp.add(j));
9648 let lo = vmovl_s8(vget_low_s8(rb));
9649 let hi = vmovl_s8(vget_high_s8(rb));
9650 for (off, half) in [(0, lo), (8, hi)] {
9651 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9652 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9653 let o = j + off;
9654 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9655 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9656 }
9657 j += 16;
9658 }
9659 while j < n {
9660 *ap.add(j) += w * (*rp.add(j)) as f32;
9661 j += 1;
9662 }
9663 }
9664}
9665
9666#[inline]
9669pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9670 #[cfg(target_arch = "aarch64")]
9671 unsafe {
9672 return dot_i8_f32_neon(w, x);
9673 }
9674 #[cfg(target_arch = "x86_64")]
9675 if avx2_enabled() {
9676 return unsafe { dot_i8_f32_avx2(w, x) };
9677 }
9678 #[allow(unreachable_code)]
9679 {
9680 let mut sum = 0.0f32;
9681 for (j, &b) in w.iter().enumerate() {
9682 sum += (b as i8) as f32 * x[j];
9683 }
9684 sum
9685 }
9686}
9687
9688#[inline]
9692fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9693 #[cfg(target_arch = "aarch64")]
9694 unsafe {
9695 return dot_i8_col_f32_neon(w, x, col);
9696 }
9697 #[allow(unreachable_code)]
9698 {
9699 let mut sum = 0.0f32;
9700 for (j, &b) in w.iter().enumerate() {
9701 sum += (b as i8) as f32 * x[j] * col[j];
9702 }
9703 sum
9704 }
9705}
9706
9707#[cfg(target_arch = "aarch64")]
9708#[target_feature(enable = "neon")]
9709unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9710 unsafe {
9712 use core::arch::aarch64::*;
9713 let n = x.len();
9714 let wp = w.as_ptr() as *const i8;
9715 let xp = x.as_ptr();
9716 let cp = col.as_ptr();
9717 let (mut a0, mut a1, mut a2, mut a3) = (
9718 vdupq_n_f32(0.0),
9719 vdupq_n_f32(0.0),
9720 vdupq_n_f32(0.0),
9721 vdupq_n_f32(0.0),
9722 );
9723 let mut j = 0usize;
9724 while j + 16 <= n {
9725 let wb = vld1q_s8(wp.add(j));
9726 let lo = vmovl_s8(vget_low_s8(wb));
9727 let hi = vmovl_s8(vget_high_s8(wb));
9728 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9729 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9730 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9731 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9732 a0 = vfmaq_f32(
9733 a0,
9734 w0,
9735 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9736 );
9737 a1 = vfmaq_f32(
9738 a1,
9739 w1,
9740 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9741 );
9742 a2 = vfmaq_f32(
9743 a2,
9744 w2,
9745 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9746 );
9747 a3 = vfmaq_f32(
9748 a3,
9749 w3,
9750 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9751 );
9752 j += 16;
9753 }
9754 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9755 while j < n {
9756 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9757 j += 1;
9758 }
9759 sum
9760 }
9761}
9762
9763#[cfg(target_arch = "aarch64")]
9764#[target_feature(enable = "neon")]
9765unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9766 unsafe {
9768 use core::arch::aarch64::*;
9769 let n = x.len();
9770 let wp = w.as_ptr() as *const i8;
9771 let xp = x.as_ptr();
9772 let (mut a0, mut a1, mut a2, mut a3) = (
9773 vdupq_n_f32(0.0),
9774 vdupq_n_f32(0.0),
9775 vdupq_n_f32(0.0),
9776 vdupq_n_f32(0.0),
9777 );
9778 let mut j = 0usize;
9779 while j + 16 <= n {
9780 let wb = vld1q_s8(wp.add(j));
9781 let lo = vmovl_s8(vget_low_s8(wb));
9782 let hi = vmovl_s8(vget_high_s8(wb));
9783 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9784 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9785 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9786 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9787 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9788 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9789 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9790 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9791 j += 16;
9792 }
9793 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9794 while j < n {
9795 sum += (*wp.add(j)) as f32 * *xp.add(j);
9796 j += 1;
9797 }
9798 sum
9799 }
9800}
9801
9802#[allow(clippy::too_many_arguments)]
9803fn qmatvec(
9804 q: &[u8],
9805 rep: &[u8],
9806 row_scale: &[f32],
9807 x: &[f32],
9808 col_field: &[f32],
9809 dtype: TensorDtype,
9810 rows: usize,
9811 cols: usize,
9812 out: &mut [f32],
9813 pool: Option<&Pool>,
9814) {
9815 debug_assert_eq!(out.len(), rows);
9816 #[cfg(not(target_arch = "aarch64"))]
9817 let _ = rep;
9818
9819 #[cfg(target_arch = "aarch64")]
9820 if sdot_enabled() {
9821 let act = if dtype == TensorDtype::Q8_2f {
9822 split_act_q8_2f(x, col_field)
9823 } else {
9824 split_act(x)
9825 };
9826 let out_addr = SendMut(out.as_mut_ptr());
9827 let run_range = |start: usize, end: usize| {
9828 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9829 };
9830 match pool {
9831 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9832 _ => run_range(0, rows),
9833 }
9834 return;
9835 }
9836 #[cfg(target_arch = "x86_64")]
9839 if avx2_a8w8_enabled() {
9840 let act = if dtype == TensorDtype::Q8_2f {
9841 split_act_q8_2f(x, col_field)
9842 } else {
9843 split_act(x)
9844 };
9845 let out_addr = SendMut(out.as_mut_ptr());
9846 let run_range = |start: usize, end: usize| {
9847 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9848 };
9849 match pool {
9850 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9851 _ => run_range(0, rows),
9852 }
9853 return;
9854 }
9855
9856 prescale_with(x, col_field, dtype, 1, |xs| {
9857 let out_addr = SendMut(out.as_mut_ptr());
9858 let run_range = move |start: usize, end: usize| {
9859 for o in start..end {
9860 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9861 unsafe { *out_addr.at(o) = v };
9863 }
9864 };
9865 match pool {
9866 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9867 _ => run_range(0, rows),
9868 }
9869 });
9870}
9871
9872#[allow(clippy::too_many_arguments)]
9873fn qmatvec2(
9874 q: &[u8],
9875 row_scale: &[f32],
9876 x1: &[f32],
9877 x2: &[f32],
9878 col_field: &[f32],
9879 dtype: TensorDtype,
9880 rows: usize,
9881 cols: usize,
9882 o1: &mut [f32],
9883 o2: &mut [f32],
9884 pool: Option<&Pool>,
9885) {
9886 #[cfg(target_arch = "aarch64")]
9887 if sdot_enabled() {
9888 let a1s = if dtype == TensorDtype::Q8_2f {
9889 split_act_q8_2f(x1, col_field)
9890 } else {
9891 split_act(x1)
9892 };
9893 let a2s = if dtype == TensorDtype::Q8_2f {
9894 split_act_q8_2f(x2, col_field)
9895 } else {
9896 split_act(x2)
9897 };
9898 let p1 = SendMut(o1.as_mut_ptr());
9899 let p2 = SendMut(o2.as_mut_ptr());
9900 let run_range = |start: usize, end: usize| {
9901 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9902 };
9903 match pool {
9904 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9905 _ => run_range(0, rows),
9906 }
9907 return;
9908 }
9909 #[cfg(target_arch = "x86_64")]
9910 if avx2_a8w8_enabled() {
9911 let a1s = if dtype == TensorDtype::Q8_2f {
9912 split_act_q8_2f(x1, col_field)
9913 } else {
9914 split_act(x1)
9915 };
9916 let a2s = if dtype == TensorDtype::Q8_2f {
9917 split_act_q8_2f(x2, col_field)
9918 } else {
9919 split_act(x2)
9920 };
9921 let p1 = SendMut(o1.as_mut_ptr());
9922 let p2 = SendMut(o2.as_mut_ptr());
9923 let run_range = |start: usize, end: usize| {
9924 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9925 };
9926 match pool {
9927 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9928 _ => run_range(0, rows),
9929 }
9930 return;
9931 }
9932
9933 prescale_with(x1, col_field, dtype, 1, |x1s| {
9934 prescale_with(x2, col_field, dtype, 2, |x2s| {
9935 let p1 = SendMut(o1.as_mut_ptr());
9936 let p2 = SendMut(o2.as_mut_ptr());
9937 let run_range = move |start: usize, end: usize| {
9938 for o in start..end {
9939 let row = &q[o * cols..(o + 1) * cols];
9940 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
9941 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
9942 unsafe {
9944 *p1.at(o) = s1;
9945 *p2.at(o) = s2;
9946 }
9947 }
9948 };
9949 match pool {
9950 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9951 _ => run_range(0, rows),
9952 }
9953 });
9954 });
9955}
9956
9957#[derive(Clone, Copy)]
9958struct SendMut(*mut f32);
9959unsafe impl Send for SendMut {}
9960unsafe impl Sync for SendMut {}
9961
9962impl SendMut {
9963 #[inline]
9964 fn at(self, i: usize) -> *mut f32 {
9965 unsafe { self.0.add(i) }
9966 }
9967}
9968
9969#[cfg(test)]
9970mod tests {
9971 use super::*;
9972
9973 #[test]
9974 fn q2tp_i8_dot_matches_exact_on_grid() {
9975 let (rows, cols) = (5, 64);
9979 let gpr = cols / GROUP_SIZE;
9980 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
9984 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
9985 .collect();
9986 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
9987 let x: Vec<f32> = (0..cols).map(|i| if i % 3 == 0 { -1.0 } else { 1.0 }).collect();
9988 let act = split_act(&x);
9989 assert!(act.outliers.is_empty(), "on-grid input must have no outliers");
9990 let gsum = q1_group_sums(&act.xq, gpr);
9991 for r in 0..rows {
9992 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
9993 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
9994 assert!(
9995 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
9996 "row {r}: exact {exact} vs i8 {fast}"
9997 );
9998 }
9999 }
10000
10001 #[test]
10002 fn q8_row_dot_fast_matches_scalar() {
10003 let cols = 96;
10006 let row: Vec<u8> = (0..cols)
10007 .map(|i| ((i as i32 * 37 % 251) - 125) as i8 as u8)
10008 .collect();
10009 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10010 let act = split_act(&x);
10011 let fast = q8_row_dot(&row, &act);
10012 let scalar = q8_row_dot_scalar(&row, &act);
10013 assert!(
10014 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10015 "fast {fast} vs scalar {scalar}"
10016 );
10017 }
10018
10019 #[test]
10020 fn f32_matvec_matches_matvec_rows_bitexact() {
10021 let (rows, cols) = (300, 40);
10022 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10023 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10024 let qt = QTensor::from_f32(w.clone(), rows, cols);
10025
10026 let mut a = vec![0.0f32; rows];
10027 matvec_rows(None, &w, &x, &mut a);
10028 let mut b = vec![0.0f32; rows];
10029 qt.matvec(&x, &mut b, None);
10030 assert_eq!(a, b);
10031 }
10032
10033 #[test]
10034 fn sdot_kernel_exact_on_grid() {
10035 eprintln!("sdot_enabled = {}", sdot_enabled());
10040 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
10042 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10043 .collect();
10044 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10045 let x: Vec<f32> = (0..cols)
10046 .map(|i| match i % 3 {
10047 0 => 1.0,
10048 1 => -1.0,
10049 _ => 0.0,
10050 })
10051 .collect();
10052 let mut a = vec![0.0f32; rows];
10053 qmatvec(
10054 &w,
10055 &[],
10056 &scales,
10057 &x,
10058 &[],
10059 TensorDtype::Q8Row,
10060 rows,
10061 cols,
10062 &mut a,
10063 None,
10064 );
10065 for o in 0..rows {
10066 let mut acc = 0.0f32;
10067 for j in 0..cols {
10068 acc += (w[o * cols + j] as i8) as f32 * x[j];
10069 }
10070 let expect = acc * scales[o];
10071 assert!(
10072 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10073 "row {o}: {} vs {expect}",
10074 a[o]
10075 );
10076 }
10077 }
10078
10079 #[test]
10080 fn q1_tbl_fast_path_matches_reference() {
10081 let (rows, cols) = (5, 256);
10086 let gpr = cols / GROUP_SIZE;
10087 let mut bytes = Vec::new();
10088 for t in 0..rows * gpr {
10089 let s = 0.007 + (t % 11) as f32 * 0.004;
10090 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10091 for j in 0..4 {
10092 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10093 }
10094 }
10095 let x: Vec<f32> = (0..cols)
10096 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10097 .collect();
10098 let mut w = vec![0.0f32; rows * cols];
10099 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10100 let mut got = vec![0.0f32; rows];
10101 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10102 for o in 0..rows {
10103 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10104 assert!(
10105 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10106 "row {o}: {} vs {expect}",
10107 got[o]
10108 );
10109 }
10110 let b = 5usize;
10113 let mut xs_all = Vec::new();
10114 for bi in 0..b {
10115 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10116 }
10117 let mut mm = vec![0.0f32; b * rows];
10118 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10119 for bi in 0..b {
10120 let mut single = vec![0.0f32; rows];
10121 q1_matvec(
10122 &bytes,
10123 &xs_all[bi * cols..(bi + 1) * cols],
10124 rows,
10125 cols,
10126 &mut single,
10127 None,
10128 );
10129 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10130 }
10131 }
10132
10133 #[test]
10134 fn q1_kernels_match_exact_reference() {
10135 let (rows, cols) = (7, 96);
10137 let gpr = cols / GROUP_SIZE;
10138 let mut bytes = Vec::new();
10139 for t in 0..rows * gpr {
10140 let s = 0.01 + (t % 13) as f32 * 0.003;
10141 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10142 for j in 0..4 {
10143 bytes.push(((t * 31 + j * 97) % 251) as u8);
10144 }
10145 }
10146 let x: Vec<f32> = (0..cols)
10148 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10149 .collect();
10150 let mut w = vec![0.0f32; rows * cols];
10152 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10153 let mut expect = vec![0.0f32; rows];
10154 for o in 0..rows {
10155 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10156 }
10157 let mut got = vec![0.0f32; rows];
10158 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10159 for o in 0..rows {
10160 assert!(
10161 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10162 "row {o}: {} vs {}",
10163 got[o],
10164 expect[o]
10165 );
10166 }
10167 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10169 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10170 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10171 assert_eq!(a1, got);
10172 let mut xs = x.clone();
10173 xs.extend_from_slice(&x2);
10174 let mut mm = vec![0.0f32; 2 * rows];
10175 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10176 assert_eq!(&mm[..rows], got.as_slice());
10177 assert_eq!(&mm[rows..], a2.as_slice());
10178 }
10179
10180 #[test]
10181 fn repack_is_bit_identical() {
10182 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
10188 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10189 .collect();
10190 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10191 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10192 let rep = q8_repack_layout(&w, rows, cols);
10193 for g in 0..rows / 4 {
10195 for c in 0..cols / 16 {
10196 for lane in 0..4 {
10197 assert_eq!(
10198 &rep[g * 4 * cols + c * 64 + lane * 16
10199 ..g * 4 * cols + c * 64 + lane * 16 + 16],
10200 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10201 );
10202 }
10203 }
10204 }
10205 let mut a = vec![0.0f32; rows];
10206 qmatvec(
10207 &w,
10208 &[],
10209 &scales,
10210 &x,
10211 &[],
10212 TensorDtype::Q8Row,
10213 rows,
10214 cols,
10215 &mut a,
10216 None,
10217 );
10218 let mut b = vec![0.0f32; rows];
10219 qmatvec(
10220 &w,
10221 &rep,
10222 &scales,
10223 &x,
10224 &[],
10225 TensorDtype::Q8Row,
10226 rows,
10227 cols,
10228 &mut b,
10229 None,
10230 );
10231 assert_eq!(a, b, "full-range repack output diverged");
10232
10233 #[cfg(target_arch = "aarch64")]
10234 if sdot_enabled() {
10235 let act = split_act(&x);
10237 let mut c1 = vec![0.0f32; rows];
10238 let mut c2 = vec![0.0f32; rows];
10239 q8_range_sdot(
10240 &w,
10241 &[],
10242 &scales,
10243 &act,
10244 cols,
10245 SendMut(c1.as_mut_ptr()),
10246 3,
10247 rows - 2,
10248 );
10249 q8_range_sdot(
10250 &w,
10251 &rep,
10252 &scales,
10253 &act,
10254 cols,
10255 SendMut(c2.as_mut_ptr()),
10256 3,
10257 rows - 2,
10258 );
10259 assert_eq!(c1, c2, "unaligned-range repack output diverged");
10260 }
10261 }
10262
10263 #[test]
10264 fn sdot_a8w8_noise_is_bounded() {
10265 let (rows, cols) = (16, 512);
10269 let w: Vec<u8> = (0..rows * cols)
10270 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10271 .collect();
10272 let scales = vec![0.01f32; rows];
10273 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10274 let mut a = vec![0.0f32; rows];
10275 qmatvec(
10276 &w,
10277 &[],
10278 &scales,
10279 &x,
10280 &[],
10281 TensorDtype::Q8Row,
10282 rows,
10283 cols,
10284 &mut a,
10285 None,
10286 );
10287 let (mut num, mut den) = (0f64, 0f64);
10288 for o in 0..rows {
10289 let mut acc = 0.0f32;
10290 for j in 0..cols {
10291 acc += (w[o * cols + j] as i8) as f32 * x[j];
10292 }
10293 let expect = acc * scales[o];
10294 num += ((a[o] - expect) as f64).powi(2);
10295 den += (expect as f64).powi(2);
10296 }
10297 let rel = (num / den.max(1e-12)).sqrt();
10298 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10299 }
10300
10301 #[test]
10302 fn i8_dot_neon_matches_scalar() {
10303 let n = 100;
10304 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10305 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10306 let mut scalar = 0.0f32;
10307 for j in 0..n {
10308 scalar += (w[j] as i8) as f32 * x[j];
10309 }
10310 let fast = dot_i8_f32(&w, &x);
10311 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10312 }
10313
10314 #[test]
10316 fn vbitmatvec_matches_full_dequant() {
10317 let (rows, cols) = (6, 64);
10318 let ng = cols / GROUP_SIZE;
10319 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10321 let mut bytes = bits.clone();
10322 for g in 0..rows * ng {
10323 let s = 0.02 + 0.001 * g as f32;
10324 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10325 }
10326 for r in 0..rows {
10327 let b = bits[r] as usize;
10328 let (mut acc, mut nb) = (0u64, 0usize);
10329 let mut rowbytes = Vec::new();
10330 for i in 0..cols {
10331 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10332 acc = (acc << b) | v;
10333 nb += b;
10334 while nb >= 8 {
10335 nb -= 8;
10336 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10337 }
10338 }
10339 if nb > 0 {
10340 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10341 }
10342 bytes.extend_from_slice(&rowbytes);
10343 }
10344 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10345
10346 let mut reference = vec![0f32; rows * cols];
10347 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10348 let mut expect = vec![0f32; rows];
10349 for r in 0..rows {
10350 expect[r] = reference[r * cols..(r + 1) * cols]
10351 .iter()
10352 .zip(&x)
10353 .map(|(w, xv)| w * xv)
10354 .sum();
10355 }
10356 let mut got = vec![0f32; rows];
10357 let offsets = vbit_row_offsets(&bytes, rows, cols);
10358 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10359 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10363 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10364 for r in 0..rows {
10365 assert!(
10366 (got[r] - expect[r]).abs() < tol * scale,
10367 "row {r}: {} vs {}",
10368 got[r],
10369 expect[r]
10370 );
10371 }
10372 }
10373
10374 #[test]
10379 #[cfg(target_arch = "x86_64")]
10380 fn vbit_matmat_blocked_matches_per_row() {
10381 let (rows, cols, b) = (64usize, 128usize, 9usize);
10382 let ng = cols / GROUP_SIZE;
10383 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10384 let mut bytes = bits.clone();
10385 for g in 0..rows * ng {
10386 let sc = 0.02 + 0.0005 * g as f32;
10387 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10388 }
10389 for r in 0..rows {
10390 let bw = bits[r] as usize;
10391 let (mut acc, mut nb) = (0u64, 0usize);
10392 let mut rowbytes = Vec::new();
10393 for i in 0..cols {
10394 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10395 acc = (acc << bw) | v;
10396 nb += bw;
10397 while nb >= 8 {
10398 nb -= 8;
10399 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10400 }
10401 }
10402 if nb > 0 {
10403 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10404 }
10405 bytes.extend_from_slice(&rowbytes);
10406 }
10407 let x: Vec<f32> = (0..b * cols)
10408 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10409 .collect();
10410 let offsets = vbit_row_offsets(&bytes, rows, cols);
10411 let mut y_a = vec![0f32; b * rows];
10412 let mut y_b = vec![0f32; b * rows];
10413 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10414 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10415 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10416 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10417 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10418 let max_d = y_a
10419 .iter()
10420 .zip(&y_b)
10421 .map(|(p, q)| (p - q).abs())
10422 .fold(0.0f32, f32::max);
10423 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10424 }
10425
10426 #[test]
10434 fn q4t_matmat_blocked_matches_per_row() {
10435 let (rows, cols, b) = (16usize, 64usize, 9usize);
10436 let gpr = cols / GROUP_SIZE;
10437 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10438 for r in 0..rows {
10439 for g in 0..gpr {
10440 let t = (r * gpr + g) * Q4_TILE;
10441 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10442 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10443 for k in 0..16 {
10444 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10445 }
10446 }
10447 }
10448 let x: Vec<f32> = (0..b * cols)
10449 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10450 .collect();
10451 let mut y_blk = vec![0f32; b * rows];
10452 let mut y_row = vec![0f32; b * rows];
10453 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10454 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10455 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10456 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10457 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10458 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10459 }
10460
10461 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10468 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10469 let gpr = cols / GROUP_SIZE;
10470 let stride = q4tp_code_stride(gpr);
10471 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10472 let mut b = vec![0u8; codes_off + rows * stride];
10473 for r in 0..rows {
10474 for g in 0..gpr {
10475 let t = (r * gpr + g) * Q4TP_NIB;
10476 for k in 0..16 {
10477 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10478 }
10479 }
10480 let lo = -6.0 - 0.03 * (r % 17) as f32;
10481 let step = 0.01 + 0.004 * (r % 11) as f32;
10482 let p = params_off + r * 4;
10483 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10484 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10485 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10486 for g in 0..gpr {
10487 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10488 }
10489 }
10490 b
10491 }
10492
10493 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10497 let gpr = cols / GROUP_SIZE;
10498 let v = Q4tpView::new(bytes, rows, cols);
10499 let mut out = vec![0u8; rows * gpr * Q4_TILE];
10500 let mut sc = vec![0f32; gpr];
10501 for r in 0..rows {
10502 v.scales_into(r, gpr, &mut sc);
10503 for g in 0..gpr {
10504 let t = (r * gpr + g) * Q4_TILE;
10505 let s = sc[g];
10506 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10507 let src = (r * gpr + g) * Q4TP_NIB;
10508 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10509 }
10510 }
10511 out
10512 }
10513
10514 #[test]
10520 fn q4tp_exact_path_matches_dequant_reference() {
10521 let (rows, cols) = (256usize, 512usize);
10522 let gpr = cols / GROUP_SIZE;
10523 let bytes = synth_q4tp(rows, cols);
10524 let mut w = vec![0f32; rows * cols];
10525 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10526
10527 let x: Vec<f32> = (0..cols)
10528 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10529 .collect();
10530 let v = Q4tpView::new(&bytes, rows, cols);
10531 let mut sc = vec![0f32; gpr];
10532 for r in 0..rows {
10533 v.scales_into(r, gpr, &mut sc);
10534 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10535 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10536 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10540 assert!(
10541 (got - want).abs() <= 1e-5 * mag,
10542 "row {r}: kernel {got} vs dequant {want}"
10543 );
10544 }
10545 }
10546
10547 #[test]
10553 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10554 let (rows, cols) = (256usize, 512usize);
10555 let bytes = synth_q4tp(rows, cols);
10556 let twin = q4tp_as_q4t(&bytes, rows, cols);
10557 let x: Vec<f32> = (0..cols)
10558 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10559 .collect();
10560
10561 let mut got = vec![0f32; rows];
10562 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10563 let mut want = vec![0f32; rows];
10564 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10565
10566 let mut w = vec![0f32; rows * cols];
10569 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10570 for r in 0..rows {
10571 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10572 assert!(
10573 (got[r] - want[r]).abs() <= 1e-3 * mag,
10574 "row {r}: q4tp {} vs q4t {}",
10575 got[r],
10576 want[r]
10577 );
10578 }
10579 }
10580
10581 #[test]
10586 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10587 let (rows, cols, b) = (256usize, 512usize, 5usize);
10588 let bytes = synth_q4tp(rows, cols);
10589 let twin = q4tp_as_q4t(&bytes, rows, cols);
10590 let xs: Vec<f32> = (0..b * cols)
10591 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10592 .collect();
10593
10594 let mut got = vec![0f32; b * rows];
10595 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10596 let mut want = vec![0f32; b * rows];
10597 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10598
10599 let mut w = vec![0f32; rows * cols];
10600 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10601 for t in 0..b {
10602 for r in 0..rows {
10603 let mag: f32 = (0..cols)
10604 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10605 .sum();
10606 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10607 assert!(
10608 (g - wa).abs() <= 1e-3 * mag,
10609 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10610 );
10611 }
10612 }
10613 }
10614
10615 #[test]
10616 fn q4tp_matvec2_matches_the_single_stream_kernel() {
10617 let (rows, cols) = (128usize, 256usize);
10618 let gpr = cols / GROUP_SIZE;
10619 let bytes = synth_q4tp(rows, cols);
10620 let xs: Vec<f32> = (0..2 * cols)
10621 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10622 .collect();
10623
10624 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10625 q4tp_matvec2(
10626 &bytes,
10627 &xs[..cols],
10628 &xs[cols..],
10629 rows,
10630 cols,
10631 &mut o1,
10632 &mut o2,
10633 None,
10634 );
10635
10636 let v = Q4tpView::new(&bytes, rows, cols);
10639 let mut sc = vec![0f32; gpr];
10640 for r in 0..rows {
10641 v.scales_into(r, gpr, &mut sc);
10642 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10643 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10644 }
10645 }
10646
10647 #[test]
10654 fn q4tp_matvec_keeps_pace_with_q4t() {
10655 let (rows, cols) = (4096usize, 3072usize);
10656 let bytes = synth_q4tp(rows, cols);
10657 let twin = q4tp_as_q4t(&bytes, rows, cols);
10658 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10659 let mut o = vec![0f32; rows];
10660 let n = 12;
10661 let mut best = (f64::MAX, f64::MAX);
10662 for _ in 0..3 {
10665 let t0 = std::time::Instant::now();
10666 for _ in 0..n {
10667 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10668 }
10669 best.0 = best.0.min(t0.elapsed().as_secs_f64());
10670 let t0 = std::time::Instant::now();
10671 for _ in 0..n {
10672 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10673 }
10674 best.1 = best.1.min(t0.elapsed().as_secs_f64());
10675 }
10676 let ratio = best.1 / best.0;
10677 println!(
10678 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10679 best.0 * 1e3 / n as f64,
10680 best.1 * 1e3 / n as f64
10681 );
10682 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10683 }
10684
10685 #[cfg(target_os = "macos")]
10686 #[test]
10687 fn q4t_matmat_accel_matches_dequant_reference() {
10688 if !accel_gemm_enabled() {
10689 return; }
10691 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
10693 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10694 for r in 0..rows {
10695 for g in 0..gpr {
10696 let t = (r * gpr + g) * Q4_TILE;
10697 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10698 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10699 for k in 0..16 {
10700 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10701 }
10702 }
10703 }
10704 let x: Vec<f32> = (0..b * cols)
10705 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10706 .collect();
10707 let mut got = vec![0f32; b * rows];
10708 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10709 let mut w = vec![0f32; rows * cols];
10711 for r in 0..rows {
10712 for g in 0..gpr {
10713 let t = (r * gpr + g) * Q4_TILE;
10714 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10715 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10716 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10717 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10718 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10719 }
10720 }
10721 }
10722 for bi in 0..b {
10723 for r in 0..rows {
10724 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10725 let d = (got[bi * rows + r] - want).abs();
10726 assert!(
10727 d <= want.abs().max(1.0) * 1e-4,
10728 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10729 got[bi * rows + r]
10730 );
10731 }
10732 }
10733 }
10734
10735 #[test]
10736 fn q4matvec_matches_full_dequant() {
10737 let (rows, cols) = (8, 64);
10738 let groups = rows * cols / GROUP_SIZE;
10739 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10741 for i in 0..groups * 16 {
10742 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10743 }
10744 for g in 0..groups {
10745 let s = 0.01 + 0.003 * g as f32;
10746 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10747 }
10748 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10749
10750 let mut reference = vec![0.0f32; rows * cols];
10751 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10752 let mut expect = vec![0.0f32; rows];
10753 for r in 0..rows {
10754 expect[r] = reference[r * cols..(r + 1) * cols]
10755 .iter()
10756 .zip(&x)
10757 .map(|(w, xv)| w * xv)
10758 .sum();
10759 }
10760
10761 let mut got = vec![0.0f32; rows];
10762 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10763 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10767 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10768 for r in 0..rows {
10769 assert!(
10770 (got[r] - expect[r]).abs() < tol * scale,
10771 "row {r}: {} vs {}",
10772 got[r],
10773 expect[r]
10774 );
10775 }
10776 }
10777
10778 #[test]
10781 fn vbitmatvec2_equals_two_singles() {
10782 let (rows, cols) = (6, 64);
10783 let ng = cols / GROUP_SIZE;
10784 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10785 let mut bytes = bits.clone();
10786 for g in 0..rows * ng {
10787 let s = 0.02 + 0.001 * g as f32;
10788 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10789 }
10790 for r in 0..rows {
10791 let b = bits[r] as usize;
10792 let (mut acc, mut nb) = (0u64, 0usize);
10793 let mut rowbytes = Vec::new();
10794 for i in 0..cols {
10795 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10796 acc = (acc << b) | v;
10797 nb += b;
10798 while nb >= 8 {
10799 nb -= 8;
10800 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10801 }
10802 }
10803 if nb > 0 {
10804 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10805 }
10806 bytes.extend_from_slice(&rowbytes);
10807 }
10808 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10809 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10810 let offsets = vbit_row_offsets(&bytes, rows, cols);
10811
10812 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10813 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10814 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10815 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10816 vbitmatvec2(
10817 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10818 );
10819 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10820 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10821 }
10822
10823 #[test]
10825 fn q4matvec2_equals_two_singles() {
10826 let (rows, cols) = (8, 128);
10827 let groups = rows * cols / GROUP_SIZE;
10828 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10829 for i in 0..groups * 16 {
10830 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10831 }
10832 for g in 0..groups {
10833 let s = 0.01 + 0.003 * g as f32;
10834 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10835 }
10836 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10839 x1[9] = 250.0;
10840 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10841
10842 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10843 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10844 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10845 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10846 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10847 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10848 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10849 }
10850
10851 #[test]
10854 fn matvec_many_equals_separate_matvecs() {
10855 use crate::pool::Pool;
10856 let (r1, r2, cols) = (300, 200, 64);
10857 let mk = |salt: usize, rows: usize| {
10858 QTensor::from_f32(
10859 (0..rows * cols)
10860 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10861 .collect(),
10862 rows,
10863 cols,
10864 )
10865 };
10866 let (a, b) = (mk(1, r1), mk(5, r2));
10867 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10868 let pool = Pool::new(3);
10869
10870 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10871 a.matvec(&x, &mut ea, Some(&pool));
10872 b.matvec(&x, &mut eb, Some(&pool));
10873 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10874 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10875 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10876 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10877 }
10878
10879 #[test]
10882 fn batched_matmat_equals_per_position_matvec() {
10883 let (rows, cols, b) = (8, 64, 5);
10884 let groups = rows * cols / GROUP_SIZE;
10886 let mut q4 = Vec::new();
10887 for i in 0..groups * 16 {
10888 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10889 }
10890 for g in 0..groups {
10891 q4.extend_from_slice(
10892 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10893 );
10894 }
10895 let ng = cols / GROUP_SIZE;
10897 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10898 let mut vb = bits.clone();
10899 for g in 0..rows * ng {
10900 vb.extend_from_slice(
10901 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10902 );
10903 }
10904 for r in 0..rows {
10905 let bw = bits[r] as usize;
10906 let (mut acc, mut nb) = (0u64, 0usize);
10907 let mut rowbytes = Vec::new();
10908 for i in 0..cols {
10909 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10910 acc = (acc << bw) | v;
10911 nb += bw;
10912 while nb >= 8 {
10913 nb -= 8;
10914 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10915 }
10916 }
10917 if nb > 0 {
10918 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10919 }
10920 vb.extend_from_slice(&rowbytes);
10921 }
10922 let offsets = vbit_row_offsets(&vb, rows, cols);
10923
10924 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10925
10926 let mut got = vec![0f32; b * rows];
10928 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
10929 for bi in 0..b {
10930 let mut expect = vec![0f32; rows];
10931 q4matvec(
10932 &q4,
10933 &xs[bi * cols..(bi + 1) * cols],
10934 rows,
10935 cols,
10936 &mut expect,
10937 None,
10938 );
10939 assert_eq!(
10940 &got[bi * rows..(bi + 1) * rows],
10941 &expect[..],
10942 "q4 batch pos {bi}"
10943 );
10944 }
10945
10946 let mut got = vec![0f32; b * rows];
10948 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
10949 for bi in 0..b {
10950 let mut expect = vec![0f32; rows];
10951 vbitmatvec(
10952 &vb,
10953 &offsets,
10954 &xs[bi * cols..(bi + 1) * cols],
10955 rows,
10956 cols,
10957 &mut expect,
10958 None,
10959 );
10960 assert_eq!(
10961 &got[bi * rows..(bi + 1) * rows],
10962 &expect[..],
10963 "vbit batch pos {bi}"
10964 );
10965 }
10966 }
10967
10968 #[test]
10972 fn q4_tiled_matches_q4_block_bitexact() {
10973 let (rows, cols, b) = (8usize, 128usize, 3usize);
10974 let groups = rows * cols / GROUP_SIZE;
10975 let mut split = Vec::with_capacity(groups * 18);
10976 for i in 0..groups * 16 {
10977 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10978 }
10979 for g in 0..groups {
10980 split.extend_from_slice(
10981 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10982 );
10983 }
10984 let (packed, scales) = split.split_at(groups * 16);
10986 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
10987 for g in 0..groups {
10988 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
10989 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
10990 }
10991
10992 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10993 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10995
10996 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
10997 q4matvec(&split, &x1, rows, cols, &mut a, None);
10998 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
10999 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11000
11001 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11002 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11003 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11004 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11005 assert_eq!(a1, t1);
11006 assert_eq!(a2, t2);
11007
11008 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11009 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11010 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11011 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11012 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11013 }
11014
11015 #[test]
11022 fn q4matvec_sdot_outlier_exact() {
11023 let (rows, cols) = (4, 128);
11024 let groups = rows * cols / GROUP_SIZE;
11025 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11026 for i in 0..groups * 16 {
11027 bytes.push(((i * 11 + 5) % 256) as u8);
11028 }
11029 for g in 0..groups {
11030 let s = 0.02 + 0.002 * g as f32;
11031 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11032 }
11033 let mut x: Vec<f32> = (0..cols)
11034 .map(|i| match i % 3 {
11035 0 => 1.0,
11036 1 => -1.0,
11037 _ => 0.0,
11038 })
11039 .collect();
11040 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
11043 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11044 let mut expect = vec![0.0f32; rows];
11045 for r in 0..rows {
11046 expect[r] = reference[r * cols..(r + 1) * cols]
11047 .iter()
11048 .zip(&x)
11049 .map(|(w, xv)| w * xv)
11050 .sum();
11051 }
11052 let mut got = vec![0.0f32; rows];
11053 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11054 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11055 for r in 0..rows {
11056 assert!(
11057 (got[r] - expect[r]).abs() < 2e-3 * scale,
11058 "row {r}: {} vs {} (outlier term must be exact)",
11059 got[r],
11060 expect[r]
11061 );
11062 }
11063 }
11064
11065 #[test]
11069 fn q1t_matvec_matches_reference() {
11070 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11071 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
11073 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11074 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11076 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11077 let mut bytes = Vec::new();
11078 for r in 0..rows {
11079 for g in 0..gpr {
11080 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11081 let mut c = [0u8; 7];
11082 for k in 0..GROUP_SIZE {
11083 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11085 0
11086 } else {
11087 ((k + r * 3 + g) % 3) as u8 };
11089 cortiq_core::quant::q1t_pack(&mut c, k, code);
11090 }
11091 bytes.extend_from_slice(&c);
11092 }
11093 }
11094 let mut row_ptr = vec![0u32; rows + 1];
11097 for &(idx, _) in &outliers {
11098 row_ptr[idx as usize / cols + 1] += 1;
11099 }
11100 for r in 0..rows {
11101 row_ptr[r + 1] += row_ptr[r];
11102 }
11103 for &p in &row_ptr {
11104 bytes.extend_from_slice(&p.to_le_bytes());
11105 }
11106 for &(idx, v) in &outliers {
11107 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11108 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11109 }
11110
11111 let mut refw = vec![0f32; rows * cols];
11112 dequant_q1t(&bytes, rows, cols, &mut refw);
11113 let x: Vec<f32> = (0..cols)
11116 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11117 .collect();
11118 let mut expect = vec![0f32; rows];
11119 for r in 0..rows {
11120 let mut a = 0.0f32;
11121 for j in 0..cols {
11122 a += refw[r * cols + j] * x[j];
11123 }
11124 expect[r] = a;
11125 }
11126 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11127 let mut got = vec![0f32; rows];
11128 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11129 for r in 0..rows {
11130 assert!(
11131 (got[r] - expect[r]).abs() < tol(expect[r]),
11132 "row {r}: {} vs {}",
11133 got[r],
11134 expect[r]
11135 );
11136 }
11137 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
11139 let mut gm = vec![0f32; 2 * rows];
11140 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11141 for r in 0..rows {
11142 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11143 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11144 }
11145 let xb: Vec<f32> = (0..cols)
11149 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11150 .collect();
11151 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11152 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11153 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11154 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11155 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11156 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11157 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11158 }
11159
11160 #[test]
11163 fn q1t_matvec2_odd_gpr_matches_singles() {
11164 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11165 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
11167 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11168 for r in 0..rows {
11169 for g in 0..gpr {
11170 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11171 let mut c = [0u8; 7];
11172 for k in 0..GROUP_SIZE {
11173 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11174 }
11175 bytes.extend_from_slice(&c);
11176 }
11177 }
11178 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11179 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11180 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11181 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11182 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11183 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11184 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11185 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11186 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11187 }
11188
11189 #[test]
11193 #[ignore]
11194 fn q1t_matvec2_speed() {
11195 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11196 use std::time::Instant;
11197 let (rows, cols) = (8192usize, 4096usize);
11198 let gpr = cols / GROUP_SIZE;
11199 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11200 for r in 0..rows {
11201 for g in 0..gpr {
11202 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11203 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11204 let mut c = [0u8; 7];
11205 for k in 0..GROUP_SIZE {
11206 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11207 }
11208 bytes.extend_from_slice(&c);
11209 }
11210 }
11211 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11212 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11213 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11214 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11215 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11217 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11218 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11219 for _ in 0..8 {
11220 let t0 = Instant::now();
11221 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11222 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11223 let t1 = Instant::now();
11224 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11225 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11226 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11227 }
11228 assert_eq!(p1, s1);
11229 assert_eq!(p2, s2);
11230 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11231 }
11232
11233 #[test]
11237 #[ignore]
11238 fn q1t_matvec_speed() {
11239 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11240 use std::time::Instant;
11241 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
11243 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11244 for r in 0..rows {
11245 for g in 0..gpr {
11246 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11247 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11248 let mut c = [0u8; 7];
11249 for k in 0..GROUP_SIZE {
11250 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11251 }
11252 bytes.extend_from_slice(&c);
11253 }
11254 }
11255 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
11257 let mut idx = 0usize;
11258 while idx < n {
11259 row_ptr[idx / cols + 1] += 1;
11260 idx += stride;
11261 }
11262 for r in 0..rows {
11263 row_ptr[r + 1] += row_ptr[r];
11264 }
11265 for &p in &row_ptr {
11266 bytes.extend_from_slice(&p.to_le_bytes());
11267 }
11268 let mut idx = 0usize;
11269 while idx < n {
11270 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11271 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11272 idx += stride;
11273 }
11274 let x: Vec<f32> = (0..cols)
11277 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11278 .collect();
11279 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11280
11281 let slow = |out: &mut [f32]| {
11283 let mut buf = vec![0f32; cols];
11284 for r in 0..rows {
11285 for g in 0..gpr {
11286 let off = (r * gpr + g) * Q1T_TILE;
11287 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11288 let codes = &bytes[off + 2..off + Q1T_TILE];
11289 for k in 0..GROUP_SIZE {
11290 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11291 1 => s,
11292 2 => -s,
11293 _ => 0.0,
11294 };
11295 }
11296 }
11297 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11298 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11299 }
11300 };
11301 let iters = 5;
11302 let mut a = vec![0f32; rows];
11303 slow(&mut a); let t = Instant::now();
11305 for _ in 0..iters {
11306 slow(&mut a);
11307 }
11308 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11309
11310 let mut b = vec![0f32; rows];
11311 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
11313 for _ in 0..iters {
11314 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11315 }
11316 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11317
11318 for r in 0..rows {
11319 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11320 }
11321 println!(
11322 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
11323 slow_ms / fast_ms
11324 );
11325 }
11326}
11327
11328
11329#[cfg(test)]
11330mod gemm_bench {
11331 #[test]
11342 #[ignore]
11343 fn q4tp_matmat_throughput() {
11344 let b: usize = std::env::var("CMF_BENCH_B")
11348 .ok()
11349 .and_then(|v| v.parse().ok())
11350 .unwrap_or(296);
11351 let (rows, cols) = (9216usize, 2304usize);
11352 let (_, _, _) = (rows, cols, b);
11353 let total = cortiq_core::quant::expected_nbytes(
11354 cortiq_core::TensorDtype::Q4TiledP,
11355 &[rows, cols],
11356 )
11357 .unwrap();
11358 let (params_off, codes_off, _) =
11363 cortiq_core::quant::q4tp_sections(rows, cols);
11364 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11365 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11366 let step = cortiq_core::quant::f32_to_f16(0.1);
11367 for r in 0..rows {
11368 let o = params_off + r * 4;
11369 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11370 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11371 }
11372 let _ = codes_off;
11373 let xs: Vec<f32> = (0..b * cols)
11374 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11375 .collect();
11376 let mut out = vec![0f32; b * rows];
11377 let pool = crate::pool::Pool::from_env();
11378 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11384 let reps: usize = std::env::var("CMF_BENCH_REPS")
11385 .ok()
11386 .and_then(|v| v.parse().ok())
11387 .unwrap_or(10);
11388 let mut best = [f64::MAX; 2];
11389 let mut sums = [0f32; 2];
11390 for _ in 0..reps {
11391 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11392 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11393 let t = std::time::Instant::now();
11394 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11395 best[k] = best[k].min(t.elapsed().as_secs_f64());
11396 sums[k] = out.iter().take(64).sum::<f32>();
11397 }
11398 }
11399 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11400 for (k, name) in ["previous", "tuned "].iter().enumerate() {
11401 println!(
11402 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11403 best[k] * 1e3,
11404 flops / best[k] / 1e9,
11405 sums[k]
11406 );
11407 }
11408 assert!(
11409 (sums[0] - sums[1]).abs() < 1e-2,
11410 "the tuned kernel changed the result: {} vs {}",
11411 sums[0],
11412 sums[1]
11413 );
11414 }
11415
11416 #[test]
11422 fn q4tp_matmat_blocked_matches_scalar() {
11423 use std::sync::atomic::Ordering::Relaxed;
11424 for &(rows, cols, b) in &[
11432 (64usize, 128usize, 7usize),
11433 (33, 96, 4),
11434 (16, 256, 9),
11435 (192, 2304, 37),
11436 ] {
11437 let total = cortiq_core::quant::expected_nbytes(
11438 cortiq_core::TensorDtype::Q4TiledP,
11439 &[rows, cols],
11440 )
11441 .unwrap();
11442 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11443 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11444 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11445 let step = cortiq_core::quant::f32_to_f16(0.1);
11446 for r in 0..rows {
11447 let o = params_off + r * 4;
11448 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11449 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11450 }
11451 let xs: Vec<f32> = (0..b * cols)
11452 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11453 .collect();
11454 let mut got = vec![0f32; b * rows];
11455 let mut want = vec![0f32; b * rows];
11456 let gpr = cols / 32;
11457 let view = super::Q4tpView::new(&bytes, rows, cols);
11458 let pool = crate::pool::Pool::from_env();
11459 super::Q4TP_ALT.store(2, Relaxed);
11460 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11461 super::Q4TP_ALT.store(1, Relaxed);
11462 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11463 super::Q4TP_ALT.store(0, Relaxed);
11464 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11471 let (mut worst, mut at) = (0f32, 0usize);
11472 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11473 if (g - w).abs() > worst {
11474 worst = (g - w).abs();
11475 at = i;
11476 }
11477 }
11478 assert!(
11479 worst <= 1e-4 * scale,
11480 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11481 (scale {scale:.3e}) at cell {at}: {} vs {}",
11482 got[at],
11483 want[at]
11484 );
11485
11486 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11493 for bi in 0..b {
11494 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11495 for r in 0..rows {
11496 let mut sc = vec![0f32; gpr];
11497 view.scales_into(r, gpr, &mut sc);
11498 let mut exact = 0f64;
11499 for j in 0..cols {
11500 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11501 exact += w as f64 * sq as f64 * act.xq[j] as f64;
11502 }
11503 exact *= act.sx as f64;
11504 for &(j, xv) in &act.outliers {
11505 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11506 exact += w as f64 * sq as f64 * xv as f64;
11507 }
11508 let i = bi * rows + r;
11509 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11510 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11511 }
11512 }
11513 println!(
11514 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11515 per-column {e_scalar:.3e}"
11516 );
11517 assert!(
11522 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11523 "{rows}x{cols} b={b}: error against f64 too large — blocked \
11524 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11525 );
11526 }
11527 }
11528
11529 #[test]
11533 #[ignore]
11534 fn q4t_matmat_throughput() {
11535 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11536 let total = cortiq_core::quant::expected_nbytes(
11537 cortiq_core::TensorDtype::Q4Tiled,
11538 &[rows, cols],
11539 )
11540 .unwrap();
11541 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11544 let sc = cortiq_core::quant::f32_to_f16(0.02);
11545 for t in bytes.chunks_mut(super::Q4_TILE) {
11546 t[..2].copy_from_slice(&sc.to_le_bytes());
11547 }
11548 let xs: Vec<f32> = (0..b * cols)
11549 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11550 .collect();
11551 let mut out = vec![0f32; b * rows];
11552 let pool = crate::pool::Pool::from_env();
11553 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11554 let reps: usize = std::env::var("CMF_BENCH_REPS")
11555 .ok()
11556 .and_then(|v| v.parse().ok())
11557 .unwrap_or(10);
11558 let mut best = f64::MAX;
11559 for _ in 0..reps {
11560 let t = std::time::Instant::now();
11561 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11562 best = best.min(t.elapsed().as_secs_f64());
11563 }
11564 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11565 println!(
11566 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11567 best * 1e3,
11568 flops / best / 1e9,
11569 out.iter().take(64).sum::<f32>()
11570 );
11571 }
11572
11573}