1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16 GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17 q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23 F32 {
24 data: Vec<f32>,
25 rows: usize,
26 cols: usize,
27 },
28 Mapped {
29 model: Arc<CmfModel>,
30 idx: usize,
32 dtype: TensorDtype,
33 rows: usize,
34 cols: usize,
35 row_scale: Vec<f32>,
37 col_field: Vec<f32>,
39 vbit_offsets: Vec<usize>,
43 repack: Vec<u8>,
51 },
52}
53
54fn repack_enabled() -> bool {
61 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62 *ON.get_or_init(|| {
63 std::env::var("CMF_REPACK")
64 .map(|v| v == "1")
65 .unwrap_or(cfg!(target_os = "android"))
66 })
67}
68
69fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73 #[cfg(target_arch = "aarch64")]
74 let arch_ok = sdot_enabled();
75 #[cfg(not(target_arch = "aarch64"))]
76 let arch_ok = false;
77 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78 return Vec::new();
79 }
80 q8_repack_layout(bytes, rows, cols)
81}
82
83fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86 let groups = rows / 4;
87 let mut rep = vec![0u8; groups * 4 * cols];
88 for g in 0..groups {
89 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90 for c in 0..cols / 16 {
91 for lane in 0..4 {
92 let src = (g * 4 + lane) * cols + c * 16;
93 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94 .copy_from_slice(&bytes[src..src + 16]);
95 }
96 }
97 }
98 rep
99}
100
101fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104 let ng = cols / GROUP_SIZE;
105 let bits = &bytes[..rows];
106 let mut offsets = Vec::with_capacity(rows + 1);
107 let mut off = rows + rows * ng * 2;
108 for r in 0..rows {
109 offsets.push(off);
110 off += (cols * bits[r] as usize).div_ceil(8);
111 }
112 offsets.push(off);
113 offsets
114}
115
116fn blocked_enabled() -> bool {
122 use std::sync::atomic::Ordering::Relaxed;
123 match BLOCKED_OVERRIDE.load(Relaxed) {
124 1 => false,
125 2 => true,
126 _ => {
127 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128 *ON.get_or_init(|| {
129 std::env::var("CMF_X86_BLOCKED")
130 .map(|v| v != "0")
131 .unwrap_or(true)
132 })
133 }
134 }
135}
136
137static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
138
139pub fn set_blocked_override(on: Option<bool>) {
146 let v = match on {
147 None => 0,
148 Some(false) => 1,
149 Some(true) => 2,
150 };
151 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
152}
153
154fn gpu_lmhead_enabled() -> bool {
155 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
156 *ON.get_or_init(|| {
157 std::env::var("CMF_GPU_LMHEAD")
158 .map(|v| v != "0")
159 .unwrap_or(true)
160 })
161}
162
163fn gpu_split_frac() -> f32 {
164 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
165 *F.get_or_init(|| {
166 std::env::var("CMF_GPU_SPLIT")
167 .ok()
168 .and_then(|v| v.parse::<f32>().ok())
169 .unwrap_or(0.5)
170 .clamp(0.0, 1.0)
171 })
172}
173
174impl QTensor {
175 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
176 debug_assert_eq!(data.len(), rows * cols);
177 Self::F32 { data, rows, cols }
178 }
179
180 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
183 let idx = model
186 .tensor_index(name)
187 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
188 let entry = &model.tensors[idx];
189 if entry.shape.len() != 2 {
190 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
191 }
192 let (rows, cols) = (entry.shape[0], entry.shape[1]);
193 let bytes = model.entry_bytes(entry);
194
195 match entry.dtype {
196 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
197 let n = rows * cols;
198 let scales_off = n;
199 let row_scale: Vec<f32> = (0..rows)
200 .map(|o| {
201 f16_to_f32(u16::from_le_bytes([
202 bytes[scales_off + o * 2],
203 bytes[scales_off + o * 2 + 1],
204 ]))
205 })
206 .collect();
207 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
208 let col_off = n + rows * 2;
209 (0..cols)
210 .map(|i| {
211 f16_to_f32(u16::from_le_bytes([
212 bytes[col_off + i * 2],
213 bytes[col_off + i * 2 + 1],
214 ]))
215 })
216 .collect()
217 } else {
218 Vec::new()
219 };
220 Ok(Self::Mapped {
221 model: model.clone(),
222 idx,
223 dtype: entry.dtype,
224 rows,
225 cols,
226 row_scale,
227 col_field,
228 vbit_offsets: Vec::new(),
229 repack: q8_repack(bytes, rows, cols),
230 })
231 }
232 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
234 model: model.clone(),
235 idx,
236 dtype: entry.dtype,
237 rows,
238 cols,
239 row_scale: Vec::new(),
240 col_field: Vec::new(),
241 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
242 repack: Vec::new(),
243 }),
244 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
248 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
249 let offsets: Vec<usize> = (0..=rows)
250 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
251 .collect();
252 Ok(Self::Mapped {
253 model: model.clone(),
254 idx,
255 dtype: entry.dtype,
256 rows,
257 cols,
258 row_scale: Vec::new(),
259 col_field: Vec::new(),
260 vbit_offsets: offsets,
261 repack: Vec::new(),
262 })
263 }
264 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
270 model: model.clone(),
271 idx,
272 dtype: entry.dtype,
273 rows,
274 cols,
275 row_scale: Vec::new(),
276 col_field: Vec::new(),
277 vbit_offsets: Vec::new(),
278 repack: Vec::new(),
279 }),
280 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
283 model: model.clone(),
284 idx,
285 dtype: entry.dtype,
286 rows,
287 cols,
288 row_scale: Vec::new(),
289 col_field: Vec::new(),
290 vbit_offsets: Vec::new(),
291 repack: Vec::new(),
292 }),
293 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
295 model: model.clone(),
296 idx,
297 dtype: entry.dtype,
298 rows,
299 cols,
300 row_scale: Vec::new(),
301 col_field: Vec::new(),
302 vbit_offsets: Vec::new(),
303 repack: Vec::new(),
304 }),
305 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
306 model: model.clone(),
307 idx,
308 dtype: entry.dtype,
309 rows,
310 cols,
311 row_scale: Vec::new(),
312 col_field: Vec::new(),
313 vbit_offsets: Vec::new(),
314 repack: Vec::new(),
315 }),
316 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
318 model: model.clone(),
319 idx,
320 dtype: entry.dtype,
321 rows,
322 cols,
323 row_scale: Vec::new(),
324 col_field: Vec::new(),
325 vbit_offsets: Vec::new(),
326 repack: Vec::new(),
327 }),
328 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
332 model: model.clone(),
333 idx,
334 dtype: entry.dtype,
335 rows,
336 cols,
337 row_scale: Vec::new(),
338 col_field: Vec::new(),
339 vbit_offsets: Vec::new(),
340 repack: Vec::new(),
341 }),
342 _ => {
344 let mut data = vec![0.0f32; rows * cols];
345 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
346 Ok(Self::from_f32(data, rows, cols))
347 }
348 }
349 }
350
351 pub(crate) fn is_q1(&self) -> bool {
354 matches!(
355 self,
356 Self::Mapped {
357 dtype: TensorDtype::Q1,
358 ..
359 }
360 )
361 }
362
363 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
366 match self {
367 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
368 _ => None,
369 }
370 }
371
372 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
379 match self {
380 #[cfg(target_os = "macos")]
381 Self::Mapped {
382 dtype: TensorDtype::Q1T,
383 ..
384 } if !crate::gpu::metal_q1t_enabled() => None,
385 Self::Mapped {
386 idx,
387 dtype:
388 TensorDtype::Q1
389 | TensorDtype::Q1T
390 | TensorDtype::Q4Block
391 | TensorDtype::Q4Tiled
392 | TensorDtype::Q4TiledP
396 | TensorDtype::Q8Row
397 | TensorDtype::Q8_2f,
398 rows,
399 cols,
400 ..
401 } => Some((*idx, *rows, *cols)),
402 _ => None,
403 }
404 }
405
406 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
412 match self {
413 Self::Mapped {
414 idx,
415 dtype: TensorDtype::Q4Tiled,
416 rows,
417 cols,
418 ..
419 } => Some((*idx, *rows, *cols)),
420 _ => None,
421 }
422 }
423
424 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
428 match self {
429 Self::Mapped {
430 idx,
431 dtype: TensorDtype::Q4TiledP,
432 rows,
433 cols,
434 ..
435 } => Some((*idx, *rows, *cols)),
436 _ => None,
437 }
438 }
439
440 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
445 match self {
446 Self::Mapped {
447 idx,
448 dtype: TensorDtype::Q8Row,
449 rows,
450 cols,
451 row_scale,
452 col_field,
453 ..
454 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
455 _ => None,
456 }
457 }
458
459 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
463 match self {
464 Self::Mapped { dtype, .. } => Some(*dtype),
465 _ => None,
466 }
467 }
468
469 pub fn model_idx(&self) -> Option<usize> {
474 match self {
475 Self::Mapped { idx, .. } => Some(*idx),
476 _ => None,
477 }
478 }
479
480 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
485 match self {
486 Self::Mapped { model, .. } => Some(model.clone()),
487 _ => None,
488 }
489 }
490
491 pub fn rows(&self) -> usize {
492 match self {
493 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
494 }
495 }
496
497 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
500 match self {
501 Self::Mapped {
502 model,
503 idx,
504 dtype: TensorDtype::Q4Tiled,
505 ..
506 } => Some((model, *idx)),
507 _ => None,
508 }
509 }
510
511 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
514 match self {
515 Self::Mapped {
516 model,
517 idx,
518 dtype: TensorDtype::Q4TiledP,
519 ..
520 } => Some((model, *idx)),
521 _ => None,
522 }
523 }
524
525 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
528 match self {
529 Self::Mapped {
530 model,
531 idx,
532 dtype: TensorDtype::Q2TiledP,
533 ..
534 } => Some((model, *idx)),
535 _ => None,
536 }
537 }
538
539 pub fn cols(&self) -> usize {
540 match self {
541 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
542 }
543 }
544
545 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
548 match self {
549 Self::Mapped {
550 model,
551 idx,
552 dtype: TensorDtype::Q1,
553 ..
554 } => Some((model, *idx)),
555 _ => None,
556 }
557 }
558
559 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
569 match self {
570 Self::Mapped {
571 model,
572 idx,
573 dtype: TensorDtype::Q8Row,
574 row_scale,
575 ..
576 } => Some((model, *idx, 0, row_scale.as_slice())),
577 Self::Mapped {
578 model,
579 idx,
580 dtype: TensorDtype::Q1,
581 ..
582 } => Some((model, *idx, 1, &[])),
583 Self::Mapped {
588 model,
589 idx,
590 dtype: TensorDtype::Q4Tiled,
591 ..
592 } => Some((model, *idx, 5, &[])),
593 Self::Mapped {
597 model,
598 idx,
599 dtype: TensorDtype::Q4TiledP,
600 ..
601 } => Some((model, *idx, 6, &[])),
602 Self::Mapped {
603 model,
604 idx,
605 dtype: TensorDtype::Q4Block,
606 ..
607 } => Some((model, *idx, 2, &[])),
608 Self::Mapped {
613 model,
614 idx,
615 dtype: TensorDtype::Q8_2f,
616 ..
617 } => Some((model, *idx, 7, &[])),
618 Self::Mapped {
619 model,
620 idx,
621 dtype: TensorDtype::Q1T,
622 ..
623 } => Some((model, *idx, 3, &[])),
624 Self::Mapped {
628 model,
629 idx,
630 dtype: TensorDtype::Q2TiledP,
631 ..
632 } => Some((model, *idx, 9, &[])),
633 _ => None,
634 }
635 }
636
637 pub fn as_f32(&self) -> Option<&[f32]> {
640 match self {
641 Self::F32 { data, .. } => Some(data),
642 Self::Mapped { .. } => None,
643 }
644 }
645
646 fn quant_bytes(&self) -> &[u8] {
647 match self {
648 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
649 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
650 }
651 }
652
653 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
655 let cols = self.cols();
656 debug_assert_eq!(dst.len(), cols);
657 match self {
658 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
659 Self::Mapped {
660 dtype,
661 row_scale,
662 col_field,
663 vbit_offsets,
664 ..
665 } => {
666 if *dtype == TensorDtype::Q4Tiled {
667 let bytes = self.quant_bytes();
668 let gpr = cols / GROUP_SIZE;
669 for gi in 0..gpr {
670 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
671 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
672 for (k, &b) in tile[2..].iter().enumerate() {
673 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
674 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
675 }
676 }
677 return;
678 }
679 if *dtype == TensorDtype::Q4TiledP {
680 let bytes = self.quant_bytes();
681 let gpr = cols / GROUP_SIZE;
682 let v = Q4tpView::new(bytes, self.rows(), cols);
683 let mut sc = vec![0f32; gpr];
684 v.scales_into(r, gpr, &mut sc);
685 for gi in 0..gpr {
686 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
687 let s = sc[gi];
688 for (k, &b) in tile.iter().enumerate() {
689 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
690 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
691 }
692 }
693 return;
694 }
695 if *dtype == TensorDtype::Q2TiledP {
696 let bytes = self.quant_bytes();
697 let gpr = cols / GROUP_SIZE;
698 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
699 let mut sc = vec![0f32; gpr];
700 v.scales_into(r, gpr, &mut sc);
701 for gi in 0..gpr {
702 let ch =
703 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
704 let s = sc[gi];
705 for (k, &b) in ch.iter().enumerate() {
706 for j in 0..4 {
707 dst[gi * GROUP_SIZE + k * 4 + j] =
708 (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
709 }
710 }
711 }
712 return;
713 }
714 if *dtype == TensorDtype::Q4Block {
715 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
716 let gpr = cols / GROUP_SIZE;
717 for gi in 0..gpr {
718 let g = r * gpr + gi;
719 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
720 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
721 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
722 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
723 }
724 }
725 return;
726 }
727 if *dtype == TensorDtype::Q1 {
728 let bytes = self.quant_bytes();
729 let gpr = cols / GROUP_SIZE;
730 for gi in 0..gpr {
731 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
732 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
733 for (j, &b) in tile[2..].iter().enumerate() {
734 for k in 0..8 {
735 dst[gi * GROUP_SIZE + j * 8 + k] =
736 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
737 }
738 }
739 }
740 return;
741 }
742 if *dtype == TensorDtype::Q1T {
743 let bytes = self.quant_bytes();
744 let gpr = cols / GROUP_SIZE;
745 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
746 for gi in 0..gpr {
747 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
748 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
749 bytes[off],
750 bytes[off + 1],
751 ]));
752 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
753 for k in 0..GROUP_SIZE {
754 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
755 {
756 1 => s,
757 2 => -s,
758 _ => 0.0,
759 };
760 }
761 }
762 let rows = self.rows();
764 let entries = base_len + (rows + 1) * 4;
765 if entries <= bytes.len() {
766 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
767 let r0 = u32::from_le_bytes([
768 ptrs[r * 4],
769 ptrs[r * 4 + 1],
770 ptrs[r * 4 + 2],
771 ptrs[r * 4 + 3],
772 ]) as usize;
773 let r1 = u32::from_le_bytes([
774 ptrs[(r + 1) * 4],
775 ptrs[(r + 1) * 4 + 1],
776 ptrs[(r + 1) * 4 + 2],
777 ptrs[(r + 1) * 4 + 3],
778 ]) as usize;
779 let off = entries + r0 * 4;
780 for i in 0..r1 - r0 {
781 let item = &bytes[off + i * 4..off + i * 4 + 4];
782 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
783 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
784 item[2], item[3],
785 ]));
786 if c < cols {
787 dst[c] = v;
788 }
789 }
790 }
791 return;
792 }
793 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
794 let bytes = self.quant_bytes();
795 let rows = self.rows();
796 let ng = cols / GROUP_SIZE;
797 let bits = &bytes[..rows];
798 let sc_off = rows;
799 let off = vbit_offsets[r];
802 let b = bits[r] as usize;
803 let l = ((1usize << (b - 1)) - 1) as f32;
804 let data = &bytes[off..];
805 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
806 for (i, d) in dst.iter_mut().enumerate() {
807 while nbits < b {
808 acc = (acc << 8) | data[idx] as u64;
809 idx += 1;
810 nbits += 8;
811 }
812 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
813 nbits -= b;
814 let so = (r * ng + i / GROUP_SIZE) * 2;
815 let sv = f16_to_f32(u16::from_le_bytes([
816 bytes[sc_off + so],
817 bytes[sc_off + so + 1],
818 ]));
819 *d = (u - l) * sv;
820 }
821 return;
822 }
823 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
824 let s = row_scale[r];
825 match dtype {
826 TensorDtype::Q8Row => {
827 for (d, &b) in dst.iter_mut().zip(q) {
828 *d = (b as i8) as f32 * s;
829 }
830 }
831 TensorDtype::Q8_2f => {
832 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
833 *d = (b as i8) as f32 * s * col_field[i];
834 }
835 }
836 _ => unreachable!(),
837 }
838 }
839 }
840 }
841
842 pub fn sparse_col_ok(&self) -> bool {
847 match self {
848 Self::F32 { .. } => true,
849 Self::Mapped { dtype, .. } => {
850 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
851 }
852 }
853 }
854
855 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
859 let inter = self.cols();
860 let hidden = self.rows();
861 debug_assert_eq!(out.len(), hidden);
862 match self {
863 Self::F32 { data, .. } => {
864 for (k, o) in out.iter_mut().enumerate() {
865 *o += w * data[k * inter + c];
866 }
867 }
868 Self::Mapped {
869 dtype,
870 row_scale,
871 col_field,
872 ..
873 } => {
874 let q = self.quant_bytes();
875 let colf = if *dtype == TensorDtype::Q8_2f {
876 col_field[c]
877 } else {
878 1.0
879 };
880 let wc = w * colf;
881 for (k, o) in out.iter_mut().enumerate() {
882 let b = q[k * inter + c] as i8 as f32;
883 *o += wc * b * row_scale[k];
884 }
885 }
886 }
887 }
888
889 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
893 let cols = self.cols();
894 match self {
895 Self::F32 { data, .. } => {
896 let row = &data[r * cols..(r + 1) * cols];
897 row.iter().zip(x).map(|(w, v)| w * v).sum()
898 }
899 Self::Mapped {
900 dtype,
901 row_scale,
902 col_field,
903 ..
904 } => match dtype {
905 TensorDtype::Q8Row => {
906 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
907 dot_i8_f32(q, x) * row_scale[r]
908 }
909 TensorDtype::Q8_2f => {
910 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
911 dot_i8_col_f32(q, x, col_field) * row_scale[r]
912 }
913 _ => {
914 self.row_f32(r, scratch);
915 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
916 }
917 },
918 }
919 }
920
921 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
924 match self {
925 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
929 Self::Mapped {
930 model,
931 idx,
932 dtype,
933 rows,
934 cols,
935 row_scale,
936 col_field,
937 vbit_offsets,
938 repack,
939 } => {
940 let _ = (model, idx);
941 assert!(
951 out.len() >= *rows && x.len() >= *cols,
952 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
953 out.len(),
954 x.len(),
955 );
956 if *dtype == TensorDtype::Q4Block {
957 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
961 let t0 = std::time::Instant::now();
962 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
963 crate::gpu::ProbeArm::Gpu => {
964 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
965 crate::gpu::probe_record(
966 crate::gpu::OpClass::Matvec,
967 true,
968 t0.elapsed(),
969 );
970 return;
971 }
972 }
973 crate::gpu::ProbeArm::CpuTimed => {
974 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
975 crate::gpu::probe_record(
976 crate::gpu::OpClass::Matvec,
977 false,
978 t0.elapsed(),
979 );
980 return;
981 }
982 crate::gpu::ProbeArm::Cpu => {}
983 }
984 }
985 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
986 return;
987 }
988 if *dtype == TensorDtype::Q4Tiled {
989 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
994 let t0 = std::time::Instant::now();
995 let cls = crate::gpu::matvec_class(*rows, *cols);
996 match crate::gpu::probe_arm(cls) {
997 crate::gpu::ProbeArm::Gpu => {
998 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
999 crate::gpu::probe_record(cls, true, t0.elapsed());
1000 return;
1001 }
1002 }
1003 crate::gpu::ProbeArm::CpuTimed => {
1004 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1005 crate::gpu::probe_record(cls, false, t0.elapsed());
1006 return;
1007 }
1008 crate::gpu::ProbeArm::Cpu => {}
1009 }
1010 }
1011 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1012 return;
1013 }
1014 if *dtype == TensorDtype::Q4TiledP {
1015 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1021 let t0 = std::time::Instant::now();
1022 let cls = crate::gpu::matvec_class(*rows, *cols);
1023 match crate::gpu::probe_arm(cls) {
1024 crate::gpu::ProbeArm::Gpu => {
1025 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1026 crate::gpu::probe_record(cls, true, t0.elapsed());
1027 return;
1028 }
1029 }
1030 crate::gpu::ProbeArm::CpuTimed => {
1031 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1032 crate::gpu::probe_record(cls, false, t0.elapsed());
1033 return;
1034 }
1035 crate::gpu::ProbeArm::Cpu => {}
1036 }
1037 }
1038 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1039 return;
1040 }
1041 if *dtype == TensorDtype::Q2TiledP {
1042 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1043 return;
1044 }
1045 if *dtype == TensorDtype::Q1 {
1046 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1051 let t0 = std::time::Instant::now();
1052 let arm = if crate::gpu::q1_force() {
1053 crate::gpu::ProbeArm::Gpu
1054 } else {
1055 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1056 };
1057 match arm {
1058 crate::gpu::ProbeArm::Gpu => {
1059 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1060 crate::gpu::probe_record(
1061 crate::gpu::OpClass::Matvec,
1062 true,
1063 t0.elapsed(),
1064 );
1065 return;
1066 }
1067 }
1068 crate::gpu::ProbeArm::CpuTimed => {
1069 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1070 crate::gpu::probe_record(
1071 crate::gpu::OpClass::Matvec,
1072 false,
1073 t0.elapsed(),
1074 );
1075 return;
1076 }
1077 crate::gpu::ProbeArm::Cpu => {}
1078 }
1079 }
1080 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1081 return;
1082 }
1083 if *dtype == TensorDtype::Q1T {
1084 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1088 let t0 = std::time::Instant::now();
1089 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1090 crate::gpu::ProbeArm::Gpu => {
1091 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1092 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1093 crate::gpu::probe_record(
1094 crate::gpu::OpClass::Matvec,
1095 true,
1096 t0.elapsed(),
1097 );
1098 return;
1099 }
1100 }
1101 crate::gpu::ProbeArm::CpuTimed => {
1102 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1103 crate::gpu::probe_record(
1104 crate::gpu::OpClass::Matvec,
1105 false,
1106 t0.elapsed(),
1107 );
1108 return;
1109 }
1110 crate::gpu::ProbeArm::Cpu => {}
1111 }
1112 }
1113 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1114 return;
1115 }
1116 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1117 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1118 return;
1119 }
1120 let xs = prescale(x, col_field, *dtype);
1121 if *rows >= crate::gpu::min_rows()
1126 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1127 && gpu_lmhead_enabled()
1128 && crate::gpu::enabled_here()
1129 {
1130 let t0 = std::time::Instant::now();
1133 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1134 crate::gpu::ProbeArm::Gpu => {}
1135 crate::gpu::ProbeArm::CpuTimed => {
1136 qmatvec(
1137 self.quant_bytes(),
1138 repack,
1139 row_scale,
1140 x,
1141 col_field,
1142 *dtype,
1143 *rows,
1144 *cols,
1145 out,
1146 pool,
1147 );
1148 crate::gpu::probe_record(
1149 crate::gpu::OpClass::Matvec,
1150 false,
1151 t0.elapsed(),
1152 );
1153 return;
1154 }
1155 crate::gpu::ProbeArm::Cpu => {
1156 qmatvec(
1157 self.quant_bytes(),
1158 repack,
1159 row_scale,
1160 x,
1161 col_field,
1162 *dtype,
1163 *rows,
1164 *cols,
1165 out,
1166 pool,
1167 );
1168 return;
1169 }
1170 }
1171 let frac = gpu_split_frac();
1172 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1173 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1174 let bytes = self.quant_bytes();
1175 let ok = std::thread::scope(|sc| {
1176 let g = sc.spawn(|| {
1177 crate::gpu::q8_matvec_range(
1178 model,
1179 *idx,
1180 cpu_rows,
1181 &row_scale[cpu_rows..],
1182 &xs,
1183 *rows - cpu_rows,
1184 *cols,
1185 out_gpu,
1186 )
1187 });
1188 if cpu_rows > 0 {
1189 let rep_cpu = if repack.is_empty() {
1192 &[][..]
1193 } else {
1194 &repack[..(cpu_rows / 4) * 4 * *cols]
1195 };
1196 qmatvec(
1197 &bytes[..cpu_rows * *cols],
1198 rep_cpu,
1199 &row_scale[..cpu_rows],
1200 x,
1201 col_field,
1202 *dtype,
1203 cpu_rows,
1204 *cols,
1205 out_cpu,
1206 pool,
1207 );
1208 }
1209 g.join().unwrap_or(false)
1210 });
1211 if ok {
1212 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1213 return;
1214 }
1215 qmatvec(
1218 &bytes[cpu_rows * *cols..(*rows) * *cols],
1219 &[],
1220 &row_scale[cpu_rows..],
1221 x,
1222 col_field,
1223 *dtype,
1224 *rows - cpu_rows,
1225 *cols,
1226 out_gpu,
1227 pool,
1228 );
1229 return;
1230 }
1231 qmatvec(
1232 self.quant_bytes(),
1233 repack,
1234 row_scale,
1235 x,
1236 col_field,
1237 *dtype,
1238 *rows,
1239 *cols,
1240 out,
1241 pool,
1242 );
1243 }
1244 }
1245 }
1246
1247 pub fn matvec2(
1249 &self,
1250 x1: &[f32],
1251 x2: &[f32],
1252 o1: &mut [f32],
1253 o2: &mut [f32],
1254 pool: Option<&Pool>,
1255 ) {
1256 match self {
1257 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1258 Self::Mapped {
1259 dtype,
1260 rows,
1261 cols,
1262 row_scale,
1263 col_field,
1264 vbit_offsets,
1265 ..
1266 } => {
1267 if *dtype == TensorDtype::Q4Block {
1268 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1269 return;
1270 }
1271 if *dtype == TensorDtype::Q4Tiled {
1272 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1273 return;
1274 }
1275 if *dtype == TensorDtype::Q4TiledP {
1276 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1277 return;
1278 }
1279 if *dtype == TensorDtype::Q2TiledP {
1280 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1281 return;
1282 }
1283 if *dtype == TensorDtype::Q1 {
1284 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1285 return;
1286 }
1287 if *dtype == TensorDtype::Q1T {
1288 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1294 return;
1295 }
1296 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1297 vbitmatvec2(
1298 self.quant_bytes(),
1299 vbit_offsets,
1300 x1,
1301 x2,
1302 *rows,
1303 *cols,
1304 o1,
1305 o2,
1306 pool,
1307 );
1308 return;
1309 }
1310 qmatvec2(
1311 self.quant_bytes(),
1312 row_scale,
1313 x1,
1314 x2,
1315 col_field,
1316 *dtype,
1317 *rows,
1318 *cols,
1319 o1,
1320 o2,
1321 pool,
1322 );
1323 }
1324 }
1325 }
1326}
1327
1328impl QTensor {
1329 pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1337 match self {
1338 Self::Mapped {
1339 model, idx, dtype, ..
1340 } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1341 _ => None,
1342 }
1343 }
1344
1345 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1346 let cols = self.cols();
1347 let rows = self.rows();
1348 debug_assert_eq!(xs_all.len(), b * cols);
1349 debug_assert_eq!(out.len(), b * rows);
1350 if crate::gptq_capture::capturing() {
1354 if let Self::Mapped { model, idx, .. } = self {
1355 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1356 }
1357 }
1358 match self {
1359 Self::F32 { data, .. } => {
1360 let out_addr = SendMut(out.as_mut_ptr());
1361 let run = |start: usize, end: usize| {
1362 for o in start..end {
1363 let row = &data[o * cols..(o + 1) * cols];
1364 for bi in 0..b {
1365 let x = &xs_all[bi * cols..(bi + 1) * cols];
1366 let mut acc = 0f32;
1367 for j in 0..cols {
1368 acc += row[j] * x[j];
1369 }
1370 unsafe { *out_addr.at(bi * rows + o) = acc };
1371 }
1372 }
1373 };
1374 dispatch_rows(pool, rows, &run);
1375 }
1376 Self::Mapped {
1377 dtype,
1378 row_scale,
1379 col_field,
1380 vbit_offsets,
1381 ..
1382 } => {
1383 if *dtype == TensorDtype::Q4Block {
1384 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1385 return;
1386 }
1387 if *dtype == TensorDtype::Q4TiledP {
1388 if b >= 32
1400 && b * rows * cols >= 128_000_000
1401 && cols % 32 == 0
1402 && !crate::gpu::mm_killed()
1403 && crate::gpu::enabled_here()
1404 {
1405 let class = if b >= 128 {
1406 crate::gpu::OpClass::MatmatWide
1407 } else {
1408 crate::gpu::OpClass::Matmat
1409 };
1410 if let Self::Mapped { model, idx, .. } = self {
1411 if crate::mm_ab::on() {
1423 let mut g = vec![0f32; b * rows];
1424 let t = std::time::Instant::now();
1425 let took = crate::gpu::q4tp_matmat(
1426 model, *idx, xs_all, b, rows, cols, &mut g,
1427 );
1428 let dg = t.elapsed();
1429 let t = std::time::Instant::now();
1430 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1431 let dc = t.elapsed();
1432 crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1433 return;
1434 }
1435 let t0 = std::time::Instant::now();
1436 let resident = crate::gpu::weight_is_resident(model, *idx);
1440 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1441 crate::gpu::ProbeArm::Gpu => {
1442 if crate::gpu::q4tp_matmat(
1443 model, *idx, xs_all, b, rows, cols, out,
1444 ) {
1445 let el = t0.elapsed();
1446 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1456 let budget = std::time::Duration::from_secs_f64(
1457 flops / 1.5e12 * 8.0 + 0.020,
1458 );
1459 crate::gpu::mm_budget_check(
1460 "q4tp matmat",
1461 el,
1462 budget,
1463 crate::gpu::probe_was_cold() || !resident,
1464 );
1465 crate::gpu::probe_record(class, true, el);
1466 return;
1467 }
1468 }
1469 crate::gpu::ProbeArm::CpuTimed => {
1470 q4tp_matmat(
1471 self.quant_bytes(),
1472 xs_all,
1473 b,
1474 rows,
1475 cols,
1476 out,
1477 pool,
1478 );
1479 crate::gpu::probe_record(class, false, t0.elapsed());
1480 return;
1481 }
1482 crate::gpu::ProbeArm::Cpu => {}
1483 }
1484 }
1485 }
1486 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1487 return;
1488 }
1489 if *dtype == TensorDtype::Q2TiledP {
1490 if b >= 32
1496 && b * rows * cols >= 128_000_000
1497 && cols % 32 == 0
1498 && !crate::gpu::mm_killed()
1499 && crate::gpu::enabled_here()
1500 {
1501 let class = if b >= 128 {
1502 crate::gpu::OpClass::MatmatWide
1503 } else {
1504 crate::gpu::OpClass::Matmat
1505 };
1506 if let Self::Mapped { model, idx, .. } = self {
1507 let t0 = std::time::Instant::now();
1508 match crate::gpu::probe_arm(class) {
1509 crate::gpu::ProbeArm::Gpu => {
1510 if crate::gpu::q2tp_matmat(
1511 model, *idx, xs_all, b, rows, cols, out,
1512 ) {
1513 crate::gpu::probe_record(class, true, t0.elapsed());
1514 return;
1515 }
1516 }
1517 crate::gpu::ProbeArm::CpuTimed => {
1518 q2tp_matmat(
1519 self.quant_bytes(),
1520 xs_all,
1521 b,
1522 rows,
1523 cols,
1524 out,
1525 pool,
1526 );
1527 crate::gpu::probe_record(class, false, t0.elapsed());
1528 return;
1529 }
1530 crate::gpu::ProbeArm::Cpu => {}
1531 }
1532 }
1533 }
1534 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1539 return;
1540 }
1541 if *dtype == TensorDtype::Q4Tiled {
1542 if b >= 32
1554 && b * rows * cols >= 128_000_000
1555 && cols % 32 == 0
1556 && !crate::gpu::mm_killed()
1557 && crate::gpu::enabled_here()
1558 {
1559 let class = if b >= 128 {
1560 crate::gpu::OpClass::MatmatWide
1561 } else {
1562 crate::gpu::OpClass::Matmat
1563 };
1564 if let Self::Mapped { model, idx, .. } = self {
1565 let t0 = std::time::Instant::now();
1566 match crate::gpu::probe_arm(class) {
1567 crate::gpu::ProbeArm::Gpu => {
1568 if crate::gpu::q4t_matmat(
1569 model, *idx, xs_all, b, rows, cols, out,
1570 ) {
1571 let el = t0.elapsed();
1572 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1582 let budget = std::time::Duration::from_secs_f64(
1583 flops / 1.5e12 * 8.0 + 0.020,
1584 );
1585 crate::gpu::mm_budget_check(
1586 "q4t matmat",
1587 el,
1588 budget,
1589 crate::gpu::probe_was_cold(),
1590 );
1591 crate::gpu::probe_record(class, true, el);
1592 return;
1593 }
1594 }
1595 crate::gpu::ProbeArm::CpuTimed => {
1596 q4t_matmat(
1597 self.quant_bytes(),
1598 xs_all,
1599 b,
1600 rows,
1601 cols,
1602 out,
1603 pool,
1604 );
1605 crate::gpu::probe_record(class, false, t0.elapsed());
1606 return;
1607 }
1608 crate::gpu::ProbeArm::Cpu => {}
1609 }
1610 }
1611 }
1612 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1613 return;
1614 }
1615 if *dtype == TensorDtype::Q1 {
1616 if b >= 32
1619 && b * rows * cols >= 128_000_000
1620 && cols % 64 == 0
1621 && crate::gpu::enabled_here()
1622 {
1623 if let Self::Mapped { model, idx, .. } = self {
1624 let t0 = std::time::Instant::now();
1625 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1626 crate::gpu::ProbeArm::Gpu => {
1627 if crate::gpu::q1_matmat(
1628 model, *idx, xs_all, b, rows, cols, out,
1629 ) {
1630 crate::gpu::probe_record(
1631 crate::gpu::OpClass::Matmat,
1632 true,
1633 t0.elapsed(),
1634 );
1635 return;
1636 }
1637 }
1638 crate::gpu::ProbeArm::CpuTimed => {
1639 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1640 crate::gpu::probe_record(
1641 crate::gpu::OpClass::Matmat,
1642 false,
1643 t0.elapsed(),
1644 );
1645 return;
1646 }
1647 crate::gpu::ProbeArm::Cpu => {}
1648 }
1649 }
1650 }
1651 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1652 return;
1653 }
1654 if *dtype == TensorDtype::Q1T {
1655 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1658 if let Self::Mapped { model, idx, .. } = self {
1659 let t0 = std::time::Instant::now();
1660 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1661 crate::gpu::ProbeArm::Gpu => {
1662 if crate::gpu::q1t_matmat(
1663 model, *idx, xs_all, b, rows, cols, out,
1664 ) {
1665 crate::gpu::probe_record(
1666 crate::gpu::OpClass::Matmat,
1667 true,
1668 t0.elapsed(),
1669 );
1670 return;
1671 }
1672 }
1673 crate::gpu::ProbeArm::CpuTimed => {
1674 q1t_matmat(
1675 self.quant_bytes(),
1676 xs_all,
1677 b,
1678 rows,
1679 cols,
1680 out,
1681 pool,
1682 );
1683 crate::gpu::probe_record(
1684 crate::gpu::OpClass::Matmat,
1685 false,
1686 t0.elapsed(),
1687 );
1688 return;
1689 }
1690 crate::gpu::ProbeArm::Cpu => {}
1691 }
1692 }
1693 }
1694 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1695 return;
1696 }
1697 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1698 vbitmatmat(
1699 self.quant_bytes(),
1700 vbit_offsets,
1701 xs_all,
1702 b,
1703 rows,
1704 cols,
1705 out,
1706 pool,
1707 );
1708 return;
1709 }
1710 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1711 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1712 .collect();
1713 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1719 if let Self::Mapped { model, idx, .. } = self {
1720 let t0 = std::time::Instant::now();
1721 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1722 crate::gpu::ProbeArm::Gpu
1723 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1724 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1725 {
1726 let q = self.quant_bytes();
1730 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1731 return;
1732 }
1733 crate::gpu::ProbeArm::Gpu => {
1734 let flat: Vec<f32> =
1735 pre.iter().flat_map(|v| v.iter().copied()).collect();
1736 if crate::gpu::q8_matmat(
1737 model, *idx, row_scale, &flat, b, rows, cols, out,
1738 ) {
1739 crate::gpu::probe_record(
1740 crate::gpu::OpClass::Matmat,
1741 true,
1742 t0.elapsed(),
1743 );
1744 return;
1745 }
1746 }
1747 crate::gpu::ProbeArm::CpuTimed => {
1748 let q = self.quant_bytes();
1749 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1750 crate::gpu::probe_record(
1751 crate::gpu::OpClass::Matmat,
1752 false,
1753 t0.elapsed(),
1754 );
1755 return;
1756 }
1757 crate::gpu::ProbeArm::Cpu => {}
1758 }
1759 }
1760 }
1761 let q = self.quant_bytes();
1762 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1763 }
1764 }
1765 }
1766}
1767
1768impl QTensor {
1769 pub fn matvec_many<const N: usize>(
1776 ts: [&QTensor; N],
1777 x: &[f32],
1778 mut outs: [&mut [f32]; N],
1779 pool: Option<&Pool>,
1780 ) {
1781 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1782 let uniform_q8 = ts.iter().all(|t| {
1783 matches!(
1784 t,
1785 Self::Mapped {
1786 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1787 ..
1788 }
1789 )
1790 });
1791 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1792 let uniform_q4 = ts.iter().all(|t| {
1793 matches!(
1794 t,
1795 Self::Mapped {
1796 dtype: TensorDtype::Q4Block,
1797 ..
1798 }
1799 )
1800 });
1801 let uniform_vbit = ts.iter().all(|t| {
1802 matches!(
1803 t,
1804 Self::Mapped {
1805 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1806 ..
1807 }
1808 )
1809 });
1810 let uniform_q1 = ts.iter().all(|t| {
1811 matches!(
1812 t,
1813 Self::Mapped {
1814 dtype: TensorDtype::Q1,
1815 ..
1816 }
1817 )
1818 });
1819 let uniform_q1t = ts.iter().all(|t| {
1820 matches!(
1821 t,
1822 Self::Mapped {
1823 dtype: TensorDtype::Q1T,
1824 ..
1825 }
1826 )
1827 });
1828 let uniform_q4tp = ts.iter().all(|t| {
1833 matches!(
1834 t,
1835 Self::Mapped {
1836 dtype: TensorDtype::Q4TiledP,
1837 ..
1838 }
1839 )
1840 }) && ts
1841 .iter()
1842 .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1843 let Some(pool) = pool else {
1844 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1845 t.matvec(x, o, None);
1846 }
1847 return;
1848 };
1849 if total_rows < 256
1850 || !(uniform_q8
1851 || uniform_f32
1852 || uniform_q4
1853 || uniform_vbit
1854 || uniform_q1
1855 || uniform_q1t
1856 || uniform_q4tp)
1857 {
1858 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1859 t.matvec(x, o, Some(pool));
1860 }
1861 return;
1862 }
1863
1864 if uniform_q4tp {
1865 let cols = ts[0].cols();
1871 let gpr = cols / GROUP_SIZE;
1872 let views: Vec<Q4tpView> = ts
1873 .iter()
1874 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1875 .collect();
1876 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1877 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1878 let locate = |flat: usize| -> (usize, usize) {
1880 let mut acc = 0;
1881 for (i, &r) in rows_of.iter().enumerate() {
1882 if flat < acc + r {
1883 return (i, flat - acc);
1884 }
1885 acc += r;
1886 }
1887 (rows_of.len() - 1, 0)
1888 };
1889 let (views, outs_addr) = (&views, &outs_addr);
1890 if a8w8_enabled() {
1891 let act = split_act(x);
1892 let act = &act;
1893 let run = |start: usize, end: usize| {
1894 let mut sc = vec![0f32; gpr];
1895 for flat in start..end {
1896 let (t, r) = locate(flat);
1897 let v = &views[t];
1898 v.scales_into(r, gpr, &mut sc);
1899 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1900 for &(j, xv) in &act.outliers {
1901 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1902 acc += w * s * xv;
1903 }
1904 unsafe { *outs_addr[t].at(r) = acc };
1906 }
1907 };
1908 pool.run_rows(total_rows, &run);
1909 } else {
1910 let run = |start: usize, end: usize| {
1911 let mut sc = vec![0f32; gpr];
1912 for flat in start..end {
1913 let (t, r) = locate(flat);
1914 let v = &views[t];
1915 v.scales_into(r, gpr, &mut sc);
1916 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1918 }
1919 };
1920 pool.run_rows(total_rows, &run);
1921 }
1922 return;
1923 }
1924
1925 if uniform_q1 {
1926 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1929 if a8w8_enabled() {
1930 let act = split_act(x);
1931 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1932 let (act, gsum) = (&act, &gsum);
1933 let closures: [_; N] = std::array::from_fn(|i| {
1934 let (bytes, gpr, out) =
1935 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1936 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
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 } else {
1942 let closures: [_; N] = std::array::from_fn(|i| {
1943 let (bytes, gpr, out) =
1944 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1945 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1946 });
1947 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1948 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1949 pool.run_many(&parts);
1950 }
1951 return;
1952 }
1953
1954 if uniform_q1t {
1955 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1959 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1960 if a8w8_enabled() {
1961 let act = split_act(x);
1962 let act = &act;
1963 let x_ref = x;
1964 let closures: [_; N] = std::array::from_fn(|i| {
1965 let bytes = ts[i].quant_bytes();
1966 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1967 let gpr = cols / GROUP_SIZE;
1968 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1969 let out = outs_addr[i];
1970 move |s: usize, e: usize| {
1971 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1972 }
1973 });
1974 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1975 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1976 pool.run_many(&parts);
1977 } else {
1978 let x_ref = x;
1979 let closures: [_; N] = std::array::from_fn(|i| {
1980 let bytes = ts[i].quant_bytes();
1981 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1982 let gpr = cols / GROUP_SIZE;
1983 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1984 let out = outs_addr[i];
1985 move |s: usize, e: usize| {
1986 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1987 }
1988 });
1989 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1990 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1991 pool.run_many(&parts);
1992 }
1993 return;
1994 }
1995
1996 if uniform_q4 || uniform_vbit {
1997 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1998 if a8w8_enabled() {
2000 let act = split_act(x);
2001 let act = &act;
2002 if uniform_q4 {
2003 let closures: [_; N] = std::array::from_fn(|i| {
2004 let (packed, scales) =
2005 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2006 let (gpr, cols, out) =
2007 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2008 move |s: usize, e: usize| {
2009 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2010 }
2011 });
2012 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2013 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2014 pool.run_many(&parts);
2015 } else {
2016 let closures: [_; N] = std::array::from_fn(|i| {
2017 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2018 unreachable!()
2019 };
2020 let (bytes, rows, cols, out) = (
2021 ts[i].quant_bytes(),
2022 ts[i].rows(),
2023 ts[i].cols(),
2024 outs_addr[i],
2025 );
2026 move |s: usize, e: usize| {
2027 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2028 }
2029 });
2030 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2031 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2032 pool.run_many(&parts);
2033 }
2034 return;
2035 }
2036 if uniform_q4 {
2037 let closures: [_; N] = std::array::from_fn(|i| {
2038 let (packed, scales) =
2039 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2040 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2041 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2042 });
2043 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2044 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2045 pool.run_many(&parts);
2046 } else {
2047 let closures: [_; N] = std::array::from_fn(|i| {
2048 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2049 unreachable!()
2050 };
2051 let (bytes, rows, cols, out) = (
2052 ts[i].quant_bytes(),
2053 ts[i].rows(),
2054 ts[i].cols(),
2055 outs_addr[i],
2056 );
2057 move |s: usize, e: usize| {
2058 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2059 }
2060 });
2061 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2062 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2063 pool.run_many(&parts);
2064 }
2065 return;
2066 }
2067
2068 if uniform_f32 {
2069 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2070 let closures: [_; N] = std::array::from_fn(|i| {
2071 let Self::F32 { data, cols, .. } = ts[i] else {
2072 unreachable!()
2073 };
2074 let out = outs_addr[i];
2075 move |start: usize, end: usize| {
2076 for o in start..end {
2077 let row = &data[o * cols..(o + 1) * cols];
2078 let mut sum = 0.0f32;
2079 for j in 0..*cols {
2080 sum += row[j] * x[j];
2081 }
2082 unsafe { *out.at(o) = sum };
2084 }
2085 }
2086 });
2087 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2088 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2089 pool.run_many(&parts);
2090 return;
2091 }
2092
2093 struct Ctx<'a> {
2096 bytes: &'a [u8],
2097 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2098 rep: &'a [u8],
2099 row_scale: &'a [f32],
2100 cols: usize,
2101 xs: std::borrow::Cow<'a, [f32]>,
2102 }
2103 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2104 let Self::Mapped {
2105 dtype,
2106 cols,
2107 row_scale,
2108 col_field,
2109 repack,
2110 ..
2111 } = ts[i]
2112 else {
2113 unreachable!()
2114 };
2115 Ctx {
2116 bytes: ts[i].quant_bytes(),
2117 rep: repack,
2118 row_scale,
2119 cols: *cols,
2120 xs: prescale(x, col_field, *dtype),
2121 }
2122 });
2123 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2124 #[cfg(target_arch = "aarch64")]
2125 if sdot_enabled() {
2126 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2127 let closures: [_; N] = std::array::from_fn(|i| {
2128 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2129 move |start: usize, end: usize| {
2130 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2131 }
2132 });
2133 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2134 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2135 pool.run_many(&parts);
2136 return;
2137 }
2138 #[cfg(target_arch = "x86_64")]
2139 if avx2_a8w8_enabled() {
2140 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2141 let closures: [_; N] = std::array::from_fn(|i| {
2142 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2143 move |start: usize, end: usize| {
2144 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2145 }
2146 });
2147 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2148 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2149 pool.run_many(&parts);
2150 return;
2151 }
2152 let closures: [_; N] = std::array::from_fn(|i| {
2153 let (c, out) = (&ctxs[i], outs_addr[i]);
2154 move |start: usize, end: usize| {
2155 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2156 }
2157 });
2158 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2159 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2160 pool.run_many(&parts);
2161 }
2162}
2163
2164impl QTensor {
2165 #[allow(clippy::needless_range_loop)]
2170 pub fn matvec2_many<const N: usize>(
2171 ts: [&QTensor; N],
2172 x1: &[f32],
2173 x2: &[f32],
2174 mut o1s: [&mut [f32]; N],
2175 mut o2s: [&mut [f32]; N],
2176 pool: Option<&Pool>,
2177 ) {
2178 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2179 let uniform_q8 = ts.iter().all(|t| {
2180 matches!(
2181 t,
2182 Self::Mapped {
2183 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2184 ..
2185 }
2186 )
2187 });
2188 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2189 let uniform_q4 = ts.iter().all(|t| {
2190 matches!(
2191 t,
2192 Self::Mapped {
2193 dtype: TensorDtype::Q4Block,
2194 ..
2195 }
2196 )
2197 });
2198 let uniform_vbit = ts.iter().all(|t| {
2199 matches!(
2200 t,
2201 Self::Mapped {
2202 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2203 ..
2204 }
2205 )
2206 });
2207 let fusable = pool.is_some()
2208 && total_rows >= 256
2209 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2210 if !fusable {
2211 for i in 0..N {
2212 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2213 }
2214 return;
2215 }
2216 let pool = pool.unwrap();
2217
2218 if uniform_q4 || uniform_vbit {
2219 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2220 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2221 if a8w8_enabled() {
2223 let a1 = split_act(x1);
2224 let a2 = split_act(x2);
2225 let (a1, a2) = (&a1, &a2);
2226 if uniform_q4 {
2227 let closures: [_; N] = std::array::from_fn(|i| {
2228 let (packed, scales) =
2229 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2230 let (gpr, cols, o1, o2) =
2231 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2232 move |s: usize, e: usize| {
2233 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2234 }
2235 });
2236 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2237 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2238 pool.run_many(&parts);
2239 } else {
2240 let closures: [_; N] = std::array::from_fn(|i| {
2241 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2242 unreachable!()
2243 };
2244 let (bytes, rows, cols, o1, o2) = (
2245 ts[i].quant_bytes(),
2246 ts[i].rows(),
2247 ts[i].cols(),
2248 p1[i],
2249 p2[i],
2250 );
2251 move |s: usize, e: usize| {
2252 vbit_range2_a8w8(
2253 bytes,
2254 vbit_offsets,
2255 x1,
2256 x2,
2257 a1,
2258 a2,
2259 rows,
2260 cols,
2261 o1,
2262 o2,
2263 s,
2264 e,
2265 )
2266 }
2267 });
2268 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2269 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2270 pool.run_many(&parts);
2271 }
2272 return;
2273 }
2274 if uniform_q4 {
2275 let closures: [_; N] = std::array::from_fn(|i| {
2276 let (packed, scales) =
2277 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2278 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2279 move |s: usize, e: usize| {
2280 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2281 }
2282 });
2283 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2284 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2285 pool.run_many(&parts);
2286 } else {
2287 let closures: [_; N] = std::array::from_fn(|i| {
2288 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2289 unreachable!()
2290 };
2291 let (bytes, rows, cols, o1, o2) = (
2292 ts[i].quant_bytes(),
2293 ts[i].rows(),
2294 ts[i].cols(),
2295 p1[i],
2296 p2[i],
2297 );
2298 move |s: usize, e: usize| {
2299 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2300 }
2301 });
2302 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2303 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2304 pool.run_many(&parts);
2305 }
2306 return;
2307 }
2308
2309 if uniform_f32 {
2310 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2311 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2312 let closures: [_; N] = std::array::from_fn(|i| {
2313 let Self::F32 { data, cols, .. } = ts[i] else {
2314 unreachable!()
2315 };
2316 let (o1, o2) = (p1[i], p2[i]);
2317 move |start: usize, end: usize| {
2318 for o in start..end {
2319 let row = &data[o * cols..(o + 1) * cols];
2320 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2321 for j in 0..*cols {
2322 s1 += row[j] * x1[j];
2323 s2 += row[j] * x2[j];
2324 }
2325 unsafe {
2327 *o1.at(o) = s1;
2328 *o2.at(o) = s2;
2329 }
2330 }
2331 }
2332 });
2333 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2334 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2335 pool.run_many(&parts);
2336 return;
2337 }
2338
2339 struct Ctx<'a> {
2340 bytes: &'a [u8],
2341 row_scale: &'a [f32],
2342 cols: usize,
2343 xs1: std::borrow::Cow<'a, [f32]>,
2344 xs2: std::borrow::Cow<'a, [f32]>,
2345 }
2346 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2347 let Self::Mapped {
2348 dtype,
2349 cols,
2350 row_scale,
2351 col_field,
2352 ..
2353 } = ts[i]
2354 else {
2355 unreachable!()
2356 };
2357 Ctx {
2358 bytes: ts[i].quant_bytes(),
2359 row_scale,
2360 cols: *cols,
2361 xs1: prescale(x1, col_field, *dtype),
2362 xs2: prescale(x2, col_field, *dtype),
2363 }
2364 });
2365 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2366 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2367 #[cfg(target_arch = "aarch64")]
2368 if sdot_enabled() {
2369 let acts: [(SplitAct, SplitAct); N] =
2370 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2371 let closures: [_; N] = std::array::from_fn(|i| {
2372 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2373 move |start: usize, end: usize| {
2374 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2375 }
2376 });
2377 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2378 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2379 pool.run_many(&parts);
2380 return;
2381 }
2382 #[cfg(target_arch = "x86_64")]
2383 if avx2_a8w8_enabled() {
2384 let acts: [(SplitAct, SplitAct); N] =
2385 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2386 let closures: [_; N] = std::array::from_fn(|i| {
2387 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2388 move |start: usize, end: usize| {
2389 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2390 }
2391 });
2392 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2393 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2394 pool.run_many(&parts);
2395 return;
2396 }
2397 let closures: [_; N] = std::array::from_fn(|i| {
2398 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2399 move |start: usize, end: usize| {
2400 q8_range2_f32(
2401 c.bytes,
2402 c.row_scale,
2403 &c.xs1,
2404 &c.xs2,
2405 c.cols,
2406 o1,
2407 o2,
2408 start,
2409 end,
2410 )
2411 }
2412 });
2413 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2414 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2415 pool.run_many(&parts);
2416 }
2417
2418 pub fn matvec_silu_mul(
2423 gate: &QTensor,
2424 up: &QTensor,
2425 x: &[f32],
2426 out: &mut [f32],
2427 pool: Option<&Pool>,
2428 ) -> bool {
2429 let inter = gate.rows();
2430 debug_assert_eq!(up.rows(), inter);
2431 debug_assert_eq!(out.len(), inter);
2432 debug_assert_eq!(gate.cols(), up.cols());
2433 if !a8w8_enabled() {
2434 return false;
2435 }
2436 let act = split_act(x);
2437 let act = &act;
2438 let x_ref = x;
2439 let out_addr = SendMut(out.as_mut_ptr());
2440
2441 match (gate, up) {
2442 (
2444 Self::Mapped {
2445 dtype: TensorDtype::Q4Block,
2446 ..
2447 },
2448 Self::Mapped {
2449 dtype: TensorDtype::Q4Block,
2450 ..
2451 },
2452 ) => {
2453 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2454 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2455 let gpr = gate.cols() / GROUP_SIZE;
2456 let cols = gate.cols();
2457 let run = move |start: usize, end: usize| {
2458 for r in start..end {
2459 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2460 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2461 for &(j, xv) in &act.outliers {
2462 let flat = r * cols + j;
2463 let gb = gp[flat / 2];
2464 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2465 let gsc = f16_to_f32(u16::from_le_bytes([
2466 gs[(flat / GROUP_SIZE) * 2],
2467 gs[(flat / GROUP_SIZE) * 2 + 1],
2468 ]));
2469 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2470 let ub = up_p[flat / 2];
2471 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2472 let usc = f16_to_f32(u16::from_le_bytes([
2473 up_s[(flat / GROUP_SIZE) * 2],
2474 up_s[(flat / GROUP_SIZE) * 2 + 1],
2475 ]));
2476 uv += ((un as i32 - 8) as f32) * usc * xv;
2477 }
2478 let silu_g = gv / (1.0 + (-gv).exp());
2479 unsafe { *out_addr.at(r) = silu_g * uv };
2481 }
2482 };
2483 dispatch_rows(pool, inter, &run);
2484 true
2485 }
2486 (
2490 Self::Mapped {
2491 dtype: TensorDtype::Q4Tiled,
2492 ..
2493 },
2494 Self::Mapped {
2495 dtype: TensorDtype::Q4Tiled,
2496 ..
2497 },
2498 ) => {
2499 let g_bytes = gate.quant_bytes();
2500 let u_bytes = up.quant_bytes();
2501 let gpr = gate.cols() / GROUP_SIZE;
2502 let run = move |start: usize, end: usize| {
2503 for r in start..end {
2504 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2505 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2506 for &(j, xv) in &act.outliers {
2507 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2508 gv += w * s * xv;
2509 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2510 uv += w * s * xv;
2511 }
2512 let silu_g = gv / (1.0 + (-gv).exp());
2513 unsafe { *out_addr.at(r) = silu_g * uv };
2515 }
2516 };
2517 dispatch_rows(pool, inter, &run);
2518 true
2519 }
2520 (
2523 Self::Mapped {
2524 dtype: TensorDtype::Q4TiledP,
2525 ..
2526 },
2527 Self::Mapped {
2528 dtype: TensorDtype::Q4TiledP,
2529 ..
2530 },
2531 ) => {
2532 let cols = gate.cols();
2533 let gpr = cols / GROUP_SIZE;
2534 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2535 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2536 let run = |start: usize, end: usize| {
2537 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2538 for r in start..end {
2539 gv_view.scales_into(r, gpr, &mut gsc);
2540 uv_view.scales_into(r, gpr, &mut usc);
2541 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2542 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2543 for &(j, xv) in &act.outliers {
2544 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2545 gv += w * s * xv;
2546 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2547 uv += w * s * xv;
2548 }
2549 let silu_g = gv / (1.0 + (-gv).exp());
2550 unsafe { *out_addr.at(r) = silu_g * uv };
2552 }
2553 };
2554 dispatch_rows(pool, inter, &run);
2555 true
2556 }
2557 (
2563 Self::Mapped {
2564 dtype: TensorDtype::Q1,
2565 ..
2566 },
2567 Self::Mapped {
2568 dtype: TensorDtype::Q1,
2569 ..
2570 },
2571 ) => {
2572 let g_bytes = gate.quant_bytes();
2573 let u_bytes = up.quant_bytes();
2574 let gpr = gate.cols() / GROUP_SIZE;
2575 let gsum = q1_group_sums(&act.xq, gpr);
2576 let gsum = &gsum;
2577 let run = move |start: usize, end: usize| {
2578 for r in start..end {
2579 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2580 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2581 for &(j, xv) in &act.outliers {
2582 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2583 gv += w * s * xv;
2584 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2585 uv += w * s * xv;
2586 }
2587 let silu_g = gv / (1.0 + (-gv).exp());
2588 unsafe { *out_addr.at(r) = silu_g * uv };
2590 }
2591 };
2592 dispatch_rows(pool, inter, &run);
2593 true
2594 }
2595 (
2599 Self::Mapped {
2600 dtype: TensorDtype::Q2TiledP,
2601 ..
2602 },
2603 Self::Mapped {
2604 dtype: TensorDtype::Q2TiledP,
2605 ..
2606 },
2607 ) => {
2608 let cols = gate.cols();
2609 let gpr = cols / GROUP_SIZE;
2610 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2611 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2612 let gsum = q1_group_sums(&act.xq, gpr);
2613 let gsum = &gsum;
2614 let run = move |start: usize, end: usize| {
2615 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2616 for r in start..end {
2617 gv_view.scales_into(r, gpr, &mut gsc);
2618 uv_view.scales_into(r, gpr, &mut usc);
2619 let mut gv =
2620 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2621 let mut uv =
2622 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2623 for &(j, xv) in &act.outliers {
2624 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2625 gv += w * s * xv;
2626 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2627 uv += w * s * xv;
2628 }
2629 let silu_g = gv / (1.0 + (-gv).exp());
2630 unsafe { *out_addr.at(r) = silu_g * uv };
2632 }
2633 };
2634 dispatch_rows(pool, inter, &run);
2635 true
2636 }
2637 (
2642 Self::Mapped {
2643 dtype: TensorDtype::Q8Row,
2644 row_scale: g_rs,
2645 ..
2646 },
2647 Self::Mapped {
2648 dtype: TensorDtype::Q8Row,
2649 row_scale: u_rs,
2650 ..
2651 },
2652 ) => {
2653 let g_bytes = gate.quant_bytes();
2654 let u_bytes = up.quant_bytes();
2655 let cols = gate.cols();
2656 let run = move |start: usize, end: usize| {
2657 for r in start..end {
2658 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2659 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2660 let silu_g = gv / (1.0 + (-gv).exp());
2661 unsafe { *out_addr.at(r) = silu_g * uv };
2663 }
2664 };
2665 dispatch_rows(pool, inter, &run);
2666 true
2667 }
2668 (
2670 Self::Mapped {
2671 dtype: TensorDtype::Q1T,
2672 ..
2673 },
2674 Self::Mapped {
2675 dtype: TensorDtype::Q1T,
2676 ..
2677 },
2678 ) => {
2679 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2680 let g_bytes = gate.quant_bytes();
2681 let u_bytes = up.quant_bytes();
2682 let gpr = gate.cols() / GROUP_SIZE;
2683 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2684 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2685 let run = move |start: usize, end: usize| {
2686 for r in start..end {
2687 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2688 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2689 for &(j, xv) in &act.outliers {
2690 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2691 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2692 }
2693 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2694 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2695 let silu_g = gv / (1.0 + (-gv).exp());
2696 unsafe { *out_addr.at(r) = silu_g * uv };
2698 }
2699 };
2700 dispatch_rows(pool, inter, &run);
2701 true
2702 }
2703 _ => false,
2704 }
2705 }
2706
2707 pub fn moe_gate_up_many(
2721 pairs: &[(&QTensor, &QTensor)],
2722 x: &[f32],
2723 outs: &mut [Vec<f32>],
2724 pool: Option<&Pool>,
2725 ) -> bool {
2726 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2727 return false;
2728 }
2729 let inter = pairs[0].0.rows();
2730 let cols = pairs[0].0.cols();
2731 if cols % GROUP_SIZE != 0 {
2732 return false;
2733 }
2734 let gpr = cols / GROUP_SIZE;
2735 let q2 = matches!(
2738 pairs[0].0,
2739 Self::Mapped {
2740 dtype: TensorDtype::Q2TiledP,
2741 ..
2742 }
2743 );
2744 let want = if q2 {
2745 TensorDtype::Q2TiledP
2746 } else {
2747 TensorDtype::Q4TiledP
2748 };
2749 let mut views = Vec::with_capacity(pairs.len() * 2);
2750 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2751 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2752 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2753 if !both
2754 || g.rows() != inter
2755 || u.rows() != inter
2756 || g.cols() != cols
2757 || u.cols() != cols
2758 || o.len() != inter
2759 {
2760 return false;
2761 }
2762 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2763 views.push(mk(g.quant_bytes(), inter, cols));
2764 views.push(mk(u.quant_bytes(), inter, cols));
2765 }
2766 let act = split_act(x);
2767 let gsum = if q2 {
2768 q1_group_sums(&act.xq, gpr)
2769 } else {
2770 Vec::new()
2771 };
2772 let (act, gsum) = (&act, &gsum);
2773 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2774 let (views, ptrs) = (&views, &ptrs);
2775 let run = |start: usize, end: usize| {
2776 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2777 for flat in start..end {
2778 let (e, r) = (flat / inter, flat % inter);
2779 let gv_view = &views[e * 2];
2780 let uv_view = &views[e * 2 + 1];
2781 gv_view.scales_into(r, gpr, &mut gsc);
2782 uv_view.scales_into(r, gpr, &mut usc);
2783 let (mut gv, mut uv) = if q2 {
2784 (
2785 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2786 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2787 )
2788 } else {
2789 (
2790 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2791 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2792 )
2793 };
2794 for &(j, xv) in &act.outliers {
2795 let (og, ou) = if q2 {
2796 (
2797 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2798 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2799 )
2800 } else {
2801 (
2802 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2803 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2804 )
2805 };
2806 gv += og.0 * og.1 * xv;
2807 uv += ou.0 * ou.1 * xv;
2808 }
2809 let silu_g = gv / (1.0 + (-gv).exp());
2810 unsafe { *ptrs[e].at(r) = silu_g * uv };
2812 }
2813 };
2814 dispatch_rows(pool, pairs.len() * inter, &run);
2815 true
2816 }
2817
2818 pub fn moe_down_many(
2827 downs: &[&QTensor],
2828 gs: &[Vec<f32>],
2829 weights: &[f32],
2830 out: &mut [f32],
2831 pool: Option<&Pool>,
2832 ) -> bool {
2833 if downs.is_empty()
2834 || downs.len() != gs.len()
2835 || downs.len() != weights.len()
2836 || !a8w8_enabled()
2837 {
2838 return false;
2839 }
2840 let rows = out.len();
2841 let cols = downs[0].cols();
2842 if cols % GROUP_SIZE != 0 {
2843 return false;
2844 }
2845 let gpr = cols / GROUP_SIZE;
2846 let mut views = Vec::with_capacity(downs.len());
2847 for (d, g) in downs.iter().zip(gs.iter()) {
2848 if !matches!(
2849 d,
2850 Self::Mapped {
2851 dtype: TensorDtype::Q4TiledP,
2852 ..
2853 }
2854 ) || d.rows() != rows
2855 || d.cols() != cols
2856 || g.len() != cols
2857 {
2858 return false;
2859 }
2860 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2861 }
2862 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2864 let out_addr = SendMut(out.as_mut_ptr());
2872 let (views, acts, weights) = (&views, &acts, &weights);
2873 let run = |start: usize, end: usize| {
2874 let mut sc = vec![0f32; gpr];
2875 for r in start..end {
2876 let mut acc = 0f32;
2877 for (e, v) in views.iter().enumerate() {
2878 v.scales_into(r, gpr, &mut sc);
2879 let a = &acts[e];
2880 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2881 for &(j, xv) in &a.outliers {
2882 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2883 d += w * s * xv;
2884 }
2885 acc += weights[e] * d;
2886 }
2887 unsafe { *out_addr.at(r) = acc };
2889 }
2890 };
2891 dispatch_rows(pool, rows, &run);
2892 true
2893 }
2894}
2895
2896#[cfg(target_os = "macos")]
2901mod accel_blas {
2902 #[link(name = "Accelerate", kind = "framework")]
2903 unsafe extern "C" {
2904 pub fn cblas_sgemm(
2905 order: i32,
2906 trans_a: i32,
2907 trans_b: i32,
2908 m: i32,
2909 n: i32,
2910 k: i32,
2911 alpha: f32,
2912 a: *const f32,
2913 lda: i32,
2914 b: *const f32,
2915 ldb: i32,
2916 beta: f32,
2917 c: *mut f32,
2918 ldc: i32,
2919 );
2920 }
2921}
2922
2923#[cfg(target_os = "macos")]
2924pub(crate) fn accel_gemm_enabled() -> bool {
2925 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2926 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2927}
2928
2929#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2932pub(crate) fn accel_gemm_enabled() -> bool {
2933 true
2934}
2935
2936#[cfg(target_arch = "aarch64")]
2943#[allow(clippy::too_many_arguments)]
2944pub(crate) fn neon_gemm_rm(
2945 m: usize,
2946 n: usize,
2947 k: usize,
2948 alpha: f32,
2949 a: &[f32],
2950 lda: usize,
2951 b_mat: &[f32],
2952 ldb: usize,
2953 b_rows_are_n: bool,
2954 c: &mut [f32],
2955 ldc: usize,
2956) {
2957 debug_assert!(a.len() >= (m - 1) * lda + k);
2958 debug_assert!(c.len() >= (m - 1) * ldc + n);
2959 unsafe {
2961 use core::arch::aarch64::*;
2962 let mut i = 0usize;
2963 while i < m {
2964 let mi = (m - i).min(4);
2965 let mut j = 0usize;
2966 while j < n {
2967 let nj = (n - j).min(8);
2968 if mi == 4 && nj == 8 {
2969 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2970 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2971 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2972 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2973 for p in 0..k {
2974 let (b0, b1) = if b_rows_are_n {
2975 let base = b_mat.as_ptr().add(j * ldb + p);
2978 let g = |o: usize| *base.add(o * ldb);
2979 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2980 } else {
2981 let base = b_mat.as_ptr().add(p * ldb + j);
2982 (
2983 [*base, *base.add(1), *base.add(2), *base.add(3)],
2984 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2985 )
2986 };
2987 let bv0 = vld1q_f32(b0.as_ptr());
2988 let bv1 = vld1q_f32(b1.as_ptr());
2989 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2990 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2991 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2992 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2993 c0a = vfmaq_f32(c0a, a0, bv0);
2994 c0b = vfmaq_f32(c0b, a0, bv1);
2995 c1a = vfmaq_f32(c1a, a1, bv0);
2996 c1b = vfmaq_f32(c1b, a1, bv1);
2997 c2a = vfmaq_f32(c2a, a2, bv0);
2998 c2b = vfmaq_f32(c2b, a2, bv1);
2999 c3a = vfmaq_f32(c3a, a3, bv0);
3000 c3b = vfmaq_f32(c3b, a3, bv1);
3001 }
3002 let al = vdupq_n_f32(alpha);
3003 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3004 .iter()
3005 .enumerate()
3006 {
3007 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3008 vst1q_f32(dst, vmulq_f32(*ca, al));
3009 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3010 }
3011 } else {
3012 for r in 0..mi {
3013 for q in 0..nj {
3014 let mut acc = 0f32;
3015 for p in 0..k {
3016 let bv = if b_rows_are_n {
3017 b_mat[(j + q) * ldb + p]
3018 } else {
3019 b_mat[p * ldb + j + q]
3020 };
3021 acc += a[(i + r) * lda + p] * bv;
3022 }
3023 c[(i + r) * ldc + j + q] = acc * alpha;
3024 }
3025 }
3026 }
3027 j += nj;
3028 }
3029 i += mi;
3030 }
3031 }
3032}
3033
3034#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3036#[allow(clippy::too_many_arguments)]
3037pub(crate) fn sgemm_rm(
3038 m: usize,
3039 n: usize,
3040 k: usize,
3041 alpha: f32,
3042 a: &[f32],
3043 lda: usize,
3044 b_mat: &[f32],
3045 ldb: usize,
3046 b_rows_are_n: bool,
3047 c: &mut [f32],
3048 ldc: usize,
3049) {
3050 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3051}
3052
3053#[allow(clippy::too_many_arguments)]
3057pub fn sgemm_public(
3058 m: usize,
3059 n: usize,
3060 k: usize,
3061 alpha: f32,
3062 a: &[f32],
3063 lda: usize,
3064 b_mat: &[f32],
3065 ldb: usize,
3066 b_rows_are_n: bool,
3067 c: &mut [f32],
3068 ldc: usize,
3069) {
3070 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3071 {
3072 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3073 }
3074 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3079 {
3080 for i in 0..m {
3081 for j in 0..n {
3082 let mut acc = 0f32;
3083 for p in 0..k {
3084 let bv = if b_rows_are_n {
3085 b_mat[j * ldb + p]
3086 } else {
3087 b_mat[p * ldb + j]
3088 };
3089 acc += a[i * lda + p] * bv;
3090 }
3091 c[i * ldc + j] = alpha * acc;
3092 }
3093 }
3094 }
3095}
3096
3097#[cfg(target_os = "macos")]
3100#[allow(clippy::too_many_arguments)]
3101pub(crate) fn sgemm_rm(
3102 m: usize,
3103 n: usize,
3104 k: usize,
3105 alpha: f32,
3106 a: &[f32],
3107 lda: usize,
3108 b_mat: &[f32],
3109 ldb: usize,
3110 b_rows_are_n: bool,
3111 c: &mut [f32],
3112 ldc: usize,
3113) {
3114 debug_assert!(a.len() >= (m - 1) * lda + k);
3115 debug_assert!(c.len() >= (m - 1) * ldc + n);
3116 #[cfg(target_arch = "aarch64")]
3121 if std::env::var("CMF_FORCE_NEON_GEMM")
3122 .map(|v| v == "1")
3123 .unwrap_or(false)
3124 {
3125 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3126 }
3127 unsafe {
3128 accel_blas::cblas_sgemm(
3129 101, 111, if b_rows_are_n { 112 } else { 111 },
3132 m as i32,
3133 n as i32,
3134 k as i32,
3135 alpha,
3136 a.as_ptr(),
3137 lda as i32,
3138 b_mat.as_ptr(),
3139 ldb as i32,
3140 0.0,
3141 c.as_mut_ptr(),
3142 ldc as i32,
3143 );
3144 }
3145}
3146
3147#[cfg(target_os = "macos")]
3154fn qmatmat_accel(
3155 q: &[u8],
3156 row_scale: &[f32],
3157 pre: &[std::borrow::Cow<'_, [f32]>],
3158 rows: usize,
3159 cols: usize,
3160 out: &mut [f32],
3161 pool: Option<&Pool>,
3162) {
3163 const TR: usize = 2048;
3168 let b = pre.len();
3169 thread_local! {
3170 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3171 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3172 }
3173 XPANEL.with(|xp| {
3174 WTILE.with(|wt| {
3175 let mut xpanel = xp.borrow_mut();
3176 xpanel.clear();
3177 for x in pre {
3178 xpanel.extend_from_slice(x);
3179 }
3180 let mut wtile = wt.borrow_mut();
3181 wtile.resize(TR * cols, 0.0);
3182 let mut r0 = 0usize;
3183 while r0 < rows {
3184 let tr = TR.min(rows - r0);
3185 let wt_addr = SendMut(wtile.as_mut_ptr());
3187 let run = |start: usize, end: usize| {
3188 for r in start..end {
3189 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3190 let s = row_scale[r0 + r];
3191 let dst =
3193 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3194 for (d, &v) in dst.iter_mut().zip(row) {
3195 *d = (v as i8) as f32 * s;
3196 }
3197 }
3198 };
3199 dispatch_rows(pool, tr, &run);
3200 unsafe {
3202 accel_blas::cblas_sgemm(
3203 101, 111, 112, b as i32,
3207 tr as i32,
3208 cols as i32,
3209 1.0,
3210 xpanel.as_ptr(),
3211 cols as i32,
3212 wtile.as_ptr(),
3213 cols as i32,
3214 0.0,
3215 out.as_mut_ptr().add(r0),
3216 rows as i32,
3217 );
3218 }
3219 r0 += tr;
3220 }
3221 })
3222 });
3223}
3224
3225fn qmatmat(
3226 q: &[u8],
3227 row_scale: &[f32],
3228 pre: &[std::borrow::Cow<'_, [f32]>],
3229 rows: usize,
3230 cols: usize,
3231 out: &mut [f32],
3232 pool: Option<&Pool>,
3233) {
3234 let b = pre.len();
3235 debug_assert_eq!(out.len(), b * rows);
3236 #[cfg(target_os = "macos")]
3241 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3242 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3243 return;
3244 }
3245 #[cfg(target_arch = "aarch64")]
3246 if sdot_enabled() {
3247 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3248 let out_addr = SendMut(out.as_mut_ptr());
3249 let blocked_ok = blocked_enabled();
3252 let use_i8mm = i8mm_enabled();
3253 if blocked_ok {
3254 let run = |start: usize, end: usize| {
3255 let mut o = start;
3256 while o < end {
3257 if o + 2 <= end {
3258 let r0 = &q[o * cols..(o + 1) * cols];
3259 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3260 let mut bi = 0usize;
3261 while bi + 4 <= acts.len() {
3262 let xs = [
3263 acts[bi].xq.as_slice(),
3264 acts[bi + 1].xq.as_slice(),
3265 acts[bi + 2].xq.as_slice(),
3266 acts[bi + 3].xq.as_slice(),
3267 ];
3268 let d = if use_i8mm {
3269 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3270 } else {
3271 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3272 };
3273 for (r, row) in [r0, r1].into_iter().enumerate() {
3274 for k in 0..4 {
3275 let act = &acts[bi + k];
3276 let mut v = d[r][k] as f32 * act.sx;
3277 for &(j, xv) in &act.outliers {
3278 v += (row[j] as i8) as f32 * xv;
3279 }
3280 unsafe {
3281 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3282 };
3283 }
3284 }
3285 bi += 4;
3286 }
3287 while bi < acts.len() {
3288 for (r, row) in [r0, r1].into_iter().enumerate() {
3289 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3290 unsafe { *out_addr.at(bi * rows + o + r) = v };
3291 }
3292 bi += 1;
3293 }
3294 o += 2;
3295 } else {
3296 let row = &q[o * cols..(o + 1) * cols];
3297 for (bi, act) in acts.iter().enumerate() {
3298 let v = row_dot_sdot(row, act) * row_scale[o];
3299 unsafe { *out_addr.at(bi * rows + o) = v };
3300 }
3301 o += 1;
3302 }
3303 }
3304 };
3305 dispatch_rows(pool, rows, &run);
3306 return;
3307 }
3308 let run = |start: usize, end: usize| {
3309 for o in start..end {
3310 let row = &q[o * cols..(o + 1) * cols];
3311 for (bi, act) in acts.iter().enumerate() {
3312 let v = row_dot_sdot(row, act) * row_scale[o];
3313 unsafe { *out_addr.at(bi * rows + o) = v };
3314 }
3315 }
3316 };
3317 dispatch_rows(pool, rows, &run);
3318 return;
3319 }
3320 #[cfg(target_arch = "x86_64")]
3325 if avx2_a8w8_enabled() {
3326 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3327 let out_addr = SendMut(out.as_mut_ptr());
3328 let blocked_ok = blocked_enabled();
3331 if !avx512vnni_enabled() && blocked_ok {
3332 let run = |start: usize, end: usize| {
3333 let mut o = start;
3334 while o < end {
3335 if o + 2 <= end {
3336 let r0 = &q[o * cols..(o + 1) * cols];
3337 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3338 let mut bi = 0usize;
3339 while bi + 4 <= acts.len() {
3340 let xs = [
3341 acts[bi].xq.as_slice(),
3342 acts[bi + 1].xq.as_slice(),
3343 acts[bi + 2].xq.as_slice(),
3344 acts[bi + 3].xq.as_slice(),
3345 ];
3346 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3347 for (r, row) in [r0, r1].into_iter().enumerate() {
3348 for k in 0..4 {
3349 let act = &acts[bi + k];
3350 let mut v = d[r][k] as f32 * act.sx;
3351 for &(j, xv) in &act.outliers {
3352 v += (row[j] as i8) as f32 * xv;
3353 }
3354 unsafe {
3355 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3356 };
3357 }
3358 }
3359 bi += 4;
3360 }
3361 while bi < acts.len() {
3362 for (r, row) in [r0, r1].into_iter().enumerate() {
3363 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3364 unsafe { *out_addr.at(bi * rows + o + r) = v };
3365 }
3366 bi += 1;
3367 }
3368 o += 2;
3369 } else {
3370 let row = &q[o * cols..(o + 1) * cols];
3371 for (bi, act) in acts.iter().enumerate() {
3372 let v = row_dot_avx2(row, act) * row_scale[o];
3373 unsafe { *out_addr.at(bi * rows + o) = v };
3374 }
3375 o += 1;
3376 }
3377 }
3378 };
3379 dispatch_rows(pool, rows, &run);
3380 return;
3381 }
3382 let run = |start: usize, end: usize| {
3383 for o in start..end {
3384 let row = &q[o * cols..(o + 1) * cols];
3385 for (bi, act) in acts.iter().enumerate() {
3386 let v = row_dot_avx2(row, act) * row_scale[o];
3387 unsafe { *out_addr.at(bi * rows + o) = v };
3388 }
3389 }
3390 };
3391 dispatch_rows(pool, rows, &run);
3392 return;
3393 }
3394 let out_addr = SendMut(out.as_mut_ptr());
3395 let run = |start: usize, end: usize| {
3396 for o in start..end {
3397 let row = &q[o * cols..(o + 1) * cols];
3398 for (bi, x) in pre.iter().enumerate() {
3399 let mut acc = 0f32;
3400 for j in 0..cols {
3401 acc += (row[j] as i8) as f32 * x[j];
3402 }
3403 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3404 }
3405 }
3406 };
3407 dispatch_rows(pool, rows, &run);
3408}
3409
3410fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3413 match pool {
3414 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3415 _ => run(0, rows),
3416 }
3417}
3418
3419fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3421 let groups = rows * cols / GROUP_SIZE;
3422 bytes.split_at(groups * 16)
3423}
3424
3425#[inline]
3430fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3431 #[cfg(target_arch = "aarch64")]
3432 unsafe {
3433 return vbit_fill4_neon(data, buf);
3434 }
3435 #[cfg(target_arch = "x86_64")]
3436 if avx2_enabled() {
3437 return unsafe { vbit_fill4_avx2(data, buf) };
3438 }
3439 #[allow(unreachable_code)]
3440 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3441 let u = unpack8::<4>(&data[blk * 4..]);
3442 for k in 0..8 {
3443 chunk[k] = (u[k] - 7) as i8 as u8;
3444 }
3445 }
3446}
3447
3448#[cfg(target_arch = "aarch64")]
3449#[target_feature(enable = "neon")]
3450unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3451 unsafe {
3454 use core::arch::aarch64::*;
3455 let n = buf.len();
3456 let mask = vdupq_n_u8(0x0F);
3457 let seven = vdupq_n_s8(7);
3458 let mut g = 0usize;
3459 while g * 32 + 32 <= n {
3460 let b = vld1q_u8(data.as_ptr().add(g * 16));
3461 let hi = vshrq_n_u8::<4>(b);
3462 let lo = vandq_u8(b, mask);
3463 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3464 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3465 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3466 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3467 g += 1;
3468 }
3469 }
3470}
3471
3472#[cfg(target_arch = "x86_64")]
3473#[target_feature(enable = "avx2")]
3474unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3475 unsafe {
3477 use core::arch::x86_64::*;
3478 let n = buf.len();
3479 let mask = _mm_set1_epi8(0x0F);
3480 let seven = _mm256_set1_epi8(7);
3481 let mut g = 0usize;
3482 while g * 32 + 32 <= n {
3483 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3484 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3485 let lo = _mm_and_si128(b, mask);
3486 let z = _mm256_sub_epi8(
3487 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3488 seven,
3489 );
3490 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3491 g += 1;
3492 }
3493 }
3494}
3495
3496#[inline(always)]
3501fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3502 let mut acc = 0u64;
3503 for i in 0..B {
3504 acc = (acc << 8) | data[i] as u64;
3505 }
3506 let mask = (1u64 << B) - 1;
3507 let mut out = [0i32; 8];
3508 for (k, o) in out.iter_mut().enumerate() {
3509 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3510 }
3511 out
3512}
3513
3514#[allow(clippy::too_many_arguments)]
3520fn vbitmatvec(
3521 bytes: &[u8],
3522 offsets: &[usize],
3523 x: &[f32],
3524 rows: usize,
3525 cols: usize,
3526 out: &mut [f32],
3527 pool: Option<&Pool>,
3528) {
3529 debug_assert_eq!(out.len(), rows);
3530 debug_assert_eq!(offsets.len(), rows + 1);
3531
3532 if a8w8_enabled() {
3536 let act = split_act(x);
3537 let out_addr = SendMut(out.as_mut_ptr());
3538 let run = move |start: usize, end: usize| {
3539 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3540 };
3541 dispatch_rows(pool, rows, &run);
3542 return;
3543 }
3544
3545 let out_addr = SendMut(out.as_mut_ptr());
3546 let run = move |start: usize, end: usize| {
3547 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3548 };
3549 dispatch_rows(pool, rows, &run);
3550}
3551
3552#[allow(clippy::too_many_arguments)]
3556fn vbit_range_a8w8(
3557 bytes: &[u8],
3558 offsets: &[usize],
3559 x: &[f32],
3560 act: &SplitAct,
3561 rows: usize,
3562 cols: usize,
3563 out: SendMut,
3564 start: usize,
3565 end: usize,
3566) {
3567 let ng = cols / GROUP_SIZE;
3568 let bits = &bytes[..rows];
3569 let sc_off = rows;
3570 let row_dot = |r: usize| -> f32 {
3571 let b = bits[r] as usize;
3572 let l = (1i32 << (b - 1)) - 1;
3573 let mask = (1u64 << b) - 1;
3574 let data = &bytes[offsets[r]..offsets[r + 1]];
3575 if b == 8 {
3576 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3578 let mut dot = 0f32;
3579 for g in 0..ng {
3580 let so = (r * ng + g) * 2;
3581 let sgf = f16_to_f32(u16::from_le_bytes([
3582 bytes[sc_off + so],
3583 bytes[sc_off + so + 1],
3584 ]));
3585 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3586 let mut gd = 0f32;
3587 for &xv in xg.iter() {
3588 if nbits < 8 {
3589 acc = (acc << 8) | data[idx] as u64;
3590 idx += 1;
3591 nbits += 8;
3592 }
3593 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3594 nbits -= 8;
3595 gd += (u - l) as f32 * xv;
3596 }
3597 dot += gd * sgf;
3598 }
3599 return dot;
3600 }
3601 thread_local! {
3605 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3606 const { std::cell::RefCell::new(Vec::new()) };
3607 }
3608 #[inline(always)]
3609 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3610 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3611 let u = unpack8::<B>(&data[blk * B..]);
3612 for k in 0..8 {
3613 chunk[k] = (u[k] - l) as i8 as u8;
3614 }
3615 }
3616 }
3617 let _ = mask;
3618 VBIT_SCRATCH.with(|scratch| {
3619 let mut buf = scratch.borrow_mut();
3620 buf.resize(cols, 0);
3621 match b {
3622 3 => fill::<3>(data, l, &mut buf),
3623 4 => vbit_fill4(data, &mut buf),
3624 5 => fill::<5>(data, l, &mut buf),
3625 6 => fill::<6>(data, l, &mut buf),
3626 _ => unreachable!(),
3627 }
3628 let mut dot = 0f32;
3629 for g in 0..ng {
3630 let so = (r * ng + g) * 2;
3631 let s = f16_to_f32(u16::from_le_bytes([
3632 bytes[sc_off + so],
3633 bytes[sc_off + so + 1],
3634 ]));
3635 let d = dot_i8_i8(
3636 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3637 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3638 ) as f32
3639 * act.sx;
3640 dot += d * s;
3641 }
3642 for &(j, xv) in &act.outliers {
3643 let so = (r * ng + j / GROUP_SIZE) * 2;
3644 let s = f16_to_f32(u16::from_le_bytes([
3645 bytes[sc_off + so],
3646 bytes[sc_off + so + 1],
3647 ]));
3648 dot += (buf[j] as i8) as f32 * s * xv;
3650 }
3651 dot
3652 })
3653 };
3654 for r in start..end {
3655 unsafe { *out.at(r) = row_dot(r) };
3657 }
3658}
3659
3660#[allow(clippy::too_many_arguments)]
3662fn vbit_range_f32(
3663 bytes: &[u8],
3664 offsets: &[usize],
3665 x: &[f32],
3666 rows: usize,
3667 cols: usize,
3668 out: SendMut,
3669 start: usize,
3670 end: usize,
3671) {
3672 let ng = cols / GROUP_SIZE;
3673 let bits = &bytes[..rows];
3674 let sc_off = rows;
3675 #[inline(always)]
3679 fn dot_row<const B: usize>(
3680 data: &[u8],
3681 bytes: &[u8],
3682 sc_off: usize,
3683 r: usize,
3684 ng: usize,
3685 x: &[f32],
3686 ) -> f32 {
3687 let l = ((1i32 << (B - 1)) - 1) as f32;
3688 let gbytes = GROUP_SIZE * B / 8;
3689 let mut dot = 0f32;
3690 for g in 0..ng {
3691 let so = (r * ng + g) * 2;
3692 let s = f16_to_f32(u16::from_le_bytes([
3693 bytes[sc_off + so],
3694 bytes[sc_off + so + 1],
3695 ]));
3696 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3697 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3698 let mut gd = 0f32;
3699 for blk in 0..GROUP_SIZE / 8 {
3700 let u = unpack8::<B>(&gd0[blk * B..]);
3701 let xb = &xg[blk * 8..blk * 8 + 8];
3702 for k in 0..8 {
3703 gd += (u[k] as f32 - l) * xb[k];
3704 }
3705 }
3706 dot += gd * s;
3707 }
3708 dot
3709 }
3710 for r in start..end {
3711 let data = &bytes[offsets[r]..offsets[r + 1]];
3712 let v = match bits[r] {
3713 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3714 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3715 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3716 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3717 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3718 b => unreachable!("vbit bit-width {b} (validated at load)"),
3719 };
3720 unsafe { *out.at(r) = v };
3722 }
3723}
3724
3725#[allow(clippy::too_many_arguments)]
3730fn vbitmatvec2(
3731 bytes: &[u8],
3732 offsets: &[usize],
3733 x1: &[f32],
3734 x2: &[f32],
3735 rows: usize,
3736 cols: usize,
3737 o1: &mut [f32],
3738 o2: &mut [f32],
3739 pool: Option<&Pool>,
3740) {
3741 debug_assert_eq!(o1.len(), rows);
3742 debug_assert_eq!(o2.len(), rows);
3743
3744 if a8w8_enabled() {
3745 let a1 = split_act(x1);
3746 let a2 = split_act(x2);
3747 let p1 = SendMut(o1.as_mut_ptr());
3748 let p2 = SendMut(o2.as_mut_ptr());
3749 let run = move |start: usize, end: usize| {
3750 vbit_range2_a8w8(
3751 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3752 )
3753 };
3754 dispatch_rows(pool, rows, &run);
3755 return;
3756 }
3757
3758 let p1 = SendMut(o1.as_mut_ptr());
3759 let p2 = SendMut(o2.as_mut_ptr());
3760 let run = move |start: usize, end: usize| {
3761 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3762 };
3763 dispatch_rows(pool, rows, &run);
3764}
3765
3766#[allow(clippy::too_many_arguments)]
3770fn vbit_range2_a8w8(
3771 bytes: &[u8],
3772 offsets: &[usize],
3773 x1: &[f32],
3774 x2: &[f32],
3775 a1: &SplitAct,
3776 a2: &SplitAct,
3777 rows: usize,
3778 cols: usize,
3779 p1: SendMut,
3780 p2: SendMut,
3781 start: usize,
3782 end: usize,
3783) {
3784 let ng = cols / GROUP_SIZE;
3785 let bits = &bytes[..rows];
3786 let sc_off = rows;
3787 let row_dots = |r: usize| -> (f32, f32) {
3788 let b = bits[r] as usize;
3789 let l = (1i32 << (b - 1)) - 1;
3790 let data = &bytes[offsets[r]..offsets[r + 1]];
3791 if b == 8 {
3792 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3795 let (mut d1, mut d2) = (0f32, 0f32);
3796 for g in 0..ng {
3797 let so = (r * ng + g) * 2;
3798 let sgf = f16_to_f32(u16::from_le_bytes([
3799 bytes[sc_off + so],
3800 bytes[sc_off + so + 1],
3801 ]));
3802 let (mut g1, mut g2) = (0f32, 0f32);
3803 for k in 0..GROUP_SIZE {
3804 if nbits < 8 {
3805 acc = (acc << 8) | data[idx] as u64;
3806 idx += 1;
3807 nbits += 8;
3808 }
3809 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3810 nbits -= 8;
3811 let w = (u - l) as f32;
3812 g1 += w * x1[g * GROUP_SIZE + k];
3813 g2 += w * x2[g * GROUP_SIZE + k];
3814 }
3815 d1 += g1 * sgf;
3816 d2 += g2 * sgf;
3817 }
3818 return (d1, d2);
3819 }
3820 thread_local! {
3821 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3822 const { std::cell::RefCell::new(Vec::new()) };
3823 }
3824 #[inline(always)]
3825 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3826 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3827 let u = unpack8::<B>(&data[blk * B..]);
3828 for k in 0..8 {
3829 chunk[k] = (u[k] - l) as i8 as u8;
3830 }
3831 }
3832 }
3833 VBIT_SCRATCH2.with(|scratch| {
3834 let mut buf = scratch.borrow_mut();
3835 buf.resize(cols, 0);
3836 match b {
3837 3 => fill::<3>(data, l, &mut buf),
3838 4 => vbit_fill4(data, &mut buf),
3839 5 => fill::<5>(data, l, &mut buf),
3840 6 => fill::<6>(data, l, &mut buf),
3841 _ => unreachable!(),
3842 }
3843 let (mut d1, mut d2) = (0f32, 0f32);
3844 for g in 0..ng {
3845 let so = (r * ng + g) * 2;
3846 let s = f16_to_f32(u16::from_le_bytes([
3847 bytes[sc_off + so],
3848 bytes[sc_off + so + 1],
3849 ]));
3850 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3851 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3852 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3853 d1 += v1 * s;
3854 d2 += v2 * s;
3855 }
3856 for &(j, xv) in &a1.outliers {
3857 let so = (r * ng + j / GROUP_SIZE) * 2;
3858 let s = f16_to_f32(u16::from_le_bytes([
3859 bytes[sc_off + so],
3860 bytes[sc_off + so + 1],
3861 ]));
3862 d1 += (buf[j] as i8) as f32 * s * xv;
3863 }
3864 for &(j, xv) in &a2.outliers {
3865 let so = (r * ng + j / GROUP_SIZE) * 2;
3866 let s = f16_to_f32(u16::from_le_bytes([
3867 bytes[sc_off + so],
3868 bytes[sc_off + so + 1],
3869 ]));
3870 d2 += (buf[j] as i8) as f32 * s * xv;
3871 }
3872 (d1, d2)
3873 })
3874 };
3875 for r in start..end {
3876 let (v1, v2) = row_dots(r);
3877 unsafe {
3879 *p1.at(r) = v1;
3880 *p2.at(r) = v2;
3881 }
3882 }
3883}
3884
3885#[allow(clippy::too_many_arguments)]
3889fn vbit_range2_f32(
3890 bytes: &[u8],
3891 offsets: &[usize],
3892 x1: &[f32],
3893 x2: &[f32],
3894 rows: usize,
3895 cols: usize,
3896 p1: SendMut,
3897 p2: SendMut,
3898 start: usize,
3899 end: usize,
3900) {
3901 let ng = cols / GROUP_SIZE;
3902 let bits = &bytes[..rows];
3903 let sc_off = rows;
3904 #[inline(always)]
3905 #[allow(clippy::too_many_arguments)]
3906 fn dot_row2<const B: usize>(
3907 data: &[u8],
3908 bytes: &[u8],
3909 sc_off: usize,
3910 r: usize,
3911 ng: usize,
3912 x1: &[f32],
3913 x2: &[f32],
3914 ) -> (f32, f32) {
3915 let l = ((1i32 << (B - 1)) - 1) as f32;
3916 let gbytes = GROUP_SIZE * B / 8;
3917 let (mut d1, mut d2) = (0f32, 0f32);
3918 for g in 0..ng {
3919 let so = (r * ng + g) * 2;
3920 let s = f16_to_f32(u16::from_le_bytes([
3921 bytes[sc_off + so],
3922 bytes[sc_off + so + 1],
3923 ]));
3924 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3925 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3926 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3927 let (mut g1, mut g2) = (0f32, 0f32);
3928 for blk in 0..GROUP_SIZE / 8 {
3929 let u = unpack8::<B>(&gd0[blk * B..]);
3930 for k in 0..8 {
3931 let w = u[k] as f32 - l;
3932 g1 += w * x1g[blk * 8 + k];
3933 g2 += w * x2g[blk * 8 + k];
3934 }
3935 }
3936 d1 += g1 * s;
3937 d2 += g2 * s;
3938 }
3939 (d1, d2)
3940 }
3941 for r in start..end {
3942 let data = &bytes[offsets[r]..offsets[r + 1]];
3943 let (v1, v2) = match bits[r] {
3944 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3945 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3946 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3947 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3948 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3949 b => unreachable!("vbit bit-width {b} (validated at load)"),
3950 };
3951 unsafe {
3953 *p1.at(r) = v1;
3954 *p2.at(r) = v2;
3955 }
3956 }
3957}
3958
3959#[inline]
3966#[allow(unreachable_code)]
3967fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3968 #[cfg(target_arch = "aarch64")]
3969 unsafe {
3970 return dot_q4t_row_sdot(bytes, r, gpr, xq);
3971 }
3972 #[cfg(target_arch = "x86_64")]
3973 unsafe {
3974 if vnni_tiles_enabled() {
3975 return dot_q4t_row_vnni(bytes, r, gpr, xq);
3976 }
3977 return dot_q4t_row_avx2(bytes, r, gpr, xq);
3978 }
3979 let mut acc = 0f32;
3980 for gi in 0..gpr {
3981 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3982 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3983 let mut d = 0i32;
3984 for (k, &b) in tile[2..].iter().enumerate() {
3985 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3986 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3987 }
3988 acc += d as f32 * s;
3989 }
3990 acc
3991}
3992
3993#[cfg(target_arch = "aarch64")]
3994#[target_feature(enable = "neon,dotprod")]
3995unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3996 unsafe {
3999 use core::arch::aarch64::*;
4000 use core::arch::asm;
4001 let lomask = vdupq_n_u8(0x0F);
4002 let eight = vdupq_n_s8(8);
4003 let mut acc = 0f32;
4004 for gi in 0..gpr {
4005 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4006 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4007 let b = vld1q_u8(t.add(2));
4008 let lo = vandq_u8(b, lomask);
4009 let hi = vshrq_n_u8::<4>(b);
4010 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4011 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4012 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4013 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4014 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4015 asm!(
4016 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4017 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4018 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4019 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4020 options(pure, nomem, nostack),
4021 );
4022 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4023 }
4024 acc
4025 }
4026}
4027
4028#[cfg(target_arch = "x86_64")]
4029#[target_feature(enable = "avx2")]
4030unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4031 unsafe {
4033 use core::arch::x86_64::*;
4034 let lomask = _mm_set1_epi8(0x0F);
4035 let eight = _mm256_set1_epi8(8);
4036 let ones = _mm256_set1_epi16(1);
4037 let mut acc = 0f32;
4038 for gi in 0..gpr {
4039 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4040 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4041 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4042 let lo = _mm_and_si128(b, lomask);
4043 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4044 let w = _mm256_sub_epi8(
4045 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4046 eight,
4047 );
4048 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4049 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4050 let d = _mm256_madd_epi16(p16, ones);
4051 let hi128 = _mm256_extracti128_si256::<1>(d);
4052 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4053 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4054 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4055 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4056 }
4057 acc
4058 }
4059}
4060
4061#[cfg(target_arch = "x86_64")]
4065#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4066unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4067 unsafe {
4069 use core::arch::x86_64::*;
4070 let lomask = _mm_set1_epi8(0x0F);
4071 let eight = _mm256_set1_epi8(8);
4072 let mut acc = 0f32;
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 b = _mm_loadu_si128(t.add(2) as *const __m128i);
4077 let lo = _mm_and_si128(b, lomask);
4078 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4079 let w = _mm256_sub_epi8(
4080 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4081 eight,
4082 );
4083 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4084 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4085 acc += d as f32 * s;
4086 }
4087 acc
4088 }
4089}
4090
4091#[cfg(target_arch = "x86_64")]
4096#[target_feature(enable = "avx2,fma")]
4101unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4102 unsafe {
4104 use core::arch::x86_64::*;
4105 let lomask = _mm_set1_epi8(0x0F);
4106 let eight = _mm256_set1_epi8(8);
4107 let ones = _mm256_set1_epi16(1);
4108 let mut f0 = _mm256_setzero_ps();
4121 let mut f1 = _mm256_setzero_ps();
4122 let mut f2 = _mm256_setzero_ps();
4123 let mut f3 = _mm256_setzero_ps();
4124 for gi in 0..gpr {
4125 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4126 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4127 let sv = _mm256_set1_ps(s);
4128 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4129 let lo = _mm_and_si128(bb, lomask);
4130 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4131 let w = _mm256_sub_epi8(
4132 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4133 eight,
4134 );
4135 let aw = _mm256_abs_epi8(w);
4136 let off = gi * GROUP_SIZE;
4137 let dot = |xq: &[i8]| {
4138 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4139 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4140 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4141 };
4142 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4143 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4144 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4145 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4146 }
4147 [
4148 hsum256_ps(f0),
4149 hsum256_ps(f1),
4150 hsum256_ps(f2),
4151 hsum256_ps(f3),
4152 ]
4153 }
4154}
4155
4156#[cfg(target_arch = "x86_64")]
4159#[target_feature(enable = "avx2")]
4160#[inline]
4161unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4162 unsafe {
4164 use core::arch::x86_64::*;
4165 let hi = _mm256_extractf128_ps::<1>(v);
4166 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4167 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4168 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4169 _mm_cvtss_f32(s)
4170 }
4171}
4172
4173#[cfg(target_arch = "x86_64")]
4175#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4176unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4177 unsafe {
4179 use core::arch::x86_64::*;
4180 let lomask = _mm_set1_epi8(0x0F);
4181 let eight = _mm256_set1_epi8(8);
4182 let mut f0 = _mm256_setzero_ps();
4185 let mut f1 = _mm256_setzero_ps();
4186 let mut f2 = _mm256_setzero_ps();
4187 let mut f3 = _mm256_setzero_ps();
4188 for gi in 0..gpr {
4189 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4190 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4191 let sv = _mm256_set1_ps(s);
4192 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4193 let lo = _mm_and_si128(bb, lomask);
4194 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4195 let w = _mm256_sub_epi8(
4196 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4197 eight,
4198 );
4199 let aw = _mm256_abs_epi8(w);
4200 let off = gi * GROUP_SIZE;
4201 let dot = |xq: &[i8]| {
4202 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4203 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4204 _mm256_setzero_si256(),
4205 aw,
4206 _mm256_sign_epi8(x, w),
4207 ))
4208 };
4209 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4210 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4211 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4212 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4213 }
4214 let acc = [
4215 hsum256_ps(f0),
4216 hsum256_ps(f1),
4217 hsum256_ps(f2),
4218 hsum256_ps(f3),
4219 ];
4220 acc
4221 }
4222}
4223
4224#[cfg(target_arch = "aarch64")]
4229#[target_feature(enable = "neon,dotprod")]
4230unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4231 unsafe {
4233 use core::arch::aarch64::*;
4234 use core::arch::asm;
4235 let lomask = vdupq_n_u8(0x0F);
4236 let eight = vdupq_n_s8(8);
4237 let mut acc = [0f32; 4];
4238 for gi in 0..gpr {
4239 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4240 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4241 let b = vld1q_u8(t.add(2));
4242 let lo = vandq_u8(b, lomask);
4243 let hi = vshrq_n_u8::<4>(b);
4244 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4245 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4246 for (k, xq) in xs.iter().enumerate() {
4247 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4248 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4249 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4250 asm!(
4251 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4252 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4253 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4254 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4255 options(pure, nomem, nostack),
4256 );
4257 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4258 }
4259 }
4260 acc
4261 }
4262}
4263
4264#[inline]
4266fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4267 let gi = j / GROUP_SIZE;
4268 let k = j % GROUP_SIZE;
4269 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4270 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4271 let byte = tile[2 + k / 2];
4272 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4273 ((nib as i32 - 8) as f32, s)
4274}
4275
4276#[inline]
4279fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4280 let mut acc = 0f32;
4281 for gi in 0..gpr {
4282 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4283 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4284 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4285 let mut ga = 0f32;
4286 for (k, &b) in tile[2..].iter().enumerate() {
4287 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4288 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4289 }
4290 acc += ga * s;
4291 }
4292 acc
4293}
4294
4295struct Q4tpView<'a> {
4299 nib: &'a [u8],
4300 params: &'a [u8],
4301 codes: &'a [u8],
4302 stride: usize,
4303 zero_rung: bool,
4305}
4306
4307impl<'a> Q4tpView<'a> {
4308 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4309 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4310 Self {
4311 nib: &bytes[..params_off],
4312 params: &bytes[params_off..codes_off],
4313 codes: &bytes[codes_off..],
4314 stride,
4315 zero_rung: false,
4316 }
4317 }
4318
4319 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4321 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4322 Self {
4323 nib: &bytes[..params_off],
4324 params: &bytes[params_off..codes_off],
4325 codes: &bytes[codes_off..],
4326 stride,
4327 zero_rung: true,
4328 }
4329 }
4330
4331 #[inline]
4347 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4348 let tab = if self.zero_rung {
4349 q2tp_ladder(self.params, r)
4350 } else {
4351 q4tp_ladder(self.params, r)
4352 };
4353 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4354 let out = &mut out[..gpr];
4355 let mut chunks = out.chunks_exact_mut(8);
4356 let mut ci = 0usize;
4357 for c in &mut chunks {
4358 let w = u64::from(codes[ci])
4359 | u64::from(codes[ci + 1]) << 8
4360 | u64::from(codes[ci + 2]) << 16
4361 | u64::from(codes[ci + 3]) << 24
4362 | u64::from(codes[ci + 4]) << 32;
4363 for (k, o) in c.iter_mut().enumerate() {
4364 *o = tab[((w >> (5 * k)) & 31) as usize];
4365 }
4366 ci += 5;
4367 }
4368 let tail = &codes[ci..];
4371 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4372 *o = tab[q4tp_code(tail, k)];
4373 }
4374 }
4375}
4376
4377#[inline]
4378fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4379 #[cfg(target_arch = "aarch64")]
4380 unsafe {
4381 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4382 }
4383 #[cfg(target_arch = "x86_64")]
4384 unsafe {
4385 if vnni_tiles_enabled() {
4386 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4387 }
4388 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4389 }
4390 #[allow(unreachable_code)]
4391 {
4392 let mut acc = 0f32;
4393 for gi in 0..gpr {
4394 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4395 let s = scales[gi];
4396 let mut d = 0i32;
4397 for (k, &b) in tile.iter().enumerate() {
4398 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4399 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4400 }
4401 acc += d as f32 * s;
4402 }
4403 acc
4404 }
4405}
4406
4407#[cfg(target_arch = "aarch64")]
4410#[target_feature(enable = "neon,dotprod")]
4411unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4412 unsafe {
4415 use core::arch::aarch64::*;
4416 use core::arch::asm;
4417 let lomask = vdupq_n_u8(0x0F);
4418 let eight = vdupq_n_s8(8);
4419 let mut acc = 0f32;
4420 for gi in 0..gpr {
4421 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4422 let s = *scales.get_unchecked(gi);
4423 let b = vld1q_u8(t);
4424 let lo = vandq_u8(b, lomask);
4425 let hi = vshrq_n_u8::<4>(b);
4426 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4427 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4428 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4429 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4430 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4431 asm!(
4432 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4433 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4434 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4435 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4436 options(pure, nomem, nostack),
4437 );
4438 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4439 }
4440 acc
4441 }
4442}
4443
4444#[cfg(target_arch = "x86_64")]
4445#[target_feature(enable = "avx2")]
4446unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4447 unsafe {
4449 use core::arch::x86_64::*;
4450 let lomask = _mm_set1_epi8(0x0F);
4451 let eight = _mm256_set1_epi8(8);
4452 let ones = _mm256_set1_epi16(1);
4453 let mut acc = 0f32;
4454 for gi in 0..gpr {
4455 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4456 let s = *scales.get_unchecked(gi);
4457 let b = _mm_loadu_si128(t as *const __m128i);
4458 let lo = _mm_and_si128(b, lomask);
4459 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4460 let w = _mm256_sub_epi8(
4461 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4462 eight,
4463 );
4464 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4465 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4466 let d = _mm256_madd_epi16(p16, ones);
4467 let hi128 = _mm256_extracti128_si256::<1>(d);
4468 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4469 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4470 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4471 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4472 }
4473 acc
4474 }
4475}
4476
4477#[cfg(target_arch = "x86_64")]
4480#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4481unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4482 unsafe {
4484 use core::arch::x86_64::*;
4485 let lomask = _mm_set1_epi8(0x0F);
4486 let eight = _mm256_set1_epi8(8);
4487 let mut acc = 0f32;
4488 for gi in 0..gpr {
4489 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4490 let s = *scales.get_unchecked(gi);
4491 let b = _mm_loadu_si128(t as *const __m128i);
4492 let lo = _mm_and_si128(b, lomask);
4493 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4494 let w = _mm256_sub_epi8(
4495 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4496 eight,
4497 );
4498 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4499 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4500 }
4501 acc
4502 }
4503}
4504
4505#[inline]
4508fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4509 let mut acc = 0f32;
4510 for gi in 0..gpr {
4511 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4512 let s = scales[gi];
4513 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4514 let mut ga = 0f32;
4515 for (k, &b) in tile.iter().enumerate() {
4516 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4517 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4518 }
4519 acc += ga * s;
4520 }
4521 acc
4522}
4523
4524#[inline]
4527fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4528 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4529 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4530 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4531 ((n as i32 - 8) as f32, scales[gi])
4532}
4533
4534fn q4tp_matvec(
4536 bytes: &[u8],
4537 x: &[f32],
4538 rows: usize,
4539 cols: usize,
4540 out: &mut [f32],
4541 pool: Option<&Pool>,
4542) {
4543 debug_assert_eq!(out.len(), rows);
4544 let gpr = cols / GROUP_SIZE;
4545 let v = Q4tpView::new(bytes, rows, cols);
4546 let out_addr = SendMut(out.as_mut_ptr());
4547 if a8w8_enabled() {
4548 let act = split_act(x);
4549 let run = |start: usize, end: usize| {
4550 with_krow(gpr, |sc| {
4552 for r in start..end {
4553 v.scales_into(r, gpr, sc);
4554 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4555 for &(j, xv) in &act.outliers {
4556 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4557 acc += w * s * xv;
4558 }
4559 unsafe { *out_addr.at(r) = acc };
4561 }
4562 })
4563 };
4564 dispatch_rows(pool, rows, &run);
4565 return;
4566 }
4567 let run = |start: usize, end: usize| {
4568 with_krow(gpr, |sc| {
4569 for r in start..end {
4570 v.scales_into(r, gpr, sc);
4571 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4573 }
4574 })
4575 };
4576 dispatch_rows(pool, rows, &run);
4577}
4578
4579#[allow(clippy::too_many_arguments)]
4582fn q4tp_matvec2(
4583 bytes: &[u8],
4584 x1: &[f32],
4585 x2: &[f32],
4586 rows: usize,
4587 cols: usize,
4588 o1: &mut [f32],
4589 o2: &mut [f32],
4590 pool: Option<&Pool>,
4591) {
4592 let gpr = cols / GROUP_SIZE;
4593 let v = Q4tpView::new(bytes, rows, cols);
4594 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4595 let run = |start: usize, end: usize| {
4596 let mut sc = vec![0f32; gpr];
4597 for r in start..end {
4598 v.scales_into(r, gpr, &mut sc);
4599 unsafe {
4601 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4602 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4603 }
4604 }
4605 };
4606 dispatch_rows(pool, rows, &run);
4607}
4608
4609#[inline]
4612fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4613 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4614 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4615 let c = (byte >> (2 * (k % 4))) & 3;
4616 (c as f32 - 1.5, scales[gi])
4617}
4618
4619#[inline]
4626fn dot_q2tp_row_i8(
4627 chunks: &[u8],
4628 r: usize,
4629 gpr: usize,
4630 xq: &[i8],
4631 gsum: &[i32],
4632 scales: &[f32],
4633) -> f32 {
4634 let mut acc = 0f32;
4635 let base = r * gpr * Q2TP_CHUNK;
4636 #[cfg(not(target_arch = "aarch64"))]
4637 let mut codes = [0i8; GROUP_SIZE];
4638 for gi in 0..gpr {
4639 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4640 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4641 #[cfg(target_arch = "aarch64")]
4642 let dot = unsafe {
4648 use core::arch::aarch64::*;
4649 let b = vld1_u8(ch.as_ptr());
4650 let three = vdup_n_u8(3);
4651 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4652 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4653 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4654 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4655 let x4 = vld4_s8(xg.as_ptr());
4656 let mut acc4 = vdupq_n_s32(0);
4657 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4658 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4659 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4660 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4661 vaddvq_s32(acc4)
4662 };
4663 #[cfg(not(target_arch = "aarch64"))]
4664 let dot: i32 = {
4665 for (k, &b) in ch.iter().enumerate() {
4666 codes[k * 4] = (b & 3) as i8;
4667 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4668 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4669 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4670 }
4671 codes
4672 .iter()
4673 .zip(xg)
4674 .map(|(&c, &x)| c as i32 * x as i32)
4675 .sum()
4676 };
4677 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4678 }
4679 acc
4680}
4681
4682fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4686 let mut acc = 0f32;
4687 for gi in 0..gpr {
4688 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4689 let s = scales[gi];
4690 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4691 let mut g = 0f32;
4692 for (k, &b) in ch.iter().enumerate() {
4693 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4694 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4695 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4696 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4697 }
4698 acc += s * g;
4699 }
4700 acc
4701}
4702
4703fn q2tp_matvec(
4704 bytes: &[u8],
4705 x: &[f32],
4706 rows: usize,
4707 cols: usize,
4708 out: &mut [f32],
4709 pool: Option<&Pool>,
4710) {
4711 debug_assert_eq!(out.len(), rows);
4712 let gpr = cols / GROUP_SIZE;
4713 let v = Q4tpView::new_q2(bytes, rows, cols);
4714 let out_addr = SendMut(out.as_mut_ptr());
4715 if a8w8_enabled() {
4720 let act = split_act(x);
4721 let gsum = q1_group_sums(&act.xq, gpr);
4722 let (act, gsum) = (&act, &gsum);
4723 let run = move |start: usize, end: usize| {
4724 with_krow(gpr, |sc| {
4725 for r in start..end {
4726 v.scales_into(r, gpr, sc);
4727 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4728 for &(j, xv) in &act.outliers {
4729 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4730 acc += w * s * xv;
4731 }
4732 unsafe { *out_addr.at(r) = acc };
4734 }
4735 })
4736 };
4737 dispatch_rows(pool, rows, &run);
4738 return;
4739 }
4740 let run = |start: usize, end: usize| {
4741 with_krow(gpr, |sc| {
4742 for r in start..end {
4743 v.scales_into(r, gpr, sc);
4744 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4746 }
4747 })
4748 };
4749 dispatch_rows(pool, rows, &run);
4750}
4751
4752#[allow(clippy::too_many_arguments)]
4754fn q2tp_matvec2(
4755 bytes: &[u8],
4756 x1: &[f32],
4757 x2: &[f32],
4758 rows: usize,
4759 cols: usize,
4760 o1: &mut [f32],
4761 o2: &mut [f32],
4762 pool: Option<&Pool>,
4763) {
4764 let gpr = cols / GROUP_SIZE;
4765 let v = Q4tpView::new_q2(bytes, rows, cols);
4766 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4767 let run = |start: usize, end: usize| {
4768 let mut sc = vec![0f32; gpr];
4769 for r in start..end {
4770 v.scales_into(r, gpr, &mut sc);
4771 unsafe {
4773 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4774 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4775 }
4776 }
4777 };
4778 dispatch_rows(pool, rows, &run);
4779}
4780
4781pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4788 let gpr = cols / GROUP_SIZE;
4793 let v = Q4tpView::new_q2(bytes, rows, cols);
4794 with_krow(gpr, |sc| {
4795 for r in 0..rows {
4796 v.scales_into(r, gpr, sc);
4797 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4798 }
4799 });
4800}
4801
4802pub fn q2tp_matmat_for_test(
4803 bytes: &[u8],
4804 xs_all: &[f32],
4805 b: usize,
4806 rows: usize,
4807 cols: usize,
4808 out: &mut [f32],
4809) {
4810 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4811}
4812
4813fn q2tp_matmat(
4814 bytes: &[u8],
4815 xs_all: &[f32],
4816 b: usize,
4817 rows: usize,
4818 cols: usize,
4819 out: &mut [f32],
4820 pool: Option<&Pool>,
4821) {
4822 debug_assert_eq!(out.len(), b * rows);
4823 let gpr = cols / GROUP_SIZE;
4824 let v = Q4tpView::new_q2(bytes, rows, cols);
4825 let out_addr = SendMut(out.as_mut_ptr());
4826 let run = |start: usize, end: usize| {
4827 let mut sc = vec![0f32; gpr];
4828 for r in start..end {
4829 v.scales_into(r, gpr, &mut sc);
4830 for bi in 0..b {
4831 let x = &xs_all[bi * cols..(bi + 1) * cols];
4832 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4834 }
4835 }
4836 };
4837 dispatch_rows(pool, rows, &run);
4838}
4839
4840#[cfg(target_arch = "aarch64")]
4844#[target_feature(enable = "neon,dotprod")]
4845unsafe fn dot_q4tp_row_1x4_sdot_v1(
4846 nib: &[u8],
4847 r: usize,
4848 gpr: usize,
4849 xs: [&[i8]; 4],
4850 scales: &[f32],
4851) -> [f32; 4] {
4852 unsafe {
4853 use core::arch::aarch64::*;
4854 use core::arch::asm;
4855 let lomask = vdupq_n_u8(0x0F);
4856 let eight = vdupq_n_s8(8);
4857 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4858 for gi in 0..gpr {
4859 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4860 let s = *scales.get_unchecked(gi);
4861 let bb = vld1q_u8(t);
4862 let lo = vandq_u8(bb, lomask);
4863 let hi = vshrq_n_u8::<4>(bb);
4864 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4865 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4866 let mut d = [0f32; 4];
4867 for (k, dk) in d.iter_mut().enumerate() {
4868 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4869 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4870 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4871 asm!(
4872 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4873 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4874 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4875 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4876 options(pure, nomem, nostack),
4877 );
4878 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4879 }
4880 f0 += d[0];
4881 f1 += d[1];
4882 f2 += d[2];
4883 f3 += d[3];
4884 }
4885 [f0, f1, f2, f3]
4886 }
4887}
4888
4889#[allow(dead_code)]
4896static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4897
4898#[cfg(target_arch = "x86_64")]
4904fn q4tp_blocked_x86() -> bool {
4905 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4906 1 => false,
4907 2 => avx512vnni_enabled(),
4912 _ => blocked_enabled() && avx512vnni_enabled(),
4916 }
4917}
4918
4919#[cfg(target_arch = "aarch64")]
4921#[allow(dead_code)]
4922fn q4tp_v1() -> bool {
4923 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4924 1 => true,
4925 2 => false,
4926 _ => {
4927 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4928 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4929 }
4930 }
4931}
4932
4933#[cfg(target_arch = "x86_64")]
4945#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4946unsafe fn dot_q4tp_2x8_avx512(
4947 nib: &[u8],
4948 r0: usize,
4949 gpr: usize,
4950 xs: [&[i8]; 8],
4951 sc0: &[f32],
4952 sc1: &[f32],
4953) -> [[f32; 8]; 2] {
4954 unsafe {
4957 use core::arch::x86_64::*;
4958 let lomask = _mm256_set1_epi8(0x0F);
4959 let eight = _mm256_set1_epi8(8);
4960 let zero = _mm512_setzero_si512();
4961 let mut v0 = [_mm512_setzero_ps(); 8];
4962 let mut v1 = [_mm512_setzero_ps(); 8];
4963 let pairs = gpr / 2;
4964 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4965 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4966 let bb = _mm256_loadu_si256(t as *const __m256i);
4967 let lo = _mm256_and_si256(bb, lomask);
4968 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4969 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4970 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4971 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4972 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4973 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
4974 };
4975 for gp in 0..pairs {
4976 let gi = gp * 2;
4977 let (wa0, neg0) = unpack(r0, gi);
4978 let (wa1, neg1) = unpack(r0 + 1, gi);
4979 let off = gi * GROUP_SIZE;
4980 let sv = |sc: &[f32]| {
4981 _mm512_insertf32x8::<1>(
4982 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
4983 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
4984 )
4985 };
4986 let s0 = sv(sc0);
4987 let s1 = sv(sc1);
4988 for k in 0..8 {
4989 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
4990 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4991 zero,
4992 wa0,
4993 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
4994 ));
4995 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4996 zero,
4997 wa1,
4998 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
4999 ));
5000 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5001 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5002 }
5003 }
5004 let mut acc = [[0f32; 8]; 2];
5005 for k in 0..8 {
5006 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5007 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5008 }
5009 if gpr % 2 == 1 {
5010 let off = (gpr - 1) * GROUP_SIZE;
5011 for j in off..off + GROUP_SIZE {
5012 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5013 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5014 for k in 0..8 {
5015 let x = *xs[k].get_unchecked(j) as f32;
5016 acc[0][k] += w0 * sa * x;
5017 acc[1][k] += w1 * sb * x;
5018 }
5019 }
5020 }
5021 acc
5022 }
5023}
5024
5025#[cfg(target_arch = "x86_64")]
5030#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5031unsafe fn dot_q4tp_row_1x8_avx512(
5032 nib: &[u8],
5033 r: usize,
5034 gpr: usize,
5035 xs: [&[i8]; 8],
5036 scales: &[f32],
5037) -> [f32; 8] {
5038 unsafe {
5040 use core::arch::x86_64::*;
5041 let lomask = _mm256_set1_epi8(0x0F);
5042 let eight = _mm256_set1_epi8(8);
5043 let zero = _mm512_setzero_si512();
5044 let (mut v0, mut v1, mut v2, mut v3) = (
5045 _mm512_setzero_ps(),
5046 _mm512_setzero_ps(),
5047 _mm512_setzero_ps(),
5048 _mm512_setzero_ps(),
5049 );
5050 let (mut v4, mut v5, mut v6, mut v7) = (
5051 _mm512_setzero_ps(),
5052 _mm512_setzero_ps(),
5053 _mm512_setzero_ps(),
5054 _mm512_setzero_ps(),
5055 );
5056 let pairs = gpr / 2;
5057 for gp in 0..pairs {
5058 let gi = gp * 2;
5059 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5060 let bb = _mm256_loadu_si256(t as *const __m256i);
5061 let lo = _mm256_and_si256(bb, lomask);
5062 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5063 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5068 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5069 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5070 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5071 let wabs = _mm512_abs_epi8(w);
5072 let neg = _mm512_movepi8_mask(w);
5073 let off = gi * GROUP_SIZE;
5074 let sv = _mm512_insertf32x8::<1>(
5075 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5076 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5077 );
5078 let dot = |x: &[i8]| -> __m512 {
5079 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5080 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5081 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5082 };
5083 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5084 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5085 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5086 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5087 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5088 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5089 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5090 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5091 }
5092 let mut acc = [
5093 _mm512_reduce_add_ps(v0),
5094 _mm512_reduce_add_ps(v1),
5095 _mm512_reduce_add_ps(v2),
5096 _mm512_reduce_add_ps(v3),
5097 _mm512_reduce_add_ps(v4),
5098 _mm512_reduce_add_ps(v5),
5099 _mm512_reduce_add_ps(v6),
5100 _mm512_reduce_add_ps(v7),
5101 ];
5102 if gpr % 2 == 1 {
5105 let off = (gpr - 1) * GROUP_SIZE;
5106 for j in off..off + GROUP_SIZE {
5107 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5108 let ws = w * s;
5109 for k in 0..8 {
5110 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5111 }
5112 }
5113 }
5114 acc
5115 }
5116}
5117
5118#[cfg(target_arch = "x86_64")]
5131#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5132unsafe fn dot_q4tp_row_1x4_avx512(
5133 nib: &[u8],
5134 r: usize,
5135 gpr: usize,
5136 xs: [&[i8]; 4],
5137 scales: &[f32],
5138) -> [f32; 4] {
5139 unsafe {
5141 use core::arch::x86_64::*;
5142 let lomask = _mm256_set1_epi8(0x0F);
5143 let eight = _mm256_set1_epi8(8);
5144 let zero = _mm512_setzero_si512();
5145 let (mut v0, mut v1, mut v2, mut v3) = (
5146 _mm512_setzero_ps(),
5147 _mm512_setzero_ps(),
5148 _mm512_setzero_ps(),
5149 _mm512_setzero_ps(),
5150 );
5151 let pairs = gpr / 2;
5152 for gp in 0..pairs {
5153 let gi = gp * 2;
5154 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5155 let bb = _mm256_loadu_si256(t as *const __m256i);
5156 let lo = _mm256_and_si256(bb, lomask);
5157 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5158 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5163 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5164 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5165 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5166 let wabs = _mm512_abs_epi8(w);
5167 let neg = _mm512_movepi8_mask(w);
5168 let off = gi * GROUP_SIZE;
5169 let sv = _mm512_insertf32x8::<1>(
5170 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5171 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5172 );
5173 let dot = |x: &[i8]| -> __m512 {
5174 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5175 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5176 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5177 };
5178 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5179 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5180 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5181 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5182 }
5183 let mut acc = [
5184 _mm512_reduce_add_ps(v0),
5185 _mm512_reduce_add_ps(v1),
5186 _mm512_reduce_add_ps(v2),
5187 _mm512_reduce_add_ps(v3),
5188 ];
5189 if gpr % 2 == 1 {
5192 let off = (gpr - 1) * GROUP_SIZE;
5193 for j in off..off + GROUP_SIZE {
5194 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5195 let ws = w * s;
5196 for k in 0..4 {
5197 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5198 }
5199 }
5200 }
5201 acc
5202 }
5203}
5204
5205#[cfg(target_arch = "aarch64")]
5209#[target_feature(enable = "neon,dotprod")]
5210unsafe fn dot_q4tp_row_1x4_sdot(
5211 nib: &[u8],
5212 r: usize,
5213 gpr: usize,
5214 xs: [&[i8]; 4],
5215 scales: &[f32],
5216) -> [f32; 4] {
5217 unsafe {
5219 use core::arch::aarch64::*;
5220 use core::arch::asm;
5221 let lomask = vdupq_n_u8(0x0F);
5222 let eight = vdupq_n_s8(8);
5223 let (mut v0, mut v1, mut v2, mut v3) = (
5239 vdupq_n_f32(0.0),
5240 vdupq_n_f32(0.0),
5241 vdupq_n_f32(0.0),
5242 vdupq_n_f32(0.0),
5243 );
5244 for gi in 0..gpr {
5245 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5246 let s = *scales.get_unchecked(gi);
5247 let bb = vld1q_u8(t);
5248 let lo = vandq_u8(bb, lomask);
5249 let hi = vshrq_n_u8::<4>(bb);
5250 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5251 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5252 let off = gi * GROUP_SIZE;
5253 let dot4 = |x: &[i8]| -> int32x4_t {
5254 let x0 = vld1q_s8(x.as_ptr().add(off));
5255 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5256 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5257 asm!(
5258 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5259 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5260 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5261 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5262 options(pure, nomem, nostack),
5263 );
5264 vaddq_s32(a0, a1)
5265 };
5266 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5267 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5268 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5269 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5270 }
5271 [
5272 vaddvq_f32(v0),
5273 vaddvq_f32(v1),
5274 vaddvq_f32(v2),
5275 vaddvq_f32(v3),
5276 ]
5277 }
5278}
5279
5280fn q4tp_matmat(
5284 bytes: &[u8],
5285 xs_all: &[f32],
5286 b: usize,
5287 rows: usize,
5288 cols: usize,
5289 out: &mut [f32],
5290 pool: Option<&Pool>,
5291) {
5292 debug_assert_eq!(out.len(), b * rows);
5293 let gpr = cols / GROUP_SIZE;
5294 let v = Q4tpView::new(bytes, rows, cols);
5295
5296 #[cfg(target_os = "macos")]
5298 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5299 dequant_matmat_accel(
5300 &|r, dst| {
5301 let mut sc = [0f32; 32];
5302 let mut scv;
5303 let s: &[f32] = if gpr <= 32 {
5304 v.scales_into(r, gpr, &mut sc);
5305 &sc[..gpr]
5306 } else {
5307 scv = vec![0f32; gpr];
5308 v.scales_into(r, gpr, &mut scv);
5309 &scv
5310 };
5311 for gi in 0..gpr {
5312 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5313 for (k, &bb) in tile.iter().enumerate() {
5314 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5315 dst[gi * GROUP_SIZE + k * 2 + 1] =
5316 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5317 }
5318 }
5319 },
5320 xs_all,
5321 b,
5322 rows,
5323 cols,
5324 out,
5325 pool,
5326 );
5327 return;
5328 }
5329
5330 let out_addr = SendMut(out.as_mut_ptr());
5331 if a8w8_enabled() {
5332 let acts: Vec<SplitAct> = (0..b)
5333 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5334 .collect();
5335 let acts = &acts;
5336 #[cfg(target_arch = "aarch64")]
5337 let blocked_ok = sdot_enabled() && blocked_enabled();
5338 #[cfg(target_arch = "x86_64")]
5345 let blocked_ok = q4tp_blocked_x86();
5346 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5347 let blocked_ok = false;
5348 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5358 .ok()
5359 .and_then(|v| v.parse().ok())
5360 .filter(|v| *v > 0)
5361 .unwrap_or(256);
5362 let run = |start: usize, end: usize| {
5363 for abase in (0..acts.len()).step_by(panel_cols) {
5364 let alen = (acts.len() - abase).min(panel_cols);
5365 let mut sc = vec![0f32; gpr];
5366 #[cfg(target_arch = "x86_64")]
5367 let mut r_lo = start;
5368 #[cfg(target_arch = "x86_64")]
5369 if blocked_ok && alen >= 8 {
5370 let mut sc1 = vec![0f32; gpr];
5371 while r_lo + 2 <= end {
5372 v.scales_into(r_lo, gpr, &mut sc);
5373 v.scales_into(r_lo + 1, gpr, &mut sc1);
5374 let mut bi = 0usize;
5375 while bi + 8 <= alen {
5376 let xs = [
5377 acts[abase + bi].xq.as_slice(),
5378 acts[abase + bi + 1].xq.as_slice(),
5379 acts[abase + bi + 2].xq.as_slice(),
5380 acts[abase + bi + 3].xq.as_slice(),
5381 acts[abase + bi + 4].xq.as_slice(),
5382 acts[abase + bi + 5].xq.as_slice(),
5383 acts[abase + bi + 6].xq.as_slice(),
5384 acts[abase + bi + 7].xq.as_slice(),
5385 ];
5386 let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5387 for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
5388 for k in 0..8 {
5389 let act = &acts[abase + bi + k];
5390 let mut acc = dr[k] * act.sx;
5391 for &(j, xv) in &act.outliers {
5392 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5393 acc += w * s * xv;
5394 }
5395 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5397 }
5398 }
5399 bi += 8;
5400 }
5401 for row in [r_lo, r_lo + 1] {
5404 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5405 for b2 in bi..alen {
5406 let act = &acts[abase + b2];
5407 let xs4 = [
5408 act.xq.as_slice(),
5409 act.xq.as_slice(),
5410 act.xq.as_slice(),
5411 act.xq.as_slice(),
5412 ];
5413 let d =
5414 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5415 let mut acc = d[0] * act.sx;
5416 for &(j, xv) in &act.outliers {
5417 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5418 acc += w * s * xv;
5419 }
5420 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5422 }
5423 }
5424 r_lo += 2;
5425 }
5426 }
5427 #[cfg(target_arch = "x86_64")]
5428 let row_start = r_lo;
5429 #[cfg(not(target_arch = "x86_64"))]
5430 let row_start = start;
5431 for r in row_start..end {
5432 v.scales_into(r, gpr, &mut sc);
5433 let mut bi = 0usize;
5434 #[cfg(target_arch = "x86_64")]
5435 if blocked_ok {
5436 while bi + 8 <= 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 acts[abase + bi + 4].xq.as_slice(),
5443 acts[abase + bi + 5].xq.as_slice(),
5444 acts[abase + bi + 6].xq.as_slice(),
5445 acts[abase + bi + 7].xq.as_slice(),
5446 ];
5447 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5448 for k in 0..8 {
5449 let act = &acts[abase + bi + k];
5450 let mut acc = d[k] * act.sx;
5451 for &(j, xv) in &act.outliers {
5452 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5453 acc += w * s * xv;
5454 }
5455 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5457 }
5458 bi += 8;
5459 }
5460 while bi + 4 <= alen {
5461 let xs = [
5462 acts[abase + bi].xq.as_slice(),
5463 acts[abase + bi + 1].xq.as_slice(),
5464 acts[abase + bi + 2].xq.as_slice(),
5465 acts[abase + bi + 3].xq.as_slice(),
5466 ];
5467 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5468 for k in 0..4 {
5469 let act = &acts[abase + bi + k];
5470 let mut acc = d[k] * act.sx;
5471 for &(j, xv) in &act.outliers {
5472 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5473 acc += w * s * xv;
5474 }
5475 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5477 }
5478 bi += 4;
5479 }
5480 }
5481 #[cfg(target_arch = "aarch64")]
5482 if blocked_ok {
5483 while bi + 4 <= alen {
5484 let xs = [
5485 acts[abase + bi].xq.as_slice(),
5486 acts[abase + bi + 1].xq.as_slice(),
5487 acts[abase + bi + 2].xq.as_slice(),
5488 acts[abase + bi + 3].xq.as_slice(),
5489 ];
5490 let d = unsafe {
5491 if q4tp_v1() {
5492 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5493 } else {
5494 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5495 }
5496 };
5497 for k in 0..4 {
5498 let act = &acts[abase + bi + k];
5499 let mut acc = d[k] * act.sx;
5500 for &(j, xv) in &act.outliers {
5501 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5502 acc += w * s * xv;
5503 }
5504 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5506 }
5507 bi += 4;
5508 }
5509 }
5510 let _ = blocked_ok;
5511 while bi < alen {
5512 let act = &acts[abase + bi];
5513 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5514 for &(j, xv) in &act.outliers {
5515 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5516 acc += w * s * xv;
5517 }
5518 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5520 bi += 1;
5521 }
5522 }
5523 }
5524 };
5525 dispatch_rows(pool, rows, &run);
5526 return;
5527 }
5528
5529 let run = |start: usize, end: usize| {
5530 let mut sc = vec![0f32; gpr];
5531 for r in start..end {
5532 v.scales_into(r, gpr, &mut sc);
5533 for bi in 0..b {
5534 let x = &xs_all[bi * cols..(bi + 1) * cols];
5535 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5537 }
5538 }
5539 };
5540 dispatch_rows(pool, rows, &run);
5541}
5542
5543fn q4t_matvec(
5545 bytes: &[u8],
5546 x: &[f32],
5547 rows: usize,
5548 cols: usize,
5549 out: &mut [f32],
5550 pool: Option<&Pool>,
5551) {
5552 debug_assert_eq!(out.len(), rows);
5553 let gpr = cols / GROUP_SIZE;
5554 let out_addr = SendMut(out.as_mut_ptr());
5555 if a8w8_enabled() {
5556 let act = split_act(x);
5557 let run = move |start: usize, end: usize| {
5558 for r in start..end {
5559 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5560 for &(j, xv) in &act.outliers {
5561 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5562 acc += w * s * xv;
5563 }
5564 unsafe { *out_addr.at(r) = acc };
5566 }
5567 };
5568 dispatch_rows(pool, rows, &run);
5569 return;
5570 }
5571 let run = move |start: usize, end: usize| {
5572 for r in start..end {
5573 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5575 }
5576 };
5577 dispatch_rows(pool, rows, &run);
5578}
5579
5580#[allow(clippy::too_many_arguments)]
5582fn q4t_matvec2(
5583 bytes: &[u8],
5584 x1: &[f32],
5585 x2: &[f32],
5586 rows: usize,
5587 cols: usize,
5588 o1: &mut [f32],
5589 o2: &mut [f32],
5590 pool: Option<&Pool>,
5591) {
5592 let gpr = cols / GROUP_SIZE;
5593 let p1 = SendMut(o1.as_mut_ptr());
5594 let p2 = SendMut(o2.as_mut_ptr());
5595 if a8w8_enabled() {
5596 let a1 = split_act(x1);
5597 let a2 = split_act(x2);
5598 let run = move |start: usize, end: usize| {
5599 for r in start..end {
5600 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5601 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5602 for &(j, xv) in &a1.outliers {
5603 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5604 v1 += w * s * xv;
5605 }
5606 for &(j, xv) in &a2.outliers {
5607 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5608 v2 += w * s * xv;
5609 }
5610 unsafe {
5612 *p1.at(r) = v1;
5613 *p2.at(r) = v2;
5614 }
5615 }
5616 };
5617 dispatch_rows(pool, rows, &run);
5618 return;
5619 }
5620 let run = move |start: usize, end: usize| {
5621 for r in start..end {
5622 unsafe {
5624 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5625 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5626 }
5627 }
5628 };
5629 dispatch_rows(pool, rows, &run);
5630}
5631
5632#[allow(clippy::too_many_arguments)]
5634#[cfg(target_os = "macos")]
5640fn dequant_matmat_accel(
5641 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5642 xs_all: &[f32],
5643 b: usize,
5644 rows: usize,
5645 cols: usize,
5646 out: &mut [f32],
5647 pool: Option<&Pool>,
5648) {
5649 const TR: usize = 2048;
5650 thread_local! {
5651 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5652 }
5653 WTILE.with(|wt| {
5654 let mut wtile = wt.borrow_mut();
5655 wtile.resize(TR * cols, 0.0);
5656 let mut r0 = 0usize;
5657 while r0 < rows {
5658 let tr = TR.min(rows - r0);
5659 let wt_addr = SendMut(wtile.as_mut_ptr());
5660 let run = |start: usize, end: usize| {
5661 for r in start..end {
5662 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5664 dequant_row(r0 + r, dst);
5665 }
5666 };
5667 dispatch_rows(pool, tr, &run);
5668 unsafe {
5669 accel_blas::cblas_sgemm(
5670 101, 111, 112, b as i32,
5674 tr as i32,
5675 cols as i32,
5676 1.0,
5677 xs_all.as_ptr(),
5678 cols as i32,
5679 wtile.as_ptr(),
5680 cols as i32,
5681 0.0,
5682 out.as_mut_ptr().add(r0),
5683 rows as i32,
5684 );
5685 }
5686 r0 += tr;
5687 }
5688 });
5689}
5690
5691fn q4t_matmat(
5692 bytes: &[u8],
5693 xs_all: &[f32],
5694 b: usize,
5695 rows: usize,
5696 cols: usize,
5697 out: &mut [f32],
5698 pool: Option<&Pool>,
5699) {
5700 debug_assert_eq!(out.len(), b * rows);
5701 let gpr = cols / GROUP_SIZE;
5702 #[cfg(target_os = "macos")]
5706 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5707 dequant_matmat_accel(
5708 &|r, dst| {
5709 for gi in 0..gpr {
5710 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5711 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5712 for (k, &bb) in tile[2..].iter().enumerate() {
5713 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5714 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5715 }
5716 }
5717 },
5718 xs_all,
5719 b,
5720 rows,
5721 cols,
5722 out,
5723 pool,
5724 );
5725 return;
5726 }
5727 let out_addr = SendMut(out.as_mut_ptr());
5728 if a8w8_enabled() {
5729 let acts: Vec<SplitAct> = (0..b)
5730 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5731 .collect();
5732 let acts = &acts;
5733 #[cfg(target_arch = "x86_64")]
5734 let blocked_ok = avx2_enabled() && blocked_enabled();
5735 #[cfg(target_arch = "aarch64")]
5736 let blocked_ok = sdot_enabled() && blocked_enabled();
5737 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5738 let blocked_ok = false;
5739 let run = move |start: usize, end: usize| {
5740 for r in start..end {
5741 let mut bi = 0usize;
5742 #[cfg(target_arch = "aarch64")]
5743 if blocked_ok {
5744 while bi + 4 <= acts.len() {
5745 let xs = [
5746 acts[bi].xq.as_slice(),
5747 acts[bi + 1].xq.as_slice(),
5748 acts[bi + 2].xq.as_slice(),
5749 acts[bi + 3].xq.as_slice(),
5750 ];
5751 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5752 for k in 0..4 {
5753 let act = &acts[bi + k];
5754 let mut acc = d[k] * act.sx;
5755 for &(j, xv) in &act.outliers {
5756 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5757 acc += w * sc * xv;
5758 }
5759 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5761 }
5762 bi += 4;
5763 }
5764 }
5765 #[cfg(target_arch = "x86_64")]
5766 if blocked_ok {
5767 while bi + 4 <= acts.len() {
5768 let xs = [
5769 acts[bi].xq.as_slice(),
5770 acts[bi + 1].xq.as_slice(),
5771 acts[bi + 2].xq.as_slice(),
5772 acts[bi + 3].xq.as_slice(),
5773 ];
5774 let d = unsafe {
5775 if vnni_tiles_enabled() {
5776 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5777 } else {
5778 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5779 }
5780 };
5781 for k in 0..4 {
5782 let act = &acts[bi + k];
5783 let mut acc = d[k] * act.sx;
5784 for &(j, xv) in &act.outliers {
5785 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5786 acc += w * sc * xv;
5787 }
5788 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5790 }
5791 bi += 4;
5792 }
5793 }
5794 let _ = blocked_ok;
5795 while bi < acts.len() {
5796 let act = &acts[bi];
5797 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5798 for &(j, xv) in &act.outliers {
5799 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5800 acc += w * s * xv;
5801 }
5802 unsafe { *out_addr.at(bi * rows + r) = acc };
5804 bi += 1;
5805 }
5806 }
5807 };
5808 dispatch_rows(pool, rows, &run);
5809 return;
5810 }
5811 let run = move |start: usize, end: usize| {
5812 for r in start..end {
5813 for bi in 0..b {
5814 let x = &xs_all[bi * cols..(bi + 1) * cols];
5815 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5817 }
5818 }
5819 };
5820 dispatch_rows(pool, rows, &run);
5821}
5822
5823fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5832 (0..gpr)
5833 .map(|gi| {
5834 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5835 .iter()
5836 .map(|&v| v as i32)
5837 .sum()
5838 })
5839 .collect()
5840}
5841
5842#[inline]
5846#[allow(unreachable_code)]
5847#[cfg(target_arch = "x86_64")]
5852#[target_feature(enable = "avx2")]
5853unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5854 unsafe {
5856 use core::arch::x86_64::*;
5857 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 ones16 = _mm256_set1_epi16(1);
5868 let mut acc = 0f32;
5869 for gi in 0..gpr {
5870 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5871 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5872 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5873 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5874 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5875 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5876 let sel = _mm256_and_si256(x, mask);
5877 let p16 = _mm256_maddubs_epi16(ones8, sel);
5879 let d32 = _mm256_madd_epi16(p16, ones16);
5880 let hi128 = _mm256_extracti128_si256::<1>(d32);
5881 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5882 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5883 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5884 let msum = _mm_cvtsi128_si32(s32);
5885 let d = 2 * msum - gsum[gi];
5888 acc += d as f32 * s;
5889 }
5890 acc
5891 }
5892}
5893
5894#[cfg(target_arch = "x86_64")]
5897#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5898unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5899 unsafe {
5901 use core::arch::x86_64::*;
5902 let expand = _mm256_setr_epi8(
5903 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,
5904 3, 3, 3,
5905 );
5906 let bitsel = _mm256_setr_epi8(
5907 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5908 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5909 );
5910 let ones8 = _mm256_set1_epi8(1);
5911 let mut acc = 0f32;
5912 for gi in 0..gpr {
5913 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5914 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5915 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5916 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5917 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5918 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5919 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5920 let d = 2 * msum - gsum[gi];
5921 acc += d as f32 * s;
5922 }
5923 acc
5924 }
5925}
5926
5927#[cfg(target_arch = "x86_64")]
5929#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5930unsafe fn dot_q1_row_1x4_vnni(
5931 bytes: &[u8],
5932 r: usize,
5933 gpr: usize,
5934 xs: [&[i8]; 4],
5935 gsums: [&[i32]; 4],
5936) -> [f32; 4] {
5937 unsafe {
5939 use core::arch::x86_64::*;
5940 let expand = _mm256_setr_epi8(
5941 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,
5942 3, 3, 3,
5943 );
5944 let bitsel = _mm256_setr_epi8(
5945 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5946 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5947 );
5948 let ones8 = _mm256_set1_epi8(1);
5949 let mut acc = [0f32; 4];
5950 for gi in 0..gpr {
5951 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5952 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5953 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5954 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5955 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5956 for (k, xq) in xs.iter().enumerate() {
5957 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5958 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5959 let d = 2 * msum - gsums[k][gi];
5960 acc[k] += d as f32 * s;
5961 }
5962 }
5963 acc
5964 }
5965}
5966
5967#[cfg(target_arch = "x86_64")]
5970#[target_feature(enable = "avx2")]
5971unsafe fn dot_q1_row_1x4_avx2(
5972 bytes: &[u8],
5973 r: usize,
5974 gpr: usize,
5975 xs: [&[i8]; 4],
5976 gsums: [&[i32]; 4],
5977) -> [f32; 4] {
5978 unsafe {
5980 use core::arch::x86_64::*;
5981 let expand = _mm256_setr_epi8(
5982 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,
5983 3, 3, 3,
5984 );
5985 let bitsel = _mm256_setr_epi8(
5986 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5987 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5988 );
5989 let ones8 = _mm256_set1_epi8(1);
5990 let ones16 = _mm256_set1_epi16(1);
5991 let mut acc = [0f32; 4];
5992 for gi in 0..gpr {
5993 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5994 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5995 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5996 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5997 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5998 for (k, xq) in xs.iter().enumerate() {
5999 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6000 let sel = _mm256_and_si256(x, mask);
6001 let p16 = _mm256_maddubs_epi16(ones8, sel);
6002 let d32 = _mm256_madd_epi16(p16, ones16);
6003 let hi128 = _mm256_extracti128_si256::<1>(d32);
6004 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6005 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6006 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6007 let msum = _mm_cvtsi128_si32(s32);
6008 let d = 2 * msum - gsums[k][gi];
6009 acc[k] += d as f32 * s;
6010 }
6011 }
6012 acc
6013 }
6014}
6015
6016#[allow(unreachable_code)]
6017fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6018 #[cfg(target_arch = "aarch64")]
6019 unsafe {
6020 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6021 }
6022 #[cfg(target_arch = "x86_64")]
6023 if avx2_enabled() {
6024 unsafe {
6025 if vnni_tiles_enabled() {
6026 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6027 }
6028 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6029 }
6030 }
6031 let _ = gsum;
6032 let mut acc = 0f32;
6033 for gi in 0..gpr {
6034 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6035 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6036 let mut d = 0i32;
6037 for (j, &b) in tile[2..].iter().enumerate() {
6038 for k in 0..8 {
6039 let w = ((b >> k) & 1) as i32 * 2 - 1;
6040 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6041 }
6042 }
6043 acc += d as f32 * s;
6044 }
6045 acc
6046}
6047
6048#[cfg(target_arch = "aarch64")]
6057#[target_feature(enable = "neon,dotprod")]
6058unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6059 unsafe {
6062 use core::arch::aarch64::*;
6063 use core::arch::asm;
6064 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6065 let m = vld1q_u8(MASKS.as_ptr());
6066 macro_rules! tile_dot {
6068 ($t:expr, $x:expr) => {{
6069 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6070 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6071 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6072 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6073 let x0 = vld1q_s8($x);
6074 let x1 = vld1q_s8($x.add(16));
6075 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6076 asm!(
6077 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6078 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6079 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6080 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6081 options(pure, nomem, nostack),
6082 );
6083 vaddq_s32(a0, a1)
6084 }};
6085 }
6086 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6095 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6096 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6097 const IW11: [u8; 16] = [
6098 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6099 ];
6100 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6101 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6102 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6103 let isc = vld1_u8(ISC.as_ptr());
6104 macro_rules! tile_dot_tbl {
6106 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6107 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6108 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6109 let x0 = vld1q_s8($x);
6110 let x1 = vld1q_s8($x.add(16));
6111 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6112 asm!(
6113 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6114 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6115 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6116 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6117 options(pure, nomem, nostack),
6118 );
6119 vaddq_s32(a0, a1)
6120 }};
6121 }
6122 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6123 let row_base = r * gpr * Q1_TILE;
6124 let abs_end = bytes.len();
6125 let xp = xq.as_ptr();
6126 let gp = gsum.as_ptr();
6127 let mut accv = vdupq_n_f32(0.0);
6128 let mut gi = 0;
6129 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6132 let t0 = base.add(gi * Q1_TILE);
6133 let ld_a = vld1q_u8(t0);
6134 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6135 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6136 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6137 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6138 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6139 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6141 let g = vld1q_s32(gp.add(gi));
6142 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6143 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6144 let scf: float32x4_t;
6145 asm!(
6146 "fcvtl {o:v}.4s, {i:v}.4h",
6147 o = out(vreg) scf, i = in(vreg) sc16,
6148 options(pure, nomem, nostack),
6149 );
6150 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6151 gi += 4;
6152 }
6153 let mut acc = vaddvq_f32(accv);
6154 while gi < gpr {
6155 let t = base.add(gi * Q1_TILE);
6156 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6157 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6158 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6159 gi += 1;
6160 }
6161 acc
6162 }
6163}
6164
6165#[cfg(target_arch = "aarch64")]
6170#[target_feature(enable = "neon,dotprod")]
6171unsafe fn dot_q1_row_1x4_sdot(
6172 bytes: &[u8],
6173 r: usize,
6174 gpr: usize,
6175 xs: [&[i8]; 4],
6176 gs: [&[i32]; 4],
6177) -> [f32; 4] {
6178 unsafe {
6180 use core::arch::aarch64::*;
6181 use core::arch::asm;
6182 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6183 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6184 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6185 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6186 const IW11: [u8; 16] = [
6187 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6188 ];
6189 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6190 let m = vld1q_u8(MASKS.as_ptr());
6191 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6192 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6193 let isc = vld1_u8(ISC.as_ptr());
6194 macro_rules! sdot2 {
6195 ($w0:expr, $w1:expr, $x:expr) => {{
6196 let x0 = vld1q_s8($x);
6197 let x1 = vld1q_s8($x.add(16));
6198 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6199 asm!(
6200 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6201 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6202 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6203 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6204 options(pure, nomem, nostack),
6205 );
6206 vaddq_s32(a0, a1)
6207 }};
6208 }
6209 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6210 let row_base = r * gpr * Q1_TILE;
6211 let abs_end = bytes.len();
6212 let mut accv = [vdupq_n_f32(0.0); 4];
6213 let mut gi = 0;
6214 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6215 let t0 = base.add(gi * Q1_TILE);
6216 let ld_a = vld1q_u8(t0);
6217 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6218 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6220 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6221 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6222 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6223 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6224 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6225 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6226 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6227 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6228 let scf: float32x4_t;
6229 asm!(
6230 "fcvtl {o:v}.4s, {i:v}.4h",
6231 o = out(vreg) scf, i = in(vreg) sc16,
6232 options(pure, nomem, nostack),
6233 );
6234 for k in 0..4 {
6235 let xp = xs[k].as_ptr();
6236 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6237 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6238 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6239 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6240 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6241 let g = vld1q_s32(gs[k].as_ptr().add(gi));
6242 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6243 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6244 }
6245 gi += 4;
6246 }
6247 let mut acc = [
6248 vaddvq_f32(accv[0]),
6249 vaddvq_f32(accv[1]),
6250 vaddvq_f32(accv[2]),
6251 vaddvq_f32(accv[3]),
6252 ];
6253 while gi < gpr {
6254 let t = base.add(gi * Q1_TILE);
6255 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6256 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6257 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6258 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6259 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6260 for k in 0..4 {
6261 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6262 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6263 }
6264 gi += 1;
6265 }
6266 acc
6267 }
6268}
6269
6270#[inline]
6272fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6273 let gi = j / GROUP_SIZE;
6274 let k = j % GROUP_SIZE;
6275 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6276 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6277 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6278 ((bit as i32 * 2 - 1) as f32, s)
6279}
6280
6281#[inline]
6283fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6284 let mut acc = 0f32;
6285 for gi in 0..gpr {
6286 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6287 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6288 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6289 let mut ga = 0f32;
6290 for (j, &b) in tile[2..].iter().enumerate() {
6291 for k in 0..8 {
6292 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6293 }
6294 }
6295 acc += ga * s;
6296 }
6297 acc
6298}
6299
6300#[allow(clippy::too_many_arguments)]
6303fn q1_range_a8w8(
6304 bytes: &[u8],
6305 gpr: usize,
6306 act: &SplitAct,
6307 gsum: &[i32],
6308 out: SendMut,
6309 start: usize,
6310 end: usize,
6311) {
6312 for r in start..end {
6313 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6314 for &(j, xv) in &act.outliers {
6315 let (w, s) = q1_outlier(bytes, r, gpr, j);
6316 acc += w * s * xv;
6317 }
6318 unsafe { *out.at(r) = acc };
6320 }
6321}
6322
6323fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6325 for r in start..end {
6326 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6328 }
6329}
6330
6331fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6336 let entries = base_len + (rows + 1) * 4;
6337 (base_len, entries, entries <= bytes.len())
6338}
6339
6340#[inline]
6342fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6343 let o = rp_off + r * 4;
6344 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6345}
6346
6347const SIGN5: [[f32; 5]; 256] = {
6351 let mut lut = [[0.0f32; 5]; 256];
6352 let pow3 = [1u16, 3, 9, 27, 81];
6353 let mut byte = 0usize;
6354 while byte < 256 {
6355 let mut i = 0usize;
6356 while i < 5 {
6357 let code = (byte as u16 / pow3[i]) % 3;
6358 lut[byte][i] = if code == 1 {
6359 1.0
6360 } else if code == 2 {
6361 -1.0
6362 } else {
6363 0.0
6364 };
6365 i += 1;
6366 }
6367 byte += 1;
6368 }
6369 lut
6370};
6371
6372const SIGN5_I8: [[i8; 5]; 256] = {
6374 let mut lut = [[0i8; 5]; 256];
6375 let pow3 = [1u16, 3, 9, 27, 81];
6376 let mut byte = 0usize;
6377 while byte < 256 {
6378 let mut i = 0usize;
6379 while i < 5 {
6380 let code = (byte as u16 / pow3[i]) % 3;
6381 lut[byte][i] = if code == 1 {
6382 1
6383 } else if code == 2 {
6384 -1
6385 } else {
6386 0
6387 };
6388 i += 1;
6389 }
6390 byte += 1;
6391 }
6392 lut
6393};
6394
6395const SIGN5_U64: [u64; 256] = {
6401 let mut lut = [0u64; 256];
6402 let pow3 = [1u16, 3, 9, 27, 81];
6403 let mut byte = 0usize;
6404 while byte < 256 {
6405 let mut v = 0u64;
6406 let mut i = 0usize;
6407 while i < 5 {
6408 let code = (byte as u16 / pow3[i]) % 3;
6409 let s: u8 = if code == 1 {
6410 1
6411 } else if code == 2 {
6412 0xFF
6413 } else {
6414 0
6415 };
6416 v |= (s as u64) << (i * 8);
6417 i += 1;
6418 }
6419 lut[byte] = v;
6420 byte += 1;
6421 }
6422 lut
6423};
6424
6425#[inline]
6430fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6431 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6432 let off = (r * gpr + j / GROUP_SIZE) * TILE;
6433 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6434 let within = j % GROUP_SIZE;
6435 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6436}
6437
6438#[cfg(target_arch = "aarch64")]
6441#[target_feature(enable = "neon,dotprod")]
6442#[inline]
6443unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6444 unsafe {
6446 use core::arch::aarch64::*;
6447 use core::arch::asm;
6448 let w0 = vld1q_s8(w);
6449 let w1 = vld1q_s8(w.add(16));
6450 let x0 = vld1q_s8(x);
6451 let x1 = vld1q_s8(x.add(16));
6452 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6453 asm!(
6454 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6455 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6456 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6457 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6458 options(pure, nomem, nostack),
6459 );
6460 vaddvq_s32(vaddq_s32(a0, a1))
6461 }
6462}
6463
6464#[cfg(target_arch = "x86_64")]
6467#[target_feature(enable = "avx2")]
6468#[inline]
6469unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6470 unsafe {
6472 use core::arch::x86_64::*;
6473 let wv = _mm256_loadu_si256(w as *const __m256i);
6474 let xv = _mm256_loadu_si256(x as *const __m256i);
6475 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6476 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6477 let hi128 = _mm256_extracti128_si256::<1>(d);
6478 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6479 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6480 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6481 _mm_cvtsi128_si32(s32)
6482 }
6483}
6484
6485#[inline]
6490fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6491 debug_assert!(dst.len() >= 40);
6492 unsafe {
6495 let p = dst.as_mut_ptr();
6496 for bi in 0..7 {
6497 core::ptr::write_unaligned(
6498 p.add(bi * 5) as *mut u64,
6499 SIGN5_U64[*codes.add(bi) as usize],
6500 );
6501 }
6502 }
6503}
6504
6505#[inline]
6510fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6511 #[cfg(target_arch = "aarch64")]
6512 unsafe {
6513 return sdot32_i8(w, x);
6514 }
6515 #[cfg(target_arch = "x86_64")]
6516 unsafe {
6517 return i8dot32_avx2(w, x);
6518 }
6519 #[allow(unreachable_code)]
6520 unsafe {
6521 let mut s = 0i32;
6522 for k in 0..GROUP_SIZE {
6523 s += *w.add(k) as i32 * *x.add(k) as i32;
6524 }
6525 s
6526 }
6527}
6528
6529#[inline]
6530unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6531 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6532 (
6533 SIGN5_U64[*codes as usize],
6534 SIGN5_U64[*codes.add(1) as usize],
6535 SIGN5_U64[*codes.add(2) as usize],
6536 SIGN5_U64[*codes.add(3) as usize],
6537 SIGN5_U64[*codes.add(4) as usize],
6538 SIGN5_U64[*codes.add(5) as usize],
6539 SIGN5_U64[*codes.add(6) as usize],
6540 )
6541 };
6542
6543 let u0 = s0 | (s1 << 40);
6544 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6545 let u2 = (s3 >> 8) | (s4 << 32);
6546 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6547
6548 (u0, u1, u2, u3)
6549}
6550
6551#[cfg(target_arch = "aarch64")]
6555#[target_feature(enable = "neon,dotprod")]
6556unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6557 use core::arch::aarch64::*;
6558 use core::arch::asm;
6559 unsafe {
6560 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6561 let mut acc = 0f32;
6562 let bytes_ptr = bytes.as_ptr();
6563 let xq_ptr = xq.as_ptr();
6564 let row_off = r * gpr * TILE;
6565
6566 let gpr2 = gpr & !1;
6567 let mut gi = 0;
6568 while gi < gpr2 {
6569 let off0 = row_off + gi * TILE;
6570 let off1 = off0 + TILE;
6571 let s0 = f16_to_f32(u16::from_le_bytes([
6572 *bytes_ptr.add(off0),
6573 *bytes_ptr.add(off0 + 1),
6574 ]));
6575 let s1 = f16_to_f32(u16::from_le_bytes([
6576 *bytes_ptr.add(off1),
6577 *bytes_ptr.add(off1 + 1),
6578 ]));
6579
6580 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6581 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6582
6583 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6584 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6585 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6586 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6587
6588 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6589 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6590 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6591 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6592
6593 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6594 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6595 asm!(
6596 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6597 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6598 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6599 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6600 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6601 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6602 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6603 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6604 options(pure, nomem, nostack),
6605 );
6606 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6607 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6608 acc += d0 as f32 * s0 + d1 as f32 * s1;
6609 gi += 2;
6610 }
6611
6612 if gi < gpr {
6613 let off = row_off + gi * TILE;
6614 let s = f16_to_f32(u16::from_le_bytes([
6615 *bytes_ptr.add(off),
6616 *bytes_ptr.add(off + 1),
6617 ]));
6618 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6619 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6620 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6621 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6622 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6623 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6624 asm!(
6625 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6626 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6627 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6628 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6629 options(pure, nomem, nostack),
6630 );
6631 let d = vaddvq_s32(vaddq_s32(a0, a1));
6632 acc += d as f32 * s;
6633 }
6634 acc
6635 }
6636}
6637
6638#[cfg(target_arch = "x86_64")]
6640#[target_feature(enable = "avx2")]
6641unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6642 use core::arch::x86_64::*;
6643 unsafe {
6644 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6645 let mut acc = 0f32;
6646 let bytes_ptr = bytes.as_ptr();
6647 let xq_ptr = xq.as_ptr();
6648 let row_off = r * gpr * TILE;
6649
6650 let ones = _mm256_set1_epi16(1);
6651 for gi in 0..gpr {
6652 let off = row_off + gi * TILE;
6653 let s = f16_to_f32(u16::from_le_bytes([
6654 *bytes_ptr.add(off),
6655 *bytes_ptr.add(off + 1),
6656 ]));
6657 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6658 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6659 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6660 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6661 let d256 = _mm256_madd_epi16(p16, ones);
6662 let d128 = _mm_add_epi32(
6663 _mm256_castsi256_si128(d256),
6664 _mm256_extracti128_si256(d256, 1),
6665 );
6666 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6667 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6668 acc += d32 as f32 * s;
6669 }
6670 acc
6671 }
6672}
6673
6674#[cfg(target_arch = "x86_64")]
6676#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6677unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6678 use core::arch::x86_64::*;
6679 unsafe {
6681 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6682 let mut acc = 0f32;
6683 let bytes_ptr = bytes.as_ptr();
6684 let xq_ptr = xq.as_ptr();
6685 let row_off = r * gpr * TILE;
6686 for gi in 0..gpr {
6687 let off = row_off + gi * TILE;
6688 let s = f16_to_f32(u16::from_le_bytes([
6689 *bytes_ptr.add(off),
6690 *bytes_ptr.add(off + 1),
6691 ]));
6692 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6693 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6694 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6695 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6696 acc += d as f32 * s;
6697 }
6698 acc
6699 }
6700}
6701
6702#[inline]
6706fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6707 #[cfg(target_arch = "aarch64")]
6708 unsafe {
6709 return q1t_dot_row_sdot(bytes, r, gpr, xq);
6710 }
6711 #[cfg(target_arch = "x86_64")]
6712 unsafe {
6713 if vnni_tiles_enabled() {
6714 return q1t_dot_row_vnni(bytes, r, gpr, xq);
6715 }
6716 return q1t_dot_row_avx2(bytes, r, gpr, xq);
6717 }
6718 #[allow(unreachable_code)]
6719 {
6720 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6721 let mut acc = 0f32;
6722 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
6724 let off = (r * gpr + gi) * TILE;
6725 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6726 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6727 let mut d = 0i32;
6728 for k in 0..GROUP_SIZE {
6729 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6730 }
6731 acc += d as f32 * s;
6732 }
6733 acc
6734 }
6735}
6736
6737fn q1t_row_outlier_correction(
6744 bytes: &[u8],
6745 r: usize,
6746 rp_off: usize,
6747 entries_off: usize,
6748 has_ov: bool,
6749 x: &[f32],
6750) -> f32 {
6751 if !has_ov {
6752 return 0.0;
6753 }
6754 let (c0, c1) = (
6755 q1t_rowptr(bytes, rp_off, r),
6756 q1t_rowptr(bytes, rp_off, r + 1),
6757 );
6758 let mut corr = 0f32;
6759 for p in c0..c1 {
6760 let e = entries_off + p * 4;
6761 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6762 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6763 corr += val * x[col];
6764 }
6765 corr
6766}
6767
6768fn q1t_dequant_row(
6772 bytes: &[u8],
6773 r: usize,
6774 gpr: usize,
6775 rp_off: usize,
6776 entries_off: usize,
6777 has_ov: bool,
6778 buf: &mut [f32],
6779) {
6780 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6781 for g in 0..gpr {
6782 let off = (r * gpr + g) * TILE;
6783 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6784 let codes = &bytes[off + 2..off + TILE];
6785 let bc = g * GROUP_SIZE;
6786 for bi in 0..6 {
6788 let lut = &SIGN5[codes[bi] as usize];
6789 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6790 for i in 0..5 {
6791 d[i] = lut[i] * s;
6792 }
6793 }
6794 let lut = &SIGN5[codes[6] as usize];
6795 buf[bc + 30] = lut[0] * s;
6796 buf[bc + 31] = lut[1] * s;
6797 }
6798 if !has_ov {
6799 return;
6800 }
6801 let (c0, c1) = (
6802 q1t_rowptr(bytes, rp_off, r),
6803 q1t_rowptr(bytes, rp_off, r + 1),
6804 );
6805 for p in c0..c1 {
6806 let e = entries_off + p * 4;
6807 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6808 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6809 }
6810}
6811
6812fn q1t_add_overlay(
6816 bytes: &[u8],
6817 x: &[f32],
6818 rows: usize,
6819 cols: usize,
6820 out: &mut [f32],
6821 pool: Option<&Pool>,
6822) {
6823 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6824 let gpr = cols / GROUP_SIZE;
6825 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6826 if !has_ov {
6827 return;
6828 }
6829 let out_addr = SendMut(out.as_mut_ptr());
6830 let run = move |start: usize, end: usize| {
6831 for r in start..end {
6832 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6833 unsafe { *out_addr.at(r) += corr };
6835 }
6836 };
6837 dispatch_rows(pool, rows, &run);
6838}
6839
6840#[allow(clippy::too_many_arguments)]
6843fn q1t_range_a8w8(
6844 bytes: &[u8],
6845 gpr: usize,
6846 rp_off: usize,
6847 ent_off: usize,
6848 has_ov: bool,
6849 act: &SplitAct,
6850 x: &[f32],
6851 out: SendMut,
6852 start: usize,
6853 end: usize,
6854) {
6855 for r in start..end {
6856 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6857 for &(j, xv) in &act.outliers {
6858 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6859 }
6860 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6861 unsafe { *out.at(r) = acc };
6863 }
6864}
6865
6866#[allow(clippy::too_many_arguments)]
6869fn q1t_range_f32_batch(
6870 bytes: &[u8],
6871 gpr: usize,
6872 rp_off: usize,
6873 ent_off: usize,
6874 has_ov: bool,
6875 x: &[f32],
6876 out: SendMut,
6877 start: usize,
6878 end: usize,
6879) {
6880 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6881 let mut sg = [0f32; GROUP_SIZE];
6882 for r in start..end {
6883 let mut acc = 0f32;
6884 for g in 0..gpr {
6885 let off = (r * gpr + g) * TILE;
6886 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6887 let codes = &bytes[off + 2..off + TILE];
6888 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6889 for bi in 0..6 {
6890 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6891 }
6892 let lut = &SIGN5[codes[6] as usize];
6893 sg[30] = lut[0];
6894 sg[31] = lut[1];
6895 let mut gsum = 0f32;
6896 for k in 0..GROUP_SIZE {
6897 gsum += sg[k] * xg[k];
6898 }
6899 acc += s * gsum;
6900 }
6901 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6902 unsafe { *out.at(r) = acc };
6904 }
6905}
6906
6907fn q1t_matvec(
6911 bytes: &[u8],
6912 x: &[f32],
6913 rows: usize,
6914 cols: usize,
6915 out: &mut [f32],
6916 pool: Option<&Pool>,
6917) {
6918 debug_assert_eq!(out.len(), rows);
6919 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6920 let gpr = cols / GROUP_SIZE;
6921 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6922 let out_addr = SendMut(out.as_mut_ptr());
6923 if a8w8_enabled() {
6927 let act = split_act(x);
6928 let act = &act;
6929 let run = move |start: usize, end: usize| {
6930 for r in start..end {
6931 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6932 for &(j, xv) in &act.outliers {
6933 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6934 }
6935 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6936 unsafe { *out_addr.at(r) = acc };
6938 }
6939 };
6940 dispatch_rows(pool, rows, &run);
6941 return;
6942 }
6943 let run = move |start: usize, end: usize| {
6944 let mut sg = [0f32; GROUP_SIZE];
6948 for r in start..end {
6949 let mut acc = 0f32;
6950 for g in 0..gpr {
6951 let off = (r * gpr + g) * TILE;
6952 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6953 let codes = &bytes[off + 2..off + TILE];
6954 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6955 for bi in 0..6 {
6956 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6957 }
6958 let lut = &SIGN5[codes[6] as usize];
6959 sg[30] = lut[0];
6960 sg[31] = lut[1];
6961 let mut gsum = 0f32;
6962 for k in 0..GROUP_SIZE {
6963 gsum += sg[k] * xg[k];
6964 }
6965 acc += s * gsum;
6966 }
6967 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6968 unsafe { *out_addr.at(r) = acc };
6969 }
6970 };
6971 dispatch_rows(pool, rows, &run);
6972}
6973
6974#[cfg(target_arch = "aarch64")]
6980#[target_feature(enable = "neon,dotprod")]
6981unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6982 use core::arch::aarch64::*;
6983 use core::arch::asm;
6984 unsafe {
6986 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6987 let bytes_ptr = bytes.as_ptr();
6988 let row_off = r * gpr * TILE;
6989 let xp = [xa.as_ptr(), xb.as_ptr()];
6990 let mut acc = [0f32; 2];
6991 macro_rules! sdot2 {
6992 ($w0:expr, $w1:expr, $x:expr) => {{
6993 let x0 = vld1q_s8($x);
6994 let x1 = vld1q_s8($x.add(16));
6995 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6996 asm!(
6997 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6998 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6999 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7000 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7001 options(pure, nomem, nostack),
7002 );
7003 vaddvq_s32(vaddq_s32(a0, a1))
7004 }};
7005 }
7006 let gpr2 = gpr & !1;
7007 let mut gi = 0;
7008 while gi < gpr2 {
7009 let off0 = row_off + gi * TILE;
7010 let off1 = off0 + TILE;
7011 let s0 = f16_to_f32(u16::from_le_bytes([
7012 *bytes_ptr.add(off0),
7013 *bytes_ptr.add(off0 + 1),
7014 ]));
7015 let s1 = f16_to_f32(u16::from_le_bytes([
7016 *bytes_ptr.add(off1),
7017 *bytes_ptr.add(off1 + 1),
7018 ]));
7019 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7020 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7021 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7022 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7023 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7024 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7025 for k in 0..2 {
7026 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7027 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7028 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7029 }
7030 gi += 2;
7031 }
7032 if gi < gpr {
7033 let off = row_off + gi * TILE;
7034 let s = f16_to_f32(u16::from_le_bytes([
7035 *bytes_ptr.add(off),
7036 *bytes_ptr.add(off + 1),
7037 ]));
7038 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7039 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7040 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7041 for k in 0..2 {
7042 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7043 acc[k] += d as f32 * s;
7044 }
7045 }
7046 acc
7047 }
7048}
7049
7050fn q1t_matvec2(
7056 bytes: &[u8],
7057 x1: &[f32],
7058 x2: &[f32],
7059 rows: usize,
7060 cols: usize,
7061 o1: &mut [f32],
7062 o2: &mut [f32],
7063 pool: Option<&Pool>,
7064) {
7065 debug_assert_eq!(o1.len(), rows);
7066 debug_assert_eq!(o2.len(), rows);
7067 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7068 let gpr = cols / GROUP_SIZE;
7069 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7070 let out1 = SendMut(o1.as_mut_ptr());
7071 let out2 = SendMut(o2.as_mut_ptr());
7072 if a8w8_enabled() {
7073 let a1 = split_act(x1);
7074 let a2 = split_act(x2);
7075 let (a1, a2) = (&a1, &a2);
7076 let run = move |start: usize, end: usize| {
7077 for r in start..end {
7078 #[cfg(target_arch = "aarch64")]
7079 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7082 #[cfg(not(target_arch = "aarch64"))]
7083 let ds = [
7084 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7085 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7086 ];
7087 let mut acc1 = ds[0] * a1.sx;
7088 for &(j, xv) in &a1.outliers {
7089 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7090 }
7091 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7092 let mut acc2 = ds[1] * a2.sx;
7093 for &(j, xv) in &a2.outliers {
7094 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7095 }
7096 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7097 unsafe {
7099 *out1.at(r) = acc1;
7100 *out2.at(r) = acc2;
7101 }
7102 }
7103 };
7104 dispatch_rows(pool, rows, &run);
7105 return;
7106 }
7107 let run = move |start: usize, end: usize| {
7108 let mut sg = [0f32; GROUP_SIZE];
7111 for r in start..end {
7112 let mut acc1 = 0f32;
7113 let mut acc2 = 0f32;
7114 for g in 0..gpr {
7115 let off = (r * gpr + g) * TILE;
7116 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7117 let codes = &bytes[off + 2..off + TILE];
7118 for bi in 0..6 {
7119 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7120 }
7121 let lut = &SIGN5[codes[6] as usize];
7122 sg[30] = lut[0];
7123 sg[31] = lut[1];
7124 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7125 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7126 let mut gsum1 = 0f32;
7127 for k in 0..GROUP_SIZE {
7128 gsum1 += sg[k] * xg1[k];
7129 }
7130 acc1 += s * gsum1;
7131 let mut gsum2 = 0f32;
7132 for k in 0..GROUP_SIZE {
7133 gsum2 += sg[k] * xg2[k];
7134 }
7135 acc2 += s * gsum2;
7136 }
7137 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7138 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7139 unsafe {
7141 *out1.at(r) = acc1;
7142 *out2.at(r) = acc2;
7143 }
7144 }
7145 };
7146 dispatch_rows(pool, rows, &run);
7147}
7148
7149fn q1t_matmat(
7152 bytes: &[u8],
7153 xs: &[f32],
7154 b: usize,
7155 rows: usize,
7156 cols: usize,
7157 out: &mut [f32],
7158 pool: Option<&Pool>,
7159) {
7160 debug_assert_eq!(out.len(), b * rows);
7161 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7162 let gpr = cols / GROUP_SIZE;
7163 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7164 let out_addr = SendMut(out.as_mut_ptr());
7165 if a8w8_enabled() {
7169 let acts: Vec<SplitAct> = (0..b)
7170 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7171 .collect();
7172 let acts = &acts;
7173 let run = move |start: usize, end: usize| {
7174 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
7178 for g in 0..gpr {
7179 let off = (r * gpr + g) * TILE;
7180 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7181 q1t_unpack_group_i8(
7182 bytes.as_ptr().wrapping_add(off + 2),
7183 &mut sg[g * GROUP_SIZE..],
7184 );
7185 }
7186 for bi in 0..b {
7187 let act = &acts[bi];
7188 let mut isum = 0f32;
7189 for g in 0..gpr {
7190 let d = q1t_i8dot32(
7191 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7192 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7193 );
7194 isum += d as f32 * sc[g];
7195 }
7196 let mut acc = isum * act.sx;
7197 for &(j, xv) in &act.outliers {
7198 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7199 }
7200 accs[bi] = acc;
7201 }
7202 if has_ov {
7206 let (c0, c1) = (
7207 q1t_rowptr(bytes, rp_off, r),
7208 q1t_rowptr(bytes, rp_off, r + 1),
7209 );
7210 for p in c0..c1 {
7211 let e = ent_off + p * 4;
7212 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7213 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7214 for bi in 0..b {
7215 accs[bi] += val * xs[bi * cols + col];
7216 }
7217 }
7218 }
7219 for bi in 0..b {
7220 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7221 }
7222 }
7223 };
7224 dispatch_rows(pool, rows, &run);
7225 return;
7226 }
7227 let run = move |start: usize, end: usize| {
7228 let mut buf = vec![0f32; cols];
7229 for r in start..end {
7230 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7231 for bi in 0..b {
7232 let xr = &xs[bi * cols..(bi + 1) * cols];
7233 let mut acc = 0f32;
7234 for j in 0..cols {
7235 acc += buf[j] * xr[j];
7236 }
7237 unsafe { *out_addr.at(bi * rows + r) = acc };
7238 }
7239 }
7240 };
7241 dispatch_rows(pool, rows, &run);
7242}
7243
7244fn q1_matvec(
7245 bytes: &[u8],
7246 x: &[f32],
7247 rows: usize,
7248 cols: usize,
7249 out: &mut [f32],
7250 pool: Option<&Pool>,
7251) {
7252 debug_assert_eq!(out.len(), rows);
7253 let gpr = cols / GROUP_SIZE;
7254 let out_addr = SendMut(out.as_mut_ptr());
7255 if a8w8_enabled() {
7256 let act = split_act(x);
7257 let gsum = q1_group_sums(&act.xq, gpr);
7258 let (act, gsum) = (&act, &gsum);
7259 let run = move |start: usize, end: usize| {
7260 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7261 };
7262 dispatch_rows(pool, rows, &run);
7263 return;
7264 }
7265 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7266 dispatch_rows(pool, rows, &run);
7267}
7268
7269#[allow(clippy::too_many_arguments)]
7271fn q1_matvec2(
7272 bytes: &[u8],
7273 x1: &[f32],
7274 x2: &[f32],
7275 rows: usize,
7276 cols: usize,
7277 o1: &mut [f32],
7278 o2: &mut [f32],
7279 pool: Option<&Pool>,
7280) {
7281 let gpr = cols / GROUP_SIZE;
7282 let p1 = SendMut(o1.as_mut_ptr());
7283 let p2 = SendMut(o2.as_mut_ptr());
7284 if a8w8_enabled() {
7285 let a1 = split_act(x1);
7286 let a2 = split_act(x2);
7287 let g1 = q1_group_sums(&a1.xq, gpr);
7288 let g2 = q1_group_sums(&a2.xq, gpr);
7289 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7290 let run = move |start: usize, end: usize| {
7291 for r in start..end {
7292 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7293 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7294 for &(j, xv) in &a1.outliers {
7295 let (w, s) = q1_outlier(bytes, r, gpr, j);
7296 v1 += w * s * xv;
7297 }
7298 for &(j, xv) in &a2.outliers {
7299 let (w, s) = q1_outlier(bytes, r, gpr, j);
7300 v2 += w * s * xv;
7301 }
7302 unsafe {
7304 *p1.at(r) = v1;
7305 *p2.at(r) = v2;
7306 }
7307 }
7308 };
7309 dispatch_rows(pool, rows, &run);
7310 return;
7311 }
7312 let run = move |start: usize, end: usize| {
7313 for r in start..end {
7314 unsafe {
7316 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7317 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7318 }
7319 }
7320 };
7321 dispatch_rows(pool, rows, &run);
7322}
7323
7324#[allow(clippy::too_many_arguments)]
7326fn q1_matmat(
7327 bytes: &[u8],
7328 xs_all: &[f32],
7329 b: usize,
7330 rows: usize,
7331 cols: usize,
7332 out: &mut [f32],
7333 pool: Option<&Pool>,
7334) {
7335 debug_assert_eq!(out.len(), b * rows);
7336 let gpr = cols / GROUP_SIZE;
7337 let out_addr = SendMut(out.as_mut_ptr());
7338 if a8w8_enabled() {
7339 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7340 .map(|bi| {
7341 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7342 let gsum = q1_group_sums(&act.xq, gpr);
7343 (act, gsum)
7344 })
7345 .collect();
7346 let acts = &acts;
7347 #[cfg(target_arch = "x86_64")]
7348 let blocked_ok = avx2_enabled() && blocked_enabled();
7349 #[cfg(target_arch = "aarch64")]
7350 let blocked_ok = sdot_enabled() && blocked_enabled();
7351 let run = move |start: usize, end: usize| {
7352 for r in start..end {
7353 let mut bi = 0usize;
7354 #[cfg(target_arch = "aarch64")]
7357 if blocked_ok {
7358 while bi + 4 <= acts.len() {
7359 let xs = [
7360 acts[bi].0.xq.as_slice(),
7361 acts[bi + 1].0.xq.as_slice(),
7362 acts[bi + 2].0.xq.as_slice(),
7363 acts[bi + 3].0.xq.as_slice(),
7364 ];
7365 let gs = [
7366 acts[bi].1.as_slice(),
7367 acts[bi + 1].1.as_slice(),
7368 acts[bi + 2].1.as_slice(),
7369 acts[bi + 3].1.as_slice(),
7370 ];
7371 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7372 for k in 0..4 {
7373 let (act, _) = &acts[bi + k];
7374 let mut acc = d[k] * act.sx;
7375 for &(j, xv) in &act.outliers {
7376 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7377 acc += w * sc * xv;
7378 }
7379 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7381 }
7382 bi += 4;
7383 }
7384 }
7385 #[cfg(target_arch = "x86_64")]
7386 if blocked_ok {
7387 while bi + 4 <= acts.len() {
7388 let xs = [
7389 acts[bi].0.xq.as_slice(),
7390 acts[bi + 1].0.xq.as_slice(),
7391 acts[bi + 2].0.xq.as_slice(),
7392 acts[bi + 3].0.xq.as_slice(),
7393 ];
7394 let gs = [
7395 acts[bi].1.as_slice(),
7396 acts[bi + 1].1.as_slice(),
7397 acts[bi + 2].1.as_slice(),
7398 acts[bi + 3].1.as_slice(),
7399 ];
7400 let d = unsafe {
7401 if vnni_tiles_enabled() {
7402 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7403 } else {
7404 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7405 }
7406 };
7407 for k in 0..4 {
7408 let (act, _) = &acts[bi + k];
7409 let mut acc = d[k] * act.sx;
7410 for &(j, xv) in &act.outliers {
7411 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7412 acc += w * sc * xv;
7413 }
7414 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7416 }
7417 bi += 4;
7418 }
7419 }
7420 while bi < acts.len() {
7421 let (act, gsum) = &acts[bi];
7422 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7423 for &(j, xv) in &act.outliers {
7424 let (w, s) = q1_outlier(bytes, r, gpr, j);
7425 acc += w * s * xv;
7426 }
7427 unsafe { *out_addr.at(bi * rows + r) = acc };
7429 bi += 1;
7430 }
7431 }
7432 };
7433 dispatch_rows(pool, rows, &run);
7434 return;
7435 }
7436 let run = move |start: usize, end: usize| {
7437 for r in start..end {
7438 for bi in 0..b {
7439 let x = &xs_all[bi * cols..(bi + 1) * cols];
7440 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7442 }
7443 }
7444 };
7445 dispatch_rows(pool, rows, &run);
7446}
7447
7448fn q4matvec(
7454 bytes: &[u8],
7455 x: &[f32],
7456 rows: usize,
7457 cols: usize,
7458 out: &mut [f32],
7459 pool: Option<&Pool>,
7460) {
7461 debug_assert_eq!(out.len(), rows);
7462 let (packed, scales) = q4_split(bytes, rows, cols);
7463 let gpr = cols / GROUP_SIZE;
7464 let out_addr = SendMut(out.as_mut_ptr());
7465
7466 if a8w8_enabled() {
7467 let act = split_act(x);
7468 let run = move |start: usize, end: usize| {
7469 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7470 };
7471 dispatch_rows(pool, rows, &run);
7472 return;
7473 }
7474
7475 let run =
7476 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7477 dispatch_rows(pool, rows, &run);
7478}
7479
7480#[inline]
7483#[allow(unreachable_code)]
7484#[cfg(target_arch = "x86_64")]
7489#[target_feature(enable = "avx2")]
7490unsafe fn dot_q4b_row_1x4_avx2(
7491 buf: &[u8],
7492 scales: &[u8],
7493 g0: usize,
7494 gpr: usize,
7495 xs: [&[i8]; 4],
7496) -> [f32; 4] {
7497 unsafe {
7499 use core::arch::x86_64::*;
7500 let ones = _mm256_set1_epi16(1);
7501 let mut acc = [0f32; 4];
7502 for gi in 0..gpr {
7503 let s = f16_to_f32(u16::from_le_bytes([
7504 scales[(g0 + gi) * 2],
7505 scales[(g0 + gi) * 2 + 1],
7506 ]));
7507 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7508 let aw = _mm256_abs_epi8(w);
7509 for (k, xq) in xs.iter().enumerate() {
7510 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7511 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7512 let d = _mm256_madd_epi16(p16, ones);
7513 let hi128 = _mm256_extracti128_si256::<1>(d);
7514 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7515 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7516 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7517 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7518 }
7519 }
7520 acc
7521 }
7522}
7523
7524#[cfg(target_arch = "x86_64")]
7526#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7527unsafe fn dot_q4b_row_1x4_vnni(
7528 buf: &[u8],
7529 scales: &[u8],
7530 g0: usize,
7531 gpr: usize,
7532 xs: [&[i8]; 4],
7533) -> [f32; 4] {
7534 unsafe {
7536 use core::arch::x86_64::*;
7537 let mut acc = [0f32; 4];
7538 for gi in 0..gpr {
7539 let s = f16_to_f32(u16::from_le_bytes([
7540 scales[(g0 + gi) * 2],
7541 scales[(g0 + gi) * 2 + 1],
7542 ]));
7543 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7544 let aw = _mm256_abs_epi8(w);
7545 for (k, xq) in xs.iter().enumerate() {
7546 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7547 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7548 acc[k] += d as f32 * s;
7549 }
7550 }
7551 acc
7552 }
7553}
7554
7555#[cfg(target_arch = "x86_64")]
7561#[target_feature(enable = "avx2")]
7562unsafe fn dot_q4b_row_1x4_sx_avx2(
7563 buf: &[u8],
7564 scales: &[u8],
7565 g0: usize,
7566 gpr: usize,
7567 xs: [&[i8]; 4],
7568 sxs: [f32; 4],
7569) -> [f32; 4] {
7570 unsafe {
7572 use core::arch::x86_64::*;
7573 let ones = _mm256_set1_epi16(1);
7574 let mut acc = [0f32; 4];
7575 for gi in 0..gpr {
7576 let s = f16_to_f32(u16::from_le_bytes([
7577 scales[(g0 + gi) * 2],
7578 scales[(g0 + gi) * 2 + 1],
7579 ]));
7580 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7581 let aw = _mm256_abs_epi8(w);
7582 for (k, xq) in xs.iter().enumerate() {
7583 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7584 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7585 let d = _mm256_madd_epi16(p16, ones);
7586 let hi128 = _mm256_extracti128_si256::<1>(d);
7587 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7588 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7589 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7590 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7591 }
7592 }
7593 acc
7594 }
7595}
7596
7597#[cfg(target_arch = "x86_64")]
7600#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7601unsafe fn dot_q4b_row_1x4_sx_vnni(
7602 buf: &[u8],
7603 scales: &[u8],
7604 g0: usize,
7605 gpr: usize,
7606 xs: [&[i8]; 4],
7607 sxs: [f32; 4],
7608) -> [f32; 4] {
7609 unsafe {
7611 use core::arch::x86_64::*;
7612 let mut acc = [0f32; 4];
7613 for gi in 0..gpr {
7614 let s = f16_to_f32(u16::from_le_bytes([
7615 scales[(g0 + gi) * 2],
7616 scales[(g0 + gi) * 2 + 1],
7617 ]));
7618 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7619 let aw = _mm256_abs_epi8(w);
7620 for (k, xq) in xs.iter().enumerate() {
7621 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7622 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7623 acc[k] += (d as f32 * sxs[k]) * s;
7624 }
7625 }
7626 acc
7627 }
7628}
7629
7630#[allow(unreachable_code)]
7631fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7632 #[cfg(target_arch = "aarch64")]
7633 unsafe {
7634 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7635 }
7636 #[cfg(target_arch = "x86_64")]
7637 unsafe {
7638 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7639 }
7640 let mut acc = 0f32;
7641 for gi in 0..gpr {
7642 let g = g0 + gi;
7643 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7644 let mut d = 0i32;
7645 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7646 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7647 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7648 }
7649 acc += d as f32 * s;
7650 }
7651 acc
7652}
7653
7654#[inline]
7656#[allow(unreachable_code)]
7657fn dot_q4_row_i8_2(
7658 packed: &[u8],
7659 scales: &[u8],
7660 g0: usize,
7661 gpr: usize,
7662 xq1: &[i8],
7663 xq2: &[i8],
7664) -> (f32, f32) {
7665 #[cfg(target_arch = "aarch64")]
7666 unsafe {
7667 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7668 }
7669 #[cfg(target_arch = "x86_64")]
7670 unsafe {
7671 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7672 }
7673 (
7674 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7675 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7676 )
7677}
7678
7679#[allow(clippy::too_many_arguments)]
7682fn q4_range_a8w8(
7683 packed: &[u8],
7684 scales: &[u8],
7685 gpr: usize,
7686 cols: usize,
7687 act: &SplitAct,
7688 out: SendMut,
7689 start: usize,
7690 end: usize,
7691) {
7692 for r in start..end {
7693 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7694 for &(j, xv) in &act.outliers {
7696 let flat = r * cols + j;
7697 let byte = packed[flat / 2];
7698 let nib = if flat & 1 == 0 {
7699 byte & 0x0F
7700 } else {
7701 byte >> 4
7702 };
7703 let s = f16_to_f32(u16::from_le_bytes([
7704 scales[(flat / GROUP_SIZE) * 2],
7705 scales[(flat / GROUP_SIZE) * 2 + 1],
7706 ]));
7707 acc += ((nib as i32 - 8) as f32) * s * xv;
7708 }
7709 unsafe { *out.at(r) = acc };
7711 }
7712}
7713
7714#[allow(clippy::too_many_arguments)]
7717fn q4_range2_a8w8(
7718 packed: &[u8],
7719 scales: &[u8],
7720 gpr: usize,
7721 cols: usize,
7722 a1: &SplitAct,
7723 a2: &SplitAct,
7724 p1: SendMut,
7725 p2: SendMut,
7726 start: usize,
7727 end: usize,
7728) {
7729 for r in start..end {
7730 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7731 let mut acc1 = s1 * a1.sx;
7732 let mut acc2 = s2 * a2.sx;
7733 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7735 for &(j, xv) in outliers {
7736 let flat = r * cols + j;
7737 let byte = packed[flat / 2];
7738 let nib = if flat & 1 == 0 {
7739 byte & 0x0F
7740 } else {
7741 byte >> 4
7742 };
7743 let s = f16_to_f32(u16::from_le_bytes([
7744 scales[(flat / GROUP_SIZE) * 2],
7745 scales[(flat / GROUP_SIZE) * 2 + 1],
7746 ]));
7747 *acc += ((nib as i32 - 8) as f32) * s * xv;
7748 }
7749 };
7750 fix(&a1.outliers, &mut acc1);
7751 fix(&a2.outliers, &mut acc2);
7752 unsafe {
7754 *p1.at(r) = acc1;
7755 *p2.at(r) = acc2;
7756 }
7757 }
7758}
7759
7760fn q4_range_f32(
7762 packed: &[u8],
7763 scales: &[u8],
7764 gpr: usize,
7765 x: &[f32],
7766 out: SendMut,
7767 start: usize,
7768 end: usize,
7769) {
7770 for r in start..end {
7771 let mut acc = 0f32;
7772 for gi in 0..gpr {
7773 let g = r * gpr + gi;
7774 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7775 let pk = &packed[g * 16..(g + 1) * 16];
7776 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7777 let mut ga = 0f32;
7778 for (k, &b) in pk.iter().enumerate() {
7779 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7780 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7781 }
7782 acc += ga * s;
7783 }
7784 unsafe { *out.at(r) = acc };
7786 }
7787}
7788
7789#[allow(clippy::too_many_arguments)]
7793fn q4matvec2(
7794 bytes: &[u8],
7795 x1: &[f32],
7796 x2: &[f32],
7797 rows: usize,
7798 cols: usize,
7799 o1: &mut [f32],
7800 o2: &mut [f32],
7801 pool: Option<&Pool>,
7802) {
7803 debug_assert_eq!(o1.len(), rows);
7804 debug_assert_eq!(o2.len(), rows);
7805 let (packed, scales) = q4_split(bytes, rows, cols);
7806 let gpr = cols / GROUP_SIZE;
7807
7808 if a8w8_enabled() {
7809 let a1 = split_act(x1);
7810 let a2 = split_act(x2);
7811 let p1 = SendMut(o1.as_mut_ptr());
7812 let p2 = SendMut(o2.as_mut_ptr());
7813 let run = move |start: usize, end: usize| {
7814 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7815 };
7816 dispatch_rows(pool, rows, &run);
7817 return;
7818 }
7819
7820 let p1 = SendMut(o1.as_mut_ptr());
7821 let p2 = SendMut(o2.as_mut_ptr());
7822 let run = move |start: usize, end: usize| {
7823 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7824 };
7825 dispatch_rows(pool, rows, &run);
7826}
7827
7828#[allow(clippy::too_many_arguments)]
7830fn q4_range2_f32(
7831 packed: &[u8],
7832 scales: &[u8],
7833 gpr: usize,
7834 x1: &[f32],
7835 x2: &[f32],
7836 p1: SendMut,
7837 p2: SendMut,
7838 start: usize,
7839 end: usize,
7840) {
7841 for r in start..end {
7842 let (mut acc1, mut acc2) = (0f32, 0f32);
7843 for gi in 0..gpr {
7844 let g = r * gpr + gi;
7845 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7846 let pk = &packed[g * 16..(g + 1) * 16];
7847 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7848 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7849 let (mut g1, mut g2) = (0f32, 0f32);
7850 for (k, &b) in pk.iter().enumerate() {
7851 let wl = (b & 0x0F) as f32 - 8.0;
7852 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7853 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7854 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7855 }
7856 acc1 += g1 * s;
7857 acc2 += g2 * s;
7858 }
7859 unsafe {
7861 *p1.at(r) = acc1;
7862 *p2.at(r) = acc2;
7863 }
7864 }
7865}
7866
7867thread_local! {
7868 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7871 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7872}
7873
7874#[allow(clippy::too_many_arguments)]
7880fn q4matmat(
7881 bytes: &[u8],
7882 xs_all: &[f32],
7883 b: usize,
7884 rows: usize,
7885 cols: usize,
7886 out: &mut [f32],
7887 pool: Option<&Pool>,
7888) {
7889 debug_assert_eq!(xs_all.len(), b * cols);
7890 debug_assert_eq!(out.len(), b * rows);
7891 let (packed, scales) = q4_split(bytes, rows, cols);
7892 let gpr = cols / GROUP_SIZE;
7893 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7894
7895 if a8w8_enabled() {
7896 let acts: Vec<SplitAct> = (0..b)
7897 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7898 .collect();
7899 let acts = &acts;
7900 let out_addr = SendMut(out.as_mut_ptr());
7901 let run = move |start: usize, end: usize| {
7902 ROW_I8.with(|rb| {
7903 let mut buf = rb.borrow_mut();
7904 buf.resize(cols, 0);
7905 for r in start..end {
7906 for gi in 0..gpr {
7910 let g = r * gpr + gi;
7911 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7912 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7913 buf[gi * GROUP_SIZE + k * 2 + 1] =
7914 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7915 }
7916 }
7917 let mut bi = 0usize;
7918 #[cfg(target_arch = "x86_64")]
7919 if avx2_enabled() && blocked_enabled() {
7920 while bi + 4 <= acts.len() {
7921 let xs = [
7922 acts[bi].xq.as_slice(),
7923 acts[bi + 1].xq.as_slice(),
7924 acts[bi + 2].xq.as_slice(),
7925 acts[bi + 3].xq.as_slice(),
7926 ];
7927 let d = unsafe {
7928 if vnni_tiles_enabled() {
7929 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7930 } else {
7931 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7932 }
7933 };
7934 for k in 0..4 {
7935 let act = &acts[bi + k];
7936 let mut acc = d[k] * act.sx;
7937 for &(j, xv) in &act.outliers {
7938 acc += (buf[j] as i8) as f32
7939 * gscale((r * cols + j) / GROUP_SIZE)
7940 * xv;
7941 }
7942 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7944 }
7945 bi += 4;
7946 }
7947 }
7948 while bi < acts.len() {
7949 let act = &acts[bi];
7950 let mut acc = 0f32;
7951 for gi in 0..gpr {
7952 let d = dot_i8_i8(
7953 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7954 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7955 );
7956 acc += d as f32 * gscale(r * gpr + gi);
7957 }
7958 acc *= act.sx;
7959 for &(j, xv) in &act.outliers {
7961 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7962 }
7963 unsafe { *out_addr.at(bi * rows + r) = acc };
7965 bi += 1;
7966 }
7967 }
7968 })
7969 };
7970 dispatch_rows(pool, rows, &run);
7971 return;
7972 }
7973
7974 let out_addr = SendMut(out.as_mut_ptr());
7975 let run = move |start: usize, end: usize| {
7976 ROW_F32.with(|rb| {
7977 let mut buf = rb.borrow_mut();
7978 buf.resize(cols, 0.0);
7979 for r in start..end {
7980 for gi in 0..gpr {
7983 let g = r * gpr + gi;
7984 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7985 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7986 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7987 }
7988 }
7989 for bi in 0..b {
7990 let x = &xs_all[bi * cols..(bi + 1) * cols];
7991 let mut acc = 0f32;
7992 for gi in 0..gpr {
7993 let mut ga = 0f32;
7994 for k in 0..GROUP_SIZE / 2 {
7999 let e = gi * GROUP_SIZE + k * 2;
8000 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8001 }
8002 acc += ga * gscale(r * gpr + gi);
8003 }
8004 unsafe { *out_addr.at(bi * rows + r) = acc };
8006 }
8007 }
8008 })
8009 };
8010 dispatch_rows(pool, rows, &run);
8011}
8012
8013#[allow(clippy::too_many_arguments)]
8018fn vbitmatmat(
8019 bytes: &[u8],
8020 offsets: &[usize],
8021 xs_all: &[f32],
8022 b: usize,
8023 rows: usize,
8024 cols: usize,
8025 out: &mut [f32],
8026 pool: Option<&Pool>,
8027) {
8028 debug_assert_eq!(xs_all.len(), b * cols);
8029 debug_assert_eq!(out.len(), b * rows);
8030 debug_assert_eq!(offsets.len(), rows + 1);
8031 let ng = cols / GROUP_SIZE;
8032 let bits = &bytes[..rows];
8033 let sc_off = rows;
8034 let gscale = |r: usize, g: usize| {
8035 let so = (r * ng + g) * 2;
8036 f16_to_f32(u16::from_le_bytes([
8037 bytes[sc_off + so],
8038 bytes[sc_off + so + 1],
8039 ]))
8040 };
8041
8042 let decode_f32 = |r: usize, dst: &mut [f32]| {
8044 let bw = bits[r] as usize;
8045 let l = ((1i32 << (bw - 1)) - 1) as f32;
8046 let data = &bytes[offsets[r]..offsets[r + 1]];
8047 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8048 for d in dst.iter_mut() {
8049 while nbits < bw {
8050 acc = (acc << 8) | data[idx] as u64;
8051 idx += 1;
8052 nbits += 8;
8053 }
8054 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8055 nbits -= bw;
8056 *d = u - l;
8057 }
8058 };
8059
8060 if a8w8_enabled() {
8061 let acts: Vec<SplitAct> = (0..b)
8062 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8063 .collect();
8064 let acts = &acts;
8065 let out_addr = SendMut(out.as_mut_ptr());
8066 let run = move |start: usize, end: usize| {
8067 for r in start..end {
8068 let bw = bits[r] as usize;
8069 if bw == 8 {
8070 ROW_F32.with(|rb| {
8073 let mut buf = rb.borrow_mut();
8074 buf.resize(cols, 0.0);
8075 decode_f32(r, &mut buf);
8076 for bi in 0..b {
8077 let x = &xs_all[bi * cols..(bi + 1) * cols];
8078 let mut dot = 0f32;
8079 for g in 0..ng {
8080 let mut gd = 0f32;
8081 for k in 0..GROUP_SIZE {
8082 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8083 }
8084 dot += gd * gscale(r, g);
8085 }
8086 unsafe { *out_addr.at(bi * rows + r) = dot };
8088 }
8089 });
8090 continue;
8091 }
8092 let l = (1i32 << (bw - 1)) - 1;
8093 let data = &bytes[offsets[r]..offsets[r + 1]];
8094 ROW_I8.with(|rb| {
8095 let mut buf = rb.borrow_mut();
8096 buf.resize(cols, 0);
8097 #[inline(always)]
8098 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8099 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8100 let u = unpack8::<B>(&data[blk * B..]);
8101 for k in 0..8 {
8102 chunk[k] = (u[k] - l) as i8 as u8;
8103 }
8104 }
8105 }
8106 match bw {
8107 3 => fill::<3>(data, l, &mut buf),
8108 4 => vbit_fill4(data, &mut buf),
8109 5 => fill::<5>(data, l, &mut buf),
8110 6 => fill::<6>(data, l, &mut buf),
8111 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8112 }
8113 let mut bi = 0usize;
8114 #[cfg(target_arch = "x86_64")]
8118 if avx2_enabled() && blocked_enabled() {
8119 while bi + 4 <= acts.len() {
8120 let xs = [
8121 acts[bi].xq.as_slice(),
8122 acts[bi + 1].xq.as_slice(),
8123 acts[bi + 2].xq.as_slice(),
8124 acts[bi + 3].xq.as_slice(),
8125 ];
8126 let sxs = [
8127 acts[bi].sx,
8128 acts[bi + 1].sx,
8129 acts[bi + 2].sx,
8130 acts[bi + 3].sx,
8131 ];
8132 let d = unsafe {
8133 if vnni_tiles_enabled() {
8134 dot_q4b_row_1x4_sx_vnni(
8135 &buf,
8136 &bytes[sc_off..],
8137 r * ng,
8138 ng,
8139 xs,
8140 sxs,
8141 )
8142 } else {
8143 dot_q4b_row_1x4_sx_avx2(
8144 &buf,
8145 &bytes[sc_off..],
8146 r * ng,
8147 ng,
8148 xs,
8149 sxs,
8150 )
8151 }
8152 };
8153 for k in 0..4 {
8154 let act = &acts[bi + k];
8155 let mut dot = d[k];
8156 for &(j, xv) in &act.outliers {
8157 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8158 }
8159 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8161 }
8162 bi += 4;
8163 }
8164 }
8165 while bi < acts.len() {
8166 let act = &acts[bi];
8167 let mut dot = 0f32;
8168 for g in 0..ng {
8169 let d = dot_i8_i8(
8170 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8171 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8172 ) as f32
8173 * act.sx;
8174 dot += d * gscale(r, g);
8175 }
8176 for &(j, xv) in &act.outliers {
8177 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8178 }
8179 unsafe { *out_addr.at(bi * rows + r) = dot };
8181 bi += 1;
8182 }
8183 });
8184 }
8185 };
8186 dispatch_rows(pool, rows, &run);
8187 return;
8188 }
8189
8190 let out_addr = SendMut(out.as_mut_ptr());
8191 let run = move |start: usize, end: usize| {
8192 ROW_F32.with(|rb| {
8193 let mut buf = rb.borrow_mut();
8194 buf.resize(cols, 0.0);
8195 for r in start..end {
8196 decode_f32(r, &mut buf);
8197 for bi in 0..b {
8198 let x = &xs_all[bi * cols..(bi + 1) * cols];
8199 let mut dot = 0f32;
8200 for g in 0..ng {
8201 let mut gd = 0f32;
8202 for k in 0..GROUP_SIZE {
8203 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8204 }
8205 dot += gd * gscale(r, g);
8206 }
8207 unsafe { *out_addr.at(bi * rows + r) = dot };
8209 }
8210 }
8211 })
8212 };
8213 dispatch_rows(pool, rows, &run);
8214}
8215
8216pub(crate) fn gpu_batch_job<'a>(
8220 t: &'a QTensor,
8221 x: &[f32],
8222) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8223 match t {
8224 QTensor::Mapped {
8225 model,
8226 idx,
8227 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8228 rows,
8229 cols,
8230 row_scale,
8231 col_field,
8232 ..
8233 } => Some((
8234 model.clone(),
8235 crate::gpu::BatchJob {
8236 idx: *idx,
8237 rows: *rows,
8238 cols: *cols,
8239 row_scale,
8240 xs: prescale(x, col_field, *dt).into_owned(),
8241 layout: crate::gpu::BatchLayout::Q8,
8242 },
8243 )),
8244 QTensor::Mapped {
8246 model,
8247 idx,
8248 dtype: TensorDtype::Q1,
8249 rows,
8250 cols,
8251 ..
8252 } => Some((
8253 model.clone(),
8254 crate::gpu::BatchJob {
8255 idx: *idx,
8256 rows: *rows,
8257 cols: *cols,
8258 row_scale: &[],
8259 xs: x.to_vec(),
8260 layout: crate::gpu::BatchLayout::Q1,
8261 },
8262 )),
8263 QTensor::Mapped {
8268 model,
8269 idx,
8270 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8271 rows,
8272 cols,
8273 ..
8274 } => Some((
8275 model.clone(),
8276 crate::gpu::BatchJob {
8277 idx: *idx,
8278 rows: *rows,
8279 cols: *cols,
8280 row_scale: &[],
8281 xs: x.to_vec(),
8282 layout: if *dt == TensorDtype::Q4Tiled {
8283 crate::gpu::BatchLayout::Q4t
8284 } else {
8285 crate::gpu::BatchLayout::Q4tp
8286 },
8287 },
8288 )),
8289 _ => None,
8290 }
8291}
8292
8293thread_local! {
8294 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8295 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8296}
8297
8298pub(crate) fn prescale<'a>(
8299 x: &'a [f32],
8300 col_field: &[f32],
8301 dtype: TensorDtype,
8302) -> std::borrow::Cow<'a, [f32]> {
8303 if dtype == TensorDtype::Q8_2f {
8304 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8305 } else {
8306 std::borrow::Cow::Borrowed(x)
8307 }
8308}
8309
8310pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8313 x: &[f32],
8314 col_field: &[f32],
8315 dtype: TensorDtype,
8316 buf_id: u8,
8317 f: F,
8318) -> R {
8319 if dtype == TensorDtype::Q8_2f {
8320 if buf_id == 1 {
8321 PRESCALE_BUF1.with(|b| {
8322 let mut buf = b.borrow_mut();
8323 buf.clear();
8324 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8325 f(&buf)
8326 })
8327 } else {
8328 PRESCALE_BUF2.with(|b| {
8329 let mut buf = b.borrow_mut();
8330 buf.clear();
8331 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8332 f(&buf)
8333 })
8334 }
8335 } else {
8336 f(x)
8337 }
8338}
8339
8340#[cfg(target_arch = "x86_64")]
8345pub(crate) fn avx2_enabled() -> bool {
8346 use std::sync::OnceLock;
8347 static ON: OnceLock<bool> = OnceLock::new();
8348 *ON.get_or_init(|| {
8349 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8350 && std::arch::is_x86_feature_detected!("avx2")
8351 && std::arch::is_x86_feature_detected!("fma")
8352 })
8353}
8354
8355#[cfg(target_arch = "x86_64")]
8360fn avx2_a8w8_enabled() -> bool {
8361 use std::sync::OnceLock;
8362 static ON: OnceLock<bool> = OnceLock::new();
8363 *ON.get_or_init(|| {
8364 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8365 })
8366}
8367
8368#[inline]
8372pub(crate) fn a8w8_enabled() -> bool {
8373 #[cfg(target_arch = "aarch64")]
8374 {
8375 sdot_enabled()
8376 }
8377 #[cfg(target_arch = "x86_64")]
8378 {
8379 avx2_a8w8_enabled()
8380 }
8381 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8382 {
8383 false
8384 }
8385}
8386
8387#[inline]
8390#[allow(unreachable_code)]
8391fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8392 #[cfg(target_arch = "aarch64")]
8393 unsafe {
8394 return dot_i8_sdot(w, xq);
8395 }
8396 #[cfg(target_arch = "x86_64")]
8397 unsafe {
8398 if avx512vnni_enabled() {
8399 return dot_i8_i8_vnni(w, xq);
8400 }
8401 return dot_i8_i8_avx2(w, xq);
8402 }
8403 w.iter()
8404 .zip(xq)
8405 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8406 .sum()
8407}
8408
8409#[cfg(target_arch = "x86_64")]
8413fn avx512vnni_enabled() -> bool {
8414 use std::sync::OnceLock;
8415 static ON: OnceLock<bool> = OnceLock::new();
8416 *ON.get_or_init(|| {
8417 std::env::var("CMF_AVX512")
8418 .map(|v| v != "0")
8419 .unwrap_or(true)
8420 && std::arch::is_x86_feature_detected!("avx512f")
8421 && std::arch::is_x86_feature_detected!("avx512bw")
8422 && std::arch::is_x86_feature_detected!("avx512vl")
8423 && std::arch::is_x86_feature_detected!("avx512vnni")
8424 })
8425}
8426
8427#[cfg(target_arch = "x86_64")]
8435fn vnni_tiles_enabled() -> bool {
8436 use std::sync::OnceLock;
8437 static ON: OnceLock<bool> = OnceLock::new();
8438 *ON.get_or_init(|| {
8439 std::env::var("CMF_VNNI_TILES")
8440 .map(|v| v != "0")
8441 .unwrap_or(true)
8442 && avx512vnni_enabled()
8443 })
8444}
8445
8446#[cfg(target_arch = "x86_64")]
8451#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8452#[inline]
8453unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8454 unsafe {
8456 use core::arch::x86_64::*;
8457 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8458 let hi128 = _mm256_extracti128_si256::<1>(d);
8459 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8460 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8461 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8462 _mm_cvtsi128_si32(s32)
8463 }
8464}
8465
8466#[cfg(target_arch = "x86_64")]
8471#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8472unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8473 unsafe {
8475 use core::arch::x86_64::*;
8476 let n = w.len();
8477 let mut j = 0usize;
8478 let mut total: i32;
8479 {
8484 #[inline(always)]
8485 unsafe fn step(
8486 w: *const u8,
8487 x: *const i8,
8488 acc: core::arch::x86_64::__m512i,
8489 ) -> core::arch::x86_64::__m512i {
8490 unsafe {
8491 use core::arch::x86_64::*;
8492 let wv = _mm512_loadu_si512(w as *const _);
8493 let xv = _mm512_loadu_si512(x as *const _);
8494 let aw = _mm512_abs_epi8(wv);
8495 let neg = _mm512_movepi8_mask(wv);
8496 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8497 _mm512_dpbusd_epi32(acc, aw, sx)
8498 }
8499 }
8500 let (mut a0, mut a1, mut a2, mut a3) = (
8501 _mm512_setzero_si512(),
8502 _mm512_setzero_si512(),
8503 _mm512_setzero_si512(),
8504 _mm512_setzero_si512(),
8505 );
8506 while j + 256 <= n {
8507 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8508 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8509 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8510 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8511 j += 256;
8512 }
8513 while j + 64 <= n {
8514 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8515 j += 64;
8516 }
8517 let s01 = _mm512_add_epi32(a0, a1);
8518 let s23 = _mm512_add_epi32(a2, a3);
8519 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8520 }
8521 if j + 32 <= n {
8523 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8524 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8525 let d = _mm256_dpbusd_epi32(
8526 _mm256_setzero_si256(),
8527 _mm256_abs_epi8(wv),
8528 _mm256_sign_epi8(xv, wv),
8529 );
8530 let hi128 = _mm256_extracti128_si256::<1>(d);
8531 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8532 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8533 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8534 total += _mm_cvtsi128_si32(s32);
8535 j += 32;
8536 }
8537 while j < n {
8538 total += (w[j] as i8) as i32 * xq[j] as i32;
8539 j += 1;
8540 }
8541 total
8542 }
8543}
8544
8545#[cfg(target_arch = "x86_64")]
8547#[target_feature(enable = "avx2,fma")]
8548unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8549 unsafe {
8551 use core::arch::x86_64::*;
8552 let n = x.len();
8553 let wp = w.as_ptr();
8554 let xp = x.as_ptr();
8555 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8556 let mut j = 0usize;
8557 while j + 16 <= n {
8558 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8559 let lo = _mm256_cvtepi8_epi32(wb);
8560 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8561 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8562 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8563 j += 16;
8564 }
8565 let acc = _mm256_add_ps(a0, a1);
8566 let hi128 = _mm256_extractf128_ps::<1>(acc);
8567 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8568 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8569 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8570 let mut sum = _mm_cvtss_f32(s32);
8571 while j < n {
8572 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8573 j += 1;
8574 }
8575 sum
8576 }
8577}
8578
8579#[cfg(target_arch = "x86_64")]
8584#[target_feature(enable = "avx2")]
8585unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8586 unsafe {
8588 use core::arch::x86_64::*;
8589 let n = w.len();
8590 let ones = _mm256_set1_epi16(1);
8591 let mut acc = _mm256_setzero_si256();
8592 let mut j = 0usize;
8593 while j + 32 <= n {
8594 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8595 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8596 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8597 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8598 j += 32;
8599 }
8600 let hi128 = _mm256_extracti128_si256::<1>(acc);
8601 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8602 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8603 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8604 let mut s = _mm_cvtsi128_si32(s32);
8605 while j < n {
8606 s += (w[j] as i8) as i32 * xq[j] as i32;
8607 j += 1;
8608 }
8609 s
8610 }
8611}
8612
8613#[cfg(target_arch = "aarch64")]
8617#[target_feature(enable = "neon,i8mm")]
8618unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8619 unsafe {
8621 use core::arch::aarch64::*;
8622 use core::arch::asm;
8623 let n = w0.len();
8624 let w0p = w0.as_ptr() as *const i8;
8625 let w1p = w1.as_ptr() as *const i8;
8626 let mut acc01 = vdupq_n_s32(0);
8629 let mut acc23 = vdupq_n_s32(0);
8630 let mut i = 0usize;
8631 while i + 8 <= n {
8632 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8633 let xb01 = vcombine_s8(
8634 vld1_s8(xs[0].as_ptr().add(i)),
8635 vld1_s8(xs[1].as_ptr().add(i)),
8636 );
8637 let xb23 = vcombine_s8(
8638 vld1_s8(xs[2].as_ptr().add(i)),
8639 vld1_s8(xs[3].as_ptr().add(i)),
8640 );
8641 asm!(
8642 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8643 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8644 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8645 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8646 options(pure, nomem, nostack),
8647 );
8648 i += 8;
8649 }
8650 let mut out = [[0i32; 4]; 2];
8651 let a01: [i32; 4] = core::mem::transmute(acc01);
8652 let a23: [i32; 4] = core::mem::transmute(acc23);
8653 out[0][0] = a01[0];
8654 out[0][1] = a01[1];
8655 out[1][0] = a01[2];
8656 out[1][1] = a01[3];
8657 out[0][2] = a23[0];
8658 out[0][3] = a23[1];
8659 out[1][2] = a23[2];
8660 out[1][3] = a23[3];
8661 if i < n {
8662 for (k, x) in xs.iter().enumerate() {
8663 for j in i..n {
8664 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8665 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8666 }
8667 }
8668 }
8669 out
8670 }
8671}
8672
8673#[cfg(target_arch = "aarch64")]
8677#[target_feature(enable = "neon,dotprod")]
8678unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8679 unsafe {
8681 use core::arch::aarch64::*;
8682 use core::arch::asm;
8683 let n = w0.len();
8684 let w0p = w0.as_ptr() as *const i8;
8685 let w1p = w1.as_ptr() as *const i8;
8686 let mut acc = [[vdupq_n_s32(0); 4]; 2];
8687 let mut i = 0usize;
8688 while i + 16 <= n {
8689 let wv0 = vld1q_s8(w0p.add(i));
8690 let wv1 = vld1q_s8(w1p.add(i));
8691 for (k, x) in xs.iter().enumerate() {
8692 let xv = vld1q_s8(x.as_ptr().add(i));
8693 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8694 asm!(
8695 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8696 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8697 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8698 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8699 options(pure, nomem, nostack),
8700 );
8701 acc[0][k] = a0;
8702 acc[1][k] = a1;
8703 }
8704 i += 16;
8705 }
8706 let mut out = [[0i32; 4]; 2];
8707 for r in 0..2 {
8708 for k in 0..4 {
8709 out[r][k] = vaddvq_s32(acc[r][k]);
8710 }
8711 }
8712 if i < n {
8713 for (k, x) in xs.iter().enumerate() {
8714 for j in i..n {
8715 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8716 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8717 }
8718 }
8719 }
8720 out
8721 }
8722}
8723
8724#[cfg(target_arch = "x86_64")]
8730#[target_feature(enable = "avx2")]
8731unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8732 unsafe {
8734 use core::arch::x86_64::*;
8735 let n = w0.len();
8736 let ones = _mm256_set1_epi16(1);
8737 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8738 let mut j = 0usize;
8739 while j + 32 <= n {
8740 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8741 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8742 let aw0 = _mm256_abs_epi8(wv0);
8743 let aw1 = _mm256_abs_epi8(wv1);
8744 for (k, x) in xs.iter().enumerate() {
8745 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8746 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8747 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8748 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8749 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8750 }
8751 j += 32;
8752 }
8753 let mut out = [[0i32; 4]; 2];
8754 for r in 0..2 {
8755 for k in 0..4 {
8756 let a = acc[r][k];
8757 let hi128 = _mm256_extracti128_si256::<1>(a);
8758 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8759 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8760 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8761 out[r][k] = _mm_cvtsi128_si32(s32);
8762 }
8763 }
8764 if j < n {
8765 for (k, x) in xs.iter().enumerate() {
8766 for i in j..n {
8767 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8768 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8769 }
8770 }
8771 }
8772 out
8773 }
8774}
8775
8776#[cfg(target_arch = "x86_64")]
8781#[inline]
8782fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8783 let dot = if avx512vnni_enabled() && row.len() >= 64 {
8784 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8785 } else {
8786 unsafe { dot_i8_i8_avx2(row, &act.xq) }
8787 };
8788 let mut acc = dot as f32 * act.sx;
8789 for &(j, xv) in &act.outliers {
8790 acc += (row[j] as i8) as f32 * xv;
8791 }
8792 acc
8793}
8794
8795#[cfg(target_arch = "x86_64")]
8799#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8800unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8801 unsafe {
8803 use core::arch::x86_64::*;
8804 let n = w.len();
8805 let flip = _mm512_set1_epi8(-128); #[inline(always)]
8807 unsafe fn step(
8808 w: *const u8,
8809 x: *const i8,
8810 flip: core::arch::x86_64::__m512i,
8811 acc: core::arch::x86_64::__m512i,
8812 ) -> core::arch::x86_64::__m512i {
8813 unsafe {
8814 use core::arch::x86_64::*;
8815 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8816 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8817 }
8818 }
8819 let (mut a0, mut a1, mut a2, mut a3) = (
8820 _mm512_setzero_si512(),
8821 _mm512_setzero_si512(),
8822 _mm512_setzero_si512(),
8823 _mm512_setzero_si512(),
8824 );
8825 let mut j = 0usize;
8826 while j + 256 <= n {
8827 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8828 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8829 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8830 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8831 j += 256;
8832 }
8833 while j + 64 <= n {
8834 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8835 j += 64;
8836 }
8837 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8838 _mm512_add_epi32(a0, a1),
8839 _mm512_add_epi32(a2, a3),
8840 ));
8841 while j < n {
8843 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8844 j += 1;
8845 }
8846 total
8847 }
8848}
8849
8850#[cfg(target_arch = "x86_64")]
8856#[target_feature(enable = "avx2")]
8857unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8858 unsafe {
8861 use core::arch::x86_64::*;
8862 let lomask = _mm_set1_epi8(0x0F);
8863 let eight = _mm256_set1_epi8(8);
8864 let ones = _mm256_set1_epi16(1);
8865 let mut acc = 0f32;
8866 for gi in 0..gpr {
8867 let g = g0 + gi;
8868 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8869 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8870 let lo = _mm_and_si128(b, lomask);
8871 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8872 let w = _mm256_sub_epi8(
8873 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8874 eight,
8875 );
8876 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8877 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8878 let d = _mm256_madd_epi16(p16, ones);
8879 let hi128 = _mm256_extracti128_si256::<1>(d);
8880 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8881 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8882 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8883 acc += _mm_cvtsi128_si32(s32) as f32 * s;
8884 }
8885 acc
8886 }
8887}
8888
8889#[cfg(target_arch = "x86_64")]
8892#[target_feature(enable = "avx2")]
8893unsafe fn dot_q4_row_avx2_2(
8894 packed: &[u8],
8895 scales: &[u8],
8896 g0: usize,
8897 gpr: usize,
8898 xq1: &[i8],
8899 xq2: &[i8],
8900) -> (f32, f32) {
8901 unsafe {
8903 use core::arch::x86_64::*;
8904 let lomask = _mm_set1_epi8(0x0F);
8905 let eight = _mm256_set1_epi8(8);
8906 let ones = _mm256_set1_epi16(1);
8907 let (mut acc1, mut acc2) = (0f32, 0f32);
8908 #[inline(always)]
8909 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8910 unsafe {
8911 use core::arch::x86_64::*;
8912 let hi128 = _mm256_extracti128_si256::<1>(d);
8913 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8914 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8915 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8916 _mm_cvtsi128_si32(s32)
8917 }
8918 }
8919 for gi in 0..gpr {
8920 let g = g0 + gi;
8921 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8922 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8923 let lo = _mm_and_si128(b, lomask);
8924 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8925 let w = _mm256_sub_epi8(
8926 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8927 eight,
8928 );
8929 let aw = _mm256_abs_epi8(w);
8930 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8931 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8932 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8933 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8934 acc1 += hsum(d1) as f32 * s;
8935 acc2 += hsum(d2) as f32 * s;
8936 }
8937 (acc1, acc2)
8938 }
8939}
8940
8941#[cfg(target_arch = "x86_64")]
8943fn q8_range_avx2(
8944 q: &[u8],
8945 row_scale: &[f32],
8946 act: &SplitAct,
8947 cols: usize,
8948 out_addr: SendMut,
8949 start: usize,
8950 end: usize,
8951) {
8952 for o in start..end {
8953 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8954 unsafe { *out_addr.at(o) = v };
8956 }
8957}
8958
8959#[cfg(target_arch = "x86_64")]
8961#[allow(clippy::too_many_arguments)]
8962fn q8_range2_avx2(
8963 q: &[u8],
8964 row_scale: &[f32],
8965 a1: &SplitAct,
8966 a2: &SplitAct,
8967 cols: usize,
8968 p1: SendMut,
8969 p2: SendMut,
8970 start: usize,
8971 end: usize,
8972) {
8973 for o in start..end {
8974 let row = &q[o * cols..(o + 1) * cols];
8975 unsafe {
8977 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8978 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8979 }
8980 }
8981}
8982
8983#[cfg(target_arch = "aarch64")]
8994fn i8mm_enabled() -> bool {
8995 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8996 *ON.get_or_init(|| {
8997 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8998 && std::arch::is_aarch64_feature_detected!("i8mm")
8999 })
9000}
9001
9002#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9006fn sdot_enabled() -> bool {
9007 use std::sync::OnceLock;
9008 static ON: OnceLock<bool> = OnceLock::new();
9009 *ON.get_or_init(|| {
9010 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9011 if !want {
9012 return false;
9013 }
9014
9015 #[cfg(target_arch = "aarch64")]
9016 {
9017 if std::arch::is_aarch64_feature_detected!("dotprod") {
9018 return true;
9019 }
9020 #[cfg(target_os = "android")]
9021 {
9022 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9023 if cpuinfo.lines().any(|l| {
9024 (l.starts_with("Features") || l.starts_with("features"))
9025 && l.contains("asimddp")
9026 }) {
9027 return true;
9028 }
9029 }
9030 }
9031 false
9032 }
9033 #[cfg(not(target_arch = "aarch64"))]
9034 {
9035 false
9036 }
9037 })
9038}
9039
9040struct SplitAct {
9045 xq: Vec<i8>,
9046 sx: f32,
9047 outliers: Vec<(usize, f32)>,
9048 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9051 xsum: i32,
9052}
9053
9054thread_local! {
9055 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9058 const { std::cell::RefCell::new(Vec::new()) };
9059}
9060
9061impl Drop for SplitAct {
9062 fn drop(&mut self) {
9063 let buf = std::mem::take(&mut self.xq);
9064 if buf.capacity() > 0 {
9065 XQ_FREE.with(|f| {
9066 let mut f = f.borrow_mut();
9067 if f.len() < 16 {
9068 f.push(buf);
9069 }
9070 });
9071 }
9072 }
9073}
9074
9075thread_local! {
9076 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9083}
9084
9085#[inline]
9088fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9089 KROW.with(|s| {
9090 let mut b = s.borrow_mut();
9091 if b.len() < n {
9092 b.resize(n, 0.0);
9093 }
9094 f(&mut b[..n])
9095 })
9096}
9097
9098fn split_act(x: &[f32]) -> SplitAct {
9099 let n = x.len();
9100 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9101 let thr = 8.0 * rms;
9102 let mut outliers: Vec<(usize, f32)> = Vec::new();
9106 let mut amax = 0f32;
9107 for (j, &v) in x.iter().enumerate() {
9108 let a = v.abs();
9109 if a > thr {
9110 outliers.push((j, v));
9111 } else if a > amax {
9112 amax = a;
9113 }
9114 }
9115 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9116 let inv = 1.0 / sx;
9117 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9118 xq.clear();
9119 xq.reserve(n);
9120 if outliers.is_empty() {
9121 xq.extend(
9122 x.iter()
9123 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9124 );
9125 } else {
9126 xq.extend(x.iter().map(|&v| {
9128 if v.abs() > thr {
9129 0
9130 } else {
9131 (v * inv).round().clamp(-127.0, 127.0) as i8
9132 }
9133 }));
9134 }
9135 let xsum = xq.iter().map(|&v| v as i32).sum();
9136 SplitAct {
9137 xq,
9138 sx,
9139 outliers,
9140 xsum,
9141 }
9142}
9143
9144fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9145 let n = x.len();
9146 let rms = (x
9147 .iter()
9148 .zip(col)
9149 .map(|(&a, &c)| {
9150 let v = a * c;
9151 (v * v) as f64
9152 })
9153 .sum::<f64>()
9154 / n.max(1) as f64)
9155 .sqrt() as f32;
9156 let thr = 8.0 * rms;
9157
9158 let mut outliers = Vec::new();
9159 let mut amax = 0f32;
9160 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9161 let v = a * c;
9162 let s = v.abs();
9163 if s > thr {
9164 outliers.push((j, v));
9165 } else if s > amax {
9166 amax = s;
9167 }
9168 }
9169
9170 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9171 let inv = 1.0 / sx;
9172 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9173 xq.clear();
9174 xq.reserve(n);
9175 if outliers.is_empty() {
9176 xq.extend(
9177 x.iter()
9178 .zip(col)
9179 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9180 );
9181 } else {
9182 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9183 let v = a * c;
9184 if v.abs() > thr {
9185 0
9186 } else {
9187 (v * inv).round().clamp(-127.0, 127.0) as i8
9188 }
9189 }));
9190 }
9191 let xsum = xq.iter().map(|&v| v as i32).sum();
9192 SplitAct {
9193 xq,
9194 sx,
9195 outliers,
9196 xsum,
9197 }
9198}
9199
9200#[cfg(target_arch = "aarch64")]
9203#[target_feature(enable = "neon,dotprod")]
9204unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9205 unsafe {
9207 use core::arch::aarch64::*;
9208 use core::arch::asm;
9209 let wp = w.as_ptr() as *const i8;
9210 let n = w.len();
9211 let (mut a0, mut a1, mut a2, mut a3) = (
9212 vdupq_n_s32(0),
9213 vdupq_n_s32(0),
9214 vdupq_n_s32(0),
9215 vdupq_n_s32(0),
9216 );
9217 let mut i = 0;
9218 while i + 64 <= n {
9219 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9220 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9221 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9222 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9223 asm!(
9224 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9225 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9226 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9227 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9228 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9229 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9230 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9231 options(pure, nomem, nostack),
9232 );
9233 i += 64;
9234 }
9235 while i + 16 <= n {
9236 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9237 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9238 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9239 i += 16;
9240 }
9241 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9242 while i < n {
9243 s += (*wp.add(i)) as i32 * xq[i] as i32;
9244 i += 1;
9245 }
9246 s
9247 }
9248}
9249
9250#[cfg(target_arch = "aarch64")]
9254#[target_feature(enable = "neon,dotprod")]
9255unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9256 unsafe {
9258 use core::arch::aarch64::*;
9259 use core::arch::asm;
9260 let n = xq.len();
9261 let px = xq.as_ptr();
9262 let (p0, p1, p2, p3) = (
9263 w0.as_ptr() as *const i8,
9264 w1.as_ptr() as *const i8,
9265 w2.as_ptr() as *const i8,
9266 w3.as_ptr() as *const i8,
9267 );
9268 let (mut a0, mut a1, mut a2, mut a3) = (
9269 vdupq_n_s32(0),
9270 vdupq_n_s32(0),
9271 vdupq_n_s32(0),
9272 vdupq_n_s32(0),
9273 );
9274 let mut i = 0;
9275 while i + 16 <= n {
9276 let x = vld1q_s8(px.add(i));
9277 let v0 = vld1q_s8(p0.add(i));
9278 let v1 = vld1q_s8(p1.add(i));
9279 let v2 = vld1q_s8(p2.add(i));
9280 let v3 = vld1q_s8(p3.add(i));
9281 asm!(
9282 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9283 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9284 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9285 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9286 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9287 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9288 options(pure, nomem, nostack),
9289 );
9290 i += 16;
9291 }
9292 let mut r = [
9293 vaddvq_s32(a0),
9294 vaddvq_s32(a1),
9295 vaddvq_s32(a2),
9296 vaddvq_s32(a3),
9297 ];
9298 while i < n {
9299 let xi = *px.add(i) as i32;
9300 r[0] += (*p0.add(i)) as i32 * xi;
9301 r[1] += (*p1.add(i)) as i32 * xi;
9302 r[2] += (*p2.add(i)) as i32 * xi;
9303 r[3] += (*p3.add(i)) as i32 * xi;
9304 i += 1;
9305 }
9306 r
9307 }
9308}
9309
9310#[cfg(target_arch = "aarch64")]
9317#[target_feature(enable = "neon,dotprod")]
9318unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9319 unsafe {
9322 use core::arch::aarch64::*;
9323 use core::arch::asm;
9324 let n = xq.len();
9325 let px = xq.as_ptr();
9326 let pg = g.as_ptr() as *const i8;
9327 let (mut a0, mut a1, mut a2, mut a3) = (
9328 vdupq_n_s32(0),
9329 vdupq_n_s32(0),
9330 vdupq_n_s32(0),
9331 vdupq_n_s32(0),
9332 );
9333 let mut i = 0;
9334 while i + 16 <= n {
9335 let x = vld1q_s8(px.add(i));
9336 let base = pg.add(4 * i);
9337 let v0 = vld1q_s8(base);
9338 let v1 = vld1q_s8(base.add(16));
9339 let v2 = vld1q_s8(base.add(32));
9340 let v3 = vld1q_s8(base.add(48));
9341 asm!(
9342 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9343 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9344 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9345 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9346 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9347 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9348 options(pure, nomem, nostack),
9349 );
9350 i += 16;
9351 }
9352 [
9353 vaddvq_s32(a0),
9354 vaddvq_s32(a1),
9355 vaddvq_s32(a2),
9356 vaddvq_s32(a3),
9357 ]
9358 }
9359}
9360
9361#[cfg(target_arch = "aarch64")]
9367fn q8_range_sdot(
9368 q: &[u8],
9369 rep: &[u8],
9370 row_scale: &[f32],
9371 act: &SplitAct,
9372 cols: usize,
9373 out_addr: SendMut,
9374 start: usize,
9375 end: usize,
9376) {
9377 let mut o = start;
9378 if !rep.is_empty() {
9381 while o < end && o % 4 != 0 {
9382 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9383 unsafe { *out_addr.at(o) = v };
9384 o += 1;
9385 }
9386 }
9387 while o + 4 <= end {
9388 let r = if rep.is_empty() {
9389 unsafe {
9390 dot_i8_sdot_4rows(
9391 &q[o * cols..(o + 1) * cols],
9392 &q[(o + 1) * cols..(o + 2) * cols],
9393 &q[(o + 2) * cols..(o + 3) * cols],
9394 &q[(o + 3) * cols..(o + 4) * cols],
9395 &act.xq,
9396 )
9397 }
9398 } else {
9399 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9400 };
9401 for k in 0..4 {
9402 let mut acc = r[k] as f32 * act.sx;
9403 for &(j, xv) in &act.outliers {
9404 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9405 }
9406 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9408 }
9409 o += 4;
9410 }
9411 while o < end {
9412 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9413 unsafe { *out_addr.at(o) = v };
9414 o += 1;
9415 }
9416}
9417
9418#[cfg(target_arch = "aarch64")]
9421#[allow(clippy::too_many_arguments)]
9422fn q8_range2_sdot(
9423 q: &[u8],
9424 row_scale: &[f32],
9425 a1: &SplitAct,
9426 a2: &SplitAct,
9427 cols: usize,
9428 p1: SendMut,
9429 p2: SendMut,
9430 start: usize,
9431 end: usize,
9432) {
9433 for o in start..end {
9434 let row = &q[o * cols..(o + 1) * cols];
9435 unsafe {
9437 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9438 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9439 }
9440 }
9441}
9442
9443#[allow(clippy::too_many_arguments)]
9445fn q8_range2_f32(
9446 q: &[u8],
9447 row_scale: &[f32],
9448 x1: &[f32],
9449 x2: &[f32],
9450 cols: usize,
9451 p1: SendMut,
9452 p2: SendMut,
9453 start: usize,
9454 end: usize,
9455) {
9456 for o in start..end {
9457 let row = &q[o * cols..(o + 1) * cols];
9458 unsafe {
9460 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9461 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9462 }
9463 }
9464}
9465
9466fn q8_range_f32(
9468 q: &[u8],
9469 row_scale: &[f32],
9470 xs: &[f32],
9471 cols: usize,
9472 out_addr: SendMut,
9473 start: usize,
9474 end: usize,
9475) {
9476 for o in start..end {
9477 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9478 unsafe { *out_addr.at(o) = v };
9480 }
9481}
9482
9483#[inline]
9487fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9488 #[cfg(target_arch = "aarch64")]
9489 return row_dot_sdot(row, act);
9490 #[cfg(target_arch = "x86_64")]
9491 return row_dot_avx2(row, act);
9492 #[allow(unreachable_code)]
9493 q8_row_dot_scalar(row, act)
9494}
9495
9496#[allow(dead_code)]
9497fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9498 let mut acc = 0i32;
9499 for (k, &b) in row.iter().enumerate() {
9500 acc += (b as i8) as i32 * act.xq[k] as i32;
9501 }
9502 let mut acc = acc as f32 * act.sx;
9503 for &(j, xv) in &act.outliers {
9504 acc += (row[j] as i8) as f32 * xv;
9505 }
9506 acc
9507}
9508
9509#[cfg(target_arch = "aarch64")]
9512#[inline]
9513fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9514 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9515 for &(j, xv) in &act.outliers {
9516 acc += (row[j] as i8) as f32 * xv;
9517 }
9518 acc
9519}
9520
9521#[cfg(target_arch = "aarch64")]
9529#[target_feature(enable = "neon,dotprod")]
9530unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9531 unsafe {
9534 use core::arch::aarch64::*;
9535 use core::arch::asm;
9536 let lomask = vdupq_n_u8(0x0F);
9537 let eight = vdupq_n_s8(8);
9538 let mut acc = 0f32;
9539 for gi in 0..gpr {
9540 let g = g0 + gi;
9541 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9542 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9543 let lo = vandq_u8(b, lomask);
9544 let hi = vshrq_n_u8::<4>(b);
9545 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9546 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9547 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9548 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9549 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9550 asm!(
9551 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9552 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9553 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9554 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9555 options(pure, nomem, nostack),
9556 );
9557 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9558 }
9559 acc
9560 }
9561}
9562
9563#[cfg(target_arch = "aarch64")]
9568#[target_feature(enable = "neon,dotprod")]
9569unsafe fn dot_q4_row_sdot2(
9570 packed: &[u8],
9571 scales: &[u8],
9572 g0: usize,
9573 gpr: usize,
9574 xq1: &[i8],
9575 xq2: &[i8],
9576) -> (f32, f32) {
9577 unsafe {
9580 use core::arch::aarch64::*;
9581 use core::arch::asm;
9582 let lomask = vdupq_n_u8(0x0F);
9583 let eight = vdupq_n_s8(8);
9584 let (mut acc1, mut acc2) = (0f32, 0f32);
9585 for gi in 0..gpr {
9586 let g = g0 + gi;
9587 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9588 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9589 let lo = vandq_u8(b, lomask);
9590 let hi = vshrq_n_u8::<4>(b);
9591 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9592 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9593 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9594 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9595 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9596 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9597 let (mut a0, mut a1, mut b0, mut b1) = (
9598 vdupq_n_s32(0),
9599 vdupq_n_s32(0),
9600 vdupq_n_s32(0),
9601 vdupq_n_s32(0),
9602 );
9603 asm!(
9604 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9605 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9606 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9607 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9608 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9609 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9610 e0 = in(vreg) e0, e1 = in(vreg) e1,
9611 x10 = in(vreg) x10, x11 = in(vreg) x11,
9612 x20 = in(vreg) x20, x21 = in(vreg) x21,
9613 options(pure, nomem, nostack),
9614 );
9615 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9616 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9617 }
9618 (acc1, acc2)
9619 }
9620}
9621
9622#[inline]
9627pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9628 #[cfg(target_arch = "aarch64")]
9629 unsafe {
9630 return axpy_i8_f32_neon(acc, row, w);
9631 }
9632 #[cfg(target_arch = "x86_64")]
9633 if avx2_enabled() {
9634 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9635 }
9636 #[allow(unreachable_code)]
9637 {
9638 for (a, &b) in acc.iter_mut().zip(row) {
9639 *a += w * b as f32;
9640 }
9641 }
9642}
9643
9644#[cfg(target_arch = "x86_64")]
9646#[target_feature(enable = "avx2,fma")]
9647unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9648 unsafe {
9650 use core::arch::x86_64::*;
9651 let n = acc.len().min(row.len());
9652 let ap = acc.as_mut_ptr();
9653 let rp = row.as_ptr();
9654 let wv = _mm256_set1_ps(w);
9655 let mut j = 0usize;
9656 while j + 16 <= n {
9657 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9658 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9659 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9660 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9661 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9662 _mm256_storeu_ps(ap.add(j), v0);
9663 _mm256_storeu_ps(ap.add(j + 8), v1);
9664 j += 16;
9665 }
9666 while j < n {
9667 *ap.add(j) += w * (*rp.add(j)) as f32;
9668 j += 1;
9669 }
9670 }
9671}
9672
9673#[cfg(target_arch = "aarch64")]
9674#[target_feature(enable = "neon")]
9675unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9676 unsafe {
9678 use core::arch::aarch64::*;
9679 let n = acc.len().min(row.len());
9680 let ap = acc.as_mut_ptr();
9681 let rp = row.as_ptr();
9682 let wv = vdupq_n_f32(w);
9683 let mut j = 0usize;
9684 while j + 16 <= n {
9685 let rb = vld1q_s8(rp.add(j));
9686 let lo = vmovl_s8(vget_low_s8(rb));
9687 let hi = vmovl_s8(vget_high_s8(rb));
9688 for (off, half) in [(0, lo), (8, hi)] {
9689 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9690 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9691 let o = j + off;
9692 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9693 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9694 }
9695 j += 16;
9696 }
9697 while j < n {
9698 *ap.add(j) += w * (*rp.add(j)) as f32;
9699 j += 1;
9700 }
9701 }
9702}
9703
9704#[inline]
9707pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9708 #[cfg(target_arch = "aarch64")]
9709 unsafe {
9710 return dot_i8_f32_neon(w, x);
9711 }
9712 #[cfg(target_arch = "x86_64")]
9713 if avx2_enabled() {
9714 return unsafe { dot_i8_f32_avx2(w, x) };
9715 }
9716 #[allow(unreachable_code)]
9717 {
9718 let mut sum = 0.0f32;
9719 for (j, &b) in w.iter().enumerate() {
9720 sum += (b as i8) as f32 * x[j];
9721 }
9722 sum
9723 }
9724}
9725
9726#[inline]
9730fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9731 #[cfg(target_arch = "aarch64")]
9732 unsafe {
9733 return dot_i8_col_f32_neon(w, x, col);
9734 }
9735 #[allow(unreachable_code)]
9736 {
9737 let mut sum = 0.0f32;
9738 for (j, &b) in w.iter().enumerate() {
9739 sum += (b as i8) as f32 * x[j] * col[j];
9740 }
9741 sum
9742 }
9743}
9744
9745#[cfg(target_arch = "aarch64")]
9746#[target_feature(enable = "neon")]
9747unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9748 unsafe {
9750 use core::arch::aarch64::*;
9751 let n = x.len();
9752 let wp = w.as_ptr() as *const i8;
9753 let xp = x.as_ptr();
9754 let cp = col.as_ptr();
9755 let (mut a0, mut a1, mut a2, mut a3) = (
9756 vdupq_n_f32(0.0),
9757 vdupq_n_f32(0.0),
9758 vdupq_n_f32(0.0),
9759 vdupq_n_f32(0.0),
9760 );
9761 let mut j = 0usize;
9762 while j + 16 <= n {
9763 let wb = vld1q_s8(wp.add(j));
9764 let lo = vmovl_s8(vget_low_s8(wb));
9765 let hi = vmovl_s8(vget_high_s8(wb));
9766 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9767 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9768 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9769 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9770 a0 = vfmaq_f32(
9771 a0,
9772 w0,
9773 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9774 );
9775 a1 = vfmaq_f32(
9776 a1,
9777 w1,
9778 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9779 );
9780 a2 = vfmaq_f32(
9781 a2,
9782 w2,
9783 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9784 );
9785 a3 = vfmaq_f32(
9786 a3,
9787 w3,
9788 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9789 );
9790 j += 16;
9791 }
9792 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9793 while j < n {
9794 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9795 j += 1;
9796 }
9797 sum
9798 }
9799}
9800
9801#[cfg(target_arch = "aarch64")]
9802#[target_feature(enable = "neon")]
9803unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9804 unsafe {
9806 use core::arch::aarch64::*;
9807 let n = x.len();
9808 let wp = w.as_ptr() as *const i8;
9809 let xp = x.as_ptr();
9810 let (mut a0, mut a1, mut a2, mut a3) = (
9811 vdupq_n_f32(0.0),
9812 vdupq_n_f32(0.0),
9813 vdupq_n_f32(0.0),
9814 vdupq_n_f32(0.0),
9815 );
9816 let mut j = 0usize;
9817 while j + 16 <= n {
9818 let wb = vld1q_s8(wp.add(j));
9819 let lo = vmovl_s8(vget_low_s8(wb));
9820 let hi = vmovl_s8(vget_high_s8(wb));
9821 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9822 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9823 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9824 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9825 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9826 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9827 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9828 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9829 j += 16;
9830 }
9831 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9832 while j < n {
9833 sum += (*wp.add(j)) as f32 * *xp.add(j);
9834 j += 1;
9835 }
9836 sum
9837 }
9838}
9839
9840#[allow(clippy::too_many_arguments)]
9841fn qmatvec(
9842 q: &[u8],
9843 rep: &[u8],
9844 row_scale: &[f32],
9845 x: &[f32],
9846 col_field: &[f32],
9847 dtype: TensorDtype,
9848 rows: usize,
9849 cols: usize,
9850 out: &mut [f32],
9851 pool: Option<&Pool>,
9852) {
9853 debug_assert_eq!(out.len(), rows);
9854 #[cfg(not(target_arch = "aarch64"))]
9855 let _ = rep;
9856
9857 #[cfg(target_arch = "aarch64")]
9858 if sdot_enabled() {
9859 let act = if dtype == TensorDtype::Q8_2f {
9860 split_act_q8_2f(x, col_field)
9861 } else {
9862 split_act(x)
9863 };
9864 let out_addr = SendMut(out.as_mut_ptr());
9865 let run_range = |start: usize, end: usize| {
9866 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9867 };
9868 match pool {
9869 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9870 _ => run_range(0, rows),
9871 }
9872 return;
9873 }
9874 #[cfg(target_arch = "x86_64")]
9877 if avx2_a8w8_enabled() {
9878 let act = if dtype == TensorDtype::Q8_2f {
9879 split_act_q8_2f(x, col_field)
9880 } else {
9881 split_act(x)
9882 };
9883 let out_addr = SendMut(out.as_mut_ptr());
9884 let run_range = |start: usize, end: usize| {
9885 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9886 };
9887 match pool {
9888 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9889 _ => run_range(0, rows),
9890 }
9891 return;
9892 }
9893
9894 prescale_with(x, col_field, dtype, 1, |xs| {
9895 let out_addr = SendMut(out.as_mut_ptr());
9896 let run_range = move |start: usize, end: usize| {
9897 for o in start..end {
9898 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9899 unsafe { *out_addr.at(o) = v };
9901 }
9902 };
9903 match pool {
9904 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9905 _ => run_range(0, rows),
9906 }
9907 });
9908}
9909
9910#[allow(clippy::too_many_arguments)]
9911fn qmatvec2(
9912 q: &[u8],
9913 row_scale: &[f32],
9914 x1: &[f32],
9915 x2: &[f32],
9916 col_field: &[f32],
9917 dtype: TensorDtype,
9918 rows: usize,
9919 cols: usize,
9920 o1: &mut [f32],
9921 o2: &mut [f32],
9922 pool: Option<&Pool>,
9923) {
9924 #[cfg(target_arch = "aarch64")]
9925 if sdot_enabled() {
9926 let a1s = if dtype == TensorDtype::Q8_2f {
9927 split_act_q8_2f(x1, col_field)
9928 } else {
9929 split_act(x1)
9930 };
9931 let a2s = if dtype == TensorDtype::Q8_2f {
9932 split_act_q8_2f(x2, col_field)
9933 } else {
9934 split_act(x2)
9935 };
9936 let p1 = SendMut(o1.as_mut_ptr());
9937 let p2 = SendMut(o2.as_mut_ptr());
9938 let run_range = |start: usize, end: usize| {
9939 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9940 };
9941 match pool {
9942 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9943 _ => run_range(0, rows),
9944 }
9945 return;
9946 }
9947 #[cfg(target_arch = "x86_64")]
9948 if avx2_a8w8_enabled() {
9949 let a1s = if dtype == TensorDtype::Q8_2f {
9950 split_act_q8_2f(x1, col_field)
9951 } else {
9952 split_act(x1)
9953 };
9954 let a2s = if dtype == TensorDtype::Q8_2f {
9955 split_act_q8_2f(x2, col_field)
9956 } else {
9957 split_act(x2)
9958 };
9959 let p1 = SendMut(o1.as_mut_ptr());
9960 let p2 = SendMut(o2.as_mut_ptr());
9961 let run_range = |start: usize, end: usize| {
9962 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9963 };
9964 match pool {
9965 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9966 _ => run_range(0, rows),
9967 }
9968 return;
9969 }
9970
9971 prescale_with(x1, col_field, dtype, 1, |x1s| {
9972 prescale_with(x2, col_field, dtype, 2, |x2s| {
9973 let p1 = SendMut(o1.as_mut_ptr());
9974 let p2 = SendMut(o2.as_mut_ptr());
9975 let run_range = move |start: usize, end: usize| {
9976 for o in start..end {
9977 let row = &q[o * cols..(o + 1) * cols];
9978 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
9979 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
9980 unsafe {
9982 *p1.at(o) = s1;
9983 *p2.at(o) = s2;
9984 }
9985 }
9986 };
9987 match pool {
9988 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9989 _ => run_range(0, rows),
9990 }
9991 });
9992 });
9993}
9994
9995#[derive(Clone, Copy)]
9996struct SendMut(*mut f32);
9997unsafe impl Send for SendMut {}
9998unsafe impl Sync for SendMut {}
9999
10000impl SendMut {
10001 #[inline]
10002 fn at(self, i: usize) -> *mut f32 {
10003 unsafe { self.0.add(i) }
10004 }
10005}
10006
10007#[cfg(test)]
10008mod tests {
10009 use super::*;
10010
10011 #[test]
10012 fn q2tp_i8_dot_matches_exact_on_grid() {
10013 let (rows, cols) = (5, 64);
10017 let gpr = cols / GROUP_SIZE;
10018 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10022 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10023 .collect();
10024 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10025 let x: Vec<f32> = (0..cols)
10026 .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10027 .collect();
10028 let act = split_act(&x);
10029 assert!(
10030 act.outliers.is_empty(),
10031 "on-grid input must have no outliers"
10032 );
10033 let gsum = q1_group_sums(&act.xq, gpr);
10034 for r in 0..rows {
10035 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10036 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10037 assert!(
10038 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10039 "row {r}: exact {exact} vs i8 {fast}"
10040 );
10041 }
10042 }
10043
10044 #[test]
10045 fn q8_row_dot_fast_matches_scalar() {
10046 let cols = 96;
10049 let row: Vec<u8> = (0..cols)
10050 .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10051 .collect();
10052 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10053 let act = split_act(&x);
10054 let fast = q8_row_dot(&row, &act);
10055 let scalar = q8_row_dot_scalar(&row, &act);
10056 assert!(
10057 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10058 "fast {fast} vs scalar {scalar}"
10059 );
10060 }
10061
10062 #[test]
10063 fn f32_matvec_matches_matvec_rows_bitexact() {
10064 let (rows, cols) = (300, 40);
10065 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10066 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10067 let qt = QTensor::from_f32(w.clone(), rows, cols);
10068
10069 let mut a = vec![0.0f32; rows];
10070 matvec_rows(None, &w, &x, &mut a);
10071 let mut b = vec![0.0f32; rows];
10072 qt.matvec(&x, &mut b, None);
10073 assert_eq!(a, b);
10074 }
10075
10076 #[test]
10077 fn sdot_kernel_exact_on_grid() {
10078 eprintln!("sdot_enabled = {}", sdot_enabled());
10083 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
10085 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10086 .collect();
10087 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10088 let x: Vec<f32> = (0..cols)
10089 .map(|i| match i % 3 {
10090 0 => 1.0,
10091 1 => -1.0,
10092 _ => 0.0,
10093 })
10094 .collect();
10095 let mut a = vec![0.0f32; rows];
10096 qmatvec(
10097 &w,
10098 &[],
10099 &scales,
10100 &x,
10101 &[],
10102 TensorDtype::Q8Row,
10103 rows,
10104 cols,
10105 &mut a,
10106 None,
10107 );
10108 for o in 0..rows {
10109 let mut acc = 0.0f32;
10110 for j in 0..cols {
10111 acc += (w[o * cols + j] as i8) as f32 * x[j];
10112 }
10113 let expect = acc * scales[o];
10114 assert!(
10115 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10116 "row {o}: {} vs {expect}",
10117 a[o]
10118 );
10119 }
10120 }
10121
10122 #[test]
10123 fn q1_tbl_fast_path_matches_reference() {
10124 let (rows, cols) = (5, 256);
10129 let gpr = cols / GROUP_SIZE;
10130 let mut bytes = Vec::new();
10131 for t in 0..rows * gpr {
10132 let s = 0.007 + (t % 11) as f32 * 0.004;
10133 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10134 for j in 0..4 {
10135 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10136 }
10137 }
10138 let x: Vec<f32> = (0..cols)
10139 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10140 .collect();
10141 let mut w = vec![0.0f32; rows * cols];
10142 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10143 let mut got = vec![0.0f32; rows];
10144 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10145 for o in 0..rows {
10146 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10147 assert!(
10148 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10149 "row {o}: {} vs {expect}",
10150 got[o]
10151 );
10152 }
10153 let b = 5usize;
10156 let mut xs_all = Vec::new();
10157 for bi in 0..b {
10158 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10159 }
10160 let mut mm = vec![0.0f32; b * rows];
10161 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10162 for bi in 0..b {
10163 let mut single = vec![0.0f32; rows];
10164 q1_matvec(
10165 &bytes,
10166 &xs_all[bi * cols..(bi + 1) * cols],
10167 rows,
10168 cols,
10169 &mut single,
10170 None,
10171 );
10172 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10173 }
10174 }
10175
10176 #[test]
10177 fn q1_kernels_match_exact_reference() {
10178 let (rows, cols) = (7, 96);
10180 let gpr = cols / GROUP_SIZE;
10181 let mut bytes = Vec::new();
10182 for t in 0..rows * gpr {
10183 let s = 0.01 + (t % 13) as f32 * 0.003;
10184 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10185 for j in 0..4 {
10186 bytes.push(((t * 31 + j * 97) % 251) as u8);
10187 }
10188 }
10189 let x: Vec<f32> = (0..cols)
10191 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10192 .collect();
10193 let mut w = vec![0.0f32; rows * cols];
10195 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10196 let mut expect = vec![0.0f32; rows];
10197 for o in 0..rows {
10198 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10199 }
10200 let mut got = vec![0.0f32; rows];
10201 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10202 for o in 0..rows {
10203 assert!(
10204 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10205 "row {o}: {} vs {}",
10206 got[o],
10207 expect[o]
10208 );
10209 }
10210 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10212 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10213 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10214 assert_eq!(a1, got);
10215 let mut xs = x.clone();
10216 xs.extend_from_slice(&x2);
10217 let mut mm = vec![0.0f32; 2 * rows];
10218 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10219 assert_eq!(&mm[..rows], got.as_slice());
10220 assert_eq!(&mm[rows..], a2.as_slice());
10221 }
10222
10223 #[test]
10224 fn repack_is_bit_identical() {
10225 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
10231 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10232 .collect();
10233 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10234 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10235 let rep = q8_repack_layout(&w, rows, cols);
10236 for g in 0..rows / 4 {
10238 for c in 0..cols / 16 {
10239 for lane in 0..4 {
10240 assert_eq!(
10241 &rep[g * 4 * cols + c * 64 + lane * 16
10242 ..g * 4 * cols + c * 64 + lane * 16 + 16],
10243 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10244 );
10245 }
10246 }
10247 }
10248 let mut a = vec![0.0f32; rows];
10249 qmatvec(
10250 &w,
10251 &[],
10252 &scales,
10253 &x,
10254 &[],
10255 TensorDtype::Q8Row,
10256 rows,
10257 cols,
10258 &mut a,
10259 None,
10260 );
10261 let mut b = vec![0.0f32; rows];
10262 qmatvec(
10263 &w,
10264 &rep,
10265 &scales,
10266 &x,
10267 &[],
10268 TensorDtype::Q8Row,
10269 rows,
10270 cols,
10271 &mut b,
10272 None,
10273 );
10274 assert_eq!(a, b, "full-range repack output diverged");
10275
10276 #[cfg(target_arch = "aarch64")]
10277 if sdot_enabled() {
10278 let act = split_act(&x);
10280 let mut c1 = vec![0.0f32; rows];
10281 let mut c2 = vec![0.0f32; rows];
10282 q8_range_sdot(
10283 &w,
10284 &[],
10285 &scales,
10286 &act,
10287 cols,
10288 SendMut(c1.as_mut_ptr()),
10289 3,
10290 rows - 2,
10291 );
10292 q8_range_sdot(
10293 &w,
10294 &rep,
10295 &scales,
10296 &act,
10297 cols,
10298 SendMut(c2.as_mut_ptr()),
10299 3,
10300 rows - 2,
10301 );
10302 assert_eq!(c1, c2, "unaligned-range repack output diverged");
10303 }
10304 }
10305
10306 #[test]
10307 fn sdot_a8w8_noise_is_bounded() {
10308 let (rows, cols) = (16, 512);
10312 let w: Vec<u8> = (0..rows * cols)
10313 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10314 .collect();
10315 let scales = vec![0.01f32; rows];
10316 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10317 let mut a = vec![0.0f32; rows];
10318 qmatvec(
10319 &w,
10320 &[],
10321 &scales,
10322 &x,
10323 &[],
10324 TensorDtype::Q8Row,
10325 rows,
10326 cols,
10327 &mut a,
10328 None,
10329 );
10330 let (mut num, mut den) = (0f64, 0f64);
10331 for o in 0..rows {
10332 let mut acc = 0.0f32;
10333 for j in 0..cols {
10334 acc += (w[o * cols + j] as i8) as f32 * x[j];
10335 }
10336 let expect = acc * scales[o];
10337 num += ((a[o] - expect) as f64).powi(2);
10338 den += (expect as f64).powi(2);
10339 }
10340 let rel = (num / den.max(1e-12)).sqrt();
10341 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10342 }
10343
10344 #[test]
10345 fn i8_dot_neon_matches_scalar() {
10346 let n = 100;
10347 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10348 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10349 let mut scalar = 0.0f32;
10350 for j in 0..n {
10351 scalar += (w[j] as i8) as f32 * x[j];
10352 }
10353 let fast = dot_i8_f32(&w, &x);
10354 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10355 }
10356
10357 #[test]
10359 fn vbitmatvec_matches_full_dequant() {
10360 let (rows, cols) = (6, 64);
10361 let ng = cols / GROUP_SIZE;
10362 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10364 let mut bytes = bits.clone();
10365 for g in 0..rows * ng {
10366 let s = 0.02 + 0.001 * g as f32;
10367 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10368 }
10369 for r in 0..rows {
10370 let b = bits[r] as usize;
10371 let (mut acc, mut nb) = (0u64, 0usize);
10372 let mut rowbytes = Vec::new();
10373 for i in 0..cols {
10374 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10375 acc = (acc << b) | v;
10376 nb += b;
10377 while nb >= 8 {
10378 nb -= 8;
10379 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10380 }
10381 }
10382 if nb > 0 {
10383 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10384 }
10385 bytes.extend_from_slice(&rowbytes);
10386 }
10387 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10388
10389 let mut reference = vec![0f32; rows * cols];
10390 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10391 let mut expect = vec![0f32; rows];
10392 for r in 0..rows {
10393 expect[r] = reference[r * cols..(r + 1) * cols]
10394 .iter()
10395 .zip(&x)
10396 .map(|(w, xv)| w * xv)
10397 .sum();
10398 }
10399 let mut got = vec![0f32; rows];
10400 let offsets = vbit_row_offsets(&bytes, rows, cols);
10401 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10402 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10406 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10407 for r in 0..rows {
10408 assert!(
10409 (got[r] - expect[r]).abs() < tol * scale,
10410 "row {r}: {} vs {}",
10411 got[r],
10412 expect[r]
10413 );
10414 }
10415 }
10416
10417 #[test]
10422 #[cfg(target_arch = "x86_64")]
10423 fn vbit_matmat_blocked_matches_per_row() {
10424 let (rows, cols, b) = (64usize, 128usize, 9usize);
10425 let ng = cols / GROUP_SIZE;
10426 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10427 let mut bytes = bits.clone();
10428 for g in 0..rows * ng {
10429 let sc = 0.02 + 0.0005 * g as f32;
10430 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10431 }
10432 for r in 0..rows {
10433 let bw = bits[r] as usize;
10434 let (mut acc, mut nb) = (0u64, 0usize);
10435 let mut rowbytes = Vec::new();
10436 for i in 0..cols {
10437 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10438 acc = (acc << bw) | v;
10439 nb += bw;
10440 while nb >= 8 {
10441 nb -= 8;
10442 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10443 }
10444 }
10445 if nb > 0 {
10446 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10447 }
10448 bytes.extend_from_slice(&rowbytes);
10449 }
10450 let x: Vec<f32> = (0..b * cols)
10451 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10452 .collect();
10453 let offsets = vbit_row_offsets(&bytes, rows, cols);
10454 let mut y_a = vec![0f32; b * rows];
10455 let mut y_b = vec![0f32; b * rows];
10456 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10457 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10458 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10459 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10460 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10461 let max_d = y_a
10462 .iter()
10463 .zip(&y_b)
10464 .map(|(p, q)| (p - q).abs())
10465 .fold(0.0f32, f32::max);
10466 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10467 }
10468
10469 #[test]
10477 fn q4t_matmat_blocked_matches_per_row() {
10478 let (rows, cols, b) = (16usize, 64usize, 9usize);
10479 let gpr = cols / GROUP_SIZE;
10480 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10481 for r in 0..rows {
10482 for g in 0..gpr {
10483 let t = (r * gpr + g) * Q4_TILE;
10484 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10485 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10486 for k in 0..16 {
10487 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10488 }
10489 }
10490 }
10491 let x: Vec<f32> = (0..b * cols)
10492 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10493 .collect();
10494 let mut y_blk = vec![0f32; b * rows];
10495 let mut y_row = vec![0f32; b * rows];
10496 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10497 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10498 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10499 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10500 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10501 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10502 }
10503
10504 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10511 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10512 let gpr = cols / GROUP_SIZE;
10513 let stride = q4tp_code_stride(gpr);
10514 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10515 let mut b = vec![0u8; codes_off + rows * stride];
10516 for r in 0..rows {
10517 for g in 0..gpr {
10518 let t = (r * gpr + g) * Q4TP_NIB;
10519 for k in 0..16 {
10520 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10521 }
10522 }
10523 let lo = -6.0 - 0.03 * (r % 17) as f32;
10524 let step = 0.01 + 0.004 * (r % 11) as f32;
10525 let p = params_off + r * 4;
10526 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10527 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10528 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10529 for g in 0..gpr {
10530 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10531 }
10532 }
10533 b
10534 }
10535
10536 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10540 let gpr = cols / GROUP_SIZE;
10541 let v = Q4tpView::new(bytes, rows, cols);
10542 let mut out = vec![0u8; rows * gpr * Q4_TILE];
10543 let mut sc = vec![0f32; gpr];
10544 for r in 0..rows {
10545 v.scales_into(r, gpr, &mut sc);
10546 for g in 0..gpr {
10547 let t = (r * gpr + g) * Q4_TILE;
10548 let s = sc[g];
10549 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10550 let src = (r * gpr + g) * Q4TP_NIB;
10551 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10552 }
10553 }
10554 out
10555 }
10556
10557 #[test]
10563 fn q4tp_exact_path_matches_dequant_reference() {
10564 let (rows, cols) = (256usize, 512usize);
10565 let gpr = cols / GROUP_SIZE;
10566 let bytes = synth_q4tp(rows, cols);
10567 let mut w = vec![0f32; rows * cols];
10568 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10569
10570 let x: Vec<f32> = (0..cols)
10571 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10572 .collect();
10573 let v = Q4tpView::new(&bytes, rows, cols);
10574 let mut sc = vec![0f32; gpr];
10575 for r in 0..rows {
10576 v.scales_into(r, gpr, &mut sc);
10577 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10578 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10579 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10583 assert!(
10584 (got - want).abs() <= 1e-5 * mag,
10585 "row {r}: kernel {got} vs dequant {want}"
10586 );
10587 }
10588 }
10589
10590 #[test]
10596 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10597 let (rows, cols) = (256usize, 512usize);
10598 let bytes = synth_q4tp(rows, cols);
10599 let twin = q4tp_as_q4t(&bytes, rows, cols);
10600 let x: Vec<f32> = (0..cols)
10601 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10602 .collect();
10603
10604 let mut got = vec![0f32; rows];
10605 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10606 let mut want = vec![0f32; rows];
10607 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10608
10609 let mut w = vec![0f32; rows * cols];
10612 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10613 for r in 0..rows {
10614 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10615 assert!(
10616 (got[r] - want[r]).abs() <= 1e-3 * mag,
10617 "row {r}: q4tp {} vs q4t {}",
10618 got[r],
10619 want[r]
10620 );
10621 }
10622 }
10623
10624 #[test]
10629 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10630 let (rows, cols, b) = (256usize, 512usize, 5usize);
10631 let bytes = synth_q4tp(rows, cols);
10632 let twin = q4tp_as_q4t(&bytes, rows, cols);
10633 let xs: Vec<f32> = (0..b * cols)
10634 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10635 .collect();
10636
10637 let mut got = vec![0f32; b * rows];
10638 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10639 let mut want = vec![0f32; b * rows];
10640 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10641
10642 let mut w = vec![0f32; rows * cols];
10643 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10644 for t in 0..b {
10645 for r in 0..rows {
10646 let mag: f32 = (0..cols)
10647 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10648 .sum();
10649 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10650 assert!(
10651 (g - wa).abs() <= 1e-3 * mag,
10652 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10653 );
10654 }
10655 }
10656 }
10657
10658 #[test]
10659 fn q4tp_matvec2_matches_the_single_stream_kernel() {
10660 let (rows, cols) = (128usize, 256usize);
10661 let gpr = cols / GROUP_SIZE;
10662 let bytes = synth_q4tp(rows, cols);
10663 let xs: Vec<f32> = (0..2 * cols)
10664 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10665 .collect();
10666
10667 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10668 q4tp_matvec2(
10669 &bytes,
10670 &xs[..cols],
10671 &xs[cols..],
10672 rows,
10673 cols,
10674 &mut o1,
10675 &mut o2,
10676 None,
10677 );
10678
10679 let v = Q4tpView::new(&bytes, rows, cols);
10682 let mut sc = vec![0f32; gpr];
10683 for r in 0..rows {
10684 v.scales_into(r, gpr, &mut sc);
10685 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10686 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10687 }
10688 }
10689
10690 #[test]
10697 fn q4tp_matvec_keeps_pace_with_q4t() {
10698 let (rows, cols) = (4096usize, 3072usize);
10699 let bytes = synth_q4tp(rows, cols);
10700 let twin = q4tp_as_q4t(&bytes, rows, cols);
10701 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10702 let mut o = vec![0f32; rows];
10703 let n = 12;
10704 let mut best = (f64::MAX, f64::MAX);
10705 for _ in 0..3 {
10708 let t0 = std::time::Instant::now();
10709 for _ in 0..n {
10710 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10711 }
10712 best.0 = best.0.min(t0.elapsed().as_secs_f64());
10713 let t0 = std::time::Instant::now();
10714 for _ in 0..n {
10715 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10716 }
10717 best.1 = best.1.min(t0.elapsed().as_secs_f64());
10718 }
10719 let ratio = best.1 / best.0;
10720 println!(
10721 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10722 best.0 * 1e3 / n as f64,
10723 best.1 * 1e3 / n as f64
10724 );
10725 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10726 }
10727
10728 #[cfg(target_os = "macos")]
10729 #[test]
10730 fn q4t_matmat_accel_matches_dequant_reference() {
10731 if !accel_gemm_enabled() {
10732 return; }
10734 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
10736 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10737 for r in 0..rows {
10738 for g in 0..gpr {
10739 let t = (r * gpr + g) * Q4_TILE;
10740 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10741 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10742 for k in 0..16 {
10743 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10744 }
10745 }
10746 }
10747 let x: Vec<f32> = (0..b * cols)
10748 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10749 .collect();
10750 let mut got = vec![0f32; b * rows];
10751 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10752 let mut w = vec![0f32; rows * cols];
10754 for r in 0..rows {
10755 for g in 0..gpr {
10756 let t = (r * gpr + g) * Q4_TILE;
10757 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10758 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10759 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10760 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10761 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10762 }
10763 }
10764 }
10765 for bi in 0..b {
10766 for r in 0..rows {
10767 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10768 let d = (got[bi * rows + r] - want).abs();
10769 assert!(
10770 d <= want.abs().max(1.0) * 1e-4,
10771 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10772 got[bi * rows + r]
10773 );
10774 }
10775 }
10776 }
10777
10778 #[test]
10779 fn q4matvec_matches_full_dequant() {
10780 let (rows, cols) = (8, 64);
10781 let groups = rows * cols / GROUP_SIZE;
10782 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10784 for i in 0..groups * 16 {
10785 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10786 }
10787 for g in 0..groups {
10788 let s = 0.01 + 0.003 * g as f32;
10789 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10790 }
10791 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10792
10793 let mut reference = vec![0.0f32; rows * cols];
10794 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10795 let mut expect = vec![0.0f32; rows];
10796 for r in 0..rows {
10797 expect[r] = reference[r * cols..(r + 1) * cols]
10798 .iter()
10799 .zip(&x)
10800 .map(|(w, xv)| w * xv)
10801 .sum();
10802 }
10803
10804 let mut got = vec![0.0f32; rows];
10805 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10806 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10810 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10811 for r in 0..rows {
10812 assert!(
10813 (got[r] - expect[r]).abs() < tol * scale,
10814 "row {r}: {} vs {}",
10815 got[r],
10816 expect[r]
10817 );
10818 }
10819 }
10820
10821 #[test]
10824 fn vbitmatvec2_equals_two_singles() {
10825 let (rows, cols) = (6, 64);
10826 let ng = cols / GROUP_SIZE;
10827 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10828 let mut bytes = bits.clone();
10829 for g in 0..rows * ng {
10830 let s = 0.02 + 0.001 * g as f32;
10831 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10832 }
10833 for r in 0..rows {
10834 let b = bits[r] as usize;
10835 let (mut acc, mut nb) = (0u64, 0usize);
10836 let mut rowbytes = Vec::new();
10837 for i in 0..cols {
10838 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10839 acc = (acc << b) | v;
10840 nb += b;
10841 while nb >= 8 {
10842 nb -= 8;
10843 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10844 }
10845 }
10846 if nb > 0 {
10847 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10848 }
10849 bytes.extend_from_slice(&rowbytes);
10850 }
10851 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10852 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10853 let offsets = vbit_row_offsets(&bytes, rows, cols);
10854
10855 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10856 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10857 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10858 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10859 vbitmatvec2(
10860 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10861 );
10862 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10863 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10864 }
10865
10866 #[test]
10868 fn q4matvec2_equals_two_singles() {
10869 let (rows, cols) = (8, 128);
10870 let groups = rows * cols / GROUP_SIZE;
10871 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10872 for i in 0..groups * 16 {
10873 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10874 }
10875 for g in 0..groups {
10876 let s = 0.01 + 0.003 * g as f32;
10877 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10878 }
10879 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10882 x1[9] = 250.0;
10883 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10884
10885 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10886 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10887 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10888 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10889 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10890 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10891 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10892 }
10893
10894 #[test]
10897 fn matvec_many_equals_separate_matvecs() {
10898 use crate::pool::Pool;
10899 let (r1, r2, cols) = (300, 200, 64);
10900 let mk = |salt: usize, rows: usize| {
10901 QTensor::from_f32(
10902 (0..rows * cols)
10903 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10904 .collect(),
10905 rows,
10906 cols,
10907 )
10908 };
10909 let (a, b) = (mk(1, r1), mk(5, r2));
10910 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10911 let pool = Pool::new(3);
10912
10913 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10914 a.matvec(&x, &mut ea, Some(&pool));
10915 b.matvec(&x, &mut eb, Some(&pool));
10916 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10917 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10918 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10919 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10920 }
10921
10922 #[test]
10925 fn batched_matmat_equals_per_position_matvec() {
10926 let (rows, cols, b) = (8, 64, 5);
10927 let groups = rows * cols / GROUP_SIZE;
10929 let mut q4 = Vec::new();
10930 for i in 0..groups * 16 {
10931 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10932 }
10933 for g in 0..groups {
10934 q4.extend_from_slice(
10935 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10936 );
10937 }
10938 let ng = cols / GROUP_SIZE;
10940 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10941 let mut vb = bits.clone();
10942 for g in 0..rows * ng {
10943 vb.extend_from_slice(
10944 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10945 );
10946 }
10947 for r in 0..rows {
10948 let bw = bits[r] as usize;
10949 let (mut acc, mut nb) = (0u64, 0usize);
10950 let mut rowbytes = Vec::new();
10951 for i in 0..cols {
10952 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10953 acc = (acc << bw) | v;
10954 nb += bw;
10955 while nb >= 8 {
10956 nb -= 8;
10957 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10958 }
10959 }
10960 if nb > 0 {
10961 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10962 }
10963 vb.extend_from_slice(&rowbytes);
10964 }
10965 let offsets = vbit_row_offsets(&vb, rows, cols);
10966
10967 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10968
10969 let mut got = vec![0f32; b * rows];
10971 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
10972 for bi in 0..b {
10973 let mut expect = vec![0f32; rows];
10974 q4matvec(
10975 &q4,
10976 &xs[bi * cols..(bi + 1) * cols],
10977 rows,
10978 cols,
10979 &mut expect,
10980 None,
10981 );
10982 assert_eq!(
10983 &got[bi * rows..(bi + 1) * rows],
10984 &expect[..],
10985 "q4 batch pos {bi}"
10986 );
10987 }
10988
10989 let mut got = vec![0f32; b * rows];
10991 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
10992 for bi in 0..b {
10993 let mut expect = vec![0f32; rows];
10994 vbitmatvec(
10995 &vb,
10996 &offsets,
10997 &xs[bi * cols..(bi + 1) * cols],
10998 rows,
10999 cols,
11000 &mut expect,
11001 None,
11002 );
11003 assert_eq!(
11004 &got[bi * rows..(bi + 1) * rows],
11005 &expect[..],
11006 "vbit batch pos {bi}"
11007 );
11008 }
11009 }
11010
11011 #[test]
11015 fn q4_tiled_matches_q4_block_bitexact() {
11016 let (rows, cols, b) = (8usize, 128usize, 3usize);
11017 let groups = rows * cols / GROUP_SIZE;
11018 let mut split = Vec::with_capacity(groups * 18);
11019 for i in 0..groups * 16 {
11020 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11021 }
11022 for g in 0..groups {
11023 split.extend_from_slice(
11024 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11025 );
11026 }
11027 let (packed, scales) = split.split_at(groups * 16);
11029 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
11030 for g in 0..groups {
11031 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
11032 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
11033 }
11034
11035 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11036 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11038
11039 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
11040 q4matvec(&split, &x1, rows, cols, &mut a, None);
11041 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
11042 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11043
11044 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11045 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11046 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11047 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11048 assert_eq!(a1, t1);
11049 assert_eq!(a2, t2);
11050
11051 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11052 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11053 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11054 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11055 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11056 }
11057
11058 #[test]
11065 fn q4matvec_sdot_outlier_exact() {
11066 let (rows, cols) = (4, 128);
11067 let groups = rows * cols / GROUP_SIZE;
11068 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11069 for i in 0..groups * 16 {
11070 bytes.push(((i * 11 + 5) % 256) as u8);
11071 }
11072 for g in 0..groups {
11073 let s = 0.02 + 0.002 * g as f32;
11074 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11075 }
11076 let mut x: Vec<f32> = (0..cols)
11077 .map(|i| match i % 3 {
11078 0 => 1.0,
11079 1 => -1.0,
11080 _ => 0.0,
11081 })
11082 .collect();
11083 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
11086 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11087 let mut expect = vec![0.0f32; rows];
11088 for r in 0..rows {
11089 expect[r] = reference[r * cols..(r + 1) * cols]
11090 .iter()
11091 .zip(&x)
11092 .map(|(w, xv)| w * xv)
11093 .sum();
11094 }
11095 let mut got = vec![0.0f32; rows];
11096 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11097 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11098 for r in 0..rows {
11099 assert!(
11100 (got[r] - expect[r]).abs() < 2e-3 * scale,
11101 "row {r}: {} vs {} (outlier term must be exact)",
11102 got[r],
11103 expect[r]
11104 );
11105 }
11106 }
11107
11108 #[test]
11112 fn q1t_matvec_matches_reference() {
11113 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11114 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
11116 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11117 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11119 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11120 let mut bytes = Vec::new();
11121 for r in 0..rows {
11122 for g in 0..gpr {
11123 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11124 let mut c = [0u8; 7];
11125 for k in 0..GROUP_SIZE {
11126 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11128 0
11129 } else {
11130 ((k + r * 3 + g) % 3) as u8 };
11132 cortiq_core::quant::q1t_pack(&mut c, k, code);
11133 }
11134 bytes.extend_from_slice(&c);
11135 }
11136 }
11137 let mut row_ptr = vec![0u32; rows + 1];
11140 for &(idx, _) in &outliers {
11141 row_ptr[idx as usize / cols + 1] += 1;
11142 }
11143 for r in 0..rows {
11144 row_ptr[r + 1] += row_ptr[r];
11145 }
11146 for &p in &row_ptr {
11147 bytes.extend_from_slice(&p.to_le_bytes());
11148 }
11149 for &(idx, v) in &outliers {
11150 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11151 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11152 }
11153
11154 let mut refw = vec![0f32; rows * cols];
11155 dequant_q1t(&bytes, rows, cols, &mut refw);
11156 let x: Vec<f32> = (0..cols)
11159 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11160 .collect();
11161 let mut expect = vec![0f32; rows];
11162 for r in 0..rows {
11163 let mut a = 0.0f32;
11164 for j in 0..cols {
11165 a += refw[r * cols + j] * x[j];
11166 }
11167 expect[r] = a;
11168 }
11169 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11170 let mut got = vec![0f32; rows];
11171 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11172 for r in 0..rows {
11173 assert!(
11174 (got[r] - expect[r]).abs() < tol(expect[r]),
11175 "row {r}: {} vs {}",
11176 got[r],
11177 expect[r]
11178 );
11179 }
11180 let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
11182 let mut gm = vec![0f32; 2 * rows];
11183 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11184 for r in 0..rows {
11185 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11186 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11187 }
11188 let xb: Vec<f32> = (0..cols)
11192 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11193 .collect();
11194 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11195 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11196 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11197 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11198 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11199 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11200 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11201 }
11202
11203 #[test]
11206 fn q1t_matvec2_odd_gpr_matches_singles() {
11207 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11208 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
11210 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11211 for r in 0..rows {
11212 for g in 0..gpr {
11213 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11214 let mut c = [0u8; 7];
11215 for k in 0..GROUP_SIZE {
11216 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11217 }
11218 bytes.extend_from_slice(&c);
11219 }
11220 }
11221 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11222 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11223 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11224 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11225 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11226 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11227 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11228 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11229 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11230 }
11231
11232 #[test]
11236 #[ignore]
11237 fn q1t_matvec2_speed() {
11238 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11239 use std::time::Instant;
11240 let (rows, cols) = (8192usize, 4096usize);
11241 let gpr = cols / GROUP_SIZE;
11242 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11243 for r in 0..rows {
11244 for g in 0..gpr {
11245 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11246 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11247 let mut c = [0u8; 7];
11248 for k in 0..GROUP_SIZE {
11249 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11250 }
11251 bytes.extend_from_slice(&c);
11252 }
11253 }
11254 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11255 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11256 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11257 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11258 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11260 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11261 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11262 for _ in 0..8 {
11263 let t0 = Instant::now();
11264 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11265 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11266 let t1 = Instant::now();
11267 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11268 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11269 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11270 }
11271 assert_eq!(p1, s1);
11272 assert_eq!(p2, s2);
11273 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11274 }
11275
11276 #[test]
11280 #[ignore]
11281 fn q1t_matvec_speed() {
11282 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11283 use std::time::Instant;
11284 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
11286 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11287 for r in 0..rows {
11288 for g in 0..gpr {
11289 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11290 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11291 let mut c = [0u8; 7];
11292 for k in 0..GROUP_SIZE {
11293 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11294 }
11295 bytes.extend_from_slice(&c);
11296 }
11297 }
11298 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
11300 let mut idx = 0usize;
11301 while idx < n {
11302 row_ptr[idx / cols + 1] += 1;
11303 idx += stride;
11304 }
11305 for r in 0..rows {
11306 row_ptr[r + 1] += row_ptr[r];
11307 }
11308 for &p in &row_ptr {
11309 bytes.extend_from_slice(&p.to_le_bytes());
11310 }
11311 let mut idx = 0usize;
11312 while idx < n {
11313 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11314 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11315 idx += stride;
11316 }
11317 let x: Vec<f32> = (0..cols)
11320 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11321 .collect();
11322 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11323
11324 let slow = |out: &mut [f32]| {
11326 let mut buf = vec![0f32; cols];
11327 for r in 0..rows {
11328 for g in 0..gpr {
11329 let off = (r * gpr + g) * Q1T_TILE;
11330 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11331 let codes = &bytes[off + 2..off + Q1T_TILE];
11332 for k in 0..GROUP_SIZE {
11333 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11334 1 => s,
11335 2 => -s,
11336 _ => 0.0,
11337 };
11338 }
11339 }
11340 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11341 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11342 }
11343 };
11344 let iters = 5;
11345 let mut a = vec![0f32; rows];
11346 slow(&mut a); let t = Instant::now();
11348 for _ in 0..iters {
11349 slow(&mut a);
11350 }
11351 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11352
11353 let mut b = vec![0f32; rows];
11354 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
11356 for _ in 0..iters {
11357 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11358 }
11359 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11360
11361 for r in 0..rows {
11362 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11363 }
11364 println!(
11365 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
11366 slow_ms / fast_ms
11367 );
11368 }
11369}
11370
11371#[cfg(test)]
11372mod gemm_bench {
11373 #[test]
11384 #[ignore]
11385 fn q4tp_matmat_throughput() {
11386 let b: usize = std::env::var("CMF_BENCH_B")
11390 .ok()
11391 .and_then(|v| v.parse().ok())
11392 .unwrap_or(296);
11393 let (rows, cols) = (9216usize, 2304usize);
11394 let (_, _, _) = (rows, cols, b);
11395 let total =
11396 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
11397 .unwrap();
11398 let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11403 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11404 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11405 let step = cortiq_core::quant::f32_to_f16(0.1);
11406 for r in 0..rows {
11407 let o = params_off + r * 4;
11408 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11409 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11410 }
11411 let _ = codes_off;
11412 let xs: Vec<f32> = (0..b * cols)
11413 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11414 .collect();
11415 let mut out = vec![0f32; b * rows];
11416 let pool = crate::pool::Pool::from_env();
11417 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11423 let reps: usize = std::env::var("CMF_BENCH_REPS")
11424 .ok()
11425 .and_then(|v| v.parse().ok())
11426 .unwrap_or(10);
11427 let mut best = [f64::MAX; 2];
11428 let mut sums = [0f32; 2];
11429 for _ in 0..reps {
11430 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11431 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11432 let t = std::time::Instant::now();
11433 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11434 best[k] = best[k].min(t.elapsed().as_secs_f64());
11435 sums[k] = out.iter().take(64).sum::<f32>();
11436 }
11437 }
11438 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11439 for (k, name) in ["previous", "tuned "].iter().enumerate() {
11440 println!(
11441 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11442 best[k] * 1e3,
11443 flops / best[k] / 1e9,
11444 sums[k]
11445 );
11446 }
11447 assert!(
11448 (sums[0] - sums[1]).abs() < 1e-2,
11449 "the tuned kernel changed the result: {} vs {}",
11450 sums[0],
11451 sums[1]
11452 );
11453 }
11454
11455 #[test]
11461 fn q4tp_matmat_blocked_matches_scalar() {
11462 use std::sync::atomic::Ordering::Relaxed;
11463 for &(rows, cols, b) in &[
11471 (64usize, 128usize, 7usize),
11472 (33, 96, 4),
11473 (16, 256, 9),
11474 (192, 2304, 37),
11475 ] {
11476 let total = cortiq_core::quant::expected_nbytes(
11477 cortiq_core::TensorDtype::Q4TiledP,
11478 &[rows, cols],
11479 )
11480 .unwrap();
11481 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11482 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11483 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11484 let step = cortiq_core::quant::f32_to_f16(0.1);
11485 for r in 0..rows {
11486 let o = params_off + r * 4;
11487 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11488 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11489 }
11490 let xs: Vec<f32> = (0..b * cols)
11491 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11492 .collect();
11493 let mut got = vec![0f32; b * rows];
11494 let mut want = vec![0f32; b * rows];
11495 let gpr = cols / 32;
11496 let view = super::Q4tpView::new(&bytes, rows, cols);
11497 let pool = crate::pool::Pool::from_env();
11498 super::Q4TP_ALT.store(2, Relaxed);
11499 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11500 super::Q4TP_ALT.store(1, Relaxed);
11501 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11502 super::Q4TP_ALT.store(0, Relaxed);
11503 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11510 let (mut worst, mut at) = (0f32, 0usize);
11511 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11512 if (g - w).abs() > worst {
11513 worst = (g - w).abs();
11514 at = i;
11515 }
11516 }
11517 assert!(
11518 worst <= 1e-4 * scale,
11519 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11520 (scale {scale:.3e}) at cell {at}: {} vs {}",
11521 got[at],
11522 want[at]
11523 );
11524
11525 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11532 for bi in 0..b {
11533 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11534 for r in 0..rows {
11535 let mut sc = vec![0f32; gpr];
11536 view.scales_into(r, gpr, &mut sc);
11537 let mut exact = 0f64;
11538 for j in 0..cols {
11539 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11540 exact += w as f64 * sq as f64 * act.xq[j] as f64;
11541 }
11542 exact *= act.sx as f64;
11543 for &(j, xv) in &act.outliers {
11544 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11545 exact += w as f64 * sq as f64 * xv as f64;
11546 }
11547 let i = bi * rows + r;
11548 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11549 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11550 }
11551 }
11552 println!(
11553 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11554 per-column {e_scalar:.3e}"
11555 );
11556 assert!(
11561 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11562 "{rows}x{cols} b={b}: error against f64 too large — blocked \
11563 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11564 );
11565 }
11566 }
11567
11568 #[test]
11572 #[ignore]
11573 fn q4t_matmat_throughput() {
11574 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11575 let total =
11576 cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
11577 .unwrap();
11578 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11581 let sc = cortiq_core::quant::f32_to_f16(0.02);
11582 for t in bytes.chunks_mut(super::Q4_TILE) {
11583 t[..2].copy_from_slice(&sc.to_le_bytes());
11584 }
11585 let xs: Vec<f32> = (0..b * cols)
11586 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11587 .collect();
11588 let mut out = vec![0f32; b * rows];
11589 let pool = crate::pool::Pool::from_env();
11590 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11591 let reps: usize = std::env::var("CMF_BENCH_REPS")
11592 .ok()
11593 .and_then(|v| v.parse().ok())
11594 .unwrap_or(10);
11595 let mut best = f64::MAX;
11596 for _ in 0..reps {
11597 let t = std::time::Instant::now();
11598 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11599 best = best.min(t.elapsed().as_secs_f64());
11600 }
11601 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11602 println!(
11603 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11604 best * 1e3,
11605 flops / best / 1e9,
11606 out.iter().take(64).sum::<f32>()
11607 );
11608 }
11609}