1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16 GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17 q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23 F32 {
24 data: Vec<f32>,
25 rows: usize,
26 cols: usize,
27 },
28 Mapped {
29 model: Arc<CmfModel>,
30 idx: usize,
32 dtype: TensorDtype,
33 rows: usize,
34 cols: usize,
35 row_scale: Vec<f32>,
37 col_field: Vec<f32>,
39 vbit_offsets: Vec<usize>,
43 repack: Vec<u8>,
51 },
52}
53
54fn repack_enabled() -> bool {
61 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62 *ON.get_or_init(|| {
63 std::env::var("CMF_REPACK")
64 .map(|v| v == "1")
65 .unwrap_or(cfg!(target_os = "android"))
66 })
67}
68
69fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73 #[cfg(target_arch = "aarch64")]
74 let arch_ok = sdot_enabled();
75 #[cfg(not(target_arch = "aarch64"))]
76 let arch_ok = false;
77 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78 return Vec::new();
79 }
80 q8_repack_layout(bytes, rows, cols)
81}
82
83fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86 let groups = rows / 4;
87 let mut rep = vec![0u8; groups * 4 * cols];
88 for g in 0..groups {
89 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90 for c in 0..cols / 16 {
91 for lane in 0..4 {
92 let src = (g * 4 + lane) * cols + c * 16;
93 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94 .copy_from_slice(&bytes[src..src + 16]);
95 }
96 }
97 }
98 rep
99}
100
101fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104 let ng = cols / GROUP_SIZE;
105 let bits = &bytes[..rows];
106 let mut offsets = Vec::with_capacity(rows + 1);
107 let mut off = rows + rows * ng * 2;
108 for r in 0..rows {
109 offsets.push(off);
110 off += (cols * bits[r] as usize).div_ceil(8);
111 }
112 offsets.push(off);
113 offsets
114}
115
116fn blocked_enabled() -> bool {
122 use std::sync::atomic::Ordering::Relaxed;
123 match BLOCKED_OVERRIDE.load(Relaxed) {
124 1 => false,
125 2 => true,
126 _ => {
127 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128 *ON.get_or_init(|| {
129 std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
130 })
131 }
132 }
133}
134
135static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
136
137pub fn set_blocked_override(on: Option<bool>) {
144 let v = match on {
145 None => 0,
146 Some(false) => 1,
147 Some(true) => 2,
148 };
149 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
150}
151
152fn gpu_lmhead_enabled() -> bool {
153 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
154 *ON.get_or_init(|| std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true))
155}
156
157fn gpu_split_frac() -> f32 {
158 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
159 *F.get_or_init(|| {
160 std::env::var("CMF_GPU_SPLIT")
161 .ok()
162 .and_then(|v| v.parse::<f32>().ok())
163 .unwrap_or(0.5)
164 .clamp(0.0, 1.0)
165 })
166}
167
168impl QTensor {
169 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
170 debug_assert_eq!(data.len(), rows * cols);
171 Self::F32 { data, rows, cols }
172 }
173
174 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
177 let idx = model
180 .tensor_index(name)
181 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
182 let entry = &model.tensors[idx];
183 if entry.shape.len() != 2 {
184 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
185 }
186 let (rows, cols) = (entry.shape[0], entry.shape[1]);
187 let bytes = model.entry_bytes(entry);
188
189 match entry.dtype {
190 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
191 let n = rows * cols;
192 let scales_off = n;
193 let row_scale: Vec<f32> = (0..rows)
194 .map(|o| {
195 f16_to_f32(u16::from_le_bytes([
196 bytes[scales_off + o * 2],
197 bytes[scales_off + o * 2 + 1],
198 ]))
199 })
200 .collect();
201 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
202 let col_off = n + rows * 2;
203 (0..cols)
204 .map(|i| {
205 f16_to_f32(u16::from_le_bytes([
206 bytes[col_off + i * 2],
207 bytes[col_off + i * 2 + 1],
208 ]))
209 })
210 .collect()
211 } else {
212 Vec::new()
213 };
214 Ok(Self::Mapped {
215 model: model.clone(),
216 idx,
217 dtype: entry.dtype,
218 rows,
219 cols,
220 row_scale,
221 col_field,
222 vbit_offsets: Vec::new(),
223 repack: q8_repack(bytes, rows, cols),
224 })
225 }
226 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
228 model: model.clone(),
229 idx,
230 dtype: entry.dtype,
231 rows,
232 cols,
233 row_scale: Vec::new(),
234 col_field: Vec::new(),
235 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
236 repack: Vec::new(),
237 }),
238 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
242 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
243 let offsets: Vec<usize> = (0..=rows)
244 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
245 .collect();
246 Ok(Self::Mapped {
247 model: model.clone(),
248 idx,
249 dtype: entry.dtype,
250 rows,
251 cols,
252 row_scale: Vec::new(),
253 col_field: Vec::new(),
254 vbit_offsets: offsets,
255 repack: Vec::new(),
256 })
257 }
258 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
264 model: model.clone(),
265 idx,
266 dtype: entry.dtype,
267 rows,
268 cols,
269 row_scale: Vec::new(),
270 col_field: Vec::new(),
271 vbit_offsets: Vec::new(),
272 repack: Vec::new(),
273 }),
274 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
277 model: model.clone(),
278 idx,
279 dtype: entry.dtype,
280 rows,
281 cols,
282 row_scale: Vec::new(),
283 col_field: Vec::new(),
284 vbit_offsets: Vec::new(),
285 repack: Vec::new(),
286 }),
287 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
289 model: model.clone(),
290 idx,
291 dtype: entry.dtype,
292 rows,
293 cols,
294 row_scale: Vec::new(),
295 col_field: Vec::new(),
296 vbit_offsets: Vec::new(),
297 repack: Vec::new(),
298 }),
299 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
300 model: model.clone(),
301 idx,
302 dtype: entry.dtype,
303 rows,
304 cols,
305 row_scale: Vec::new(),
306 col_field: Vec::new(),
307 vbit_offsets: Vec::new(),
308 repack: Vec::new(),
309 }),
310 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
312 model: model.clone(),
313 idx,
314 dtype: entry.dtype,
315 rows,
316 cols,
317 row_scale: Vec::new(),
318 col_field: Vec::new(),
319 vbit_offsets: Vec::new(),
320 repack: Vec::new(),
321 }),
322 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
326 model: model.clone(),
327 idx,
328 dtype: entry.dtype,
329 rows,
330 cols,
331 row_scale: Vec::new(),
332 col_field: Vec::new(),
333 vbit_offsets: Vec::new(),
334 repack: Vec::new(),
335 }),
336 _ => {
338 let mut data = vec![0.0f32; rows * cols];
339 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
340 Ok(Self::from_f32(data, rows, cols))
341 }
342 }
343 }
344
345 pub(crate) fn is_q1(&self) -> bool {
348 matches!(
349 self,
350 Self::Mapped {
351 dtype: TensorDtype::Q1,
352 ..
353 }
354 )
355 }
356
357 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
360 match self {
361 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
362 _ => None,
363 }
364 }
365
366 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
373 match self {
374 #[cfg(target_os = "macos")]
375 Self::Mapped {
376 dtype: TensorDtype::Q1T,
377 ..
378 } if !crate::gpu::metal_q1t_enabled() => None,
379 Self::Mapped {
380 idx,
381 dtype:
382 TensorDtype::Q1
383 | TensorDtype::Q1T
384 | TensorDtype::Q4Block
385 | TensorDtype::Q4Tiled
386 | TensorDtype::Q4TiledP
387 | TensorDtype::Q2TiledP
388 | TensorDtype::Q8Row
389 | TensorDtype::Q8_2f,
390 rows,
391 cols,
392 ..
393 } => Some((*idx, *rows, *cols)),
394 _ => None,
395 }
396 }
397
398 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
404 match self {
405 Self::Mapped {
406 idx,
407 dtype: TensorDtype::Q4Tiled,
408 rows,
409 cols,
410 ..
411 } => Some((*idx, *rows, *cols)),
412 _ => None,
413 }
414 }
415
416 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
420 match self {
421 Self::Mapped {
422 idx,
423 dtype: TensorDtype::Q4TiledP,
424 rows,
425 cols,
426 ..
427 } => Some((*idx, *rows, *cols)),
428 _ => None,
429 }
430 }
431
432 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
437 match self {
438 Self::Mapped {
439 idx,
440 dtype: TensorDtype::Q8Row,
441 rows,
442 cols,
443 row_scale,
444 col_field,
445 ..
446 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
447 _ => None,
448 }
449 }
450
451 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
455 match self {
456 Self::Mapped { dtype, .. } => Some(*dtype),
457 _ => None,
458 }
459 }
460
461 pub fn model_idx(&self) -> Option<usize> {
466 match self {
467 Self::Mapped { idx, .. } => Some(*idx),
468 _ => None,
469 }
470 }
471
472 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
477 match self {
478 Self::Mapped { model, .. } => Some(model.clone()),
479 _ => None,
480 }
481 }
482
483 pub fn rows(&self) -> usize {
484 match self {
485 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
486 }
487 }
488
489 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
492 match self {
493 Self::Mapped {
494 model,
495 idx,
496 dtype: TensorDtype::Q4Tiled,
497 ..
498 } => Some((model, *idx)),
499 _ => None,
500 }
501 }
502
503 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
506 match self {
507 Self::Mapped {
508 model,
509 idx,
510 dtype: TensorDtype::Q4TiledP,
511 ..
512 } => Some((model, *idx)),
513 _ => None,
514 }
515 }
516
517 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
520 match self {
521 Self::Mapped {
522 model,
523 idx,
524 dtype: TensorDtype::Q2TiledP,
525 ..
526 } => Some((model, *idx)),
527 _ => None,
528 }
529 }
530
531 pub fn cols(&self) -> usize {
532 match self {
533 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
534 }
535 }
536
537 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
540 match self {
541 Self::Mapped {
542 model,
543 idx,
544 dtype: TensorDtype::Q1,
545 ..
546 } => Some((model, *idx)),
547 _ => None,
548 }
549 }
550
551 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
555 match self {
556 Self::Mapped {
557 model,
558 idx,
559 dtype: TensorDtype::Q8Row,
560 row_scale,
561 ..
562 } => Some((model, *idx, 0, row_scale.as_slice())),
563 Self::Mapped {
564 model,
565 idx,
566 dtype: TensorDtype::Q1,
567 ..
568 } => Some((model, *idx, 1, &[])),
569 Self::Mapped {
574 model,
575 idx,
576 dtype: TensorDtype::Q4Tiled,
577 ..
578 } => Some((model, *idx, 5, &[])),
579 Self::Mapped {
583 model,
584 idx,
585 dtype: TensorDtype::Q4TiledP,
586 ..
587 } => Some((model, *idx, 6, &[])),
588 Self::Mapped {
589 model,
590 idx,
591 dtype: TensorDtype::Q4Block,
592 ..
593 } => Some((model, *idx, 2, &[])),
594 Self::Mapped {
595 model,
596 idx,
597 dtype: TensorDtype::Q1T,
598 ..
599 } => Some((model, *idx, 3, &[])),
600 _ => None,
601 }
602 }
603
604 pub fn as_f32(&self) -> Option<&[f32]> {
607 match self {
608 Self::F32 { data, .. } => Some(data),
609 Self::Mapped { .. } => None,
610 }
611 }
612
613 fn quant_bytes(&self) -> &[u8] {
614 match self {
615 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
616 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
617 }
618 }
619
620 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
622 let cols = self.cols();
623 debug_assert_eq!(dst.len(), cols);
624 match self {
625 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
626 Self::Mapped {
627 dtype,
628 row_scale,
629 col_field,
630 vbit_offsets,
631 ..
632 } => {
633 if *dtype == TensorDtype::Q4Tiled {
634 let bytes = self.quant_bytes();
635 let gpr = cols / GROUP_SIZE;
636 for gi in 0..gpr {
637 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
638 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
639 for (k, &b) in tile[2..].iter().enumerate() {
640 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
641 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
642 }
643 }
644 return;
645 }
646 if *dtype == TensorDtype::Q4TiledP {
647 let bytes = self.quant_bytes();
648 let gpr = cols / GROUP_SIZE;
649 let v = Q4tpView::new(bytes, self.rows(), cols);
650 let mut sc = vec![0f32; gpr];
651 v.scales_into(r, gpr, &mut sc);
652 for gi in 0..gpr {
653 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
654 let s = sc[gi];
655 for (k, &b) in tile.iter().enumerate() {
656 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
657 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
658 }
659 }
660 return;
661 }
662 if *dtype == TensorDtype::Q2TiledP {
663 let bytes = self.quant_bytes();
664 let gpr = cols / GROUP_SIZE;
665 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
666 let mut sc = vec![0f32; gpr];
667 v.scales_into(r, gpr, &mut sc);
668 for gi in 0..gpr {
669 let ch =
670 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
671 let s = sc[gi];
672 for (k, &b) in ch.iter().enumerate() {
673 for j in 0..4 {
674 dst[gi * GROUP_SIZE + k * 4 + j] =
675 (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
676 }
677 }
678 }
679 return;
680 }
681 if *dtype == TensorDtype::Q4Block {
682 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
683 let gpr = cols / GROUP_SIZE;
684 for gi in 0..gpr {
685 let g = r * gpr + gi;
686 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
687 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
688 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
689 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
690 }
691 }
692 return;
693 }
694 if *dtype == TensorDtype::Q1 {
695 let bytes = self.quant_bytes();
696 let gpr = cols / GROUP_SIZE;
697 for gi in 0..gpr {
698 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
699 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
700 for (j, &b) in tile[2..].iter().enumerate() {
701 for k in 0..8 {
702 dst[gi * GROUP_SIZE + j * 8 + k] =
703 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
704 }
705 }
706 }
707 return;
708 }
709 if *dtype == TensorDtype::Q1T {
710 let bytes = self.quant_bytes();
711 let gpr = cols / GROUP_SIZE;
712 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
713 for gi in 0..gpr {
714 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
715 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
716 bytes[off],
717 bytes[off + 1],
718 ]));
719 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
720 for k in 0..GROUP_SIZE {
721 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
722 {
723 1 => s,
724 2 => -s,
725 _ => 0.0,
726 };
727 }
728 }
729 let rows = self.rows();
731 let entries = base_len + (rows + 1) * 4;
732 if entries <= bytes.len() {
733 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
734 let r0 = u32::from_le_bytes([
735 ptrs[r * 4],
736 ptrs[r * 4 + 1],
737 ptrs[r * 4 + 2],
738 ptrs[r * 4 + 3],
739 ]) as usize;
740 let r1 = u32::from_le_bytes([
741 ptrs[(r + 1) * 4],
742 ptrs[(r + 1) * 4 + 1],
743 ptrs[(r + 1) * 4 + 2],
744 ptrs[(r + 1) * 4 + 3],
745 ]) as usize;
746 let off = entries + r0 * 4;
747 for i in 0..r1 - r0 {
748 let item = &bytes[off + i * 4..off + i * 4 + 4];
749 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
750 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
751 item[2], item[3],
752 ]));
753 if c < cols {
754 dst[c] = v;
755 }
756 }
757 }
758 return;
759 }
760 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
761 let bytes = self.quant_bytes();
762 let rows = self.rows();
763 let ng = cols / GROUP_SIZE;
764 let bits = &bytes[..rows];
765 let sc_off = rows;
766 let off = vbit_offsets[r];
769 let b = bits[r] as usize;
770 let l = ((1usize << (b - 1)) - 1) as f32;
771 let data = &bytes[off..];
772 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
773 for (i, d) in dst.iter_mut().enumerate() {
774 while nbits < b {
775 acc = (acc << 8) | data[idx] as u64;
776 idx += 1;
777 nbits += 8;
778 }
779 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
780 nbits -= b;
781 let so = (r * ng + i / GROUP_SIZE) * 2;
782 let sv = f16_to_f32(u16::from_le_bytes([
783 bytes[sc_off + so],
784 bytes[sc_off + so + 1],
785 ]));
786 *d = (u - l) * sv;
787 }
788 return;
789 }
790 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
791 let s = row_scale[r];
792 match dtype {
793 TensorDtype::Q8Row => {
794 for (d, &b) in dst.iter_mut().zip(q) {
795 *d = (b as i8) as f32 * s;
796 }
797 }
798 TensorDtype::Q8_2f => {
799 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
800 *d = (b as i8) as f32 * s * col_field[i];
801 }
802 }
803 _ => unreachable!(),
804 }
805 }
806 }
807 }
808
809 pub fn sparse_col_ok(&self) -> bool {
814 match self {
815 Self::F32 { .. } => true,
816 Self::Mapped { dtype, .. } => {
817 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
818 }
819 }
820 }
821
822 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
826 let inter = self.cols();
827 let hidden = self.rows();
828 debug_assert_eq!(out.len(), hidden);
829 match self {
830 Self::F32 { data, .. } => {
831 for (k, o) in out.iter_mut().enumerate() {
832 *o += w * data[k * inter + c];
833 }
834 }
835 Self::Mapped {
836 dtype,
837 row_scale,
838 col_field,
839 ..
840 } => {
841 let q = self.quant_bytes();
842 let colf = if *dtype == TensorDtype::Q8_2f {
843 col_field[c]
844 } else {
845 1.0
846 };
847 let wc = w * colf;
848 for (k, o) in out.iter_mut().enumerate() {
849 let b = q[k * inter + c] as i8 as f32;
850 *o += wc * b * row_scale[k];
851 }
852 }
853 }
854 }
855
856 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
860 let cols = self.cols();
861 match self {
862 Self::F32 { data, .. } => {
863 let row = &data[r * cols..(r + 1) * cols];
864 row.iter().zip(x).map(|(w, v)| w * v).sum()
865 }
866 Self::Mapped {
867 dtype,
868 row_scale,
869 col_field,
870 ..
871 } => match dtype {
872 TensorDtype::Q8Row => {
873 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
874 dot_i8_f32(q, x) * row_scale[r]
875 }
876 TensorDtype::Q8_2f => {
877 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
878 dot_i8_col_f32(q, x, col_field) * row_scale[r]
879 }
880 _ => {
881 self.row_f32(r, scratch);
882 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
883 }
884 },
885 }
886 }
887
888 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
891 match self {
892 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
896 Self::Mapped {
897 model,
898 idx,
899 dtype,
900 rows,
901 cols,
902 row_scale,
903 col_field,
904 vbit_offsets,
905 repack,
906 } => {
907 let _ = (model, idx);
908 assert!(
918 out.len() >= *rows && x.len() >= *cols,
919 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
920 out.len(),
921 x.len(),
922 );
923 if *dtype == TensorDtype::Q4Block {
924 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
928 let t0 = std::time::Instant::now();
929 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
930 crate::gpu::ProbeArm::Gpu => {
931 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
932 crate::gpu::probe_record(
933 crate::gpu::OpClass::Matvec,
934 true,
935 t0.elapsed(),
936 );
937 return;
938 }
939 }
940 crate::gpu::ProbeArm::CpuTimed => {
941 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
942 crate::gpu::probe_record(
943 crate::gpu::OpClass::Matvec,
944 false,
945 t0.elapsed(),
946 );
947 return;
948 }
949 crate::gpu::ProbeArm::Cpu => {}
950 }
951 }
952 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
953 return;
954 }
955 if *dtype == TensorDtype::Q4Tiled {
956 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
957 return;
958 }
959 if *dtype == TensorDtype::Q4TiledP {
960 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
966 let t0 = std::time::Instant::now();
967 let cls = crate::gpu::matvec_class(*rows, *cols);
968 match crate::gpu::probe_arm(cls) {
969 crate::gpu::ProbeArm::Gpu => {
970 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
971 crate::gpu::probe_record(cls, true, t0.elapsed());
972 return;
973 }
974 }
975 crate::gpu::ProbeArm::CpuTimed => {
976 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
977 crate::gpu::probe_record(cls, false, t0.elapsed());
978 return;
979 }
980 crate::gpu::ProbeArm::Cpu => {}
981 }
982 }
983 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
984 return;
985 }
986 if *dtype == TensorDtype::Q2TiledP {
987 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
988 return;
989 }
990 if *dtype == TensorDtype::Q1 {
991 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
996 let t0 = std::time::Instant::now();
997 let arm = if crate::gpu::q1_force() {
998 crate::gpu::ProbeArm::Gpu
999 } else {
1000 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1001 };
1002 match arm {
1003 crate::gpu::ProbeArm::Gpu => {
1004 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1005 crate::gpu::probe_record(
1006 crate::gpu::OpClass::Matvec,
1007 true,
1008 t0.elapsed(),
1009 );
1010 return;
1011 }
1012 }
1013 crate::gpu::ProbeArm::CpuTimed => {
1014 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1015 crate::gpu::probe_record(
1016 crate::gpu::OpClass::Matvec,
1017 false,
1018 t0.elapsed(),
1019 );
1020 return;
1021 }
1022 crate::gpu::ProbeArm::Cpu => {}
1023 }
1024 }
1025 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1026 return;
1027 }
1028 if *dtype == TensorDtype::Q1T {
1029 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1033 let t0 = std::time::Instant::now();
1034 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1035 crate::gpu::ProbeArm::Gpu => {
1036 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1037 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1038 crate::gpu::probe_record(
1039 crate::gpu::OpClass::Matvec,
1040 true,
1041 t0.elapsed(),
1042 );
1043 return;
1044 }
1045 }
1046 crate::gpu::ProbeArm::CpuTimed => {
1047 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1048 crate::gpu::probe_record(
1049 crate::gpu::OpClass::Matvec,
1050 false,
1051 t0.elapsed(),
1052 );
1053 return;
1054 }
1055 crate::gpu::ProbeArm::Cpu => {}
1056 }
1057 }
1058 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1059 return;
1060 }
1061 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1062 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1063 return;
1064 }
1065 let xs = prescale(x, col_field, *dtype);
1066 if *rows >= crate::gpu::min_rows()
1071 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1072 && gpu_lmhead_enabled()
1073 && crate::gpu::enabled_here()
1074 {
1075 let t0 = std::time::Instant::now();
1078 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1079 crate::gpu::ProbeArm::Gpu => {}
1080 crate::gpu::ProbeArm::CpuTimed => {
1081 qmatvec(
1082 self.quant_bytes(),
1083 repack,
1084 row_scale,
1085 x,
1086 col_field,
1087 *dtype,
1088 *rows,
1089 *cols,
1090 out,
1091 pool,
1092 );
1093 crate::gpu::probe_record(
1094 crate::gpu::OpClass::Matvec,
1095 false,
1096 t0.elapsed(),
1097 );
1098 return;
1099 }
1100 crate::gpu::ProbeArm::Cpu => {
1101 qmatvec(
1102 self.quant_bytes(),
1103 repack,
1104 row_scale,
1105 x,
1106 col_field,
1107 *dtype,
1108 *rows,
1109 *cols,
1110 out,
1111 pool,
1112 );
1113 return;
1114 }
1115 }
1116 let frac = gpu_split_frac();
1117 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1118 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1119 let bytes = self.quant_bytes();
1120 let ok = std::thread::scope(|sc| {
1121 let g = sc.spawn(|| {
1122 crate::gpu::q8_matvec_range(
1123 model,
1124 *idx,
1125 cpu_rows,
1126 &row_scale[cpu_rows..],
1127 &xs,
1128 *rows - cpu_rows,
1129 *cols,
1130 out_gpu,
1131 )
1132 });
1133 if cpu_rows > 0 {
1134 let rep_cpu = if repack.is_empty() {
1137 &[][..]
1138 } else {
1139 &repack[..(cpu_rows / 4) * 4 * *cols]
1140 };
1141 qmatvec(
1142 &bytes[..cpu_rows * *cols],
1143 rep_cpu,
1144 &row_scale[..cpu_rows],
1145 x,
1146 col_field,
1147 *dtype,
1148 cpu_rows,
1149 *cols,
1150 out_cpu,
1151 pool,
1152 );
1153 }
1154 g.join().unwrap_or(false)
1155 });
1156 if ok {
1157 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1158 return;
1159 }
1160 qmatvec(
1163 &bytes[cpu_rows * *cols..(*rows) * *cols],
1164 &[],
1165 &row_scale[cpu_rows..],
1166 x,
1167 col_field,
1168 *dtype,
1169 *rows - cpu_rows,
1170 *cols,
1171 out_gpu,
1172 pool,
1173 );
1174 return;
1175 }
1176 qmatvec(
1177 self.quant_bytes(),
1178 repack,
1179 row_scale,
1180 x,
1181 col_field,
1182 *dtype,
1183 *rows,
1184 *cols,
1185 out,
1186 pool,
1187 );
1188 }
1189 }
1190 }
1191
1192 pub fn matvec2(
1194 &self,
1195 x1: &[f32],
1196 x2: &[f32],
1197 o1: &mut [f32],
1198 o2: &mut [f32],
1199 pool: Option<&Pool>,
1200 ) {
1201 match self {
1202 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1203 Self::Mapped {
1204 dtype,
1205 rows,
1206 cols,
1207 row_scale,
1208 col_field,
1209 vbit_offsets,
1210 ..
1211 } => {
1212 if *dtype == TensorDtype::Q4Block {
1213 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1214 return;
1215 }
1216 if *dtype == TensorDtype::Q4Tiled {
1217 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1218 return;
1219 }
1220 if *dtype == TensorDtype::Q4TiledP {
1221 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1222 return;
1223 }
1224 if *dtype == TensorDtype::Q2TiledP {
1225 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1226 return;
1227 }
1228 if *dtype == TensorDtype::Q1 {
1229 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1230 return;
1231 }
1232 if *dtype == TensorDtype::Q1T {
1233 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1239 return;
1240 }
1241 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1242 vbitmatvec2(
1243 self.quant_bytes(),
1244 vbit_offsets,
1245 x1,
1246 x2,
1247 *rows,
1248 *cols,
1249 o1,
1250 o2,
1251 pool,
1252 );
1253 return;
1254 }
1255 qmatvec2(
1256 self.quant_bytes(),
1257 row_scale,
1258 x1,
1259 x2,
1260 col_field,
1261 *dtype,
1262 *rows,
1263 *cols,
1264 o1,
1265 o2,
1266 pool,
1267 );
1268 }
1269 }
1270 }
1271}
1272
1273impl QTensor {
1274 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1279 let cols = self.cols();
1280 let rows = self.rows();
1281 debug_assert_eq!(xs_all.len(), b * cols);
1282 debug_assert_eq!(out.len(), b * rows);
1283 if crate::gptq_capture::capturing() {
1287 if let Self::Mapped { model, idx, .. } = self {
1288 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1289 }
1290 }
1291 match self {
1292 Self::F32 { data, .. } => {
1293 let out_addr = SendMut(out.as_mut_ptr());
1294 let run = |start: usize, end: usize| {
1295 for o in start..end {
1296 let row = &data[o * cols..(o + 1) * cols];
1297 for bi in 0..b {
1298 let x = &xs_all[bi * cols..(bi + 1) * cols];
1299 let mut acc = 0f32;
1300 for j in 0..cols {
1301 acc += row[j] * x[j];
1302 }
1303 unsafe { *out_addr.at(bi * rows + o) = acc };
1304 }
1305 }
1306 };
1307 dispatch_rows(pool, rows, &run);
1308 }
1309 Self::Mapped {
1310 dtype,
1311 row_scale,
1312 col_field,
1313 vbit_offsets,
1314 ..
1315 } => {
1316 if *dtype == TensorDtype::Q4Block {
1317 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1318 return;
1319 }
1320 if *dtype == TensorDtype::Q4TiledP {
1321 if b >= 32
1333 && b * rows * cols >= 128_000_000
1334 && cols % 32 == 0
1335 && !crate::gpu::mm_killed()
1336 && crate::gpu::enabled_here()
1337 {
1338 let class = if b >= 128 {
1339 crate::gpu::OpClass::MatmatWide
1340 } else {
1341 crate::gpu::OpClass::Matmat
1342 };
1343 if let Self::Mapped { model, idx, .. } = self {
1344 let t0 = std::time::Instant::now();
1345 match crate::gpu::probe_arm(class) {
1346 crate::gpu::ProbeArm::Gpu => {
1347 if crate::gpu::q4tp_matmat(
1348 model, *idx, xs_all, b, rows, cols, out,
1349 ) {
1350 let el = t0.elapsed();
1351 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1361 let budget = std::time::Duration::from_secs_f64(
1362 flops / 1.5e12 * 8.0 + 0.020,
1363 );
1364 if el > budget && !crate::gpu::probe_was_cold() {
1365 tracing::warn!(
1366 "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1367 device contended, CPU for the rest of the process"
1368 );
1369 crate::gpu::mm_kill();
1370 }
1371 crate::gpu::probe_record(class, true, el);
1372 return;
1373 }
1374 }
1375 crate::gpu::ProbeArm::CpuTimed => {
1376 q4tp_matmat(
1377 self.quant_bytes(),
1378 xs_all,
1379 b,
1380 rows,
1381 cols,
1382 out,
1383 pool,
1384 );
1385 crate::gpu::probe_record(class, false, t0.elapsed());
1386 return;
1387 }
1388 crate::gpu::ProbeArm::Cpu => {}
1389 }
1390 }
1391 }
1392 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1393 return;
1394 }
1395 if *dtype == TensorDtype::Q2TiledP {
1396 if b >= 32
1402 && b * rows * cols >= 128_000_000
1403 && cols % 32 == 0
1404 && !crate::gpu::mm_killed()
1405 && crate::gpu::enabled_here()
1406 {
1407 let class = if b >= 128 {
1408 crate::gpu::OpClass::MatmatWide
1409 } else {
1410 crate::gpu::OpClass::Matmat
1411 };
1412 if let Self::Mapped { model, idx, .. } = self {
1413 let t0 = std::time::Instant::now();
1414 match crate::gpu::probe_arm(class) {
1415 crate::gpu::ProbeArm::Gpu => {
1416 if crate::gpu::q2tp_matmat(
1417 model, *idx, xs_all, b, rows, cols, out,
1418 ) {
1419 crate::gpu::probe_record(class, true, t0.elapsed());
1420 return;
1421 }
1422 }
1423 crate::gpu::ProbeArm::CpuTimed => {
1424 q2tp_matmat(
1425 self.quant_bytes(),
1426 xs_all,
1427 b,
1428 rows,
1429 cols,
1430 out,
1431 pool,
1432 );
1433 crate::gpu::probe_record(class, false, t0.elapsed());
1434 return;
1435 }
1436 crate::gpu::ProbeArm::Cpu => {}
1437 }
1438 }
1439 }
1440 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1445 return;
1446 }
1447 if *dtype == TensorDtype::Q4Tiled {
1448 if b >= 32
1460 && b * rows * cols >= 128_000_000
1461 && cols % 32 == 0
1462 && !crate::gpu::mm_killed()
1463 && crate::gpu::enabled_here()
1464 {
1465 let class = if b >= 128 {
1466 crate::gpu::OpClass::MatmatWide
1467 } else {
1468 crate::gpu::OpClass::Matmat
1469 };
1470 if let Self::Mapped { model, idx, .. } = self {
1471 let t0 = std::time::Instant::now();
1472 match crate::gpu::probe_arm(class) {
1473 crate::gpu::ProbeArm::Gpu => {
1474 if crate::gpu::q4t_matmat(
1475 model, *idx, xs_all, b, rows, cols, out,
1476 ) {
1477 let el = t0.elapsed();
1478 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1488 let budget = std::time::Duration::from_secs_f64(
1489 flops / 1.5e12 * 8.0 + 0.020,
1490 );
1491 if el > budget && !crate::gpu::probe_was_cold() {
1492 tracing::warn!(
1493 "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1494 device contended, CPU for the rest of the process"
1495 );
1496 crate::gpu::mm_kill();
1497 }
1498 crate::gpu::probe_record(class, true, el);
1499 return;
1500 }
1501 }
1502 crate::gpu::ProbeArm::CpuTimed => {
1503 q4t_matmat(
1504 self.quant_bytes(),
1505 xs_all,
1506 b,
1507 rows,
1508 cols,
1509 out,
1510 pool,
1511 );
1512 crate::gpu::probe_record(class, false, t0.elapsed());
1513 return;
1514 }
1515 crate::gpu::ProbeArm::Cpu => {}
1516 }
1517 }
1518 }
1519 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1520 return;
1521 }
1522 if *dtype == TensorDtype::Q1 {
1523 if b >= 32
1526 && b * rows * cols >= 128_000_000
1527 && cols % 64 == 0
1528 && crate::gpu::enabled_here()
1529 {
1530 if let Self::Mapped { model, idx, .. } = self {
1531 let t0 = std::time::Instant::now();
1532 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1533 crate::gpu::ProbeArm::Gpu => {
1534 if crate::gpu::q1_matmat(
1535 model, *idx, xs_all, b, rows, cols, out,
1536 ) {
1537 crate::gpu::probe_record(
1538 crate::gpu::OpClass::Matmat,
1539 true,
1540 t0.elapsed(),
1541 );
1542 return;
1543 }
1544 }
1545 crate::gpu::ProbeArm::CpuTimed => {
1546 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1547 crate::gpu::probe_record(
1548 crate::gpu::OpClass::Matmat,
1549 false,
1550 t0.elapsed(),
1551 );
1552 return;
1553 }
1554 crate::gpu::ProbeArm::Cpu => {}
1555 }
1556 }
1557 }
1558 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1559 return;
1560 }
1561 if *dtype == TensorDtype::Q1T {
1562 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1565 if let Self::Mapped { model, idx, .. } = self {
1566 let t0 = std::time::Instant::now();
1567 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1568 crate::gpu::ProbeArm::Gpu => {
1569 if crate::gpu::q1t_matmat(
1570 model, *idx, xs_all, b, rows, cols, out,
1571 ) {
1572 crate::gpu::probe_record(
1573 crate::gpu::OpClass::Matmat,
1574 true,
1575 t0.elapsed(),
1576 );
1577 return;
1578 }
1579 }
1580 crate::gpu::ProbeArm::CpuTimed => {
1581 q1t_matmat(
1582 self.quant_bytes(),
1583 xs_all,
1584 b,
1585 rows,
1586 cols,
1587 out,
1588 pool,
1589 );
1590 crate::gpu::probe_record(
1591 crate::gpu::OpClass::Matmat,
1592 false,
1593 t0.elapsed(),
1594 );
1595 return;
1596 }
1597 crate::gpu::ProbeArm::Cpu => {}
1598 }
1599 }
1600 }
1601 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1602 return;
1603 }
1604 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1605 vbitmatmat(
1606 self.quant_bytes(),
1607 vbit_offsets,
1608 xs_all,
1609 b,
1610 rows,
1611 cols,
1612 out,
1613 pool,
1614 );
1615 return;
1616 }
1617 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1618 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1619 .collect();
1620 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1626 if let Self::Mapped { model, idx, .. } = self {
1627 let t0 = std::time::Instant::now();
1628 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1629 crate::gpu::ProbeArm::Gpu
1630 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1631 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1632 {
1633 let q = self.quant_bytes();
1637 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1638 return;
1639 }
1640 crate::gpu::ProbeArm::Gpu => {
1641 let flat: Vec<f32> =
1642 pre.iter().flat_map(|v| v.iter().copied()).collect();
1643 if crate::gpu::q8_matmat(
1644 model, *idx, row_scale, &flat, b, rows, cols, out,
1645 ) {
1646 crate::gpu::probe_record(
1647 crate::gpu::OpClass::Matmat,
1648 true,
1649 t0.elapsed(),
1650 );
1651 return;
1652 }
1653 }
1654 crate::gpu::ProbeArm::CpuTimed => {
1655 let q = self.quant_bytes();
1656 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1657 crate::gpu::probe_record(
1658 crate::gpu::OpClass::Matmat,
1659 false,
1660 t0.elapsed(),
1661 );
1662 return;
1663 }
1664 crate::gpu::ProbeArm::Cpu => {}
1665 }
1666 }
1667 }
1668 let q = self.quant_bytes();
1669 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1670 }
1671 }
1672 }
1673}
1674
1675impl QTensor {
1676 pub fn matvec_many<const N: usize>(
1683 ts: [&QTensor; N],
1684 x: &[f32],
1685 mut outs: [&mut [f32]; N],
1686 pool: Option<&Pool>,
1687 ) {
1688 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1689 let uniform_q8 = ts.iter().all(|t| {
1690 matches!(
1691 t,
1692 Self::Mapped {
1693 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1694 ..
1695 }
1696 )
1697 });
1698 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1699 let uniform_q4 = ts.iter().all(|t| {
1700 matches!(
1701 t,
1702 Self::Mapped {
1703 dtype: TensorDtype::Q4Block,
1704 ..
1705 }
1706 )
1707 });
1708 let uniform_vbit = ts.iter().all(|t| {
1709 matches!(
1710 t,
1711 Self::Mapped {
1712 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1713 ..
1714 }
1715 )
1716 });
1717 let uniform_q1 = ts.iter().all(|t| {
1718 matches!(
1719 t,
1720 Self::Mapped {
1721 dtype: TensorDtype::Q1,
1722 ..
1723 }
1724 )
1725 });
1726 let uniform_q1t = ts.iter().all(|t| {
1727 matches!(
1728 t,
1729 Self::Mapped {
1730 dtype: TensorDtype::Q1T,
1731 ..
1732 }
1733 )
1734 });
1735 let uniform_q4tp = ts.iter().all(|t| {
1740 matches!(
1741 t,
1742 Self::Mapped {
1743 dtype: TensorDtype::Q4TiledP,
1744 ..
1745 }
1746 )
1747 }) && ts.iter().all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1748 let Some(pool) = pool else {
1749 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1750 t.matvec(x, o, None);
1751 }
1752 return;
1753 };
1754 if total_rows < 256
1755 || !(uniform_q8
1756 || uniform_f32
1757 || uniform_q4
1758 || uniform_vbit
1759 || uniform_q1
1760 || uniform_q1t
1761 || uniform_q4tp)
1762 {
1763 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1764 t.matvec(x, o, Some(pool));
1765 }
1766 return;
1767 }
1768
1769 if uniform_q4tp {
1770 let cols = ts[0].cols();
1776 let gpr = cols / GROUP_SIZE;
1777 let views: Vec<Q4tpView> = ts
1778 .iter()
1779 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1780 .collect();
1781 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1782 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1783 let locate = |flat: usize| -> (usize, usize) {
1785 let mut acc = 0;
1786 for (i, &r) in rows_of.iter().enumerate() {
1787 if flat < acc + r {
1788 return (i, flat - acc);
1789 }
1790 acc += r;
1791 }
1792 (rows_of.len() - 1, 0)
1793 };
1794 let (views, outs_addr) = (&views, &outs_addr);
1795 if a8w8_enabled() {
1796 let act = split_act(x);
1797 let act = &act;
1798 let run = |start: usize, end: usize| {
1799 let mut sc = vec![0f32; gpr];
1800 for flat in start..end {
1801 let (t, r) = locate(flat);
1802 let v = &views[t];
1803 v.scales_into(r, gpr, &mut sc);
1804 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1805 for &(j, xv) in &act.outliers {
1806 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1807 acc += w * s * xv;
1808 }
1809 unsafe { *outs_addr[t].at(r) = acc };
1811 }
1812 };
1813 pool.run_rows(total_rows, &run);
1814 } else {
1815 let run = |start: usize, end: usize| {
1816 let mut sc = vec![0f32; gpr];
1817 for flat in start..end {
1818 let (t, r) = locate(flat);
1819 let v = &views[t];
1820 v.scales_into(r, gpr, &mut sc);
1821 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1823 }
1824 };
1825 pool.run_rows(total_rows, &run);
1826 }
1827 return;
1828 }
1829
1830 if uniform_q1 {
1831 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1834 if a8w8_enabled() {
1835 let act = split_act(x);
1836 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1837 let (act, gsum) = (&act, &gsum);
1838 let closures: [_; N] = std::array::from_fn(|i| {
1839 let (bytes, gpr, out) =
1840 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1841 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1842 });
1843 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1844 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1845 pool.run_many(&parts);
1846 } else {
1847 let closures: [_; N] = std::array::from_fn(|i| {
1848 let (bytes, gpr, out) =
1849 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1850 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1851 });
1852 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1853 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1854 pool.run_many(&parts);
1855 }
1856 return;
1857 }
1858
1859 if uniform_q1t {
1860 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1864 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1865 if a8w8_enabled() {
1866 let act = split_act(x);
1867 let act = &act;
1868 let x_ref = x;
1869 let closures: [_; N] = std::array::from_fn(|i| {
1870 let bytes = ts[i].quant_bytes();
1871 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1872 let gpr = cols / GROUP_SIZE;
1873 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1874 let out = outs_addr[i];
1875 move |s: usize, e: usize| {
1876 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1877 }
1878 });
1879 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1880 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1881 pool.run_many(&parts);
1882 } else {
1883 let x_ref = x;
1884 let closures: [_; N] = std::array::from_fn(|i| {
1885 let bytes = ts[i].quant_bytes();
1886 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1887 let gpr = cols / GROUP_SIZE;
1888 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1889 let out = outs_addr[i];
1890 move |s: usize, e: usize| {
1891 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1892 }
1893 });
1894 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1895 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1896 pool.run_many(&parts);
1897 }
1898 return;
1899 }
1900
1901 if uniform_q4 || uniform_vbit {
1902 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1903 if a8w8_enabled() {
1905 let act = split_act(x);
1906 let act = &act;
1907 if uniform_q4 {
1908 let closures: [_; N] = std::array::from_fn(|i| {
1909 let (packed, scales) =
1910 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1911 let (gpr, cols, out) =
1912 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1913 move |s: usize, e: usize| {
1914 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1915 }
1916 });
1917 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1918 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1919 pool.run_many(&parts);
1920 } else {
1921 let closures: [_; N] = std::array::from_fn(|i| {
1922 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1923 unreachable!()
1924 };
1925 let (bytes, rows, cols, out) = (
1926 ts[i].quant_bytes(),
1927 ts[i].rows(),
1928 ts[i].cols(),
1929 outs_addr[i],
1930 );
1931 move |s: usize, e: usize| {
1932 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1933 }
1934 });
1935 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1936 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1937 pool.run_many(&parts);
1938 }
1939 return;
1940 }
1941 if uniform_q4 {
1942 let closures: [_; N] = std::array::from_fn(|i| {
1943 let (packed, scales) =
1944 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1945 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1946 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1947 });
1948 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1949 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1950 pool.run_many(&parts);
1951 } else {
1952 let closures: [_; N] = std::array::from_fn(|i| {
1953 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1954 unreachable!()
1955 };
1956 let (bytes, rows, cols, out) = (
1957 ts[i].quant_bytes(),
1958 ts[i].rows(),
1959 ts[i].cols(),
1960 outs_addr[i],
1961 );
1962 move |s: usize, e: usize| {
1963 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1964 }
1965 });
1966 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1967 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1968 pool.run_many(&parts);
1969 }
1970 return;
1971 }
1972
1973 if uniform_f32 {
1974 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1975 let closures: [_; N] = std::array::from_fn(|i| {
1976 let Self::F32 { data, cols, .. } = ts[i] else {
1977 unreachable!()
1978 };
1979 let out = outs_addr[i];
1980 move |start: usize, end: usize| {
1981 for o in start..end {
1982 let row = &data[o * cols..(o + 1) * cols];
1983 let mut sum = 0.0f32;
1984 for j in 0..*cols {
1985 sum += row[j] * x[j];
1986 }
1987 unsafe { *out.at(o) = sum };
1989 }
1990 }
1991 });
1992 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1993 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1994 pool.run_many(&parts);
1995 return;
1996 }
1997
1998 struct Ctx<'a> {
2001 bytes: &'a [u8],
2002 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2003 rep: &'a [u8],
2004 row_scale: &'a [f32],
2005 cols: usize,
2006 xs: std::borrow::Cow<'a, [f32]>,
2007 }
2008 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2009 let Self::Mapped {
2010 dtype,
2011 cols,
2012 row_scale,
2013 col_field,
2014 repack,
2015 ..
2016 } = ts[i]
2017 else {
2018 unreachable!()
2019 };
2020 Ctx {
2021 bytes: ts[i].quant_bytes(),
2022 rep: repack,
2023 row_scale,
2024 cols: *cols,
2025 xs: prescale(x, col_field, *dtype),
2026 }
2027 });
2028 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2029 #[cfg(target_arch = "aarch64")]
2030 if sdot_enabled() {
2031 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2032 let closures: [_; N] = std::array::from_fn(|i| {
2033 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2034 move |start: usize, end: usize| {
2035 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2036 }
2037 });
2038 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2039 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2040 pool.run_many(&parts);
2041 return;
2042 }
2043 #[cfg(target_arch = "x86_64")]
2044 if avx2_a8w8_enabled() {
2045 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2046 let closures: [_; N] = std::array::from_fn(|i| {
2047 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2048 move |start: usize, end: usize| {
2049 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2050 }
2051 });
2052 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2053 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2054 pool.run_many(&parts);
2055 return;
2056 }
2057 let closures: [_; N] = std::array::from_fn(|i| {
2058 let (c, out) = (&ctxs[i], outs_addr[i]);
2059 move |start: usize, end: usize| {
2060 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2061 }
2062 });
2063 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2064 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2065 pool.run_many(&parts);
2066 }
2067}
2068
2069impl QTensor {
2070 #[allow(clippy::needless_range_loop)]
2075 pub fn matvec2_many<const N: usize>(
2076 ts: [&QTensor; N],
2077 x1: &[f32],
2078 x2: &[f32],
2079 mut o1s: [&mut [f32]; N],
2080 mut o2s: [&mut [f32]; N],
2081 pool: Option<&Pool>,
2082 ) {
2083 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2084 let uniform_q8 = ts.iter().all(|t| {
2085 matches!(
2086 t,
2087 Self::Mapped {
2088 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2089 ..
2090 }
2091 )
2092 });
2093 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2094 let uniform_q4 = ts.iter().all(|t| {
2095 matches!(
2096 t,
2097 Self::Mapped {
2098 dtype: TensorDtype::Q4Block,
2099 ..
2100 }
2101 )
2102 });
2103 let uniform_vbit = ts.iter().all(|t| {
2104 matches!(
2105 t,
2106 Self::Mapped {
2107 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2108 ..
2109 }
2110 )
2111 });
2112 let fusable = pool.is_some()
2113 && total_rows >= 256
2114 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2115 if !fusable {
2116 for i in 0..N {
2117 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2118 }
2119 return;
2120 }
2121 let pool = pool.unwrap();
2122
2123 if uniform_q4 || uniform_vbit {
2124 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2125 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2126 if a8w8_enabled() {
2128 let a1 = split_act(x1);
2129 let a2 = split_act(x2);
2130 let (a1, a2) = (&a1, &a2);
2131 if uniform_q4 {
2132 let closures: [_; N] = std::array::from_fn(|i| {
2133 let (packed, scales) =
2134 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2135 let (gpr, cols, o1, o2) =
2136 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2137 move |s: usize, e: usize| {
2138 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2139 }
2140 });
2141 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2142 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2143 pool.run_many(&parts);
2144 } else {
2145 let closures: [_; N] = std::array::from_fn(|i| {
2146 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2147 unreachable!()
2148 };
2149 let (bytes, rows, cols, o1, o2) = (
2150 ts[i].quant_bytes(),
2151 ts[i].rows(),
2152 ts[i].cols(),
2153 p1[i],
2154 p2[i],
2155 );
2156 move |s: usize, e: usize| {
2157 vbit_range2_a8w8(
2158 bytes,
2159 vbit_offsets,
2160 x1,
2161 x2,
2162 a1,
2163 a2,
2164 rows,
2165 cols,
2166 o1,
2167 o2,
2168 s,
2169 e,
2170 )
2171 }
2172 });
2173 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2174 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2175 pool.run_many(&parts);
2176 }
2177 return;
2178 }
2179 if uniform_q4 {
2180 let closures: [_; N] = std::array::from_fn(|i| {
2181 let (packed, scales) =
2182 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2183 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2184 move |s: usize, e: usize| {
2185 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2186 }
2187 });
2188 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2189 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2190 pool.run_many(&parts);
2191 } else {
2192 let closures: [_; N] = std::array::from_fn(|i| {
2193 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2194 unreachable!()
2195 };
2196 let (bytes, rows, cols, o1, o2) = (
2197 ts[i].quant_bytes(),
2198 ts[i].rows(),
2199 ts[i].cols(),
2200 p1[i],
2201 p2[i],
2202 );
2203 move |s: usize, e: usize| {
2204 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2205 }
2206 });
2207 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2208 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2209 pool.run_many(&parts);
2210 }
2211 return;
2212 }
2213
2214 if uniform_f32 {
2215 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2216 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2217 let closures: [_; N] = std::array::from_fn(|i| {
2218 let Self::F32 { data, cols, .. } = ts[i] else {
2219 unreachable!()
2220 };
2221 let (o1, o2) = (p1[i], p2[i]);
2222 move |start: usize, end: usize| {
2223 for o in start..end {
2224 let row = &data[o * cols..(o + 1) * cols];
2225 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2226 for j in 0..*cols {
2227 s1 += row[j] * x1[j];
2228 s2 += row[j] * x2[j];
2229 }
2230 unsafe {
2232 *o1.at(o) = s1;
2233 *o2.at(o) = s2;
2234 }
2235 }
2236 }
2237 });
2238 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2239 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2240 pool.run_many(&parts);
2241 return;
2242 }
2243
2244 struct Ctx<'a> {
2245 bytes: &'a [u8],
2246 row_scale: &'a [f32],
2247 cols: usize,
2248 xs1: std::borrow::Cow<'a, [f32]>,
2249 xs2: std::borrow::Cow<'a, [f32]>,
2250 }
2251 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2252 let Self::Mapped {
2253 dtype,
2254 cols,
2255 row_scale,
2256 col_field,
2257 ..
2258 } = ts[i]
2259 else {
2260 unreachable!()
2261 };
2262 Ctx {
2263 bytes: ts[i].quant_bytes(),
2264 row_scale,
2265 cols: *cols,
2266 xs1: prescale(x1, col_field, *dtype),
2267 xs2: prescale(x2, col_field, *dtype),
2268 }
2269 });
2270 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2271 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2272 #[cfg(target_arch = "aarch64")]
2273 if sdot_enabled() {
2274 let acts: [(SplitAct, SplitAct); N] =
2275 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2276 let closures: [_; N] = std::array::from_fn(|i| {
2277 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2278 move |start: usize, end: usize| {
2279 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2280 }
2281 });
2282 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2283 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2284 pool.run_many(&parts);
2285 return;
2286 }
2287 #[cfg(target_arch = "x86_64")]
2288 if avx2_a8w8_enabled() {
2289 let acts: [(SplitAct, SplitAct); N] =
2290 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2291 let closures: [_; N] = std::array::from_fn(|i| {
2292 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2293 move |start: usize, end: usize| {
2294 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2295 }
2296 });
2297 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2298 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2299 pool.run_many(&parts);
2300 return;
2301 }
2302 let closures: [_; N] = std::array::from_fn(|i| {
2303 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2304 move |start: usize, end: usize| {
2305 q8_range2_f32(
2306 c.bytes,
2307 c.row_scale,
2308 &c.xs1,
2309 &c.xs2,
2310 c.cols,
2311 o1,
2312 o2,
2313 start,
2314 end,
2315 )
2316 }
2317 });
2318 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2319 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2320 pool.run_many(&parts);
2321 }
2322
2323 pub fn matvec_silu_mul(
2328 gate: &QTensor,
2329 up: &QTensor,
2330 x: &[f32],
2331 out: &mut [f32],
2332 pool: Option<&Pool>,
2333 ) -> bool {
2334 let inter = gate.rows();
2335 debug_assert_eq!(up.rows(), inter);
2336 debug_assert_eq!(out.len(), inter);
2337 debug_assert_eq!(gate.cols(), up.cols());
2338 if !a8w8_enabled() {
2339 return false;
2340 }
2341 let act = split_act(x);
2342 let act = &act;
2343 let x_ref = x;
2344 let out_addr = SendMut(out.as_mut_ptr());
2345
2346 match (gate, up) {
2347 (
2349 Self::Mapped {
2350 dtype: TensorDtype::Q4Block,
2351 ..
2352 },
2353 Self::Mapped {
2354 dtype: TensorDtype::Q4Block,
2355 ..
2356 },
2357 ) => {
2358 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2359 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2360 let gpr = gate.cols() / GROUP_SIZE;
2361 let cols = gate.cols();
2362 let run = move |start: usize, end: usize| {
2363 for r in start..end {
2364 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2365 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2366 for &(j, xv) in &act.outliers {
2367 let flat = r * cols + j;
2368 let gb = gp[flat / 2];
2369 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2370 let gsc = f16_to_f32(u16::from_le_bytes([
2371 gs[(flat / GROUP_SIZE) * 2],
2372 gs[(flat / GROUP_SIZE) * 2 + 1],
2373 ]));
2374 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2375 let ub = up_p[flat / 2];
2376 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2377 let usc = f16_to_f32(u16::from_le_bytes([
2378 up_s[(flat / GROUP_SIZE) * 2],
2379 up_s[(flat / GROUP_SIZE) * 2 + 1],
2380 ]));
2381 uv += ((un as i32 - 8) as f32) * usc * xv;
2382 }
2383 let silu_g = gv / (1.0 + (-gv).exp());
2384 unsafe { *out_addr.at(r) = silu_g * uv };
2386 }
2387 };
2388 dispatch_rows(pool, inter, &run);
2389 true
2390 }
2391 (
2395 Self::Mapped {
2396 dtype: TensorDtype::Q4Tiled,
2397 ..
2398 },
2399 Self::Mapped {
2400 dtype: TensorDtype::Q4Tiled,
2401 ..
2402 },
2403 ) => {
2404 let g_bytes = gate.quant_bytes();
2405 let u_bytes = up.quant_bytes();
2406 let gpr = gate.cols() / GROUP_SIZE;
2407 let run = move |start: usize, end: usize| {
2408 for r in start..end {
2409 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2410 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2411 for &(j, xv) in &act.outliers {
2412 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2413 gv += w * s * xv;
2414 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2415 uv += w * s * xv;
2416 }
2417 let silu_g = gv / (1.0 + (-gv).exp());
2418 unsafe { *out_addr.at(r) = silu_g * uv };
2420 }
2421 };
2422 dispatch_rows(pool, inter, &run);
2423 true
2424 }
2425 (
2428 Self::Mapped {
2429 dtype: TensorDtype::Q4TiledP,
2430 ..
2431 },
2432 Self::Mapped {
2433 dtype: TensorDtype::Q4TiledP,
2434 ..
2435 },
2436 ) => {
2437 let cols = gate.cols();
2438 let gpr = cols / GROUP_SIZE;
2439 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2440 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2441 let run = |start: usize, end: usize| {
2442 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2443 for r in start..end {
2444 gv_view.scales_into(r, gpr, &mut gsc);
2445 uv_view.scales_into(r, gpr, &mut usc);
2446 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2447 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2448 for &(j, xv) in &act.outliers {
2449 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2450 gv += w * s * xv;
2451 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2452 uv += w * s * xv;
2453 }
2454 let silu_g = gv / (1.0 + (-gv).exp());
2455 unsafe { *out_addr.at(r) = silu_g * uv };
2457 }
2458 };
2459 dispatch_rows(pool, inter, &run);
2460 true
2461 }
2462 (
2464 Self::Mapped {
2465 dtype: TensorDtype::Q1T,
2466 ..
2467 },
2468 Self::Mapped {
2469 dtype: TensorDtype::Q1T,
2470 ..
2471 },
2472 ) => {
2473 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2474 let g_bytes = gate.quant_bytes();
2475 let u_bytes = up.quant_bytes();
2476 let gpr = gate.cols() / GROUP_SIZE;
2477 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2478 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2479 let run = move |start: usize, end: usize| {
2480 for r in start..end {
2481 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2482 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2483 for &(j, xv) in &act.outliers {
2484 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2485 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2486 }
2487 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2488 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2489 let silu_g = gv / (1.0 + (-gv).exp());
2490 unsafe { *out_addr.at(r) = silu_g * uv };
2492 }
2493 };
2494 dispatch_rows(pool, inter, &run);
2495 true
2496 }
2497 _ => false,
2498 }
2499 }
2500
2501 pub fn moe_gate_up_many(
2515 pairs: &[(&QTensor, &QTensor)],
2516 x: &[f32],
2517 outs: &mut [Vec<f32>],
2518 pool: Option<&Pool>,
2519 ) -> bool {
2520 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2521 return false;
2522 }
2523 let inter = pairs[0].0.rows();
2524 let cols = pairs[0].0.cols();
2525 if cols % GROUP_SIZE != 0 {
2526 return false;
2527 }
2528 let gpr = cols / GROUP_SIZE;
2529 let mut views = Vec::with_capacity(pairs.len() * 2);
2530 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2531 let both_q4tp = matches!(
2532 g,
2533 Self::Mapped {
2534 dtype: TensorDtype::Q4TiledP,
2535 ..
2536 }
2537 ) && matches!(
2538 u,
2539 Self::Mapped {
2540 dtype: TensorDtype::Q4TiledP,
2541 ..
2542 }
2543 );
2544 if !both_q4tp
2545 || g.rows() != inter
2546 || u.rows() != inter
2547 || g.cols() != cols
2548 || u.cols() != cols
2549 || o.len() != inter
2550 {
2551 return false;
2552 }
2553 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
2554 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
2555 }
2556 let act = split_act(x);
2557 let act = &act;
2558 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2559 let (views, ptrs) = (&views, &ptrs);
2560 let run = |start: usize, end: usize| {
2561 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2562 for flat in start..end {
2563 let (e, r) = (flat / inter, flat % inter);
2564 let gv_view = &views[e * 2];
2565 let uv_view = &views[e * 2 + 1];
2566 gv_view.scales_into(r, gpr, &mut gsc);
2567 uv_view.scales_into(r, gpr, &mut usc);
2568 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2569 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2570 for &(j, xv) in &act.outliers {
2571 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2572 gv += w * s * xv;
2573 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2574 uv += w * s * xv;
2575 }
2576 let silu_g = gv / (1.0 + (-gv).exp());
2577 unsafe { *ptrs[e].at(r) = silu_g * uv };
2579 }
2580 };
2581 dispatch_rows(pool, pairs.len() * inter, &run);
2582 true
2583 }
2584
2585 pub fn moe_down_many(
2594 downs: &[&QTensor],
2595 gs: &[Vec<f32>],
2596 weights: &[f32],
2597 out: &mut [f32],
2598 pool: Option<&Pool>,
2599 ) -> bool {
2600 if downs.is_empty()
2601 || downs.len() != gs.len()
2602 || downs.len() != weights.len()
2603 || !a8w8_enabled()
2604 {
2605 return false;
2606 }
2607 let rows = out.len();
2608 let cols = downs[0].cols();
2609 if cols % GROUP_SIZE != 0 {
2610 return false;
2611 }
2612 let gpr = cols / GROUP_SIZE;
2613 let mut views = Vec::with_capacity(downs.len());
2614 for (d, g) in downs.iter().zip(gs.iter()) {
2615 if !matches!(
2616 d,
2617 Self::Mapped {
2618 dtype: TensorDtype::Q4TiledP,
2619 ..
2620 }
2621 ) || d.rows() != rows
2622 || d.cols() != cols
2623 || g.len() != cols
2624 {
2625 return false;
2626 }
2627 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2628 }
2629 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2631 let out_addr = SendMut(out.as_mut_ptr());
2639 let (views, acts, weights) = (&views, &acts, &weights);
2640 let run = |start: usize, end: usize| {
2641 let mut sc = vec![0f32; gpr];
2642 for r in start..end {
2643 let mut acc = 0f32;
2644 for (e, v) in views.iter().enumerate() {
2645 v.scales_into(r, gpr, &mut sc);
2646 let a = &acts[e];
2647 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2648 for &(j, xv) in &a.outliers {
2649 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2650 d += w * s * xv;
2651 }
2652 acc += weights[e] * d;
2653 }
2654 unsafe { *out_addr.at(r) = acc };
2656 }
2657 };
2658 dispatch_rows(pool, rows, &run);
2659 true
2660 }
2661}
2662
2663#[cfg(target_os = "macos")]
2668mod accel_blas {
2669 #[link(name = "Accelerate", kind = "framework")]
2670 unsafe extern "C" {
2671 pub fn cblas_sgemm(
2672 order: i32,
2673 trans_a: i32,
2674 trans_b: i32,
2675 m: i32,
2676 n: i32,
2677 k: i32,
2678 alpha: f32,
2679 a: *const f32,
2680 lda: i32,
2681 b: *const f32,
2682 ldb: i32,
2683 beta: f32,
2684 c: *mut f32,
2685 ldc: i32,
2686 );
2687 }
2688}
2689
2690#[cfg(target_os = "macos")]
2691pub(crate) fn accel_gemm_enabled() -> bool {
2692 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2693 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2694}
2695
2696#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2699pub(crate) fn accel_gemm_enabled() -> bool {
2700 true
2701}
2702
2703#[cfg(target_arch = "aarch64")]
2710#[allow(clippy::too_many_arguments)]
2711pub(crate) fn neon_gemm_rm(
2712 m: usize,
2713 n: usize,
2714 k: usize,
2715 alpha: f32,
2716 a: &[f32],
2717 lda: usize,
2718 b_mat: &[f32],
2719 ldb: usize,
2720 b_rows_are_n: bool,
2721 c: &mut [f32],
2722 ldc: usize,
2723) {
2724 debug_assert!(a.len() >= (m - 1) * lda + k);
2725 debug_assert!(c.len() >= (m - 1) * ldc + n);
2726 unsafe {
2728 use core::arch::aarch64::*;
2729 let mut i = 0usize;
2730 while i < m {
2731 let mi = (m - i).min(4);
2732 let mut j = 0usize;
2733 while j < n {
2734 let nj = (n - j).min(8);
2735 if mi == 4 && nj == 8 {
2736 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2737 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2738 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2739 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2740 for p in 0..k {
2741 let (b0, b1) = if b_rows_are_n {
2742 let base = b_mat.as_ptr().add(j * ldb + p);
2745 let g = |o: usize| *base.add(o * ldb);
2746 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2747 } else {
2748 let base = b_mat.as_ptr().add(p * ldb + j);
2749 (
2750 [*base, *base.add(1), *base.add(2), *base.add(3)],
2751 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2752 )
2753 };
2754 let bv0 = vld1q_f32(b0.as_ptr());
2755 let bv1 = vld1q_f32(b1.as_ptr());
2756 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2757 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2758 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2759 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2760 c0a = vfmaq_f32(c0a, a0, bv0);
2761 c0b = vfmaq_f32(c0b, a0, bv1);
2762 c1a = vfmaq_f32(c1a, a1, bv0);
2763 c1b = vfmaq_f32(c1b, a1, bv1);
2764 c2a = vfmaq_f32(c2a, a2, bv0);
2765 c2b = vfmaq_f32(c2b, a2, bv1);
2766 c3a = vfmaq_f32(c3a, a3, bv0);
2767 c3b = vfmaq_f32(c3b, a3, bv1);
2768 }
2769 let al = vdupq_n_f32(alpha);
2770 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2771 .iter()
2772 .enumerate()
2773 {
2774 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2775 vst1q_f32(dst, vmulq_f32(*ca, al));
2776 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2777 }
2778 } else {
2779 for r in 0..mi {
2780 for q in 0..nj {
2781 let mut acc = 0f32;
2782 for p in 0..k {
2783 let bv = if b_rows_are_n {
2784 b_mat[(j + q) * ldb + p]
2785 } else {
2786 b_mat[p * ldb + j + q]
2787 };
2788 acc += a[(i + r) * lda + p] * bv;
2789 }
2790 c[(i + r) * ldc + j + q] = acc * alpha;
2791 }
2792 }
2793 }
2794 j += nj;
2795 }
2796 i += mi;
2797 }
2798 }
2799}
2800
2801#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2803#[allow(clippy::too_many_arguments)]
2804pub(crate) fn sgemm_rm(
2805 m: usize,
2806 n: usize,
2807 k: usize,
2808 alpha: f32,
2809 a: &[f32],
2810 lda: usize,
2811 b_mat: &[f32],
2812 ldb: usize,
2813 b_rows_are_n: bool,
2814 c: &mut [f32],
2815 ldc: usize,
2816) {
2817 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2818}
2819
2820#[allow(clippy::too_many_arguments)]
2824pub fn sgemm_public(
2825 m: usize,
2826 n: usize,
2827 k: usize,
2828 alpha: f32,
2829 a: &[f32],
2830 lda: usize,
2831 b_mat: &[f32],
2832 ldb: usize,
2833 b_rows_are_n: bool,
2834 c: &mut [f32],
2835 ldc: usize,
2836) {
2837 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
2838 {
2839 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2840 }
2841 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
2846 {
2847 for i in 0..m {
2848 for j in 0..n {
2849 let mut acc = 0f32;
2850 for p in 0..k {
2851 let bv = if b_rows_are_n {
2852 b_mat[j * ldb + p]
2853 } else {
2854 b_mat[p * ldb + j]
2855 };
2856 acc += a[i * lda + p] * bv;
2857 }
2858 c[i * ldc + j] = alpha * acc;
2859 }
2860 }
2861 }
2862}
2863
2864#[cfg(target_os = "macos")]
2867#[allow(clippy::too_many_arguments)]
2868pub(crate) fn sgemm_rm(
2869 m: usize,
2870 n: usize,
2871 k: usize,
2872 alpha: f32,
2873 a: &[f32],
2874 lda: usize,
2875 b_mat: &[f32],
2876 ldb: usize,
2877 b_rows_are_n: bool,
2878 c: &mut [f32],
2879 ldc: usize,
2880) {
2881 debug_assert!(a.len() >= (m - 1) * lda + k);
2882 debug_assert!(c.len() >= (m - 1) * ldc + n);
2883 #[cfg(target_arch = "aarch64")]
2888 if std::env::var("CMF_FORCE_NEON_GEMM")
2889 .map(|v| v == "1")
2890 .unwrap_or(false)
2891 {
2892 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2893 }
2894 unsafe {
2895 accel_blas::cblas_sgemm(
2896 101, 111, if b_rows_are_n { 112 } else { 111 },
2899 m as i32,
2900 n as i32,
2901 k as i32,
2902 alpha,
2903 a.as_ptr(),
2904 lda as i32,
2905 b_mat.as_ptr(),
2906 ldb as i32,
2907 0.0,
2908 c.as_mut_ptr(),
2909 ldc as i32,
2910 );
2911 }
2912}
2913
2914#[cfg(target_os = "macos")]
2921fn qmatmat_accel(
2922 q: &[u8],
2923 row_scale: &[f32],
2924 pre: &[std::borrow::Cow<'_, [f32]>],
2925 rows: usize,
2926 cols: usize,
2927 out: &mut [f32],
2928 pool: Option<&Pool>,
2929) {
2930 const TR: usize = 2048;
2935 let b = pre.len();
2936 thread_local! {
2937 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2938 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2939 }
2940 XPANEL.with(|xp| {
2941 WTILE.with(|wt| {
2942 let mut xpanel = xp.borrow_mut();
2943 xpanel.clear();
2944 for x in pre {
2945 xpanel.extend_from_slice(x);
2946 }
2947 let mut wtile = wt.borrow_mut();
2948 wtile.resize(TR * cols, 0.0);
2949 let mut r0 = 0usize;
2950 while r0 < rows {
2951 let tr = TR.min(rows - r0);
2952 let wt_addr = SendMut(wtile.as_mut_ptr());
2954 let run = |start: usize, end: usize| {
2955 for r in start..end {
2956 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2957 let s = row_scale[r0 + r];
2958 let dst =
2960 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2961 for (d, &v) in dst.iter_mut().zip(row) {
2962 *d = (v as i8) as f32 * s;
2963 }
2964 }
2965 };
2966 dispatch_rows(pool, tr, &run);
2967 unsafe {
2969 accel_blas::cblas_sgemm(
2970 101, 111, 112, b as i32,
2974 tr as i32,
2975 cols as i32,
2976 1.0,
2977 xpanel.as_ptr(),
2978 cols as i32,
2979 wtile.as_ptr(),
2980 cols as i32,
2981 0.0,
2982 out.as_mut_ptr().add(r0),
2983 rows as i32,
2984 );
2985 }
2986 r0 += tr;
2987 }
2988 })
2989 });
2990}
2991
2992fn qmatmat(
2993 q: &[u8],
2994 row_scale: &[f32],
2995 pre: &[std::borrow::Cow<'_, [f32]>],
2996 rows: usize,
2997 cols: usize,
2998 out: &mut [f32],
2999 pool: Option<&Pool>,
3000) {
3001 let b = pre.len();
3002 debug_assert_eq!(out.len(), b * rows);
3003 #[cfg(target_os = "macos")]
3008 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3009 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3010 return;
3011 }
3012 #[cfg(target_arch = "aarch64")]
3013 if sdot_enabled() {
3014 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3015 let out_addr = SendMut(out.as_mut_ptr());
3016 let blocked_ok = blocked_enabled();
3019 let use_i8mm = i8mm_enabled();
3020 if blocked_ok {
3021 let run = |start: usize, end: usize| {
3022 let mut o = start;
3023 while o < end {
3024 if o + 2 <= end {
3025 let r0 = &q[o * cols..(o + 1) * cols];
3026 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3027 let mut bi = 0usize;
3028 while bi + 4 <= acts.len() {
3029 let xs = [
3030 acts[bi].xq.as_slice(),
3031 acts[bi + 1].xq.as_slice(),
3032 acts[bi + 2].xq.as_slice(),
3033 acts[bi + 3].xq.as_slice(),
3034 ];
3035 let d = if use_i8mm {
3036 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3037 } else {
3038 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3039 };
3040 for (r, row) in [r0, r1].into_iter().enumerate() {
3041 for k in 0..4 {
3042 let act = &acts[bi + k];
3043 let mut v = d[r][k] as f32 * act.sx;
3044 for &(j, xv) in &act.outliers {
3045 v += (row[j] as i8) as f32 * xv;
3046 }
3047 unsafe {
3048 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3049 };
3050 }
3051 }
3052 bi += 4;
3053 }
3054 while bi < acts.len() {
3055 for (r, row) in [r0, r1].into_iter().enumerate() {
3056 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3057 unsafe { *out_addr.at(bi * rows + o + r) = v };
3058 }
3059 bi += 1;
3060 }
3061 o += 2;
3062 } else {
3063 let row = &q[o * cols..(o + 1) * cols];
3064 for (bi, act) in acts.iter().enumerate() {
3065 let v = row_dot_sdot(row, act) * row_scale[o];
3066 unsafe { *out_addr.at(bi * rows + o) = v };
3067 }
3068 o += 1;
3069 }
3070 }
3071 };
3072 dispatch_rows(pool, rows, &run);
3073 return;
3074 }
3075 let run = |start: usize, end: usize| {
3076 for o in start..end {
3077 let row = &q[o * cols..(o + 1) * cols];
3078 for (bi, act) in acts.iter().enumerate() {
3079 let v = row_dot_sdot(row, act) * row_scale[o];
3080 unsafe { *out_addr.at(bi * rows + o) = v };
3081 }
3082 }
3083 };
3084 dispatch_rows(pool, rows, &run);
3085 return;
3086 }
3087 #[cfg(target_arch = "x86_64")]
3092 if avx2_a8w8_enabled() {
3093 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3094 let out_addr = SendMut(out.as_mut_ptr());
3095 let blocked_ok = blocked_enabled();
3098 if !avx512vnni_enabled() && blocked_ok {
3099 let run = |start: usize, end: usize| {
3100 let mut o = start;
3101 while o < end {
3102 if o + 2 <= end {
3103 let r0 = &q[o * cols..(o + 1) * cols];
3104 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3105 let mut bi = 0usize;
3106 while bi + 4 <= acts.len() {
3107 let xs = [
3108 acts[bi].xq.as_slice(),
3109 acts[bi + 1].xq.as_slice(),
3110 acts[bi + 2].xq.as_slice(),
3111 acts[bi + 3].xq.as_slice(),
3112 ];
3113 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3114 for (r, row) in [r0, r1].into_iter().enumerate() {
3115 for k in 0..4 {
3116 let act = &acts[bi + k];
3117 let mut v = d[r][k] as f32 * act.sx;
3118 for &(j, xv) in &act.outliers {
3119 v += (row[j] as i8) as f32 * xv;
3120 }
3121 unsafe {
3122 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3123 };
3124 }
3125 }
3126 bi += 4;
3127 }
3128 while bi < acts.len() {
3129 for (r, row) in [r0, r1].into_iter().enumerate() {
3130 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3131 unsafe { *out_addr.at(bi * rows + o + r) = v };
3132 }
3133 bi += 1;
3134 }
3135 o += 2;
3136 } else {
3137 let row = &q[o * cols..(o + 1) * cols];
3138 for (bi, act) in acts.iter().enumerate() {
3139 let v = row_dot_avx2(row, act) * row_scale[o];
3140 unsafe { *out_addr.at(bi * rows + o) = v };
3141 }
3142 o += 1;
3143 }
3144 }
3145 };
3146 dispatch_rows(pool, rows, &run);
3147 return;
3148 }
3149 let run = |start: usize, end: usize| {
3150 for o in start..end {
3151 let row = &q[o * cols..(o + 1) * cols];
3152 for (bi, act) in acts.iter().enumerate() {
3153 let v = row_dot_avx2(row, act) * row_scale[o];
3154 unsafe { *out_addr.at(bi * rows + o) = v };
3155 }
3156 }
3157 };
3158 dispatch_rows(pool, rows, &run);
3159 return;
3160 }
3161 let out_addr = SendMut(out.as_mut_ptr());
3162 let run = |start: usize, end: usize| {
3163 for o in start..end {
3164 let row = &q[o * cols..(o + 1) * cols];
3165 for (bi, x) in pre.iter().enumerate() {
3166 let mut acc = 0f32;
3167 for j in 0..cols {
3168 acc += (row[j] as i8) as f32 * x[j];
3169 }
3170 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3171 }
3172 }
3173 };
3174 dispatch_rows(pool, rows, &run);
3175}
3176
3177fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3180 match pool {
3181 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3182 _ => run(0, rows),
3183 }
3184}
3185
3186fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3188 let groups = rows * cols / GROUP_SIZE;
3189 bytes.split_at(groups * 16)
3190}
3191
3192#[inline]
3197fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3198 #[cfg(target_arch = "aarch64")]
3199 unsafe {
3200 return vbit_fill4_neon(data, buf);
3201 }
3202 #[cfg(target_arch = "x86_64")]
3203 if avx2_enabled() {
3204 return unsafe { vbit_fill4_avx2(data, buf) };
3205 }
3206 #[allow(unreachable_code)]
3207 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3208 let u = unpack8::<4>(&data[blk * 4..]);
3209 for k in 0..8 {
3210 chunk[k] = (u[k] - 7) as i8 as u8;
3211 }
3212 }
3213}
3214
3215#[cfg(target_arch = "aarch64")]
3216#[target_feature(enable = "neon")]
3217unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3218 unsafe {
3221 use core::arch::aarch64::*;
3222 let n = buf.len();
3223 let mask = vdupq_n_u8(0x0F);
3224 let seven = vdupq_n_s8(7);
3225 let mut g = 0usize;
3226 while g * 32 + 32 <= n {
3227 let b = vld1q_u8(data.as_ptr().add(g * 16));
3228 let hi = vshrq_n_u8::<4>(b);
3229 let lo = vandq_u8(b, mask);
3230 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3231 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3232 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3233 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3234 g += 1;
3235 }
3236 }
3237}
3238
3239#[cfg(target_arch = "x86_64")]
3240#[target_feature(enable = "avx2")]
3241unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3242 unsafe {
3244 use core::arch::x86_64::*;
3245 let n = buf.len();
3246 let mask = _mm_set1_epi8(0x0F);
3247 let seven = _mm256_set1_epi8(7);
3248 let mut g = 0usize;
3249 while g * 32 + 32 <= n {
3250 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3251 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3252 let lo = _mm_and_si128(b, mask);
3253 let z = _mm256_sub_epi8(
3254 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3255 seven,
3256 );
3257 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3258 g += 1;
3259 }
3260 }
3261}
3262
3263#[inline(always)]
3268fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3269 let mut acc = 0u64;
3270 for i in 0..B {
3271 acc = (acc << 8) | data[i] as u64;
3272 }
3273 let mask = (1u64 << B) - 1;
3274 let mut out = [0i32; 8];
3275 for (k, o) in out.iter_mut().enumerate() {
3276 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3277 }
3278 out
3279}
3280
3281#[allow(clippy::too_many_arguments)]
3287fn vbitmatvec(
3288 bytes: &[u8],
3289 offsets: &[usize],
3290 x: &[f32],
3291 rows: usize,
3292 cols: usize,
3293 out: &mut [f32],
3294 pool: Option<&Pool>,
3295) {
3296 debug_assert_eq!(out.len(), rows);
3297 debug_assert_eq!(offsets.len(), rows + 1);
3298
3299 if a8w8_enabled() {
3303 let act = split_act(x);
3304 let out_addr = SendMut(out.as_mut_ptr());
3305 let run = move |start: usize, end: usize| {
3306 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3307 };
3308 dispatch_rows(pool, rows, &run);
3309 return;
3310 }
3311
3312 let out_addr = SendMut(out.as_mut_ptr());
3313 let run = move |start: usize, end: usize| {
3314 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3315 };
3316 dispatch_rows(pool, rows, &run);
3317}
3318
3319#[allow(clippy::too_many_arguments)]
3323fn vbit_range_a8w8(
3324 bytes: &[u8],
3325 offsets: &[usize],
3326 x: &[f32],
3327 act: &SplitAct,
3328 rows: usize,
3329 cols: usize,
3330 out: SendMut,
3331 start: usize,
3332 end: usize,
3333) {
3334 let ng = cols / GROUP_SIZE;
3335 let bits = &bytes[..rows];
3336 let sc_off = rows;
3337 let row_dot = |r: usize| -> f32 {
3338 let b = bits[r] as usize;
3339 let l = (1i32 << (b - 1)) - 1;
3340 let mask = (1u64 << b) - 1;
3341 let data = &bytes[offsets[r]..offsets[r + 1]];
3342 if b == 8 {
3343 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3345 let mut dot = 0f32;
3346 for g in 0..ng {
3347 let so = (r * ng + g) * 2;
3348 let sgf = f16_to_f32(u16::from_le_bytes([
3349 bytes[sc_off + so],
3350 bytes[sc_off + so + 1],
3351 ]));
3352 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3353 let mut gd = 0f32;
3354 for &xv in xg.iter() {
3355 if nbits < 8 {
3356 acc = (acc << 8) | data[idx] as u64;
3357 idx += 1;
3358 nbits += 8;
3359 }
3360 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3361 nbits -= 8;
3362 gd += (u - l) as f32 * xv;
3363 }
3364 dot += gd * sgf;
3365 }
3366 return dot;
3367 }
3368 thread_local! {
3372 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3373 const { std::cell::RefCell::new(Vec::new()) };
3374 }
3375 #[inline(always)]
3376 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3377 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3378 let u = unpack8::<B>(&data[blk * B..]);
3379 for k in 0..8 {
3380 chunk[k] = (u[k] - l) as i8 as u8;
3381 }
3382 }
3383 }
3384 let _ = mask;
3385 VBIT_SCRATCH.with(|scratch| {
3386 let mut buf = scratch.borrow_mut();
3387 buf.resize(cols, 0);
3388 match b {
3389 3 => fill::<3>(data, l, &mut buf),
3390 4 => vbit_fill4(data, &mut buf),
3391 5 => fill::<5>(data, l, &mut buf),
3392 6 => fill::<6>(data, l, &mut buf),
3393 _ => unreachable!(),
3394 }
3395 let mut dot = 0f32;
3396 for g in 0..ng {
3397 let so = (r * ng + g) * 2;
3398 let s = f16_to_f32(u16::from_le_bytes([
3399 bytes[sc_off + so],
3400 bytes[sc_off + so + 1],
3401 ]));
3402 let d = dot_i8_i8(
3403 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3404 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3405 ) as f32
3406 * act.sx;
3407 dot += d * s;
3408 }
3409 for &(j, xv) in &act.outliers {
3410 let so = (r * ng + j / GROUP_SIZE) * 2;
3411 let s = f16_to_f32(u16::from_le_bytes([
3412 bytes[sc_off + so],
3413 bytes[sc_off + so + 1],
3414 ]));
3415 dot += (buf[j] as i8) as f32 * s * xv;
3417 }
3418 dot
3419 })
3420 };
3421 for r in start..end {
3422 unsafe { *out.at(r) = row_dot(r) };
3424 }
3425}
3426
3427#[allow(clippy::too_many_arguments)]
3429fn vbit_range_f32(
3430 bytes: &[u8],
3431 offsets: &[usize],
3432 x: &[f32],
3433 rows: usize,
3434 cols: usize,
3435 out: SendMut,
3436 start: usize,
3437 end: usize,
3438) {
3439 let ng = cols / GROUP_SIZE;
3440 let bits = &bytes[..rows];
3441 let sc_off = rows;
3442 #[inline(always)]
3446 fn dot_row<const B: usize>(
3447 data: &[u8],
3448 bytes: &[u8],
3449 sc_off: usize,
3450 r: usize,
3451 ng: usize,
3452 x: &[f32],
3453 ) -> f32 {
3454 let l = ((1i32 << (B - 1)) - 1) as f32;
3455 let gbytes = GROUP_SIZE * B / 8;
3456 let mut dot = 0f32;
3457 for g in 0..ng {
3458 let so = (r * ng + g) * 2;
3459 let s = f16_to_f32(u16::from_le_bytes([
3460 bytes[sc_off + so],
3461 bytes[sc_off + so + 1],
3462 ]));
3463 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3464 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3465 let mut gd = 0f32;
3466 for blk in 0..GROUP_SIZE / 8 {
3467 let u = unpack8::<B>(&gd0[blk * B..]);
3468 let xb = &xg[blk * 8..blk * 8 + 8];
3469 for k in 0..8 {
3470 gd += (u[k] as f32 - l) * xb[k];
3471 }
3472 }
3473 dot += gd * s;
3474 }
3475 dot
3476 }
3477 for r in start..end {
3478 let data = &bytes[offsets[r]..offsets[r + 1]];
3479 let v = match bits[r] {
3480 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3481 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3482 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3483 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3484 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3485 b => unreachable!("vbit bit-width {b} (validated at load)"),
3486 };
3487 unsafe { *out.at(r) = v };
3489 }
3490}
3491
3492#[allow(clippy::too_many_arguments)]
3497fn vbitmatvec2(
3498 bytes: &[u8],
3499 offsets: &[usize],
3500 x1: &[f32],
3501 x2: &[f32],
3502 rows: usize,
3503 cols: usize,
3504 o1: &mut [f32],
3505 o2: &mut [f32],
3506 pool: Option<&Pool>,
3507) {
3508 debug_assert_eq!(o1.len(), rows);
3509 debug_assert_eq!(o2.len(), rows);
3510
3511 if a8w8_enabled() {
3512 let a1 = split_act(x1);
3513 let a2 = split_act(x2);
3514 let p1 = SendMut(o1.as_mut_ptr());
3515 let p2 = SendMut(o2.as_mut_ptr());
3516 let run = move |start: usize, end: usize| {
3517 vbit_range2_a8w8(
3518 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3519 )
3520 };
3521 dispatch_rows(pool, rows, &run);
3522 return;
3523 }
3524
3525 let p1 = SendMut(o1.as_mut_ptr());
3526 let p2 = SendMut(o2.as_mut_ptr());
3527 let run = move |start: usize, end: usize| {
3528 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3529 };
3530 dispatch_rows(pool, rows, &run);
3531}
3532
3533#[allow(clippy::too_many_arguments)]
3537fn vbit_range2_a8w8(
3538 bytes: &[u8],
3539 offsets: &[usize],
3540 x1: &[f32],
3541 x2: &[f32],
3542 a1: &SplitAct,
3543 a2: &SplitAct,
3544 rows: usize,
3545 cols: usize,
3546 p1: SendMut,
3547 p2: SendMut,
3548 start: usize,
3549 end: usize,
3550) {
3551 let ng = cols / GROUP_SIZE;
3552 let bits = &bytes[..rows];
3553 let sc_off = rows;
3554 let row_dots = |r: usize| -> (f32, f32) {
3555 let b = bits[r] as usize;
3556 let l = (1i32 << (b - 1)) - 1;
3557 let data = &bytes[offsets[r]..offsets[r + 1]];
3558 if b == 8 {
3559 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3562 let (mut d1, mut d2) = (0f32, 0f32);
3563 for g in 0..ng {
3564 let so = (r * ng + g) * 2;
3565 let sgf = f16_to_f32(u16::from_le_bytes([
3566 bytes[sc_off + so],
3567 bytes[sc_off + so + 1],
3568 ]));
3569 let (mut g1, mut g2) = (0f32, 0f32);
3570 for k in 0..GROUP_SIZE {
3571 if nbits < 8 {
3572 acc = (acc << 8) | data[idx] as u64;
3573 idx += 1;
3574 nbits += 8;
3575 }
3576 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3577 nbits -= 8;
3578 let w = (u - l) as f32;
3579 g1 += w * x1[g * GROUP_SIZE + k];
3580 g2 += w * x2[g * GROUP_SIZE + k];
3581 }
3582 d1 += g1 * sgf;
3583 d2 += g2 * sgf;
3584 }
3585 return (d1, d2);
3586 }
3587 thread_local! {
3588 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3589 const { std::cell::RefCell::new(Vec::new()) };
3590 }
3591 #[inline(always)]
3592 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3593 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3594 let u = unpack8::<B>(&data[blk * B..]);
3595 for k in 0..8 {
3596 chunk[k] = (u[k] - l) as i8 as u8;
3597 }
3598 }
3599 }
3600 VBIT_SCRATCH2.with(|scratch| {
3601 let mut buf = scratch.borrow_mut();
3602 buf.resize(cols, 0);
3603 match b {
3604 3 => fill::<3>(data, l, &mut buf),
3605 4 => vbit_fill4(data, &mut buf),
3606 5 => fill::<5>(data, l, &mut buf),
3607 6 => fill::<6>(data, l, &mut buf),
3608 _ => unreachable!(),
3609 }
3610 let (mut d1, mut d2) = (0f32, 0f32);
3611 for g in 0..ng {
3612 let so = (r * ng + g) * 2;
3613 let s = f16_to_f32(u16::from_le_bytes([
3614 bytes[sc_off + so],
3615 bytes[sc_off + so + 1],
3616 ]));
3617 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3618 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3619 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3620 d1 += v1 * s;
3621 d2 += v2 * s;
3622 }
3623 for &(j, xv) in &a1.outliers {
3624 let so = (r * ng + j / GROUP_SIZE) * 2;
3625 let s = f16_to_f32(u16::from_le_bytes([
3626 bytes[sc_off + so],
3627 bytes[sc_off + so + 1],
3628 ]));
3629 d1 += (buf[j] as i8) as f32 * s * xv;
3630 }
3631 for &(j, xv) in &a2.outliers {
3632 let so = (r * ng + j / GROUP_SIZE) * 2;
3633 let s = f16_to_f32(u16::from_le_bytes([
3634 bytes[sc_off + so],
3635 bytes[sc_off + so + 1],
3636 ]));
3637 d2 += (buf[j] as i8) as f32 * s * xv;
3638 }
3639 (d1, d2)
3640 })
3641 };
3642 for r in start..end {
3643 let (v1, v2) = row_dots(r);
3644 unsafe {
3646 *p1.at(r) = v1;
3647 *p2.at(r) = v2;
3648 }
3649 }
3650}
3651
3652#[allow(clippy::too_many_arguments)]
3656fn vbit_range2_f32(
3657 bytes: &[u8],
3658 offsets: &[usize],
3659 x1: &[f32],
3660 x2: &[f32],
3661 rows: usize,
3662 cols: usize,
3663 p1: SendMut,
3664 p2: SendMut,
3665 start: usize,
3666 end: usize,
3667) {
3668 let ng = cols / GROUP_SIZE;
3669 let bits = &bytes[..rows];
3670 let sc_off = rows;
3671 #[inline(always)]
3672 #[allow(clippy::too_many_arguments)]
3673 fn dot_row2<const B: usize>(
3674 data: &[u8],
3675 bytes: &[u8],
3676 sc_off: usize,
3677 r: usize,
3678 ng: usize,
3679 x1: &[f32],
3680 x2: &[f32],
3681 ) -> (f32, f32) {
3682 let l = ((1i32 << (B - 1)) - 1) as f32;
3683 let gbytes = GROUP_SIZE * B / 8;
3684 let (mut d1, mut d2) = (0f32, 0f32);
3685 for g in 0..ng {
3686 let so = (r * ng + g) * 2;
3687 let s = f16_to_f32(u16::from_le_bytes([
3688 bytes[sc_off + so],
3689 bytes[sc_off + so + 1],
3690 ]));
3691 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3692 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3693 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3694 let (mut g1, mut g2) = (0f32, 0f32);
3695 for blk in 0..GROUP_SIZE / 8 {
3696 let u = unpack8::<B>(&gd0[blk * B..]);
3697 for k in 0..8 {
3698 let w = u[k] as f32 - l;
3699 g1 += w * x1g[blk * 8 + k];
3700 g2 += w * x2g[blk * 8 + k];
3701 }
3702 }
3703 d1 += g1 * s;
3704 d2 += g2 * s;
3705 }
3706 (d1, d2)
3707 }
3708 for r in start..end {
3709 let data = &bytes[offsets[r]..offsets[r + 1]];
3710 let (v1, v2) = match bits[r] {
3711 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3712 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3713 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3714 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3715 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3716 b => unreachable!("vbit bit-width {b} (validated at load)"),
3717 };
3718 unsafe {
3720 *p1.at(r) = v1;
3721 *p2.at(r) = v2;
3722 }
3723 }
3724}
3725
3726#[inline]
3733#[allow(unreachable_code)]
3734fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3735 #[cfg(target_arch = "aarch64")]
3736 unsafe {
3737 return dot_q4t_row_sdot(bytes, r, gpr, xq);
3738 }
3739 #[cfg(target_arch = "x86_64")]
3740 unsafe {
3741 if vnni_tiles_enabled() {
3742 return dot_q4t_row_vnni(bytes, r, gpr, xq);
3743 }
3744 return dot_q4t_row_avx2(bytes, r, gpr, xq);
3745 }
3746 let mut acc = 0f32;
3747 for gi in 0..gpr {
3748 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3749 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3750 let mut d = 0i32;
3751 for (k, &b) in tile[2..].iter().enumerate() {
3752 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3753 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3754 }
3755 acc += d as f32 * s;
3756 }
3757 acc
3758}
3759
3760#[cfg(target_arch = "aarch64")]
3761#[target_feature(enable = "neon,dotprod")]
3762unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3763 unsafe {
3766 use core::arch::aarch64::*;
3767 use core::arch::asm;
3768 let lomask = vdupq_n_u8(0x0F);
3769 let eight = vdupq_n_s8(8);
3770 let mut acc = 0f32;
3771 for gi in 0..gpr {
3772 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3773 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3774 let b = vld1q_u8(t.add(2));
3775 let lo = vandq_u8(b, lomask);
3776 let hi = vshrq_n_u8::<4>(b);
3777 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3778 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3779 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3780 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3781 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3782 asm!(
3783 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3784 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3785 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3786 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3787 options(pure, nomem, nostack),
3788 );
3789 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3790 }
3791 acc
3792 }
3793}
3794
3795#[cfg(target_arch = "x86_64")]
3796#[target_feature(enable = "avx2")]
3797unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3798 unsafe {
3800 use core::arch::x86_64::*;
3801 let lomask = _mm_set1_epi8(0x0F);
3802 let eight = _mm256_set1_epi8(8);
3803 let ones = _mm256_set1_epi16(1);
3804 let mut acc = 0f32;
3805 for gi in 0..gpr {
3806 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3807 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3808 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3809 let lo = _mm_and_si128(b, lomask);
3810 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3811 let w = _mm256_sub_epi8(
3812 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3813 eight,
3814 );
3815 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3816 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3817 let d = _mm256_madd_epi16(p16, ones);
3818 let hi128 = _mm256_extracti128_si256::<1>(d);
3819 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3820 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3821 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3822 acc += _mm_cvtsi128_si32(s32) as f32 * s;
3823 }
3824 acc
3825 }
3826}
3827
3828#[cfg(target_arch = "x86_64")]
3832#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3833unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3834 unsafe {
3836 use core::arch::x86_64::*;
3837 let lomask = _mm_set1_epi8(0x0F);
3838 let eight = _mm256_set1_epi8(8);
3839 let mut acc = 0f32;
3840 for gi in 0..gpr {
3841 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3842 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3843 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3844 let lo = _mm_and_si128(b, lomask);
3845 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3846 let w = _mm256_sub_epi8(
3847 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3848 eight,
3849 );
3850 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3851 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3852 acc += d as f32 * s;
3853 }
3854 acc
3855 }
3856}
3857
3858#[cfg(target_arch = "x86_64")]
3863#[target_feature(enable = "avx2,fma")]
3868unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3869 unsafe {
3871 use core::arch::x86_64::*;
3872 let lomask = _mm_set1_epi8(0x0F);
3873 let eight = _mm256_set1_epi8(8);
3874 let ones = _mm256_set1_epi16(1);
3875 let mut f0 = _mm256_setzero_ps();
3888 let mut f1 = _mm256_setzero_ps();
3889 let mut f2 = _mm256_setzero_ps();
3890 let mut f3 = _mm256_setzero_ps();
3891 for gi in 0..gpr {
3892 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3893 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3894 let sv = _mm256_set1_ps(s);
3895 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3896 let lo = _mm_and_si128(bb, lomask);
3897 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3898 let w = _mm256_sub_epi8(
3899 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3900 eight,
3901 );
3902 let aw = _mm256_abs_epi8(w);
3903 let off = gi * GROUP_SIZE;
3904 let dot = |xq: &[i8]| {
3905 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3906 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3907 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
3908 };
3909 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3910 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3911 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3912 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3913 }
3914 [
3915 hsum256_ps(f0),
3916 hsum256_ps(f1),
3917 hsum256_ps(f2),
3918 hsum256_ps(f3),
3919 ]
3920 }
3921}
3922
3923#[cfg(target_arch = "x86_64")]
3926#[target_feature(enable = "avx2")]
3927#[inline]
3928unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
3929 unsafe {
3931 use core::arch::x86_64::*;
3932 let hi = _mm256_extractf128_ps::<1>(v);
3933 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
3934 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
3935 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
3936 _mm_cvtss_f32(s)
3937 }
3938}
3939
3940#[cfg(target_arch = "x86_64")]
3942#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
3943unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3944 unsafe {
3946 use core::arch::x86_64::*;
3947 let lomask = _mm_set1_epi8(0x0F);
3948 let eight = _mm256_set1_epi8(8);
3949 let mut f0 = _mm256_setzero_ps();
3952 let mut f1 = _mm256_setzero_ps();
3953 let mut f2 = _mm256_setzero_ps();
3954 let mut f3 = _mm256_setzero_ps();
3955 for gi in 0..gpr {
3956 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3957 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3958 let sv = _mm256_set1_ps(s);
3959 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3960 let lo = _mm_and_si128(bb, lomask);
3961 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3962 let w = _mm256_sub_epi8(
3963 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3964 eight,
3965 );
3966 let aw = _mm256_abs_epi8(w);
3967 let off = gi * GROUP_SIZE;
3968 let dot = |xq: &[i8]| {
3969 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3970 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
3971 _mm256_setzero_si256(),
3972 aw,
3973 _mm256_sign_epi8(x, w),
3974 ))
3975 };
3976 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3977 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3978 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3979 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3980 }
3981 let acc = [
3982 hsum256_ps(f0),
3983 hsum256_ps(f1),
3984 hsum256_ps(f2),
3985 hsum256_ps(f3),
3986 ];
3987 acc
3988 }
3989}
3990
3991#[cfg(target_arch = "aarch64")]
3996#[target_feature(enable = "neon,dotprod")]
3997unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3998 unsafe {
4000 use core::arch::aarch64::*;
4001 use core::arch::asm;
4002 let lomask = vdupq_n_u8(0x0F);
4003 let eight = vdupq_n_s8(8);
4004 let mut acc = [0f32; 4];
4005 for gi in 0..gpr {
4006 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4007 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4008 let b = vld1q_u8(t.add(2));
4009 let lo = vandq_u8(b, lomask);
4010 let hi = vshrq_n_u8::<4>(b);
4011 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4012 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4013 for (k, xq) in xs.iter().enumerate() {
4014 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4015 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4016 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4017 asm!(
4018 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4019 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4020 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4021 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4022 options(pure, nomem, nostack),
4023 );
4024 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4025 }
4026 }
4027 acc
4028 }
4029}
4030
4031#[inline]
4033fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4034 let gi = j / GROUP_SIZE;
4035 let k = j % GROUP_SIZE;
4036 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4037 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4038 let byte = tile[2 + k / 2];
4039 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4040 ((nib as i32 - 8) as f32, s)
4041}
4042
4043#[inline]
4046fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4047 let mut acc = 0f32;
4048 for gi in 0..gpr {
4049 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4050 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4051 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4052 let mut ga = 0f32;
4053 for (k, &b) in tile[2..].iter().enumerate() {
4054 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4055 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4056 }
4057 acc += ga * s;
4058 }
4059 acc
4060}
4061
4062struct Q4tpView<'a> {
4066 nib: &'a [u8],
4067 params: &'a [u8],
4068 codes: &'a [u8],
4069 stride: usize,
4070 zero_rung: bool,
4072}
4073
4074impl<'a> Q4tpView<'a> {
4075 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4076 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4077 Self {
4078 nib: &bytes[..params_off],
4079 params: &bytes[params_off..codes_off],
4080 codes: &bytes[codes_off..],
4081 stride,
4082 zero_rung: false,
4083 }
4084 }
4085
4086 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4088 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4089 Self {
4090 nib: &bytes[..params_off],
4091 params: &bytes[params_off..codes_off],
4092 codes: &bytes[codes_off..],
4093 stride,
4094 zero_rung: true,
4095 }
4096 }
4097
4098 #[inline]
4114 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4115 let tab = if self.zero_rung {
4116 q2tp_ladder(self.params, r)
4117 } else {
4118 q4tp_ladder(self.params, r)
4119 };
4120 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4121 let out = &mut out[..gpr];
4122 let mut chunks = out.chunks_exact_mut(8);
4123 let mut ci = 0usize;
4124 for c in &mut chunks {
4125 let w = u64::from(codes[ci])
4126 | u64::from(codes[ci + 1]) << 8
4127 | u64::from(codes[ci + 2]) << 16
4128 | u64::from(codes[ci + 3]) << 24
4129 | u64::from(codes[ci + 4]) << 32;
4130 for (k, o) in c.iter_mut().enumerate() {
4131 *o = tab[((w >> (5 * k)) & 31) as usize];
4132 }
4133 ci += 5;
4134 }
4135 let tail = &codes[ci..];
4138 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4139 *o = tab[q4tp_code(tail, k)];
4140 }
4141 }
4142}
4143
4144#[inline]
4145fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4146 #[cfg(target_arch = "aarch64")]
4147 unsafe {
4148 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4149 }
4150 #[cfg(target_arch = "x86_64")]
4151 unsafe {
4152 if vnni_tiles_enabled() {
4153 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4154 }
4155 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4156 }
4157 #[allow(unreachable_code)]
4158 {
4159 let mut acc = 0f32;
4160 for gi in 0..gpr {
4161 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4162 let s = scales[gi];
4163 let mut d = 0i32;
4164 for (k, &b) in tile.iter().enumerate() {
4165 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4166 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4167 }
4168 acc += d as f32 * s;
4169 }
4170 acc
4171 }
4172}
4173
4174#[cfg(target_arch = "aarch64")]
4177#[target_feature(enable = "neon,dotprod")]
4178unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4179 unsafe {
4182 use core::arch::aarch64::*;
4183 use core::arch::asm;
4184 let lomask = vdupq_n_u8(0x0F);
4185 let eight = vdupq_n_s8(8);
4186 let mut acc = 0f32;
4187 for gi in 0..gpr {
4188 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4189 let s = *scales.get_unchecked(gi);
4190 let b = vld1q_u8(t);
4191 let lo = vandq_u8(b, lomask);
4192 let hi = vshrq_n_u8::<4>(b);
4193 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4194 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4195 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4196 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4197 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4198 asm!(
4199 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4200 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4201 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4202 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4203 options(pure, nomem, nostack),
4204 );
4205 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4206 }
4207 acc
4208 }
4209}
4210
4211#[cfg(target_arch = "x86_64")]
4212#[target_feature(enable = "avx2")]
4213unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4214 unsafe {
4216 use core::arch::x86_64::*;
4217 let lomask = _mm_set1_epi8(0x0F);
4218 let eight = _mm256_set1_epi8(8);
4219 let ones = _mm256_set1_epi16(1);
4220 let mut acc = 0f32;
4221 for gi in 0..gpr {
4222 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4223 let s = *scales.get_unchecked(gi);
4224 let b = _mm_loadu_si128(t as *const __m128i);
4225 let lo = _mm_and_si128(b, lomask);
4226 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4227 let w = _mm256_sub_epi8(
4228 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4229 eight,
4230 );
4231 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4232 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4233 let d = _mm256_madd_epi16(p16, ones);
4234 let hi128 = _mm256_extracti128_si256::<1>(d);
4235 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4236 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4237 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4238 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4239 }
4240 acc
4241 }
4242}
4243
4244
4245#[cfg(target_arch = "x86_64")]
4248#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4249unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4250 unsafe {
4252 use core::arch::x86_64::*;
4253 let lomask = _mm_set1_epi8(0x0F);
4254 let eight = _mm256_set1_epi8(8);
4255 let mut acc = 0f32;
4256 for gi in 0..gpr {
4257 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4258 let s = *scales.get_unchecked(gi);
4259 let b = _mm_loadu_si128(t as *const __m128i);
4260 let lo = _mm_and_si128(b, lomask);
4261 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4262 let w = _mm256_sub_epi8(
4263 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4264 eight,
4265 );
4266 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4267 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4268 }
4269 acc
4270 }
4271}
4272
4273#[inline]
4276fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4277 let mut acc = 0f32;
4278 for gi in 0..gpr {
4279 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4280 let s = scales[gi];
4281 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4282 let mut ga = 0f32;
4283 for (k, &b) in tile.iter().enumerate() {
4284 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4285 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4286 }
4287 acc += ga * s;
4288 }
4289 acc
4290}
4291
4292#[inline]
4295fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4296 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4297 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4298 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4299 ((n as i32 - 8) as f32, scales[gi])
4300}
4301
4302fn q4tp_matvec(
4304 bytes: &[u8],
4305 x: &[f32],
4306 rows: usize,
4307 cols: usize,
4308 out: &mut [f32],
4309 pool: Option<&Pool>,
4310) {
4311 debug_assert_eq!(out.len(), rows);
4312 let gpr = cols / GROUP_SIZE;
4313 let v = Q4tpView::new(bytes, rows, cols);
4314 let out_addr = SendMut(out.as_mut_ptr());
4315 if a8w8_enabled() {
4316 let act = split_act(x);
4317 let run = |start: usize, end: usize| {
4318 with_krow(gpr, |sc| {
4320 for r in start..end {
4321 v.scales_into(r, gpr, sc);
4322 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4323 for &(j, xv) in &act.outliers {
4324 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4325 acc += w * s * xv;
4326 }
4327 unsafe { *out_addr.at(r) = acc };
4329 }
4330 })
4331 };
4332 dispatch_rows(pool, rows, &run);
4333 return;
4334 }
4335 let run = |start: usize, end: usize| {
4336 with_krow(gpr, |sc| {
4337 for r in start..end {
4338 v.scales_into(r, gpr, sc);
4339 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4341 }
4342 })
4343 };
4344 dispatch_rows(pool, rows, &run);
4345}
4346
4347#[allow(clippy::too_many_arguments)]
4350fn q4tp_matvec2(
4351 bytes: &[u8],
4352 x1: &[f32],
4353 x2: &[f32],
4354 rows: usize,
4355 cols: usize,
4356 o1: &mut [f32],
4357 o2: &mut [f32],
4358 pool: Option<&Pool>,
4359) {
4360 let gpr = cols / GROUP_SIZE;
4361 let v = Q4tpView::new(bytes, rows, cols);
4362 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4363 let run = |start: usize, end: usize| {
4364 let mut sc = vec![0f32; gpr];
4365 for r in start..end {
4366 v.scales_into(r, gpr, &mut sc);
4367 unsafe {
4369 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4370 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4371 }
4372 }
4373 };
4374 dispatch_rows(pool, rows, &run);
4375}
4376
4377fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4381 let mut acc = 0f32;
4382 for gi in 0..gpr {
4383 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4384 let s = scales[gi];
4385 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4386 let mut g = 0f32;
4387 for (k, &b) in ch.iter().enumerate() {
4388 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4389 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4390 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4391 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4392 }
4393 acc += s * g;
4394 }
4395 acc
4396}
4397
4398fn q2tp_matvec(
4399 bytes: &[u8],
4400 x: &[f32],
4401 rows: usize,
4402 cols: usize,
4403 out: &mut [f32],
4404 pool: Option<&Pool>,
4405) {
4406 debug_assert_eq!(out.len(), rows);
4407 let gpr = cols / GROUP_SIZE;
4408 let v = Q4tpView::new_q2(bytes, rows, cols);
4409 let out_addr = SendMut(out.as_mut_ptr());
4410 let run = |start: usize, end: usize| {
4411 with_krow(gpr, |sc| {
4412 for r in start..end {
4413 v.scales_into(r, gpr, sc);
4414 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4416 }
4417 })
4418 };
4419 dispatch_rows(pool, rows, &run);
4420}
4421
4422#[allow(clippy::too_many_arguments)]
4424fn q2tp_matvec2(
4425 bytes: &[u8],
4426 x1: &[f32],
4427 x2: &[f32],
4428 rows: usize,
4429 cols: usize,
4430 o1: &mut [f32],
4431 o2: &mut [f32],
4432 pool: Option<&Pool>,
4433) {
4434 let gpr = cols / GROUP_SIZE;
4435 let v = Q4tpView::new_q2(bytes, rows, cols);
4436 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4437 let run = |start: usize, end: usize| {
4438 let mut sc = vec![0f32; gpr];
4439 for r in start..end {
4440 v.scales_into(r, gpr, &mut sc);
4441 unsafe {
4443 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4444 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4445 }
4446 }
4447 };
4448 dispatch_rows(pool, rows, &run);
4449}
4450
4451pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4458 q2tp_matvec(bytes, x, rows, cols, out, None);
4459}
4460
4461pub fn q2tp_matmat_for_test(
4462 bytes: &[u8],
4463 xs_all: &[f32],
4464 b: usize,
4465 rows: usize,
4466 cols: usize,
4467 out: &mut [f32],
4468) {
4469 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4470}
4471
4472fn q2tp_matmat(
4473 bytes: &[u8],
4474 xs_all: &[f32],
4475 b: usize,
4476 rows: usize,
4477 cols: usize,
4478 out: &mut [f32],
4479 pool: Option<&Pool>,
4480) {
4481 debug_assert_eq!(out.len(), b * rows);
4482 let gpr = cols / GROUP_SIZE;
4483 let v = Q4tpView::new_q2(bytes, rows, cols);
4484 let out_addr = SendMut(out.as_mut_ptr());
4485 let run = |start: usize, end: usize| {
4486 let mut sc = vec![0f32; gpr];
4487 for r in start..end {
4488 v.scales_into(r, gpr, &mut sc);
4489 for bi in 0..b {
4490 let x = &xs_all[bi * cols..(bi + 1) * cols];
4491 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4493 }
4494 }
4495 };
4496 dispatch_rows(pool, rows, &run);
4497}
4498
4499#[cfg(target_arch = "aarch64")]
4503#[target_feature(enable = "neon,dotprod")]
4504unsafe fn dot_q4tp_row_1x4_sdot_v1(
4505 nib: &[u8],
4506 r: usize,
4507 gpr: usize,
4508 xs: [&[i8]; 4],
4509 scales: &[f32],
4510) -> [f32; 4] {
4511 unsafe {
4512 use core::arch::aarch64::*;
4513 use core::arch::asm;
4514 let lomask = vdupq_n_u8(0x0F);
4515 let eight = vdupq_n_s8(8);
4516 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4517 for gi in 0..gpr {
4518 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4519 let s = *scales.get_unchecked(gi);
4520 let bb = vld1q_u8(t);
4521 let lo = vandq_u8(bb, lomask);
4522 let hi = vshrq_n_u8::<4>(bb);
4523 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4524 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4525 let mut d = [0f32; 4];
4526 for (k, dk) in d.iter_mut().enumerate() {
4527 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4528 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4529 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4530 asm!(
4531 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4532 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4533 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4534 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4535 options(pure, nomem, nostack),
4536 );
4537 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4538 }
4539 f0 += d[0];
4540 f1 += d[1];
4541 f2 += d[2];
4542 f3 += d[3];
4543 }
4544 [f0, f1, f2, f3]
4545 }
4546}
4547
4548#[allow(dead_code)]
4555static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4556
4557#[cfg(target_arch = "x86_64")]
4563fn q4tp_blocked_x86() -> bool {
4564 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4565 1 => false,
4566 2 => avx512vnni_enabled(),
4571 _ => blocked_enabled() && avx512vnni_enabled(),
4575 }
4576}
4577
4578#[cfg(target_arch = "aarch64")]
4580#[allow(dead_code)]
4581fn q4tp_v1() -> bool {
4582 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4583 1 => true,
4584 2 => false,
4585 _ => {
4586 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4587 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4588 }
4589 }
4590}
4591
4592#[cfg(target_arch = "x86_64")]
4604#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4605unsafe fn dot_q4tp_2x8_avx512(
4606 nib: &[u8],
4607 r0: usize,
4608 gpr: usize,
4609 xs: [&[i8]; 8],
4610 sc0: &[f32],
4611 sc1: &[f32],
4612) -> [[f32; 8]; 2] {
4613 unsafe {
4616 use core::arch::x86_64::*;
4617 let lomask = _mm256_set1_epi8(0x0F);
4618 let eight = _mm256_set1_epi8(8);
4619 let zero = _mm512_setzero_si512();
4620 let mut v0 = [_mm512_setzero_ps(); 8];
4621 let mut v1 = [_mm512_setzero_ps(); 8];
4622 let pairs = gpr / 2;
4623 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4624 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4625 let bb = _mm256_loadu_si256(t as *const __m256i);
4626 let lo = _mm256_and_si256(bb, lomask);
4627 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4628 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4629 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4630 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4631 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4632 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
4633 };
4634 for gp in 0..pairs {
4635 let gi = gp * 2;
4636 let (wa0, neg0) = unpack(r0, gi);
4637 let (wa1, neg1) = unpack(r0 + 1, gi);
4638 let off = gi * GROUP_SIZE;
4639 let sv = |sc: &[f32]| {
4640 _mm512_insertf32x8::<1>(
4641 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
4642 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
4643 )
4644 };
4645 let s0 = sv(sc0);
4646 let s1 = sv(sc1);
4647 for k in 0..8 {
4648 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
4649 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4650 zero,
4651 wa0,
4652 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
4653 ));
4654 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4655 zero,
4656 wa1,
4657 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
4658 ));
4659 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
4660 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
4661 }
4662 }
4663 let mut acc = [[0f32; 8]; 2];
4664 for k in 0..8 {
4665 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
4666 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
4667 }
4668 if gpr % 2 == 1 {
4669 let off = (gpr - 1) * GROUP_SIZE;
4670 for j in off..off + GROUP_SIZE {
4671 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
4672 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
4673 for k in 0..8 {
4674 let x = *xs[k].get_unchecked(j) as f32;
4675 acc[0][k] += w0 * sa * x;
4676 acc[1][k] += w1 * sb * x;
4677 }
4678 }
4679 }
4680 acc
4681 }
4682}
4683
4684#[cfg(target_arch = "x86_64")]
4689#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4690unsafe fn dot_q4tp_row_1x8_avx512(
4691 nib: &[u8],
4692 r: usize,
4693 gpr: usize,
4694 xs: [&[i8]; 8],
4695 scales: &[f32],
4696) -> [f32; 8] {
4697 unsafe {
4699 use core::arch::x86_64::*;
4700 let lomask = _mm256_set1_epi8(0x0F);
4701 let eight = _mm256_set1_epi8(8);
4702 let zero = _mm512_setzero_si512();
4703 let (mut v0, mut v1, mut v2, mut v3) = (
4704 _mm512_setzero_ps(),
4705 _mm512_setzero_ps(),
4706 _mm512_setzero_ps(),
4707 _mm512_setzero_ps(),
4708 );
4709 let (mut v4, mut v5, mut v6, mut v7) = (
4710 _mm512_setzero_ps(),
4711 _mm512_setzero_ps(),
4712 _mm512_setzero_ps(),
4713 _mm512_setzero_ps(),
4714 );
4715 let pairs = gpr / 2;
4716 for gp in 0..pairs {
4717 let gi = gp * 2;
4718 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4719 let bb = _mm256_loadu_si256(t as *const __m256i);
4720 let lo = _mm256_and_si256(bb, lomask);
4721 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4722 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4727 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4728 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4729 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4730 let wabs = _mm512_abs_epi8(w);
4731 let neg = _mm512_movepi8_mask(w);
4732 let off = gi * GROUP_SIZE;
4733 let sv = _mm512_insertf32x8::<1>(
4734 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
4735 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
4736 );
4737 let dot = |x: &[i8]| -> __m512 {
4738 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
4739 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
4740 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
4741 };
4742 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
4743 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
4744 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
4745 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
4746 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
4747 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
4748 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
4749 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
4750 }
4751 let mut acc = [
4752 _mm512_reduce_add_ps(v0),
4753 _mm512_reduce_add_ps(v1),
4754 _mm512_reduce_add_ps(v2),
4755 _mm512_reduce_add_ps(v3),
4756 _mm512_reduce_add_ps(v4),
4757 _mm512_reduce_add_ps(v5),
4758 _mm512_reduce_add_ps(v6),
4759 _mm512_reduce_add_ps(v7),
4760 ];
4761 if gpr % 2 == 1 {
4764 let off = (gpr - 1) * GROUP_SIZE;
4765 for j in off..off + GROUP_SIZE {
4766 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
4767 let ws = w * s;
4768 for k in 0..8 {
4769 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
4770 }
4771 }
4772 }
4773 acc
4774 }
4775}
4776
4777#[cfg(target_arch = "x86_64")]
4790#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4791unsafe fn dot_q4tp_row_1x4_avx512(
4792 nib: &[u8],
4793 r: usize,
4794 gpr: usize,
4795 xs: [&[i8]; 4],
4796 scales: &[f32],
4797) -> [f32; 4] {
4798 unsafe {
4800 use core::arch::x86_64::*;
4801 let lomask = _mm256_set1_epi8(0x0F);
4802 let eight = _mm256_set1_epi8(8);
4803 let zero = _mm512_setzero_si512();
4804 let (mut v0, mut v1, mut v2, mut v3) = (
4805 _mm512_setzero_ps(),
4806 _mm512_setzero_ps(),
4807 _mm512_setzero_ps(),
4808 _mm512_setzero_ps(),
4809 );
4810 let pairs = gpr / 2;
4811 for gp in 0..pairs {
4812 let gi = gp * 2;
4813 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4814 let bb = _mm256_loadu_si256(t as *const __m256i);
4815 let lo = _mm256_and_si256(bb, lomask);
4816 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4817 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4822 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4823 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4824 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4825 let wabs = _mm512_abs_epi8(w);
4826 let neg = _mm512_movepi8_mask(w);
4827 let off = gi * GROUP_SIZE;
4828 let sv = _mm512_insertf32x8::<1>(
4829 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
4830 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
4831 );
4832 let dot = |x: &[i8]| -> __m512 {
4833 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
4834 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
4835 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
4836 };
4837 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
4838 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
4839 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
4840 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
4841 }
4842 let mut acc = [
4843 _mm512_reduce_add_ps(v0),
4844 _mm512_reduce_add_ps(v1),
4845 _mm512_reduce_add_ps(v2),
4846 _mm512_reduce_add_ps(v3),
4847 ];
4848 if gpr % 2 == 1 {
4851 let off = (gpr - 1) * GROUP_SIZE;
4852 for j in off..off + GROUP_SIZE {
4853 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
4854 let ws = w * s;
4855 for k in 0..4 {
4856 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
4857 }
4858 }
4859 }
4860 acc
4861 }
4862}
4863
4864#[cfg(target_arch = "aarch64")]
4868#[target_feature(enable = "neon,dotprod")]
4869unsafe fn dot_q4tp_row_1x4_sdot(
4870 nib: &[u8],
4871 r: usize,
4872 gpr: usize,
4873 xs: [&[i8]; 4],
4874 scales: &[f32],
4875) -> [f32; 4] {
4876 unsafe {
4878 use core::arch::aarch64::*;
4879 use core::arch::asm;
4880 let lomask = vdupq_n_u8(0x0F);
4881 let eight = vdupq_n_s8(8);
4882 let (mut v0, mut v1, mut v2, mut v3) = (
4898 vdupq_n_f32(0.0),
4899 vdupq_n_f32(0.0),
4900 vdupq_n_f32(0.0),
4901 vdupq_n_f32(0.0),
4902 );
4903 for gi in 0..gpr {
4904 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4905 let s = *scales.get_unchecked(gi);
4906 let bb = vld1q_u8(t);
4907 let lo = vandq_u8(bb, lomask);
4908 let hi = vshrq_n_u8::<4>(bb);
4909 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4910 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4911 let off = gi * GROUP_SIZE;
4912 let dot4 = |x: &[i8]| -> int32x4_t {
4913 let x0 = vld1q_s8(x.as_ptr().add(off));
4914 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
4915 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4916 asm!(
4917 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4918 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4919 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4920 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4921 options(pure, nomem, nostack),
4922 );
4923 vaddq_s32(a0, a1)
4924 };
4925 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
4926 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
4927 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
4928 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
4929 }
4930 [
4931 vaddvq_f32(v0),
4932 vaddvq_f32(v1),
4933 vaddvq_f32(v2),
4934 vaddvq_f32(v3),
4935 ]
4936 }
4937}
4938
4939fn q4tp_matmat(
4943 bytes: &[u8],
4944 xs_all: &[f32],
4945 b: usize,
4946 rows: usize,
4947 cols: usize,
4948 out: &mut [f32],
4949 pool: Option<&Pool>,
4950) {
4951 debug_assert_eq!(out.len(), b * rows);
4952 let gpr = cols / GROUP_SIZE;
4953 let v = Q4tpView::new(bytes, rows, cols);
4954
4955 #[cfg(target_os = "macos")]
4957 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4958 dequant_matmat_accel(
4959 &|r, dst| {
4960 let mut sc = [0f32; 32];
4961 let mut scv;
4962 let s: &[f32] = if gpr <= 32 {
4963 v.scales_into(r, gpr, &mut sc);
4964 &sc[..gpr]
4965 } else {
4966 scv = vec![0f32; gpr];
4967 v.scales_into(r, gpr, &mut scv);
4968 &scv
4969 };
4970 for gi in 0..gpr {
4971 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4972 for (k, &bb) in tile.iter().enumerate() {
4973 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
4974 dst[gi * GROUP_SIZE + k * 2 + 1] =
4975 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
4976 }
4977 }
4978 },
4979 xs_all,
4980 b,
4981 rows,
4982 cols,
4983 out,
4984 pool,
4985 );
4986 return;
4987 }
4988
4989 let out_addr = SendMut(out.as_mut_ptr());
4990 if a8w8_enabled() {
4991 let acts: Vec<SplitAct> = (0..b)
4992 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4993 .collect();
4994 let acts = &acts;
4995 #[cfg(target_arch = "aarch64")]
4996 let blocked_ok = sdot_enabled() && blocked_enabled();
4997 #[cfg(target_arch = "x86_64")]
5004 let blocked_ok = q4tp_blocked_x86();
5005 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5006 let blocked_ok = false;
5007 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5017 .ok()
5018 .and_then(|v| v.parse().ok())
5019 .filter(|v| *v > 0)
5020 .unwrap_or(256);
5021 let run = |start: usize, end: usize| {
5022 for abase in (0..acts.len()).step_by(panel_cols) {
5023 let alen = (acts.len() - abase).min(panel_cols);
5024 let mut sc = vec![0f32; gpr];
5025 #[cfg(target_arch = "x86_64")]
5026 let mut r_lo = start;
5027 #[cfg(target_arch = "x86_64")]
5028 if blocked_ok && alen >= 8 {
5029 let mut sc1 = vec![0f32; gpr];
5030 while r_lo + 2 <= end {
5031 v.scales_into(r_lo, gpr, &mut sc);
5032 v.scales_into(r_lo + 1, gpr, &mut sc1);
5033 let mut bi = 0usize;
5034 while bi + 8 <= alen {
5035 let xs = [
5036 acts[abase + bi].xq.as_slice(),
5037 acts[abase + bi + 1].xq.as_slice(),
5038 acts[abase + bi + 2].xq.as_slice(),
5039 acts[abase + bi + 3].xq.as_slice(),
5040 acts[abase + bi + 4].xq.as_slice(),
5041 acts[abase + bi + 5].xq.as_slice(),
5042 acts[abase + bi + 6].xq.as_slice(),
5043 acts[abase + bi + 7].xq.as_slice(),
5044 ];
5045 let d =
5046 unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5047 for (row, dr, scr) in
5048 [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)]
5049 {
5050 for k in 0..8 {
5051 let act = &acts[abase + bi + k];
5052 let mut acc = dr[k] * act.sx;
5053 for &(j, xv) in &act.outliers {
5054 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5055 acc += w * s * xv;
5056 }
5057 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5059 }
5060 }
5061 bi += 8;
5062 }
5063 for row in [r_lo, r_lo + 1] {
5066 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5067 for b2 in bi..alen {
5068 let act = &acts[abase + b2];
5069 let xs4 = [
5070 act.xq.as_slice(),
5071 act.xq.as_slice(),
5072 act.xq.as_slice(),
5073 act.xq.as_slice(),
5074 ];
5075 let d =
5076 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5077 let mut acc = d[0] * act.sx;
5078 for &(j, xv) in &act.outliers {
5079 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5080 acc += w * s * xv;
5081 }
5082 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5084 }
5085 }
5086 r_lo += 2;
5087 }
5088 }
5089 #[cfg(target_arch = "x86_64")]
5090 let row_start = r_lo;
5091 #[cfg(not(target_arch = "x86_64"))]
5092 let row_start = start;
5093 for r in row_start..end {
5094 v.scales_into(r, gpr, &mut sc);
5095 let mut bi = 0usize;
5096 #[cfg(target_arch = "x86_64")]
5097 if blocked_ok {
5098 while bi + 8 <= alen {
5099 let xs = [
5100 acts[abase + bi].xq.as_slice(),
5101 acts[abase + bi + 1].xq.as_slice(),
5102 acts[abase + bi + 2].xq.as_slice(),
5103 acts[abase + bi + 3].xq.as_slice(),
5104 acts[abase + bi + 4].xq.as_slice(),
5105 acts[abase + bi + 5].xq.as_slice(),
5106 acts[abase + bi + 6].xq.as_slice(),
5107 acts[abase + bi + 7].xq.as_slice(),
5108 ];
5109 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5110 for k in 0..8 {
5111 let act = &acts[abase + bi + k];
5112 let mut acc = d[k] * act.sx;
5113 for &(j, xv) in &act.outliers {
5114 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5115 acc += w * s * xv;
5116 }
5117 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5119 }
5120 bi += 8;
5121 }
5122 while bi + 4 <= alen {
5123 let xs = [
5124 acts[abase + bi].xq.as_slice(),
5125 acts[abase + bi + 1].xq.as_slice(),
5126 acts[abase + bi + 2].xq.as_slice(),
5127 acts[abase + bi + 3].xq.as_slice(),
5128 ];
5129 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5130 for k in 0..4 {
5131 let act = &acts[abase + bi + k];
5132 let mut acc = d[k] * act.sx;
5133 for &(j, xv) in &act.outliers {
5134 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5135 acc += w * s * xv;
5136 }
5137 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5139 }
5140 bi += 4;
5141 }
5142 }
5143 #[cfg(target_arch = "aarch64")]
5144 if blocked_ok {
5145 while bi + 4 <= alen {
5146 let xs = [
5147 acts[abase + bi].xq.as_slice(),
5148 acts[abase + bi + 1].xq.as_slice(),
5149 acts[abase + bi + 2].xq.as_slice(),
5150 acts[abase + bi + 3].xq.as_slice(),
5151 ];
5152 let d = unsafe {
5153 if q4tp_v1() {
5154 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5155 } else {
5156 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5157 }
5158 };
5159 for k in 0..4 {
5160 let act = &acts[abase + bi + k];
5161 let mut acc = d[k] * act.sx;
5162 for &(j, xv) in &act.outliers {
5163 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5164 acc += w * s * xv;
5165 }
5166 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5168 }
5169 bi += 4;
5170 }
5171 }
5172 let _ = blocked_ok;
5173 while bi < alen {
5174 let act = &acts[abase + bi];
5175 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5176 for &(j, xv) in &act.outliers {
5177 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5178 acc += w * s * xv;
5179 }
5180 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5182 bi += 1;
5183 }
5184 }
5185
5186 }
5187 };
5188 dispatch_rows(pool, rows, &run);
5189 return;
5190 }
5191
5192 let run = |start: usize, end: usize| {
5193 let mut sc = vec![0f32; gpr];
5194 for r in start..end {
5195 v.scales_into(r, gpr, &mut sc);
5196 for bi in 0..b {
5197 let x = &xs_all[bi * cols..(bi + 1) * cols];
5198 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5200 }
5201 }
5202 };
5203 dispatch_rows(pool, rows, &run);
5204}
5205
5206fn q4t_matvec(
5208 bytes: &[u8],
5209 x: &[f32],
5210 rows: usize,
5211 cols: usize,
5212 out: &mut [f32],
5213 pool: Option<&Pool>,
5214) {
5215 debug_assert_eq!(out.len(), rows);
5216 let gpr = cols / GROUP_SIZE;
5217 let out_addr = SendMut(out.as_mut_ptr());
5218 if a8w8_enabled() {
5219 let act = split_act(x);
5220 let run = move |start: usize, end: usize| {
5221 for r in start..end {
5222 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5223 for &(j, xv) in &act.outliers {
5224 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5225 acc += w * s * xv;
5226 }
5227 unsafe { *out_addr.at(r) = acc };
5229 }
5230 };
5231 dispatch_rows(pool, rows, &run);
5232 return;
5233 }
5234 let run = move |start: usize, end: usize| {
5235 for r in start..end {
5236 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5238 }
5239 };
5240 dispatch_rows(pool, rows, &run);
5241}
5242
5243#[allow(clippy::too_many_arguments)]
5245fn q4t_matvec2(
5246 bytes: &[u8],
5247 x1: &[f32],
5248 x2: &[f32],
5249 rows: usize,
5250 cols: usize,
5251 o1: &mut [f32],
5252 o2: &mut [f32],
5253 pool: Option<&Pool>,
5254) {
5255 let gpr = cols / GROUP_SIZE;
5256 let p1 = SendMut(o1.as_mut_ptr());
5257 let p2 = SendMut(o2.as_mut_ptr());
5258 if a8w8_enabled() {
5259 let a1 = split_act(x1);
5260 let a2 = split_act(x2);
5261 let run = move |start: usize, end: usize| {
5262 for r in start..end {
5263 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5264 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5265 for &(j, xv) in &a1.outliers {
5266 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5267 v1 += w * s * xv;
5268 }
5269 for &(j, xv) in &a2.outliers {
5270 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5271 v2 += w * s * xv;
5272 }
5273 unsafe {
5275 *p1.at(r) = v1;
5276 *p2.at(r) = v2;
5277 }
5278 }
5279 };
5280 dispatch_rows(pool, rows, &run);
5281 return;
5282 }
5283 let run = move |start: usize, end: usize| {
5284 for r in start..end {
5285 unsafe {
5287 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5288 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5289 }
5290 }
5291 };
5292 dispatch_rows(pool, rows, &run);
5293}
5294
5295#[allow(clippy::too_many_arguments)]
5297#[cfg(target_os = "macos")]
5303fn dequant_matmat_accel(
5304 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5305 xs_all: &[f32],
5306 b: usize,
5307 rows: usize,
5308 cols: usize,
5309 out: &mut [f32],
5310 pool: Option<&Pool>,
5311) {
5312 const TR: usize = 2048;
5313 thread_local! {
5314 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5315 }
5316 WTILE.with(|wt| {
5317 let mut wtile = wt.borrow_mut();
5318 wtile.resize(TR * cols, 0.0);
5319 let mut r0 = 0usize;
5320 while r0 < rows {
5321 let tr = TR.min(rows - r0);
5322 let wt_addr = SendMut(wtile.as_mut_ptr());
5323 let run = |start: usize, end: usize| {
5324 for r in start..end {
5325 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5327 dequant_row(r0 + r, dst);
5328 }
5329 };
5330 dispatch_rows(pool, tr, &run);
5331 unsafe {
5332 accel_blas::cblas_sgemm(
5333 101, 111, 112, b as i32,
5337 tr as i32,
5338 cols as i32,
5339 1.0,
5340 xs_all.as_ptr(),
5341 cols as i32,
5342 wtile.as_ptr(),
5343 cols as i32,
5344 0.0,
5345 out.as_mut_ptr().add(r0),
5346 rows as i32,
5347 );
5348 }
5349 r0 += tr;
5350 }
5351 });
5352}
5353
5354fn q4t_matmat(
5355 bytes: &[u8],
5356 xs_all: &[f32],
5357 b: usize,
5358 rows: usize,
5359 cols: usize,
5360 out: &mut [f32],
5361 pool: Option<&Pool>,
5362) {
5363 debug_assert_eq!(out.len(), b * rows);
5364 let gpr = cols / GROUP_SIZE;
5365 #[cfg(target_os = "macos")]
5369 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5370 dequant_matmat_accel(
5371 &|r, dst| {
5372 for gi in 0..gpr {
5373 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5374 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5375 for (k, &bb) in tile[2..].iter().enumerate() {
5376 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5377 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5378 }
5379 }
5380 },
5381 xs_all,
5382 b,
5383 rows,
5384 cols,
5385 out,
5386 pool,
5387 );
5388 return;
5389 }
5390 let out_addr = SendMut(out.as_mut_ptr());
5391 if a8w8_enabled() {
5392 let acts: Vec<SplitAct> = (0..b)
5393 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5394 .collect();
5395 let acts = &acts;
5396 #[cfg(target_arch = "x86_64")]
5397 let blocked_ok = avx2_enabled()
5398 && blocked_enabled();
5399 #[cfg(target_arch = "aarch64")]
5400 let blocked_ok = sdot_enabled()
5401 && blocked_enabled();
5402 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5403 let blocked_ok = false;
5404 let run = move |start: usize, end: usize| {
5405 for r in start..end {
5406 let mut bi = 0usize;
5407 #[cfg(target_arch = "aarch64")]
5408 if blocked_ok {
5409 while bi + 4 <= acts.len() {
5410 let xs = [
5411 acts[bi].xq.as_slice(),
5412 acts[bi + 1].xq.as_slice(),
5413 acts[bi + 2].xq.as_slice(),
5414 acts[bi + 3].xq.as_slice(),
5415 ];
5416 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5417 for k in 0..4 {
5418 let act = &acts[bi + k];
5419 let mut acc = d[k] * act.sx;
5420 for &(j, xv) in &act.outliers {
5421 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5422 acc += w * sc * xv;
5423 }
5424 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5426 }
5427 bi += 4;
5428 }
5429 }
5430 #[cfg(target_arch = "x86_64")]
5431 if blocked_ok {
5432 while bi + 4 <= acts.len() {
5433 let xs = [
5434 acts[bi].xq.as_slice(),
5435 acts[bi + 1].xq.as_slice(),
5436 acts[bi + 2].xq.as_slice(),
5437 acts[bi + 3].xq.as_slice(),
5438 ];
5439 let d = unsafe {
5440 if vnni_tiles_enabled() {
5441 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5442 } else {
5443 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5444 }
5445 };
5446 for k in 0..4 {
5447 let act = &acts[bi + k];
5448 let mut acc = d[k] * act.sx;
5449 for &(j, xv) in &act.outliers {
5450 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5451 acc += w * sc * xv;
5452 }
5453 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5455 }
5456 bi += 4;
5457 }
5458 }
5459 let _ = blocked_ok;
5460 while bi < acts.len() {
5461 let act = &acts[bi];
5462 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5463 for &(j, xv) in &act.outliers {
5464 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5465 acc += w * s * xv;
5466 }
5467 unsafe { *out_addr.at(bi * rows + r) = acc };
5469 bi += 1;
5470 }
5471 }
5472 };
5473 dispatch_rows(pool, rows, &run);
5474 return;
5475 }
5476 let run = move |start: usize, end: usize| {
5477 for r in start..end {
5478 for bi in 0..b {
5479 let x = &xs_all[bi * cols..(bi + 1) * cols];
5480 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5482 }
5483 }
5484 };
5485 dispatch_rows(pool, rows, &run);
5486}
5487
5488fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5497 (0..gpr)
5498 .map(|gi| {
5499 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5500 .iter()
5501 .map(|&v| v as i32)
5502 .sum()
5503 })
5504 .collect()
5505}
5506
5507#[inline]
5511#[allow(unreachable_code)]
5512#[cfg(target_arch = "x86_64")]
5517#[target_feature(enable = "avx2")]
5518unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5519 unsafe {
5521 use core::arch::x86_64::*;
5522 let expand = _mm256_setr_epi8(
5524 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,
5525 3, 3, 3,
5526 );
5527 let bitsel = _mm256_setr_epi8(
5528 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5529 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5530 );
5531 let ones8 = _mm256_set1_epi8(1);
5532 let ones16 = _mm256_set1_epi16(1);
5533 let mut acc = 0f32;
5534 for gi in 0..gpr {
5535 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5536 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5537 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5538 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5539 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5540 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5541 let sel = _mm256_and_si256(x, mask);
5542 let p16 = _mm256_maddubs_epi16(ones8, sel);
5544 let d32 = _mm256_madd_epi16(p16, ones16);
5545 let hi128 = _mm256_extracti128_si256::<1>(d32);
5546 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5547 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5548 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5549 let msum = _mm_cvtsi128_si32(s32);
5550 let d = 2 * msum - gsum[gi];
5553 acc += d as f32 * s;
5554 }
5555 acc
5556 }
5557}
5558
5559#[cfg(target_arch = "x86_64")]
5562#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5563unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5564 unsafe {
5566 use core::arch::x86_64::*;
5567 let expand = _mm256_setr_epi8(
5568 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,
5569 3, 3, 3,
5570 );
5571 let bitsel = _mm256_setr_epi8(
5572 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5573 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5574 );
5575 let ones8 = _mm256_set1_epi8(1);
5576 let mut acc = 0f32;
5577 for gi in 0..gpr {
5578 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5579 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5580 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5581 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5582 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5583 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5584 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5585 let d = 2 * msum - gsum[gi];
5586 acc += d as f32 * s;
5587 }
5588 acc
5589 }
5590}
5591
5592#[cfg(target_arch = "x86_64")]
5594#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5595unsafe fn dot_q1_row_1x4_vnni(
5596 bytes: &[u8],
5597 r: usize,
5598 gpr: usize,
5599 xs: [&[i8]; 4],
5600 gsums: [&[i32]; 4],
5601) -> [f32; 4] {
5602 unsafe {
5604 use core::arch::x86_64::*;
5605 let expand = _mm256_setr_epi8(
5606 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,
5607 3, 3, 3,
5608 );
5609 let bitsel = _mm256_setr_epi8(
5610 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5611 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5612 );
5613 let ones8 = _mm256_set1_epi8(1);
5614 let mut acc = [0f32; 4];
5615 for gi in 0..gpr {
5616 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5617 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5618 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5619 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5620 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5621 for (k, xq) in xs.iter().enumerate() {
5622 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5623 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5624 let d = 2 * msum - gsums[k][gi];
5625 acc[k] += d as f32 * s;
5626 }
5627 }
5628 acc
5629 }
5630}
5631
5632#[cfg(target_arch = "x86_64")]
5635#[target_feature(enable = "avx2")]
5636unsafe fn dot_q1_row_1x4_avx2(
5637 bytes: &[u8],
5638 r: usize,
5639 gpr: usize,
5640 xs: [&[i8]; 4],
5641 gsums: [&[i32]; 4],
5642) -> [f32; 4] {
5643 unsafe {
5645 use core::arch::x86_64::*;
5646 let expand = _mm256_setr_epi8(
5647 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,
5648 3, 3, 3,
5649 );
5650 let bitsel = _mm256_setr_epi8(
5651 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5652 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5653 );
5654 let ones8 = _mm256_set1_epi8(1);
5655 let ones16 = _mm256_set1_epi16(1);
5656 let mut acc = [0f32; 4];
5657 for gi in 0..gpr {
5658 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5659 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5660 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5661 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5662 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5663 for (k, xq) in xs.iter().enumerate() {
5664 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5665 let sel = _mm256_and_si256(x, mask);
5666 let p16 = _mm256_maddubs_epi16(ones8, sel);
5667 let d32 = _mm256_madd_epi16(p16, ones16);
5668 let hi128 = _mm256_extracti128_si256::<1>(d32);
5669 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5670 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5671 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5672 let msum = _mm_cvtsi128_si32(s32);
5673 let d = 2 * msum - gsums[k][gi];
5674 acc[k] += d as f32 * s;
5675 }
5676 }
5677 acc
5678 }
5679}
5680
5681#[allow(unreachable_code)]
5682fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5683 #[cfg(target_arch = "aarch64")]
5684 unsafe {
5685 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
5686 }
5687 #[cfg(target_arch = "x86_64")]
5688 if avx2_enabled() {
5689 unsafe {
5690 if vnni_tiles_enabled() {
5691 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
5692 }
5693 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
5694 }
5695 }
5696 let _ = gsum;
5697 let mut acc = 0f32;
5698 for gi in 0..gpr {
5699 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5700 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5701 let mut d = 0i32;
5702 for (j, &b) in tile[2..].iter().enumerate() {
5703 for k in 0..8 {
5704 let w = ((b >> k) & 1) as i32 * 2 - 1;
5705 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
5706 }
5707 }
5708 acc += d as f32 * s;
5709 }
5710 acc
5711}
5712
5713#[cfg(target_arch = "aarch64")]
5722#[target_feature(enable = "neon,dotprod")]
5723unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5724 unsafe {
5727 use core::arch::aarch64::*;
5728 use core::arch::asm;
5729 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
5730 let m = vld1q_u8(MASKS.as_ptr());
5731 macro_rules! tile_dot {
5733 ($t:expr, $x:expr) => {{
5734 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
5735 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
5736 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
5737 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
5738 let x0 = vld1q_s8($x);
5739 let x1 = vld1q_s8($x.add(16));
5740 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5741 asm!(
5742 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5743 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5744 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5745 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5746 options(pure, nomem, nostack),
5747 );
5748 vaddq_s32(a0, a1)
5749 }};
5750 }
5751 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
5760 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
5761 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
5762 const IW11: [u8; 16] = [
5763 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
5764 ];
5765 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
5766 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
5767 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
5768 let isc = vld1_u8(ISC.as_ptr());
5769 macro_rules! tile_dot_tbl {
5771 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
5772 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
5773 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
5774 let x0 = vld1q_s8($x);
5775 let x1 = vld1q_s8($x.add(16));
5776 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5777 asm!(
5778 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5779 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5780 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5781 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5782 options(pure, nomem, nostack),
5783 );
5784 vaddq_s32(a0, a1)
5785 }};
5786 }
5787 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
5788 let row_base = r * gpr * Q1_TILE;
5789 let abs_end = bytes.len();
5790 let xp = xq.as_ptr();
5791 let gp = gsum.as_ptr();
5792 let mut accv = vdupq_n_f32(0.0);
5793 let mut gi = 0;
5794 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
5797 let t0 = base.add(gi * Q1_TILE);
5798 let ld_a = vld1q_u8(t0);
5799 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
5800 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
5801 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
5802 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
5803 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
5804 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
5806 let g = vld1q_s32(gp.add(gi));
5807 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
5808 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
5809 let scf: float32x4_t;
5810 asm!(
5811 "fcvtl {o:v}.4s, {i:v}.4h",
5812 o = out(vreg) scf, i = in(vreg) sc16,
5813 options(pure, nomem, nostack),
5814 );
5815 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
5816 gi += 4;
5817 }
5818 let mut acc = vaddvq_f32(accv);
5819 while gi < gpr {
5820 let t = base.add(gi * Q1_TILE);
5821 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5822 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
5823 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
5824 gi += 1;
5825 }
5826 acc
5827 }
5828}
5829
5830#[cfg(target_arch = "aarch64")]
5835#[target_feature(enable = "neon,dotprod")]
5836unsafe fn dot_q1_row_1x4_sdot(
5837 bytes: &[u8],
5838 r: usize,
5839 gpr: usize,
5840 xs: [&[i8]; 4],
5841 gs: [&[i32]; 4],
5842) -> [f32; 4] {
5843 unsafe {
5845 use core::arch::aarch64::*;
5846 use core::arch::asm;
5847 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
5848 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
5849 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
5850 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
5851 const IW11: [u8; 16] = [
5852 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
5853 ];
5854 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
5855 let m = vld1q_u8(MASKS.as_ptr());
5856 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
5857 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
5858 let isc = vld1_u8(ISC.as_ptr());
5859 macro_rules! sdot2 {
5860 ($w0:expr, $w1:expr, $x:expr) => {{
5861 let x0 = vld1q_s8($x);
5862 let x1 = vld1q_s8($x.add(16));
5863 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5864 asm!(
5865 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5866 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5867 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5868 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
5869 options(pure, nomem, nostack),
5870 );
5871 vaddq_s32(a0, a1)
5872 }};
5873 }
5874 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
5875 let row_base = r * gpr * Q1_TILE;
5876 let abs_end = bytes.len();
5877 let mut accv = [vdupq_n_f32(0.0); 4];
5878 let mut gi = 0;
5879 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
5880 let t0 = base.add(gi * Q1_TILE);
5881 let ld_a = vld1q_u8(t0);
5882 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
5883 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
5885 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
5886 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
5887 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
5888 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
5889 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
5890 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
5891 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
5892 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
5893 let scf: float32x4_t;
5894 asm!(
5895 "fcvtl {o:v}.4s, {i:v}.4h",
5896 o = out(vreg) scf, i = in(vreg) sc16,
5897 options(pure, nomem, nostack),
5898 );
5899 for k in 0..4 {
5900 let xp = xs[k].as_ptr();
5901 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
5902 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
5903 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
5904 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
5905 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
5906 let g = vld1q_s32(gs[k].as_ptr().add(gi));
5907 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
5908 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
5909 }
5910 gi += 4;
5911 }
5912 let mut acc = [
5913 vaddvq_f32(accv[0]),
5914 vaddvq_f32(accv[1]),
5915 vaddvq_f32(accv[2]),
5916 vaddvq_f32(accv[3]),
5917 ];
5918 while gi < gpr {
5919 let t = base.add(gi * Q1_TILE);
5920 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5921 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
5922 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
5923 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
5924 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
5925 for k in 0..4 {
5926 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
5927 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
5928 }
5929 gi += 1;
5930 }
5931 acc
5932 }
5933}
5934
5935#[inline]
5937fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5938 let gi = j / GROUP_SIZE;
5939 let k = j % GROUP_SIZE;
5940 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5941 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5942 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
5943 ((bit as i32 * 2 - 1) as f32, s)
5944}
5945
5946#[inline]
5948fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5949 let mut acc = 0f32;
5950 for gi in 0..gpr {
5951 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5952 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5953 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5954 let mut ga = 0f32;
5955 for (j, &b) in tile[2..].iter().enumerate() {
5956 for k in 0..8 {
5957 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
5958 }
5959 }
5960 acc += ga * s;
5961 }
5962 acc
5963}
5964
5965#[allow(clippy::too_many_arguments)]
5968fn q1_range_a8w8(
5969 bytes: &[u8],
5970 gpr: usize,
5971 act: &SplitAct,
5972 gsum: &[i32],
5973 out: SendMut,
5974 start: usize,
5975 end: usize,
5976) {
5977 for r in start..end {
5978 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5979 for &(j, xv) in &act.outliers {
5980 let (w, s) = q1_outlier(bytes, r, gpr, j);
5981 acc += w * s * xv;
5982 }
5983 unsafe { *out.at(r) = acc };
5985 }
5986}
5987
5988fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
5990 for r in start..end {
5991 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
5993 }
5994}
5995
5996fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6001 let entries = base_len + (rows + 1) * 4;
6002 (base_len, entries, entries <= bytes.len())
6003}
6004
6005#[inline]
6007fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6008 let o = rp_off + r * 4;
6009 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6010}
6011
6012const SIGN5: [[f32; 5]; 256] = {
6016 let mut lut = [[0.0f32; 5]; 256];
6017 let pow3 = [1u16, 3, 9, 27, 81];
6018 let mut byte = 0usize;
6019 while byte < 256 {
6020 let mut i = 0usize;
6021 while i < 5 {
6022 let code = (byte as u16 / pow3[i]) % 3;
6023 lut[byte][i] = if code == 1 {
6024 1.0
6025 } else if code == 2 {
6026 -1.0
6027 } else {
6028 0.0
6029 };
6030 i += 1;
6031 }
6032 byte += 1;
6033 }
6034 lut
6035};
6036
6037const SIGN5_I8: [[i8; 5]; 256] = {
6039 let mut lut = [[0i8; 5]; 256];
6040 let pow3 = [1u16, 3, 9, 27, 81];
6041 let mut byte = 0usize;
6042 while byte < 256 {
6043 let mut i = 0usize;
6044 while i < 5 {
6045 let code = (byte as u16 / pow3[i]) % 3;
6046 lut[byte][i] = if code == 1 {
6047 1
6048 } else if code == 2 {
6049 -1
6050 } else {
6051 0
6052 };
6053 i += 1;
6054 }
6055 byte += 1;
6056 }
6057 lut
6058};
6059
6060const SIGN5_U64: [u64; 256] = {
6066 let mut lut = [0u64; 256];
6067 let pow3 = [1u16, 3, 9, 27, 81];
6068 let mut byte = 0usize;
6069 while byte < 256 {
6070 let mut v = 0u64;
6071 let mut i = 0usize;
6072 while i < 5 {
6073 let code = (byte as u16 / pow3[i]) % 3;
6074 let s: u8 = if code == 1 {
6075 1
6076 } else if code == 2 {
6077 0xFF
6078 } else {
6079 0
6080 };
6081 v |= (s as u64) << (i * 8);
6082 i += 1;
6083 }
6084 lut[byte] = v;
6085 byte += 1;
6086 }
6087 lut
6088};
6089
6090#[inline]
6095fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6096 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6097 let off = (r * gpr + j / GROUP_SIZE) * TILE;
6098 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6099 let within = j % GROUP_SIZE;
6100 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6101}
6102
6103#[cfg(target_arch = "aarch64")]
6106#[target_feature(enable = "neon,dotprod")]
6107#[inline]
6108unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6109 unsafe {
6111 use core::arch::aarch64::*;
6112 use core::arch::asm;
6113 let w0 = vld1q_s8(w);
6114 let w1 = vld1q_s8(w.add(16));
6115 let x0 = vld1q_s8(x);
6116 let x1 = vld1q_s8(x.add(16));
6117 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6118 asm!(
6119 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6120 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6121 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6122 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6123 options(pure, nomem, nostack),
6124 );
6125 vaddvq_s32(vaddq_s32(a0, a1))
6126 }
6127}
6128
6129#[cfg(target_arch = "x86_64")]
6132#[target_feature(enable = "avx2")]
6133#[inline]
6134unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6135 unsafe {
6137 use core::arch::x86_64::*;
6138 let wv = _mm256_loadu_si256(w as *const __m256i);
6139 let xv = _mm256_loadu_si256(x as *const __m256i);
6140 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6141 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6142 let hi128 = _mm256_extracti128_si256::<1>(d);
6143 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6144 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6145 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6146 _mm_cvtsi128_si32(s32)
6147 }
6148}
6149
6150#[inline]
6155fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6156 debug_assert!(dst.len() >= 40);
6157 unsafe {
6160 let p = dst.as_mut_ptr();
6161 for bi in 0..7 {
6162 core::ptr::write_unaligned(
6163 p.add(bi * 5) as *mut u64,
6164 SIGN5_U64[*codes.add(bi) as usize],
6165 );
6166 }
6167 }
6168}
6169
6170#[inline]
6175fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6176 #[cfg(target_arch = "aarch64")]
6177 unsafe {
6178 return sdot32_i8(w, x);
6179 }
6180 #[cfg(target_arch = "x86_64")]
6181 unsafe {
6182 return i8dot32_avx2(w, x);
6183 }
6184 #[allow(unreachable_code)]
6185 unsafe {
6186 let mut s = 0i32;
6187 for k in 0..GROUP_SIZE {
6188 s += *w.add(k) as i32 * *x.add(k) as i32;
6189 }
6190 s
6191 }
6192}
6193
6194#[inline]
6195unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6196 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6197 (
6198 SIGN5_U64[*codes as usize],
6199 SIGN5_U64[*codes.add(1) as usize],
6200 SIGN5_U64[*codes.add(2) as usize],
6201 SIGN5_U64[*codes.add(3) as usize],
6202 SIGN5_U64[*codes.add(4) as usize],
6203 SIGN5_U64[*codes.add(5) as usize],
6204 SIGN5_U64[*codes.add(6) as usize],
6205 )
6206 };
6207
6208 let u0 = s0 | (s1 << 40);
6209 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6210 let u2 = (s3 >> 8) | (s4 << 32);
6211 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6212
6213 (u0, u1, u2, u3)
6214}
6215
6216#[cfg(target_arch = "aarch64")]
6220#[target_feature(enable = "neon,dotprod")]
6221unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6222 use core::arch::aarch64::*;
6223 use core::arch::asm;
6224 unsafe {
6225 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6226 let mut acc = 0f32;
6227 let bytes_ptr = bytes.as_ptr();
6228 let xq_ptr = xq.as_ptr();
6229 let row_off = r * gpr * TILE;
6230
6231 let gpr2 = gpr & !1;
6232 let mut gi = 0;
6233 while gi < gpr2 {
6234 let off0 = row_off + gi * TILE;
6235 let off1 = off0 + TILE;
6236 let s0 = f16_to_f32(u16::from_le_bytes([
6237 *bytes_ptr.add(off0),
6238 *bytes_ptr.add(off0 + 1),
6239 ]));
6240 let s1 = f16_to_f32(u16::from_le_bytes([
6241 *bytes_ptr.add(off1),
6242 *bytes_ptr.add(off1 + 1),
6243 ]));
6244
6245 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6246 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6247
6248 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6249 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6250 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6251 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6252
6253 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6254 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6255 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6256 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6257
6258 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6259 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6260 asm!(
6261 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6262 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6263 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6264 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6265 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6266 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6267 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6268 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6269 options(pure, nomem, nostack),
6270 );
6271 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6272 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6273 acc += d0 as f32 * s0 + d1 as f32 * s1;
6274 gi += 2;
6275 }
6276
6277 if gi < gpr {
6278 let off = row_off + gi * TILE;
6279 let s = f16_to_f32(u16::from_le_bytes([
6280 *bytes_ptr.add(off),
6281 *bytes_ptr.add(off + 1),
6282 ]));
6283 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6284 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6285 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6286 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6287 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6288 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6289 asm!(
6290 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6291 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6292 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6293 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6294 options(pure, nomem, nostack),
6295 );
6296 let d = vaddvq_s32(vaddq_s32(a0, a1));
6297 acc += d as f32 * s;
6298 }
6299 acc
6300 }
6301}
6302
6303#[cfg(target_arch = "x86_64")]
6305#[target_feature(enable = "avx2")]
6306unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6307 use core::arch::x86_64::*;
6308 unsafe {
6309 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6310 let mut acc = 0f32;
6311 let bytes_ptr = bytes.as_ptr();
6312 let xq_ptr = xq.as_ptr();
6313 let row_off = r * gpr * TILE;
6314
6315 let ones = _mm256_set1_epi16(1);
6316 for gi in 0..gpr {
6317 let off = row_off + gi * TILE;
6318 let s = f16_to_f32(u16::from_le_bytes([
6319 *bytes_ptr.add(off),
6320 *bytes_ptr.add(off + 1),
6321 ]));
6322 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6323 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6324 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6325 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6326 let d256 = _mm256_madd_epi16(p16, ones);
6327 let d128 = _mm_add_epi32(
6328 _mm256_castsi256_si128(d256),
6329 _mm256_extracti128_si256(d256, 1),
6330 );
6331 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6332 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6333 acc += d32 as f32 * s;
6334 }
6335 acc
6336 }
6337}
6338
6339#[cfg(target_arch = "x86_64")]
6341#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6342unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6343 use core::arch::x86_64::*;
6344 unsafe {
6346 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6347 let mut acc = 0f32;
6348 let bytes_ptr = bytes.as_ptr();
6349 let xq_ptr = xq.as_ptr();
6350 let row_off = r * gpr * TILE;
6351 for gi in 0..gpr {
6352 let off = row_off + gi * TILE;
6353 let s = f16_to_f32(u16::from_le_bytes([
6354 *bytes_ptr.add(off),
6355 *bytes_ptr.add(off + 1),
6356 ]));
6357 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6358 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6359 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6360 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6361 acc += d as f32 * s;
6362 }
6363 acc
6364 }
6365}
6366
6367#[inline]
6371fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6372 #[cfg(target_arch = "aarch64")]
6373 unsafe {
6374 return q1t_dot_row_sdot(bytes, r, gpr, xq);
6375 }
6376 #[cfg(target_arch = "x86_64")]
6377 unsafe {
6378 if vnni_tiles_enabled() {
6379 return q1t_dot_row_vnni(bytes, r, gpr, xq);
6380 }
6381 return q1t_dot_row_avx2(bytes, r, gpr, xq);
6382 }
6383 #[allow(unreachable_code)]
6384 {
6385 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6386 let mut acc = 0f32;
6387 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
6389 let off = (r * gpr + gi) * TILE;
6390 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6391 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6392 let mut d = 0i32;
6393 for k in 0..GROUP_SIZE {
6394 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6395 }
6396 acc += d as f32 * s;
6397 }
6398 acc
6399 }
6400}
6401
6402fn q1t_row_outlier_correction(
6409 bytes: &[u8],
6410 r: usize,
6411 rp_off: usize,
6412 entries_off: usize,
6413 has_ov: bool,
6414 x: &[f32],
6415) -> f32 {
6416 if !has_ov {
6417 return 0.0;
6418 }
6419 let (c0, c1) = (
6420 q1t_rowptr(bytes, rp_off, r),
6421 q1t_rowptr(bytes, rp_off, r + 1),
6422 );
6423 let mut corr = 0f32;
6424 for p in c0..c1 {
6425 let e = entries_off + p * 4;
6426 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6427 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6428 corr += val * x[col];
6429 }
6430 corr
6431}
6432
6433fn q1t_dequant_row(
6437 bytes: &[u8],
6438 r: usize,
6439 gpr: usize,
6440 rp_off: usize,
6441 entries_off: usize,
6442 has_ov: bool,
6443 buf: &mut [f32],
6444) {
6445 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6446 for g in 0..gpr {
6447 let off = (r * gpr + g) * TILE;
6448 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6449 let codes = &bytes[off + 2..off + TILE];
6450 let bc = g * GROUP_SIZE;
6451 for bi in 0..6 {
6453 let lut = &SIGN5[codes[bi] as usize];
6454 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6455 for i in 0..5 {
6456 d[i] = lut[i] * s;
6457 }
6458 }
6459 let lut = &SIGN5[codes[6] as usize];
6460 buf[bc + 30] = lut[0] * s;
6461 buf[bc + 31] = lut[1] * s;
6462 }
6463 if !has_ov {
6464 return;
6465 }
6466 let (c0, c1) = (
6467 q1t_rowptr(bytes, rp_off, r),
6468 q1t_rowptr(bytes, rp_off, r + 1),
6469 );
6470 for p in c0..c1 {
6471 let e = entries_off + p * 4;
6472 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6473 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6474 }
6475}
6476
6477fn q1t_add_overlay(
6481 bytes: &[u8],
6482 x: &[f32],
6483 rows: usize,
6484 cols: usize,
6485 out: &mut [f32],
6486 pool: Option<&Pool>,
6487) {
6488 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6489 let gpr = cols / GROUP_SIZE;
6490 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6491 if !has_ov {
6492 return;
6493 }
6494 let out_addr = SendMut(out.as_mut_ptr());
6495 let run = move |start: usize, end: usize| {
6496 for r in start..end {
6497 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6498 unsafe { *out_addr.at(r) += corr };
6500 }
6501 };
6502 dispatch_rows(pool, rows, &run);
6503}
6504
6505#[allow(clippy::too_many_arguments)]
6508fn q1t_range_a8w8(
6509 bytes: &[u8],
6510 gpr: usize,
6511 rp_off: usize,
6512 ent_off: usize,
6513 has_ov: bool,
6514 act: &SplitAct,
6515 x: &[f32],
6516 out: SendMut,
6517 start: usize,
6518 end: usize,
6519) {
6520 for r in start..end {
6521 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6522 for &(j, xv) in &act.outliers {
6523 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6524 }
6525 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6526 unsafe { *out.at(r) = acc };
6528 }
6529}
6530
6531#[allow(clippy::too_many_arguments)]
6534fn q1t_range_f32_batch(
6535 bytes: &[u8],
6536 gpr: usize,
6537 rp_off: usize,
6538 ent_off: usize,
6539 has_ov: bool,
6540 x: &[f32],
6541 out: SendMut,
6542 start: usize,
6543 end: usize,
6544) {
6545 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6546 let mut sg = [0f32; GROUP_SIZE];
6547 for r in start..end {
6548 let mut acc = 0f32;
6549 for g in 0..gpr {
6550 let off = (r * gpr + g) * TILE;
6551 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6552 let codes = &bytes[off + 2..off + TILE];
6553 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6554 for bi in 0..6 {
6555 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6556 }
6557 let lut = &SIGN5[codes[6] as usize];
6558 sg[30] = lut[0];
6559 sg[31] = lut[1];
6560 let mut gsum = 0f32;
6561 for k in 0..GROUP_SIZE {
6562 gsum += sg[k] * xg[k];
6563 }
6564 acc += s * gsum;
6565 }
6566 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6567 unsafe { *out.at(r) = acc };
6569 }
6570}
6571
6572fn q1t_matvec(
6576 bytes: &[u8],
6577 x: &[f32],
6578 rows: usize,
6579 cols: usize,
6580 out: &mut [f32],
6581 pool: Option<&Pool>,
6582) {
6583 debug_assert_eq!(out.len(), rows);
6584 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6585 let gpr = cols / GROUP_SIZE;
6586 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6587 let out_addr = SendMut(out.as_mut_ptr());
6588 if a8w8_enabled() {
6592 let act = split_act(x);
6593 let act = &act;
6594 let run = move |start: usize, end: usize| {
6595 for r in start..end {
6596 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6597 for &(j, xv) in &act.outliers {
6598 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6599 }
6600 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6601 unsafe { *out_addr.at(r) = acc };
6603 }
6604 };
6605 dispatch_rows(pool, rows, &run);
6606 return;
6607 }
6608 let run = move |start: usize, end: usize| {
6609 let mut sg = [0f32; GROUP_SIZE];
6613 for r in start..end {
6614 let mut acc = 0f32;
6615 for g in 0..gpr {
6616 let off = (r * gpr + g) * TILE;
6617 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6618 let codes = &bytes[off + 2..off + TILE];
6619 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6620 for bi in 0..6 {
6621 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6622 }
6623 let lut = &SIGN5[codes[6] as usize];
6624 sg[30] = lut[0];
6625 sg[31] = lut[1];
6626 let mut gsum = 0f32;
6627 for k in 0..GROUP_SIZE {
6628 gsum += sg[k] * xg[k];
6629 }
6630 acc += s * gsum;
6631 }
6632 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6633 unsafe { *out_addr.at(r) = acc };
6634 }
6635 };
6636 dispatch_rows(pool, rows, &run);
6637}
6638
6639#[cfg(target_arch = "aarch64")]
6645#[target_feature(enable = "neon,dotprod")]
6646unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6647 use core::arch::aarch64::*;
6648 use core::arch::asm;
6649 unsafe {
6651 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6652 let bytes_ptr = bytes.as_ptr();
6653 let row_off = r * gpr * TILE;
6654 let xp = [xa.as_ptr(), xb.as_ptr()];
6655 let mut acc = [0f32; 2];
6656 macro_rules! sdot2 {
6657 ($w0:expr, $w1:expr, $x:expr) => {{
6658 let x0 = vld1q_s8($x);
6659 let x1 = vld1q_s8($x.add(16));
6660 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6661 asm!(
6662 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6663 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6664 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6665 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6666 options(pure, nomem, nostack),
6667 );
6668 vaddvq_s32(vaddq_s32(a0, a1))
6669 }};
6670 }
6671 let gpr2 = gpr & !1;
6672 let mut gi = 0;
6673 while gi < gpr2 {
6674 let off0 = row_off + gi * TILE;
6675 let off1 = off0 + TILE;
6676 let s0 = f16_to_f32(u16::from_le_bytes([
6677 *bytes_ptr.add(off0),
6678 *bytes_ptr.add(off0 + 1),
6679 ]));
6680 let s1 = f16_to_f32(u16::from_le_bytes([
6681 *bytes_ptr.add(off1),
6682 *bytes_ptr.add(off1 + 1),
6683 ]));
6684 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6685 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6686 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6687 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6688 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6689 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6690 for k in 0..2 {
6691 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
6692 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
6693 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
6694 }
6695 gi += 2;
6696 }
6697 if gi < gpr {
6698 let off = row_off + gi * TILE;
6699 let s = f16_to_f32(u16::from_le_bytes([
6700 *bytes_ptr.add(off),
6701 *bytes_ptr.add(off + 1),
6702 ]));
6703 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6704 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6705 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6706 for k in 0..2 {
6707 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
6708 acc[k] += d as f32 * s;
6709 }
6710 }
6711 acc
6712 }
6713}
6714
6715fn q1t_matvec2(
6721 bytes: &[u8],
6722 x1: &[f32],
6723 x2: &[f32],
6724 rows: usize,
6725 cols: usize,
6726 o1: &mut [f32],
6727 o2: &mut [f32],
6728 pool: Option<&Pool>,
6729) {
6730 debug_assert_eq!(o1.len(), rows);
6731 debug_assert_eq!(o2.len(), rows);
6732 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6733 let gpr = cols / GROUP_SIZE;
6734 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6735 let out1 = SendMut(o1.as_mut_ptr());
6736 let out2 = SendMut(o2.as_mut_ptr());
6737 if a8w8_enabled() {
6738 let a1 = split_act(x1);
6739 let a2 = split_act(x2);
6740 let (a1, a2) = (&a1, &a2);
6741 let run = move |start: usize, end: usize| {
6742 for r in start..end {
6743 #[cfg(target_arch = "aarch64")]
6744 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
6747 #[cfg(not(target_arch = "aarch64"))]
6748 let ds = [
6749 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
6750 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
6751 ];
6752 let mut acc1 = ds[0] * a1.sx;
6753 for &(j, xv) in &a1.outliers {
6754 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
6755 }
6756 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
6757 let mut acc2 = ds[1] * a2.sx;
6758 for &(j, xv) in &a2.outliers {
6759 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
6760 }
6761 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
6762 unsafe {
6764 *out1.at(r) = acc1;
6765 *out2.at(r) = acc2;
6766 }
6767 }
6768 };
6769 dispatch_rows(pool, rows, &run);
6770 return;
6771 }
6772 let run = move |start: usize, end: usize| {
6773 let mut sg = [0f32; GROUP_SIZE];
6776 for r in start..end {
6777 let mut acc1 = 0f32;
6778 let mut acc2 = 0f32;
6779 for g in 0..gpr {
6780 let off = (r * gpr + g) * TILE;
6781 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6782 let codes = &bytes[off + 2..off + TILE];
6783 for bi in 0..6 {
6784 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6785 }
6786 let lut = &SIGN5[codes[6] as usize];
6787 sg[30] = lut[0];
6788 sg[31] = lut[1];
6789 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6790 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6791 let mut gsum1 = 0f32;
6792 for k in 0..GROUP_SIZE {
6793 gsum1 += sg[k] * xg1[k];
6794 }
6795 acc1 += s * gsum1;
6796 let mut gsum2 = 0f32;
6797 for k in 0..GROUP_SIZE {
6798 gsum2 += sg[k] * xg2[k];
6799 }
6800 acc2 += s * gsum2;
6801 }
6802 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
6803 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
6804 unsafe {
6806 *out1.at(r) = acc1;
6807 *out2.at(r) = acc2;
6808 }
6809 }
6810 };
6811 dispatch_rows(pool, rows, &run);
6812}
6813
6814fn q1t_matmat(
6817 bytes: &[u8],
6818 xs: &[f32],
6819 b: usize,
6820 rows: usize,
6821 cols: usize,
6822 out: &mut [f32],
6823 pool: Option<&Pool>,
6824) {
6825 debug_assert_eq!(out.len(), b * rows);
6826 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6827 let gpr = cols / GROUP_SIZE;
6828 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6829 let out_addr = SendMut(out.as_mut_ptr());
6830 if a8w8_enabled() {
6834 let acts: Vec<SplitAct> = (0..b)
6835 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
6836 .collect();
6837 let acts = &acts;
6838 let run = move |start: usize, end: usize| {
6839 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
6843 for g in 0..gpr {
6844 let off = (r * gpr + g) * TILE;
6845 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6846 q1t_unpack_group_i8(
6847 bytes.as_ptr().wrapping_add(off + 2),
6848 &mut sg[g * GROUP_SIZE..],
6849 );
6850 }
6851 for bi in 0..b {
6852 let act = &acts[bi];
6853 let mut isum = 0f32;
6854 for g in 0..gpr {
6855 let d = q1t_i8dot32(
6856 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
6857 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
6858 );
6859 isum += d as f32 * sc[g];
6860 }
6861 let mut acc = isum * act.sx;
6862 for &(j, xv) in &act.outliers {
6863 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6864 }
6865 accs[bi] = acc;
6866 }
6867 if has_ov {
6871 let (c0, c1) = (
6872 q1t_rowptr(bytes, rp_off, r),
6873 q1t_rowptr(bytes, rp_off, r + 1),
6874 );
6875 for p in c0..c1 {
6876 let e = ent_off + p * 4;
6877 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6878 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6879 for bi in 0..b {
6880 accs[bi] += val * xs[bi * cols + col];
6881 }
6882 }
6883 }
6884 for bi in 0..b {
6885 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
6886 }
6887 }
6888 };
6889 dispatch_rows(pool, rows, &run);
6890 return;
6891 }
6892 let run = move |start: usize, end: usize| {
6893 let mut buf = vec![0f32; cols];
6894 for r in start..end {
6895 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
6896 for bi in 0..b {
6897 let xr = &xs[bi * cols..(bi + 1) * cols];
6898 let mut acc = 0f32;
6899 for j in 0..cols {
6900 acc += buf[j] * xr[j];
6901 }
6902 unsafe { *out_addr.at(bi * rows + r) = acc };
6903 }
6904 }
6905 };
6906 dispatch_rows(pool, rows, &run);
6907}
6908
6909fn q1_matvec(
6910 bytes: &[u8],
6911 x: &[f32],
6912 rows: usize,
6913 cols: usize,
6914 out: &mut [f32],
6915 pool: Option<&Pool>,
6916) {
6917 debug_assert_eq!(out.len(), rows);
6918 let gpr = cols / GROUP_SIZE;
6919 let out_addr = SendMut(out.as_mut_ptr());
6920 if a8w8_enabled() {
6921 let act = split_act(x);
6922 let gsum = q1_group_sums(&act.xq, gpr);
6923 let (act, gsum) = (&act, &gsum);
6924 let run = move |start: usize, end: usize| {
6925 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
6926 };
6927 dispatch_rows(pool, rows, &run);
6928 return;
6929 }
6930 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
6931 dispatch_rows(pool, rows, &run);
6932}
6933
6934#[allow(clippy::too_many_arguments)]
6936fn q1_matvec2(
6937 bytes: &[u8],
6938 x1: &[f32],
6939 x2: &[f32],
6940 rows: usize,
6941 cols: usize,
6942 o1: &mut [f32],
6943 o2: &mut [f32],
6944 pool: Option<&Pool>,
6945) {
6946 let gpr = cols / GROUP_SIZE;
6947 let p1 = SendMut(o1.as_mut_ptr());
6948 let p2 = SendMut(o2.as_mut_ptr());
6949 if a8w8_enabled() {
6950 let a1 = split_act(x1);
6951 let a2 = split_act(x2);
6952 let g1 = q1_group_sums(&a1.xq, gpr);
6953 let g2 = q1_group_sums(&a2.xq, gpr);
6954 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
6955 let run = move |start: usize, end: usize| {
6956 for r in start..end {
6957 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
6958 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
6959 for &(j, xv) in &a1.outliers {
6960 let (w, s) = q1_outlier(bytes, r, gpr, j);
6961 v1 += w * s * xv;
6962 }
6963 for &(j, xv) in &a2.outliers {
6964 let (w, s) = q1_outlier(bytes, r, gpr, j);
6965 v2 += w * s * xv;
6966 }
6967 unsafe {
6969 *p1.at(r) = v1;
6970 *p2.at(r) = v2;
6971 }
6972 }
6973 };
6974 dispatch_rows(pool, rows, &run);
6975 return;
6976 }
6977 let run = move |start: usize, end: usize| {
6978 for r in start..end {
6979 unsafe {
6981 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
6982 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
6983 }
6984 }
6985 };
6986 dispatch_rows(pool, rows, &run);
6987}
6988
6989#[allow(clippy::too_many_arguments)]
6991fn q1_matmat(
6992 bytes: &[u8],
6993 xs_all: &[f32],
6994 b: usize,
6995 rows: usize,
6996 cols: usize,
6997 out: &mut [f32],
6998 pool: Option<&Pool>,
6999) {
7000 debug_assert_eq!(out.len(), b * rows);
7001 let gpr = cols / GROUP_SIZE;
7002 let out_addr = SendMut(out.as_mut_ptr());
7003 if a8w8_enabled() {
7004 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7005 .map(|bi| {
7006 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7007 let gsum = q1_group_sums(&act.xq, gpr);
7008 (act, gsum)
7009 })
7010 .collect();
7011 let acts = &acts;
7012 #[cfg(target_arch = "x86_64")]
7013 let blocked_ok = avx2_enabled()
7014 && blocked_enabled();
7015 #[cfg(target_arch = "aarch64")]
7016 let blocked_ok = sdot_enabled()
7017 && blocked_enabled();
7018 let run = move |start: usize, end: usize| {
7019 for r in start..end {
7020 let mut bi = 0usize;
7021 #[cfg(target_arch = "aarch64")]
7024 if blocked_ok {
7025 while bi + 4 <= acts.len() {
7026 let xs = [
7027 acts[bi].0.xq.as_slice(),
7028 acts[bi + 1].0.xq.as_slice(),
7029 acts[bi + 2].0.xq.as_slice(),
7030 acts[bi + 3].0.xq.as_slice(),
7031 ];
7032 let gs = [
7033 acts[bi].1.as_slice(),
7034 acts[bi + 1].1.as_slice(),
7035 acts[bi + 2].1.as_slice(),
7036 acts[bi + 3].1.as_slice(),
7037 ];
7038 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7039 for k in 0..4 {
7040 let (act, _) = &acts[bi + k];
7041 let mut acc = d[k] * act.sx;
7042 for &(j, xv) in &act.outliers {
7043 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7044 acc += w * sc * xv;
7045 }
7046 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7048 }
7049 bi += 4;
7050 }
7051 }
7052 #[cfg(target_arch = "x86_64")]
7053 if blocked_ok {
7054 while bi + 4 <= acts.len() {
7055 let xs = [
7056 acts[bi].0.xq.as_slice(),
7057 acts[bi + 1].0.xq.as_slice(),
7058 acts[bi + 2].0.xq.as_slice(),
7059 acts[bi + 3].0.xq.as_slice(),
7060 ];
7061 let gs = [
7062 acts[bi].1.as_slice(),
7063 acts[bi + 1].1.as_slice(),
7064 acts[bi + 2].1.as_slice(),
7065 acts[bi + 3].1.as_slice(),
7066 ];
7067 let d = unsafe {
7068 if vnni_tiles_enabled() {
7069 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7070 } else {
7071 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7072 }
7073 };
7074 for k in 0..4 {
7075 let (act, _) = &acts[bi + k];
7076 let mut acc = d[k] * act.sx;
7077 for &(j, xv) in &act.outliers {
7078 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7079 acc += w * sc * xv;
7080 }
7081 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7083 }
7084 bi += 4;
7085 }
7086 }
7087 while bi < acts.len() {
7088 let (act, gsum) = &acts[bi];
7089 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7090 for &(j, xv) in &act.outliers {
7091 let (w, s) = q1_outlier(bytes, r, gpr, j);
7092 acc += w * s * xv;
7093 }
7094 unsafe { *out_addr.at(bi * rows + r) = acc };
7096 bi += 1;
7097 }
7098 }
7099 };
7100 dispatch_rows(pool, rows, &run);
7101 return;
7102 }
7103 let run = move |start: usize, end: usize| {
7104 for r in start..end {
7105 for bi in 0..b {
7106 let x = &xs_all[bi * cols..(bi + 1) * cols];
7107 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7109 }
7110 }
7111 };
7112 dispatch_rows(pool, rows, &run);
7113}
7114
7115fn q4matvec(
7121 bytes: &[u8],
7122 x: &[f32],
7123 rows: usize,
7124 cols: usize,
7125 out: &mut [f32],
7126 pool: Option<&Pool>,
7127) {
7128 debug_assert_eq!(out.len(), rows);
7129 let (packed, scales) = q4_split(bytes, rows, cols);
7130 let gpr = cols / GROUP_SIZE;
7131 let out_addr = SendMut(out.as_mut_ptr());
7132
7133 if a8w8_enabled() {
7134 let act = split_act(x);
7135 let run = move |start: usize, end: usize| {
7136 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7137 };
7138 dispatch_rows(pool, rows, &run);
7139 return;
7140 }
7141
7142 let run =
7143 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7144 dispatch_rows(pool, rows, &run);
7145}
7146
7147#[inline]
7150#[allow(unreachable_code)]
7151#[cfg(target_arch = "x86_64")]
7156#[target_feature(enable = "avx2")]
7157unsafe fn dot_q4b_row_1x4_avx2(
7158 buf: &[u8],
7159 scales: &[u8],
7160 g0: usize,
7161 gpr: usize,
7162 xs: [&[i8]; 4],
7163) -> [f32; 4] {
7164 unsafe {
7166 use core::arch::x86_64::*;
7167 let ones = _mm256_set1_epi16(1);
7168 let mut acc = [0f32; 4];
7169 for gi in 0..gpr {
7170 let s = f16_to_f32(u16::from_le_bytes([
7171 scales[(g0 + gi) * 2],
7172 scales[(g0 + gi) * 2 + 1],
7173 ]));
7174 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7175 let aw = _mm256_abs_epi8(w);
7176 for (k, xq) in xs.iter().enumerate() {
7177 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7178 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7179 let d = _mm256_madd_epi16(p16, ones);
7180 let hi128 = _mm256_extracti128_si256::<1>(d);
7181 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7182 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7183 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7184 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7185 }
7186 }
7187 acc
7188 }
7189}
7190
7191#[cfg(target_arch = "x86_64")]
7193#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7194unsafe fn dot_q4b_row_1x4_vnni(
7195 buf: &[u8],
7196 scales: &[u8],
7197 g0: usize,
7198 gpr: usize,
7199 xs: [&[i8]; 4],
7200) -> [f32; 4] {
7201 unsafe {
7203 use core::arch::x86_64::*;
7204 let mut acc = [0f32; 4];
7205 for gi in 0..gpr {
7206 let s = f16_to_f32(u16::from_le_bytes([
7207 scales[(g0 + gi) * 2],
7208 scales[(g0 + gi) * 2 + 1],
7209 ]));
7210 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7211 let aw = _mm256_abs_epi8(w);
7212 for (k, xq) in xs.iter().enumerate() {
7213 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7214 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7215 acc[k] += d as f32 * s;
7216 }
7217 }
7218 acc
7219 }
7220}
7221
7222#[cfg(target_arch = "x86_64")]
7228#[target_feature(enable = "avx2")]
7229unsafe fn dot_q4b_row_1x4_sx_avx2(
7230 buf: &[u8],
7231 scales: &[u8],
7232 g0: usize,
7233 gpr: usize,
7234 xs: [&[i8]; 4],
7235 sxs: [f32; 4],
7236) -> [f32; 4] {
7237 unsafe {
7239 use core::arch::x86_64::*;
7240 let ones = _mm256_set1_epi16(1);
7241 let mut acc = [0f32; 4];
7242 for gi in 0..gpr {
7243 let s = f16_to_f32(u16::from_le_bytes([
7244 scales[(g0 + gi) * 2],
7245 scales[(g0 + gi) * 2 + 1],
7246 ]));
7247 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7248 let aw = _mm256_abs_epi8(w);
7249 for (k, xq) in xs.iter().enumerate() {
7250 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7251 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7252 let d = _mm256_madd_epi16(p16, ones);
7253 let hi128 = _mm256_extracti128_si256::<1>(d);
7254 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7255 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7256 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7257 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7258 }
7259 }
7260 acc
7261 }
7262}
7263
7264#[cfg(target_arch = "x86_64")]
7267#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7268unsafe fn dot_q4b_row_1x4_sx_vnni(
7269 buf: &[u8],
7270 scales: &[u8],
7271 g0: usize,
7272 gpr: usize,
7273 xs: [&[i8]; 4],
7274 sxs: [f32; 4],
7275) -> [f32; 4] {
7276 unsafe {
7278 use core::arch::x86_64::*;
7279 let mut acc = [0f32; 4];
7280 for gi in 0..gpr {
7281 let s = f16_to_f32(u16::from_le_bytes([
7282 scales[(g0 + gi) * 2],
7283 scales[(g0 + gi) * 2 + 1],
7284 ]));
7285 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7286 let aw = _mm256_abs_epi8(w);
7287 for (k, xq) in xs.iter().enumerate() {
7288 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7289 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7290 acc[k] += (d as f32 * sxs[k]) * s;
7291 }
7292 }
7293 acc
7294 }
7295}
7296
7297#[allow(unreachable_code)]
7298fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7299 #[cfg(target_arch = "aarch64")]
7300 unsafe {
7301 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7302 }
7303 #[cfg(target_arch = "x86_64")]
7304 unsafe {
7305 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7306 }
7307 let mut acc = 0f32;
7308 for gi in 0..gpr {
7309 let g = g0 + gi;
7310 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7311 let mut d = 0i32;
7312 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7313 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7314 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7315 }
7316 acc += d as f32 * s;
7317 }
7318 acc
7319}
7320
7321#[inline]
7323#[allow(unreachable_code)]
7324fn dot_q4_row_i8_2(
7325 packed: &[u8],
7326 scales: &[u8],
7327 g0: usize,
7328 gpr: usize,
7329 xq1: &[i8],
7330 xq2: &[i8],
7331) -> (f32, f32) {
7332 #[cfg(target_arch = "aarch64")]
7333 unsafe {
7334 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7335 }
7336 #[cfg(target_arch = "x86_64")]
7337 unsafe {
7338 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7339 }
7340 (
7341 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7342 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7343 )
7344}
7345
7346#[allow(clippy::too_many_arguments)]
7349fn q4_range_a8w8(
7350 packed: &[u8],
7351 scales: &[u8],
7352 gpr: usize,
7353 cols: usize,
7354 act: &SplitAct,
7355 out: SendMut,
7356 start: usize,
7357 end: usize,
7358) {
7359 for r in start..end {
7360 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7361 for &(j, xv) in &act.outliers {
7363 let flat = r * cols + j;
7364 let byte = packed[flat / 2];
7365 let nib = if flat & 1 == 0 {
7366 byte & 0x0F
7367 } else {
7368 byte >> 4
7369 };
7370 let s = f16_to_f32(u16::from_le_bytes([
7371 scales[(flat / GROUP_SIZE) * 2],
7372 scales[(flat / GROUP_SIZE) * 2 + 1],
7373 ]));
7374 acc += ((nib as i32 - 8) as f32) * s * xv;
7375 }
7376 unsafe { *out.at(r) = acc };
7378 }
7379}
7380
7381#[allow(clippy::too_many_arguments)]
7384fn q4_range2_a8w8(
7385 packed: &[u8],
7386 scales: &[u8],
7387 gpr: usize,
7388 cols: usize,
7389 a1: &SplitAct,
7390 a2: &SplitAct,
7391 p1: SendMut,
7392 p2: SendMut,
7393 start: usize,
7394 end: usize,
7395) {
7396 for r in start..end {
7397 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7398 let mut acc1 = s1 * a1.sx;
7399 let mut acc2 = s2 * a2.sx;
7400 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7402 for &(j, xv) in outliers {
7403 let flat = r * cols + j;
7404 let byte = packed[flat / 2];
7405 let nib = if flat & 1 == 0 {
7406 byte & 0x0F
7407 } else {
7408 byte >> 4
7409 };
7410 let s = f16_to_f32(u16::from_le_bytes([
7411 scales[(flat / GROUP_SIZE) * 2],
7412 scales[(flat / GROUP_SIZE) * 2 + 1],
7413 ]));
7414 *acc += ((nib as i32 - 8) as f32) * s * xv;
7415 }
7416 };
7417 fix(&a1.outliers, &mut acc1);
7418 fix(&a2.outliers, &mut acc2);
7419 unsafe {
7421 *p1.at(r) = acc1;
7422 *p2.at(r) = acc2;
7423 }
7424 }
7425}
7426
7427fn q4_range_f32(
7429 packed: &[u8],
7430 scales: &[u8],
7431 gpr: usize,
7432 x: &[f32],
7433 out: SendMut,
7434 start: usize,
7435 end: usize,
7436) {
7437 for r in start..end {
7438 let mut acc = 0f32;
7439 for gi in 0..gpr {
7440 let g = r * gpr + gi;
7441 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7442 let pk = &packed[g * 16..(g + 1) * 16];
7443 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7444 let mut ga = 0f32;
7445 for (k, &b) in pk.iter().enumerate() {
7446 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7447 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7448 }
7449 acc += ga * s;
7450 }
7451 unsafe { *out.at(r) = acc };
7453 }
7454}
7455
7456#[allow(clippy::too_many_arguments)]
7460fn q4matvec2(
7461 bytes: &[u8],
7462 x1: &[f32],
7463 x2: &[f32],
7464 rows: usize,
7465 cols: usize,
7466 o1: &mut [f32],
7467 o2: &mut [f32],
7468 pool: Option<&Pool>,
7469) {
7470 debug_assert_eq!(o1.len(), rows);
7471 debug_assert_eq!(o2.len(), rows);
7472 let (packed, scales) = q4_split(bytes, rows, cols);
7473 let gpr = cols / GROUP_SIZE;
7474
7475 if a8w8_enabled() {
7476 let a1 = split_act(x1);
7477 let a2 = split_act(x2);
7478 let p1 = SendMut(o1.as_mut_ptr());
7479 let p2 = SendMut(o2.as_mut_ptr());
7480 let run = move |start: usize, end: usize| {
7481 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7482 };
7483 dispatch_rows(pool, rows, &run);
7484 return;
7485 }
7486
7487 let p1 = SendMut(o1.as_mut_ptr());
7488 let p2 = SendMut(o2.as_mut_ptr());
7489 let run = move |start: usize, end: usize| {
7490 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7491 };
7492 dispatch_rows(pool, rows, &run);
7493}
7494
7495#[allow(clippy::too_many_arguments)]
7497fn q4_range2_f32(
7498 packed: &[u8],
7499 scales: &[u8],
7500 gpr: usize,
7501 x1: &[f32],
7502 x2: &[f32],
7503 p1: SendMut,
7504 p2: SendMut,
7505 start: usize,
7506 end: usize,
7507) {
7508 for r in start..end {
7509 let (mut acc1, mut acc2) = (0f32, 0f32);
7510 for gi in 0..gpr {
7511 let g = r * gpr + gi;
7512 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7513 let pk = &packed[g * 16..(g + 1) * 16];
7514 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7515 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7516 let (mut g1, mut g2) = (0f32, 0f32);
7517 for (k, &b) in pk.iter().enumerate() {
7518 let wl = (b & 0x0F) as f32 - 8.0;
7519 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7520 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7521 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7522 }
7523 acc1 += g1 * s;
7524 acc2 += g2 * s;
7525 }
7526 unsafe {
7528 *p1.at(r) = acc1;
7529 *p2.at(r) = acc2;
7530 }
7531 }
7532}
7533
7534thread_local! {
7535 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7538 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7539}
7540
7541#[allow(clippy::too_many_arguments)]
7547fn q4matmat(
7548 bytes: &[u8],
7549 xs_all: &[f32],
7550 b: usize,
7551 rows: usize,
7552 cols: usize,
7553 out: &mut [f32],
7554 pool: Option<&Pool>,
7555) {
7556 debug_assert_eq!(xs_all.len(), b * cols);
7557 debug_assert_eq!(out.len(), b * rows);
7558 let (packed, scales) = q4_split(bytes, rows, cols);
7559 let gpr = cols / GROUP_SIZE;
7560 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7561
7562 if a8w8_enabled() {
7563 let acts: Vec<SplitAct> = (0..b)
7564 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7565 .collect();
7566 let acts = &acts;
7567 let out_addr = SendMut(out.as_mut_ptr());
7568 let run = move |start: usize, end: usize| {
7569 ROW_I8.with(|rb| {
7570 let mut buf = rb.borrow_mut();
7571 buf.resize(cols, 0);
7572 for r in start..end {
7573 for gi in 0..gpr {
7577 let g = r * gpr + gi;
7578 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7579 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7580 buf[gi * GROUP_SIZE + k * 2 + 1] =
7581 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7582 }
7583 }
7584 let mut bi = 0usize;
7585 #[cfg(target_arch = "x86_64")]
7586 if avx2_enabled()
7587 && blocked_enabled()
7588 {
7589 while bi + 4 <= acts.len() {
7590 let xs = [
7591 acts[bi].xq.as_slice(),
7592 acts[bi + 1].xq.as_slice(),
7593 acts[bi + 2].xq.as_slice(),
7594 acts[bi + 3].xq.as_slice(),
7595 ];
7596 let d = unsafe {
7597 if vnni_tiles_enabled() {
7598 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7599 } else {
7600 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7601 }
7602 };
7603 for k in 0..4 {
7604 let act = &acts[bi + k];
7605 let mut acc = d[k] * act.sx;
7606 for &(j, xv) in &act.outliers {
7607 acc += (buf[j] as i8) as f32
7608 * gscale((r * cols + j) / GROUP_SIZE)
7609 * xv;
7610 }
7611 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7613 }
7614 bi += 4;
7615 }
7616 }
7617 while bi < acts.len() {
7618 let act = &acts[bi];
7619 let mut acc = 0f32;
7620 for gi in 0..gpr {
7621 let d = dot_i8_i8(
7622 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7623 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7624 );
7625 acc += d as f32 * gscale(r * gpr + gi);
7626 }
7627 acc *= act.sx;
7628 for &(j, xv) in &act.outliers {
7630 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7631 }
7632 unsafe { *out_addr.at(bi * rows + r) = acc };
7634 bi += 1;
7635 }
7636 }
7637 })
7638 };
7639 dispatch_rows(pool, rows, &run);
7640 return;
7641 }
7642
7643 let out_addr = SendMut(out.as_mut_ptr());
7644 let run = move |start: usize, end: usize| {
7645 ROW_F32.with(|rb| {
7646 let mut buf = rb.borrow_mut();
7647 buf.resize(cols, 0.0);
7648 for r in start..end {
7649 for gi in 0..gpr {
7652 let g = r * gpr + gi;
7653 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7654 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7655 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7656 }
7657 }
7658 for bi in 0..b {
7659 let x = &xs_all[bi * cols..(bi + 1) * cols];
7660 let mut acc = 0f32;
7661 for gi in 0..gpr {
7662 let mut ga = 0f32;
7663 for k in 0..GROUP_SIZE / 2 {
7668 let e = gi * GROUP_SIZE + k * 2;
7669 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
7670 }
7671 acc += ga * gscale(r * gpr + gi);
7672 }
7673 unsafe { *out_addr.at(bi * rows + r) = acc };
7675 }
7676 }
7677 })
7678 };
7679 dispatch_rows(pool, rows, &run);
7680}
7681
7682#[allow(clippy::too_many_arguments)]
7687fn vbitmatmat(
7688 bytes: &[u8],
7689 offsets: &[usize],
7690 xs_all: &[f32],
7691 b: usize,
7692 rows: usize,
7693 cols: usize,
7694 out: &mut [f32],
7695 pool: Option<&Pool>,
7696) {
7697 debug_assert_eq!(xs_all.len(), b * cols);
7698 debug_assert_eq!(out.len(), b * rows);
7699 debug_assert_eq!(offsets.len(), rows + 1);
7700 let ng = cols / GROUP_SIZE;
7701 let bits = &bytes[..rows];
7702 let sc_off = rows;
7703 let gscale = |r: usize, g: usize| {
7704 let so = (r * ng + g) * 2;
7705 f16_to_f32(u16::from_le_bytes([
7706 bytes[sc_off + so],
7707 bytes[sc_off + so + 1],
7708 ]))
7709 };
7710
7711 let decode_f32 = |r: usize, dst: &mut [f32]| {
7713 let bw = bits[r] as usize;
7714 let l = ((1i32 << (bw - 1)) - 1) as f32;
7715 let data = &bytes[offsets[r]..offsets[r + 1]];
7716 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
7717 for d in dst.iter_mut() {
7718 while nbits < bw {
7719 acc = (acc << 8) | data[idx] as u64;
7720 idx += 1;
7721 nbits += 8;
7722 }
7723 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
7724 nbits -= bw;
7725 *d = u - l;
7726 }
7727 };
7728
7729 if a8w8_enabled() {
7730 let acts: Vec<SplitAct> = (0..b)
7731 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7732 .collect();
7733 let acts = &acts;
7734 let out_addr = SendMut(out.as_mut_ptr());
7735 let run = move |start: usize, end: usize| {
7736 for r in start..end {
7737 let bw = bits[r] as usize;
7738 if bw == 8 {
7739 ROW_F32.with(|rb| {
7742 let mut buf = rb.borrow_mut();
7743 buf.resize(cols, 0.0);
7744 decode_f32(r, &mut buf);
7745 for bi in 0..b {
7746 let x = &xs_all[bi * cols..(bi + 1) * cols];
7747 let mut dot = 0f32;
7748 for g in 0..ng {
7749 let mut gd = 0f32;
7750 for k in 0..GROUP_SIZE {
7751 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
7752 }
7753 dot += gd * gscale(r, g);
7754 }
7755 unsafe { *out_addr.at(bi * rows + r) = dot };
7757 }
7758 });
7759 continue;
7760 }
7761 let l = (1i32 << (bw - 1)) - 1;
7762 let data = &bytes[offsets[r]..offsets[r + 1]];
7763 ROW_I8.with(|rb| {
7764 let mut buf = rb.borrow_mut();
7765 buf.resize(cols, 0);
7766 #[inline(always)]
7767 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
7768 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
7769 let u = unpack8::<B>(&data[blk * B..]);
7770 for k in 0..8 {
7771 chunk[k] = (u[k] - l) as i8 as u8;
7772 }
7773 }
7774 }
7775 match bw {
7776 3 => fill::<3>(data, l, &mut buf),
7777 4 => vbit_fill4(data, &mut buf),
7778 5 => fill::<5>(data, l, &mut buf),
7779 6 => fill::<6>(data, l, &mut buf),
7780 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
7781 }
7782 let mut bi = 0usize;
7783 #[cfg(target_arch = "x86_64")]
7787 if avx2_enabled()
7788 && blocked_enabled()
7789 {
7790 while bi + 4 <= acts.len() {
7791 let xs = [
7792 acts[bi].xq.as_slice(),
7793 acts[bi + 1].xq.as_slice(),
7794 acts[bi + 2].xq.as_slice(),
7795 acts[bi + 3].xq.as_slice(),
7796 ];
7797 let sxs = [
7798 acts[bi].sx,
7799 acts[bi + 1].sx,
7800 acts[bi + 2].sx,
7801 acts[bi + 3].sx,
7802 ];
7803 let d = unsafe {
7804 if vnni_tiles_enabled() {
7805 dot_q4b_row_1x4_sx_vnni(
7806 &buf,
7807 &bytes[sc_off..],
7808 r * ng,
7809 ng,
7810 xs,
7811 sxs,
7812 )
7813 } else {
7814 dot_q4b_row_1x4_sx_avx2(
7815 &buf,
7816 &bytes[sc_off..],
7817 r * ng,
7818 ng,
7819 xs,
7820 sxs,
7821 )
7822 }
7823 };
7824 for k in 0..4 {
7825 let act = &acts[bi + k];
7826 let mut dot = d[k];
7827 for &(j, xv) in &act.outliers {
7828 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
7829 }
7830 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
7832 }
7833 bi += 4;
7834 }
7835 }
7836 while bi < acts.len() {
7837 let act = &acts[bi];
7838 let mut dot = 0f32;
7839 for g in 0..ng {
7840 let d = dot_i8_i8(
7841 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
7842 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
7843 ) as f32
7844 * act.sx;
7845 dot += d * gscale(r, g);
7846 }
7847 for &(j, xv) in &act.outliers {
7848 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
7849 }
7850 unsafe { *out_addr.at(bi * rows + r) = dot };
7852 bi += 1;
7853 }
7854 });
7855 }
7856 };
7857 dispatch_rows(pool, rows, &run);
7858 return;
7859 }
7860
7861 let out_addr = SendMut(out.as_mut_ptr());
7862 let run = move |start: usize, end: usize| {
7863 ROW_F32.with(|rb| {
7864 let mut buf = rb.borrow_mut();
7865 buf.resize(cols, 0.0);
7866 for r in start..end {
7867 decode_f32(r, &mut buf);
7868 for bi in 0..b {
7869 let x = &xs_all[bi * cols..(bi + 1) * cols];
7870 let mut dot = 0f32;
7871 for g in 0..ng {
7872 let mut gd = 0f32;
7873 for k in 0..GROUP_SIZE {
7874 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
7875 }
7876 dot += gd * gscale(r, g);
7877 }
7878 unsafe { *out_addr.at(bi * rows + r) = dot };
7880 }
7881 }
7882 })
7883 };
7884 dispatch_rows(pool, rows, &run);
7885}
7886
7887pub(crate) fn gpu_batch_job<'a>(
7891 t: &'a QTensor,
7892 x: &[f32],
7893) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
7894 match t {
7895 QTensor::Mapped {
7896 model,
7897 idx,
7898 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
7899 rows,
7900 cols,
7901 row_scale,
7902 col_field,
7903 ..
7904 } => Some((
7905 model.clone(),
7906 crate::gpu::BatchJob {
7907 idx: *idx,
7908 rows: *rows,
7909 cols: *cols,
7910 row_scale,
7911 xs: prescale(x, col_field, *dt).into_owned(),
7912 layout: crate::gpu::BatchLayout::Q8,
7913 },
7914 )),
7915 QTensor::Mapped {
7917 model,
7918 idx,
7919 dtype: TensorDtype::Q1,
7920 rows,
7921 cols,
7922 ..
7923 } => Some((
7924 model.clone(),
7925 crate::gpu::BatchJob {
7926 idx: *idx,
7927 rows: *rows,
7928 cols: *cols,
7929 row_scale: &[],
7930 xs: x.to_vec(),
7931 layout: crate::gpu::BatchLayout::Q1,
7932 },
7933 )),
7934 _ => None,
7935 }
7936}
7937
7938thread_local! {
7939 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7940 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7941}
7942
7943pub(crate) fn prescale<'a>(
7944 x: &'a [f32],
7945 col_field: &[f32],
7946 dtype: TensorDtype,
7947) -> std::borrow::Cow<'a, [f32]> {
7948 if dtype == TensorDtype::Q8_2f {
7949 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
7950 } else {
7951 std::borrow::Cow::Borrowed(x)
7952 }
7953}
7954
7955pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
7958 x: &[f32],
7959 col_field: &[f32],
7960 dtype: TensorDtype,
7961 buf_id: u8,
7962 f: F,
7963) -> R {
7964 if dtype == TensorDtype::Q8_2f {
7965 if buf_id == 1 {
7966 PRESCALE_BUF1.with(|b| {
7967 let mut buf = b.borrow_mut();
7968 buf.clear();
7969 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7970 f(&buf)
7971 })
7972 } else {
7973 PRESCALE_BUF2.with(|b| {
7974 let mut buf = b.borrow_mut();
7975 buf.clear();
7976 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7977 f(&buf)
7978 })
7979 }
7980 } else {
7981 f(x)
7982 }
7983}
7984
7985#[cfg(target_arch = "x86_64")]
7990pub(crate) fn avx2_enabled() -> bool {
7991 use std::sync::OnceLock;
7992 static ON: OnceLock<bool> = OnceLock::new();
7993 *ON.get_or_init(|| {
7994 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
7995 && std::arch::is_x86_feature_detected!("avx2")
7996 && std::arch::is_x86_feature_detected!("fma")
7997 })
7998}
7999
8000#[cfg(target_arch = "x86_64")]
8005fn avx2_a8w8_enabled() -> bool {
8006 use std::sync::OnceLock;
8007 static ON: OnceLock<bool> = OnceLock::new();
8008 *ON.get_or_init(|| {
8009 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8010 })
8011}
8012
8013#[inline]
8017pub(crate) fn a8w8_enabled() -> bool {
8018 #[cfg(target_arch = "aarch64")]
8019 {
8020 sdot_enabled()
8021 }
8022 #[cfg(target_arch = "x86_64")]
8023 {
8024 avx2_a8w8_enabled()
8025 }
8026 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8027 {
8028 false
8029 }
8030}
8031
8032#[inline]
8035#[allow(unreachable_code)]
8036fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8037 #[cfg(target_arch = "aarch64")]
8038 unsafe {
8039 return dot_i8_sdot(w, xq);
8040 }
8041 #[cfg(target_arch = "x86_64")]
8042 unsafe {
8043 if avx512vnni_enabled() {
8044 return dot_i8_i8_vnni(w, xq);
8045 }
8046 return dot_i8_i8_avx2(w, xq);
8047 }
8048 w.iter()
8049 .zip(xq)
8050 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8051 .sum()
8052}
8053
8054#[cfg(target_arch = "x86_64")]
8058fn avx512vnni_enabled() -> bool {
8059 use std::sync::OnceLock;
8060 static ON: OnceLock<bool> = OnceLock::new();
8061 *ON.get_or_init(|| {
8062 std::env::var("CMF_AVX512")
8063 .map(|v| v != "0")
8064 .unwrap_or(true)
8065 && std::arch::is_x86_feature_detected!("avx512f")
8066 && std::arch::is_x86_feature_detected!("avx512bw")
8067 && std::arch::is_x86_feature_detected!("avx512vl")
8068 && std::arch::is_x86_feature_detected!("avx512vnni")
8069 })
8070}
8071
8072#[cfg(target_arch = "x86_64")]
8080fn vnni_tiles_enabled() -> bool {
8081 use std::sync::OnceLock;
8082 static ON: OnceLock<bool> = OnceLock::new();
8083 *ON.get_or_init(|| {
8084 std::env::var("CMF_VNNI_TILES")
8085 .map(|v| v != "0")
8086 .unwrap_or(true)
8087 && avx512vnni_enabled()
8088 })
8089}
8090
8091#[cfg(target_arch = "x86_64")]
8096#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8097#[inline]
8098unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8099 unsafe {
8101 use core::arch::x86_64::*;
8102 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8103 let hi128 = _mm256_extracti128_si256::<1>(d);
8104 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8105 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8106 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8107 _mm_cvtsi128_si32(s32)
8108 }
8109}
8110
8111#[cfg(target_arch = "x86_64")]
8116#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8117unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8118 unsafe {
8120 use core::arch::x86_64::*;
8121 let n = w.len();
8122 let mut j = 0usize;
8123 let mut total: i32;
8124 {
8129 #[inline(always)]
8130 unsafe fn step(
8131 w: *const u8,
8132 x: *const i8,
8133 acc: core::arch::x86_64::__m512i,
8134 ) -> core::arch::x86_64::__m512i {
8135 unsafe {
8136 use core::arch::x86_64::*;
8137 let wv = _mm512_loadu_si512(w as *const _);
8138 let xv = _mm512_loadu_si512(x as *const _);
8139 let aw = _mm512_abs_epi8(wv);
8140 let neg = _mm512_movepi8_mask(wv);
8141 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8142 _mm512_dpbusd_epi32(acc, aw, sx)
8143 }
8144 }
8145 let (mut a0, mut a1, mut a2, mut a3) = (
8146 _mm512_setzero_si512(),
8147 _mm512_setzero_si512(),
8148 _mm512_setzero_si512(),
8149 _mm512_setzero_si512(),
8150 );
8151 while j + 256 <= n {
8152 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8153 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8154 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8155 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8156 j += 256;
8157 }
8158 while j + 64 <= n {
8159 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8160 j += 64;
8161 }
8162 let s01 = _mm512_add_epi32(a0, a1);
8163 let s23 = _mm512_add_epi32(a2, a3);
8164 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8165 }
8166 if j + 32 <= n {
8168 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8169 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8170 let d = _mm256_dpbusd_epi32(
8171 _mm256_setzero_si256(),
8172 _mm256_abs_epi8(wv),
8173 _mm256_sign_epi8(xv, wv),
8174 );
8175 let hi128 = _mm256_extracti128_si256::<1>(d);
8176 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8177 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8178 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8179 total += _mm_cvtsi128_si32(s32);
8180 j += 32;
8181 }
8182 while j < n {
8183 total += (w[j] as i8) as i32 * xq[j] as i32;
8184 j += 1;
8185 }
8186 total
8187 }
8188}
8189
8190#[cfg(target_arch = "x86_64")]
8192#[target_feature(enable = "avx2,fma")]
8193unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8194 unsafe {
8196 use core::arch::x86_64::*;
8197 let n = x.len();
8198 let wp = w.as_ptr();
8199 let xp = x.as_ptr();
8200 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8201 let mut j = 0usize;
8202 while j + 16 <= n {
8203 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8204 let lo = _mm256_cvtepi8_epi32(wb);
8205 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8206 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8207 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8208 j += 16;
8209 }
8210 let acc = _mm256_add_ps(a0, a1);
8211 let hi128 = _mm256_extractf128_ps::<1>(acc);
8212 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8213 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8214 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8215 let mut sum = _mm_cvtss_f32(s32);
8216 while j < n {
8217 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8218 j += 1;
8219 }
8220 sum
8221 }
8222}
8223
8224#[cfg(target_arch = "x86_64")]
8229#[target_feature(enable = "avx2")]
8230unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8231 unsafe {
8233 use core::arch::x86_64::*;
8234 let n = w.len();
8235 let ones = _mm256_set1_epi16(1);
8236 let mut acc = _mm256_setzero_si256();
8237 let mut j = 0usize;
8238 while j + 32 <= n {
8239 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8240 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8241 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8242 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8243 j += 32;
8244 }
8245 let hi128 = _mm256_extracti128_si256::<1>(acc);
8246 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8247 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8248 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8249 let mut s = _mm_cvtsi128_si32(s32);
8250 while j < n {
8251 s += (w[j] as i8) as i32 * xq[j] as i32;
8252 j += 1;
8253 }
8254 s
8255 }
8256}
8257
8258#[cfg(target_arch = "aarch64")]
8262#[target_feature(enable = "neon,i8mm")]
8263unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8264 unsafe {
8266 use core::arch::aarch64::*;
8267 use core::arch::asm;
8268 let n = w0.len();
8269 let w0p = w0.as_ptr() as *const i8;
8270 let w1p = w1.as_ptr() as *const i8;
8271 let mut acc01 = vdupq_n_s32(0);
8274 let mut acc23 = vdupq_n_s32(0);
8275 let mut i = 0usize;
8276 while i + 8 <= n {
8277 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8278 let xb01 = vcombine_s8(
8279 vld1_s8(xs[0].as_ptr().add(i)),
8280 vld1_s8(xs[1].as_ptr().add(i)),
8281 );
8282 let xb23 = vcombine_s8(
8283 vld1_s8(xs[2].as_ptr().add(i)),
8284 vld1_s8(xs[3].as_ptr().add(i)),
8285 );
8286 asm!(
8287 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8288 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8289 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8290 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8291 options(pure, nomem, nostack),
8292 );
8293 i += 8;
8294 }
8295 let mut out = [[0i32; 4]; 2];
8296 let a01: [i32; 4] = core::mem::transmute(acc01);
8297 let a23: [i32; 4] = core::mem::transmute(acc23);
8298 out[0][0] = a01[0];
8299 out[0][1] = a01[1];
8300 out[1][0] = a01[2];
8301 out[1][1] = a01[3];
8302 out[0][2] = a23[0];
8303 out[0][3] = a23[1];
8304 out[1][2] = a23[2];
8305 out[1][3] = a23[3];
8306 if i < n {
8307 for (k, x) in xs.iter().enumerate() {
8308 for j in i..n {
8309 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8310 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8311 }
8312 }
8313 }
8314 out
8315 }
8316}
8317
8318#[cfg(target_arch = "aarch64")]
8322#[target_feature(enable = "neon,dotprod")]
8323unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8324 unsafe {
8326 use core::arch::aarch64::*;
8327 use core::arch::asm;
8328 let n = w0.len();
8329 let w0p = w0.as_ptr() as *const i8;
8330 let w1p = w1.as_ptr() as *const i8;
8331 let mut acc = [[vdupq_n_s32(0); 4]; 2];
8332 let mut i = 0usize;
8333 while i + 16 <= n {
8334 let wv0 = vld1q_s8(w0p.add(i));
8335 let wv1 = vld1q_s8(w1p.add(i));
8336 for (k, x) in xs.iter().enumerate() {
8337 let xv = vld1q_s8(x.as_ptr().add(i));
8338 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8339 asm!(
8340 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8341 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8342 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8343 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8344 options(pure, nomem, nostack),
8345 );
8346 acc[0][k] = a0;
8347 acc[1][k] = a1;
8348 }
8349 i += 16;
8350 }
8351 let mut out = [[0i32; 4]; 2];
8352 for r in 0..2 {
8353 for k in 0..4 {
8354 out[r][k] = vaddvq_s32(acc[r][k]);
8355 }
8356 }
8357 if i < n {
8358 for (k, x) in xs.iter().enumerate() {
8359 for j in i..n {
8360 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8361 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8362 }
8363 }
8364 }
8365 out
8366 }
8367}
8368
8369#[cfg(target_arch = "x86_64")]
8375#[target_feature(enable = "avx2")]
8376unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8377 unsafe {
8379 use core::arch::x86_64::*;
8380 let n = w0.len();
8381 let ones = _mm256_set1_epi16(1);
8382 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8383 let mut j = 0usize;
8384 while j + 32 <= n {
8385 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8386 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8387 let aw0 = _mm256_abs_epi8(wv0);
8388 let aw1 = _mm256_abs_epi8(wv1);
8389 for (k, x) in xs.iter().enumerate() {
8390 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8391 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8392 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8393 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8394 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8395 }
8396 j += 32;
8397 }
8398 let mut out = [[0i32; 4]; 2];
8399 for r in 0..2 {
8400 for k in 0..4 {
8401 let a = acc[r][k];
8402 let hi128 = _mm256_extracti128_si256::<1>(a);
8403 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8404 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8405 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8406 out[r][k] = _mm_cvtsi128_si32(s32);
8407 }
8408 }
8409 if j < n {
8410 for (k, x) in xs.iter().enumerate() {
8411 for i in j..n {
8412 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8413 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8414 }
8415 }
8416 }
8417 out
8418 }
8419}
8420
8421#[cfg(target_arch = "x86_64")]
8426#[inline]
8427fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8428 let dot = if avx512vnni_enabled() && row.len() >= 64 {
8429 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8430 } else {
8431 unsafe { dot_i8_i8_avx2(row, &act.xq) }
8432 };
8433 let mut acc = dot as f32 * act.sx;
8434 for &(j, xv) in &act.outliers {
8435 acc += (row[j] as i8) as f32 * xv;
8436 }
8437 acc
8438}
8439
8440#[cfg(target_arch = "x86_64")]
8444#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8445unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8446 unsafe {
8448 use core::arch::x86_64::*;
8449 let n = w.len();
8450 let flip = _mm512_set1_epi8(-128); #[inline(always)]
8452 unsafe fn step(
8453 w: *const u8,
8454 x: *const i8,
8455 flip: core::arch::x86_64::__m512i,
8456 acc: core::arch::x86_64::__m512i,
8457 ) -> core::arch::x86_64::__m512i {
8458 unsafe {
8459 use core::arch::x86_64::*;
8460 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8461 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8462 }
8463 }
8464 let (mut a0, mut a1, mut a2, mut a3) = (
8465 _mm512_setzero_si512(),
8466 _mm512_setzero_si512(),
8467 _mm512_setzero_si512(),
8468 _mm512_setzero_si512(),
8469 );
8470 let mut j = 0usize;
8471 while j + 256 <= n {
8472 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8473 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8474 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8475 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8476 j += 256;
8477 }
8478 while j + 64 <= n {
8479 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8480 j += 64;
8481 }
8482 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8483 _mm512_add_epi32(a0, a1),
8484 _mm512_add_epi32(a2, a3),
8485 ));
8486 while j < n {
8488 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8489 j += 1;
8490 }
8491 total
8492 }
8493}
8494
8495#[cfg(target_arch = "x86_64")]
8501#[target_feature(enable = "avx2")]
8502unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8503 unsafe {
8506 use core::arch::x86_64::*;
8507 let lomask = _mm_set1_epi8(0x0F);
8508 let eight = _mm256_set1_epi8(8);
8509 let ones = _mm256_set1_epi16(1);
8510 let mut acc = 0f32;
8511 for gi in 0..gpr {
8512 let g = g0 + gi;
8513 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8514 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8515 let lo = _mm_and_si128(b, lomask);
8516 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8517 let w = _mm256_sub_epi8(
8518 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8519 eight,
8520 );
8521 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8522 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8523 let d = _mm256_madd_epi16(p16, ones);
8524 let hi128 = _mm256_extracti128_si256::<1>(d);
8525 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8526 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8527 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8528 acc += _mm_cvtsi128_si32(s32) as f32 * s;
8529 }
8530 acc
8531 }
8532}
8533
8534#[cfg(target_arch = "x86_64")]
8537#[target_feature(enable = "avx2")]
8538unsafe fn dot_q4_row_avx2_2(
8539 packed: &[u8],
8540 scales: &[u8],
8541 g0: usize,
8542 gpr: usize,
8543 xq1: &[i8],
8544 xq2: &[i8],
8545) -> (f32, f32) {
8546 unsafe {
8548 use core::arch::x86_64::*;
8549 let lomask = _mm_set1_epi8(0x0F);
8550 let eight = _mm256_set1_epi8(8);
8551 let ones = _mm256_set1_epi16(1);
8552 let (mut acc1, mut acc2) = (0f32, 0f32);
8553 #[inline(always)]
8554 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8555 unsafe {
8556 use core::arch::x86_64::*;
8557 let hi128 = _mm256_extracti128_si256::<1>(d);
8558 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8559 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8560 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8561 _mm_cvtsi128_si32(s32)
8562 }
8563 }
8564 for gi in 0..gpr {
8565 let g = g0 + gi;
8566 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8567 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8568 let lo = _mm_and_si128(b, lomask);
8569 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8570 let w = _mm256_sub_epi8(
8571 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8572 eight,
8573 );
8574 let aw = _mm256_abs_epi8(w);
8575 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8576 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8577 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8578 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8579 acc1 += hsum(d1) as f32 * s;
8580 acc2 += hsum(d2) as f32 * s;
8581 }
8582 (acc1, acc2)
8583 }
8584}
8585
8586#[cfg(target_arch = "x86_64")]
8588fn q8_range_avx2(
8589 q: &[u8],
8590 row_scale: &[f32],
8591 act: &SplitAct,
8592 cols: usize,
8593 out_addr: SendMut,
8594 start: usize,
8595 end: usize,
8596) {
8597 for o in start..end {
8598 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8599 unsafe { *out_addr.at(o) = v };
8601 }
8602}
8603
8604#[cfg(target_arch = "x86_64")]
8606#[allow(clippy::too_many_arguments)]
8607fn q8_range2_avx2(
8608 q: &[u8],
8609 row_scale: &[f32],
8610 a1: &SplitAct,
8611 a2: &SplitAct,
8612 cols: usize,
8613 p1: SendMut,
8614 p2: SendMut,
8615 start: usize,
8616 end: usize,
8617) {
8618 for o in start..end {
8619 let row = &q[o * cols..(o + 1) * cols];
8620 unsafe {
8622 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8623 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8624 }
8625 }
8626}
8627
8628#[cfg(target_arch = "aarch64")]
8639fn i8mm_enabled() -> bool {
8640 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8641 *ON.get_or_init(|| {
8642 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8643 && std::arch::is_aarch64_feature_detected!("i8mm")
8644 })
8645}
8646
8647#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
8651fn sdot_enabled() -> bool {
8652 use std::sync::OnceLock;
8653 static ON: OnceLock<bool> = OnceLock::new();
8654 *ON.get_or_init(|| {
8655 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
8656 if !want {
8657 return false;
8658 }
8659
8660 #[cfg(target_arch = "aarch64")]
8661 {
8662 if std::arch::is_aarch64_feature_detected!("dotprod") {
8663 return true;
8664 }
8665 #[cfg(target_os = "android")]
8666 {
8667 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
8668 if cpuinfo.lines().any(|l| {
8669 (l.starts_with("Features") || l.starts_with("features"))
8670 && l.contains("asimddp")
8671 }) {
8672 return true;
8673 }
8674 }
8675 }
8676 false
8677 }
8678 #[cfg(not(target_arch = "aarch64"))]
8679 {
8680 false
8681 }
8682 })
8683}
8684
8685struct SplitAct {
8690 xq: Vec<i8>,
8691 sx: f32,
8692 outliers: Vec<(usize, f32)>,
8693 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
8696 xsum: i32,
8697}
8698
8699thread_local! {
8700 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
8703 const { std::cell::RefCell::new(Vec::new()) };
8704}
8705
8706impl Drop for SplitAct {
8707 fn drop(&mut self) {
8708 let buf = std::mem::take(&mut self.xq);
8709 if buf.capacity() > 0 {
8710 XQ_FREE.with(|f| {
8711 let mut f = f.borrow_mut();
8712 if f.len() < 16 {
8713 f.push(buf);
8714 }
8715 });
8716 }
8717 }
8718}
8719
8720thread_local! {
8721 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8728}
8729
8730#[inline]
8733fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
8734 KROW.with(|s| {
8735 let mut b = s.borrow_mut();
8736 if b.len() < n {
8737 b.resize(n, 0.0);
8738 }
8739 f(&mut b[..n])
8740 })
8741}
8742
8743fn split_act(x: &[f32]) -> SplitAct {
8744 let n = x.len();
8745 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
8746 let thr = 8.0 * rms;
8747 let mut outliers: Vec<(usize, f32)> = Vec::new();
8751 let mut amax = 0f32;
8752 for (j, &v) in x.iter().enumerate() {
8753 let a = v.abs();
8754 if a > thr {
8755 outliers.push((j, v));
8756 } else if a > amax {
8757 amax = a;
8758 }
8759 }
8760 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
8761 let inv = 1.0 / sx;
8762 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
8763 xq.clear();
8764 xq.reserve(n);
8765 if outliers.is_empty() {
8766 xq.extend(
8767 x.iter()
8768 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
8769 );
8770 } else {
8771 xq.extend(x.iter().map(|&v| {
8773 if v.abs() > thr {
8774 0
8775 } else {
8776 (v * inv).round().clamp(-127.0, 127.0) as i8
8777 }
8778 }));
8779 }
8780 let xsum = xq.iter().map(|&v| v as i32).sum();
8781 SplitAct {
8782 xq,
8783 sx,
8784 outliers,
8785 xsum,
8786 }
8787}
8788
8789fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
8790 let n = x.len();
8791 let rms = (x
8792 .iter()
8793 .zip(col)
8794 .map(|(&a, &c)| {
8795 let v = a * c;
8796 (v * v) as f64
8797 })
8798 .sum::<f64>()
8799 / n.max(1) as f64)
8800 .sqrt() as f32;
8801 let thr = 8.0 * rms;
8802
8803 let mut outliers = Vec::new();
8804 let mut amax = 0f32;
8805 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
8806 let v = a * c;
8807 let s = v.abs();
8808 if s > thr {
8809 outliers.push((j, v));
8810 } else if s > amax {
8811 amax = s;
8812 }
8813 }
8814
8815 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
8816 let inv = 1.0 / sx;
8817 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
8818 xq.clear();
8819 xq.reserve(n);
8820 if outliers.is_empty() {
8821 xq.extend(
8822 x.iter()
8823 .zip(col)
8824 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
8825 );
8826 } else {
8827 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
8828 let v = a * c;
8829 if v.abs() > thr {
8830 0
8831 } else {
8832 (v * inv).round().clamp(-127.0, 127.0) as i8
8833 }
8834 }));
8835 }
8836 let xsum = xq.iter().map(|&v| v as i32).sum();
8837 SplitAct {
8838 xq,
8839 sx,
8840 outliers,
8841 xsum,
8842 }
8843}
8844
8845#[cfg(target_arch = "aarch64")]
8848#[target_feature(enable = "neon,dotprod")]
8849unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
8850 unsafe {
8852 use core::arch::aarch64::*;
8853 use core::arch::asm;
8854 let wp = w.as_ptr() as *const i8;
8855 let n = w.len();
8856 let (mut a0, mut a1, mut a2, mut a3) = (
8857 vdupq_n_s32(0),
8858 vdupq_n_s32(0),
8859 vdupq_n_s32(0),
8860 vdupq_n_s32(0),
8861 );
8862 let mut i = 0;
8863 while i + 64 <= n {
8864 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
8865 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
8866 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
8867 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
8868 asm!(
8869 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8870 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8871 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
8872 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
8873 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8874 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8875 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
8876 options(pure, nomem, nostack),
8877 );
8878 i += 64;
8879 }
8880 while i + 16 <= n {
8881 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
8882 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
8883 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
8884 i += 16;
8885 }
8886 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
8887 while i < n {
8888 s += (*wp.add(i)) as i32 * xq[i] as i32;
8889 i += 1;
8890 }
8891 s
8892 }
8893}
8894
8895#[cfg(target_arch = "aarch64")]
8899#[target_feature(enable = "neon,dotprod")]
8900unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
8901 unsafe {
8903 use core::arch::aarch64::*;
8904 use core::arch::asm;
8905 let n = xq.len();
8906 let px = xq.as_ptr();
8907 let (p0, p1, p2, p3) = (
8908 w0.as_ptr() as *const i8,
8909 w1.as_ptr() as *const i8,
8910 w2.as_ptr() as *const i8,
8911 w3.as_ptr() as *const i8,
8912 );
8913 let (mut a0, mut a1, mut a2, mut a3) = (
8914 vdupq_n_s32(0),
8915 vdupq_n_s32(0),
8916 vdupq_n_s32(0),
8917 vdupq_n_s32(0),
8918 );
8919 let mut i = 0;
8920 while i + 16 <= n {
8921 let x = vld1q_s8(px.add(i));
8922 let v0 = vld1q_s8(p0.add(i));
8923 let v1 = vld1q_s8(p1.add(i));
8924 let v2 = vld1q_s8(p2.add(i));
8925 let v3 = vld1q_s8(p3.add(i));
8926 asm!(
8927 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8928 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8929 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8930 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8931 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8932 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8933 options(pure, nomem, nostack),
8934 );
8935 i += 16;
8936 }
8937 let mut r = [
8938 vaddvq_s32(a0),
8939 vaddvq_s32(a1),
8940 vaddvq_s32(a2),
8941 vaddvq_s32(a3),
8942 ];
8943 while i < n {
8944 let xi = *px.add(i) as i32;
8945 r[0] += (*p0.add(i)) as i32 * xi;
8946 r[1] += (*p1.add(i)) as i32 * xi;
8947 r[2] += (*p2.add(i)) as i32 * xi;
8948 r[3] += (*p3.add(i)) as i32 * xi;
8949 i += 1;
8950 }
8951 r
8952 }
8953}
8954
8955#[cfg(target_arch = "aarch64")]
8962#[target_feature(enable = "neon,dotprod")]
8963unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
8964 unsafe {
8967 use core::arch::aarch64::*;
8968 use core::arch::asm;
8969 let n = xq.len();
8970 let px = xq.as_ptr();
8971 let pg = g.as_ptr() as *const i8;
8972 let (mut a0, mut a1, mut a2, mut a3) = (
8973 vdupq_n_s32(0),
8974 vdupq_n_s32(0),
8975 vdupq_n_s32(0),
8976 vdupq_n_s32(0),
8977 );
8978 let mut i = 0;
8979 while i + 16 <= n {
8980 let x = vld1q_s8(px.add(i));
8981 let base = pg.add(4 * i);
8982 let v0 = vld1q_s8(base);
8983 let v1 = vld1q_s8(base.add(16));
8984 let v2 = vld1q_s8(base.add(32));
8985 let v3 = vld1q_s8(base.add(48));
8986 asm!(
8987 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8988 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8989 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8990 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8991 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8992 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8993 options(pure, nomem, nostack),
8994 );
8995 i += 16;
8996 }
8997 [
8998 vaddvq_s32(a0),
8999 vaddvq_s32(a1),
9000 vaddvq_s32(a2),
9001 vaddvq_s32(a3),
9002 ]
9003 }
9004}
9005
9006#[cfg(target_arch = "aarch64")]
9012fn q8_range_sdot(
9013 q: &[u8],
9014 rep: &[u8],
9015 row_scale: &[f32],
9016 act: &SplitAct,
9017 cols: usize,
9018 out_addr: SendMut,
9019 start: usize,
9020 end: usize,
9021) {
9022 let mut o = start;
9023 if !rep.is_empty() {
9026 while o < end && o % 4 != 0 {
9027 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9028 unsafe { *out_addr.at(o) = v };
9029 o += 1;
9030 }
9031 }
9032 while o + 4 <= end {
9033 let r = if rep.is_empty() {
9034 unsafe {
9035 dot_i8_sdot_4rows(
9036 &q[o * cols..(o + 1) * cols],
9037 &q[(o + 1) * cols..(o + 2) * cols],
9038 &q[(o + 2) * cols..(o + 3) * cols],
9039 &q[(o + 3) * cols..(o + 4) * cols],
9040 &act.xq,
9041 )
9042 }
9043 } else {
9044 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9045 };
9046 for k in 0..4 {
9047 let mut acc = r[k] as f32 * act.sx;
9048 for &(j, xv) in &act.outliers {
9049 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9050 }
9051 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9053 }
9054 o += 4;
9055 }
9056 while o < end {
9057 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9058 unsafe { *out_addr.at(o) = v };
9059 o += 1;
9060 }
9061}
9062
9063#[cfg(target_arch = "aarch64")]
9066#[allow(clippy::too_many_arguments)]
9067fn q8_range2_sdot(
9068 q: &[u8],
9069 row_scale: &[f32],
9070 a1: &SplitAct,
9071 a2: &SplitAct,
9072 cols: usize,
9073 p1: SendMut,
9074 p2: SendMut,
9075 start: usize,
9076 end: usize,
9077) {
9078 for o in start..end {
9079 let row = &q[o * cols..(o + 1) * cols];
9080 unsafe {
9082 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9083 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9084 }
9085 }
9086}
9087
9088#[allow(clippy::too_many_arguments)]
9090fn q8_range2_f32(
9091 q: &[u8],
9092 row_scale: &[f32],
9093 x1: &[f32],
9094 x2: &[f32],
9095 cols: usize,
9096 p1: SendMut,
9097 p2: SendMut,
9098 start: usize,
9099 end: usize,
9100) {
9101 for o in start..end {
9102 let row = &q[o * cols..(o + 1) * cols];
9103 unsafe {
9105 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9106 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9107 }
9108 }
9109}
9110
9111fn q8_range_f32(
9113 q: &[u8],
9114 row_scale: &[f32],
9115 xs: &[f32],
9116 cols: usize,
9117 out_addr: SendMut,
9118 start: usize,
9119 end: usize,
9120) {
9121 for o in start..end {
9122 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9123 unsafe { *out_addr.at(o) = v };
9125 }
9126}
9127
9128#[cfg(target_arch = "aarch64")]
9131#[inline]
9132fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9133 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9134 for &(j, xv) in &act.outliers {
9135 acc += (row[j] as i8) as f32 * xv;
9136 }
9137 acc
9138}
9139
9140#[cfg(target_arch = "aarch64")]
9148#[target_feature(enable = "neon,dotprod")]
9149unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9150 unsafe {
9153 use core::arch::aarch64::*;
9154 use core::arch::asm;
9155 let lomask = vdupq_n_u8(0x0F);
9156 let eight = vdupq_n_s8(8);
9157 let mut acc = 0f32;
9158 for gi in 0..gpr {
9159 let g = g0 + gi;
9160 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9161 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9162 let lo = vandq_u8(b, lomask);
9163 let hi = vshrq_n_u8::<4>(b);
9164 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9165 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9166 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9167 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9168 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9169 asm!(
9170 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9171 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9172 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9173 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9174 options(pure, nomem, nostack),
9175 );
9176 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9177 }
9178 acc
9179 }
9180}
9181
9182#[cfg(target_arch = "aarch64")]
9187#[target_feature(enable = "neon,dotprod")]
9188unsafe fn dot_q4_row_sdot2(
9189 packed: &[u8],
9190 scales: &[u8],
9191 g0: usize,
9192 gpr: usize,
9193 xq1: &[i8],
9194 xq2: &[i8],
9195) -> (f32, f32) {
9196 unsafe {
9199 use core::arch::aarch64::*;
9200 use core::arch::asm;
9201 let lomask = vdupq_n_u8(0x0F);
9202 let eight = vdupq_n_s8(8);
9203 let (mut acc1, mut acc2) = (0f32, 0f32);
9204 for gi in 0..gpr {
9205 let g = g0 + gi;
9206 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9207 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9208 let lo = vandq_u8(b, lomask);
9209 let hi = vshrq_n_u8::<4>(b);
9210 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9211 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9212 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9213 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9214 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9215 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9216 let (mut a0, mut a1, mut b0, mut b1) = (
9217 vdupq_n_s32(0),
9218 vdupq_n_s32(0),
9219 vdupq_n_s32(0),
9220 vdupq_n_s32(0),
9221 );
9222 asm!(
9223 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9224 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9225 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9226 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9227 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9228 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9229 e0 = in(vreg) e0, e1 = in(vreg) e1,
9230 x10 = in(vreg) x10, x11 = in(vreg) x11,
9231 x20 = in(vreg) x20, x21 = in(vreg) x21,
9232 options(pure, nomem, nostack),
9233 );
9234 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9235 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9236 }
9237 (acc1, acc2)
9238 }
9239}
9240
9241#[inline]
9246pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9247 #[cfg(target_arch = "aarch64")]
9248 unsafe {
9249 return axpy_i8_f32_neon(acc, row, w);
9250 }
9251 #[cfg(target_arch = "x86_64")]
9252 if avx2_enabled() {
9253 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9254 }
9255 #[allow(unreachable_code)]
9256 {
9257 for (a, &b) in acc.iter_mut().zip(row) {
9258 *a += w * b as f32;
9259 }
9260 }
9261}
9262
9263#[cfg(target_arch = "x86_64")]
9265#[target_feature(enable = "avx2,fma")]
9266unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9267 unsafe {
9269 use core::arch::x86_64::*;
9270 let n = acc.len().min(row.len());
9271 let ap = acc.as_mut_ptr();
9272 let rp = row.as_ptr();
9273 let wv = _mm256_set1_ps(w);
9274 let mut j = 0usize;
9275 while j + 16 <= n {
9276 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9277 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9278 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9279 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9280 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9281 _mm256_storeu_ps(ap.add(j), v0);
9282 _mm256_storeu_ps(ap.add(j + 8), v1);
9283 j += 16;
9284 }
9285 while j < n {
9286 *ap.add(j) += w * (*rp.add(j)) as f32;
9287 j += 1;
9288 }
9289 }
9290}
9291
9292#[cfg(target_arch = "aarch64")]
9293#[target_feature(enable = "neon")]
9294unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9295 unsafe {
9297 use core::arch::aarch64::*;
9298 let n = acc.len().min(row.len());
9299 let ap = acc.as_mut_ptr();
9300 let rp = row.as_ptr();
9301 let wv = vdupq_n_f32(w);
9302 let mut j = 0usize;
9303 while j + 16 <= n {
9304 let rb = vld1q_s8(rp.add(j));
9305 let lo = vmovl_s8(vget_low_s8(rb));
9306 let hi = vmovl_s8(vget_high_s8(rb));
9307 for (off, half) in [(0, lo), (8, hi)] {
9308 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9309 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9310 let o = j + off;
9311 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9312 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9313 }
9314 j += 16;
9315 }
9316 while j < n {
9317 *ap.add(j) += w * (*rp.add(j)) as f32;
9318 j += 1;
9319 }
9320 }
9321}
9322
9323#[inline]
9326pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9327 #[cfg(target_arch = "aarch64")]
9328 unsafe {
9329 return dot_i8_f32_neon(w, x);
9330 }
9331 #[cfg(target_arch = "x86_64")]
9332 if avx2_enabled() {
9333 return unsafe { dot_i8_f32_avx2(w, x) };
9334 }
9335 #[allow(unreachable_code)]
9336 {
9337 let mut sum = 0.0f32;
9338 for (j, &b) in w.iter().enumerate() {
9339 sum += (b as i8) as f32 * x[j];
9340 }
9341 sum
9342 }
9343}
9344
9345#[inline]
9349fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9350 #[cfg(target_arch = "aarch64")]
9351 unsafe {
9352 return dot_i8_col_f32_neon(w, x, col);
9353 }
9354 #[allow(unreachable_code)]
9355 {
9356 let mut sum = 0.0f32;
9357 for (j, &b) in w.iter().enumerate() {
9358 sum += (b as i8) as f32 * x[j] * col[j];
9359 }
9360 sum
9361 }
9362}
9363
9364#[cfg(target_arch = "aarch64")]
9365#[target_feature(enable = "neon")]
9366unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9367 unsafe {
9369 use core::arch::aarch64::*;
9370 let n = x.len();
9371 let wp = w.as_ptr() as *const i8;
9372 let xp = x.as_ptr();
9373 let cp = col.as_ptr();
9374 let (mut a0, mut a1, mut a2, mut a3) = (
9375 vdupq_n_f32(0.0),
9376 vdupq_n_f32(0.0),
9377 vdupq_n_f32(0.0),
9378 vdupq_n_f32(0.0),
9379 );
9380 let mut j = 0usize;
9381 while j + 16 <= n {
9382 let wb = vld1q_s8(wp.add(j));
9383 let lo = vmovl_s8(vget_low_s8(wb));
9384 let hi = vmovl_s8(vget_high_s8(wb));
9385 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9386 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9387 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9388 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9389 a0 = vfmaq_f32(
9390 a0,
9391 w0,
9392 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9393 );
9394 a1 = vfmaq_f32(
9395 a1,
9396 w1,
9397 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9398 );
9399 a2 = vfmaq_f32(
9400 a2,
9401 w2,
9402 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9403 );
9404 a3 = vfmaq_f32(
9405 a3,
9406 w3,
9407 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9408 );
9409 j += 16;
9410 }
9411 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9412 while j < n {
9413 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9414 j += 1;
9415 }
9416 sum
9417 }
9418}
9419
9420#[cfg(target_arch = "aarch64")]
9421#[target_feature(enable = "neon")]
9422unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9423 unsafe {
9425 use core::arch::aarch64::*;
9426 let n = x.len();
9427 let wp = w.as_ptr() as *const i8;
9428 let xp = x.as_ptr();
9429 let (mut a0, mut a1, mut a2, mut a3) = (
9430 vdupq_n_f32(0.0),
9431 vdupq_n_f32(0.0),
9432 vdupq_n_f32(0.0),
9433 vdupq_n_f32(0.0),
9434 );
9435 let mut j = 0usize;
9436 while j + 16 <= n {
9437 let wb = vld1q_s8(wp.add(j));
9438 let lo = vmovl_s8(vget_low_s8(wb));
9439 let hi = vmovl_s8(vget_high_s8(wb));
9440 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9441 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9442 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9443 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9444 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9445 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9446 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9447 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9448 j += 16;
9449 }
9450 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9451 while j < n {
9452 sum += (*wp.add(j)) as f32 * *xp.add(j);
9453 j += 1;
9454 }
9455 sum
9456 }
9457}
9458
9459#[allow(clippy::too_many_arguments)]
9460fn qmatvec(
9461 q: &[u8],
9462 rep: &[u8],
9463 row_scale: &[f32],
9464 x: &[f32],
9465 col_field: &[f32],
9466 dtype: TensorDtype,
9467 rows: usize,
9468 cols: usize,
9469 out: &mut [f32],
9470 pool: Option<&Pool>,
9471) {
9472 debug_assert_eq!(out.len(), rows);
9473 #[cfg(not(target_arch = "aarch64"))]
9474 let _ = rep;
9475
9476 #[cfg(target_arch = "aarch64")]
9477 if sdot_enabled() {
9478 let act = if dtype == TensorDtype::Q8_2f {
9479 split_act_q8_2f(x, col_field)
9480 } else {
9481 split_act(x)
9482 };
9483 let out_addr = SendMut(out.as_mut_ptr());
9484 let run_range = |start: usize, end: usize| {
9485 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9486 };
9487 match pool {
9488 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9489 _ => run_range(0, rows),
9490 }
9491 return;
9492 }
9493 #[cfg(target_arch = "x86_64")]
9496 if avx2_a8w8_enabled() {
9497 let act = if dtype == TensorDtype::Q8_2f {
9498 split_act_q8_2f(x, col_field)
9499 } else {
9500 split_act(x)
9501 };
9502 let out_addr = SendMut(out.as_mut_ptr());
9503 let run_range = |start: usize, end: usize| {
9504 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9505 };
9506 match pool {
9507 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9508 _ => run_range(0, rows),
9509 }
9510 return;
9511 }
9512
9513 prescale_with(x, col_field, dtype, 1, |xs| {
9514 let out_addr = SendMut(out.as_mut_ptr());
9515 let run_range = move |start: usize, end: usize| {
9516 for o in start..end {
9517 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9518 unsafe { *out_addr.at(o) = v };
9520 }
9521 };
9522 match pool {
9523 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9524 _ => run_range(0, rows),
9525 }
9526 });
9527}
9528
9529#[allow(clippy::too_many_arguments)]
9530fn qmatvec2(
9531 q: &[u8],
9532 row_scale: &[f32],
9533 x1: &[f32],
9534 x2: &[f32],
9535 col_field: &[f32],
9536 dtype: TensorDtype,
9537 rows: usize,
9538 cols: usize,
9539 o1: &mut [f32],
9540 o2: &mut [f32],
9541 pool: Option<&Pool>,
9542) {
9543 #[cfg(target_arch = "aarch64")]
9544 if sdot_enabled() {
9545 let a1s = if dtype == TensorDtype::Q8_2f {
9546 split_act_q8_2f(x1, col_field)
9547 } else {
9548 split_act(x1)
9549 };
9550 let a2s = if dtype == TensorDtype::Q8_2f {
9551 split_act_q8_2f(x2, col_field)
9552 } else {
9553 split_act(x2)
9554 };
9555 let p1 = SendMut(o1.as_mut_ptr());
9556 let p2 = SendMut(o2.as_mut_ptr());
9557 let run_range = |start: usize, end: usize| {
9558 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9559 };
9560 match pool {
9561 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9562 _ => run_range(0, rows),
9563 }
9564 return;
9565 }
9566 #[cfg(target_arch = "x86_64")]
9567 if avx2_a8w8_enabled() {
9568 let a1s = if dtype == TensorDtype::Q8_2f {
9569 split_act_q8_2f(x1, col_field)
9570 } else {
9571 split_act(x1)
9572 };
9573 let a2s = if dtype == TensorDtype::Q8_2f {
9574 split_act_q8_2f(x2, col_field)
9575 } else {
9576 split_act(x2)
9577 };
9578 let p1 = SendMut(o1.as_mut_ptr());
9579 let p2 = SendMut(o2.as_mut_ptr());
9580 let run_range = |start: usize, end: usize| {
9581 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9582 };
9583 match pool {
9584 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9585 _ => run_range(0, rows),
9586 }
9587 return;
9588 }
9589
9590 prescale_with(x1, col_field, dtype, 1, |x1s| {
9591 prescale_with(x2, col_field, dtype, 2, |x2s| {
9592 let p1 = SendMut(o1.as_mut_ptr());
9593 let p2 = SendMut(o2.as_mut_ptr());
9594 let run_range = move |start: usize, end: usize| {
9595 for o in start..end {
9596 let row = &q[o * cols..(o + 1) * cols];
9597 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
9598 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
9599 unsafe {
9601 *p1.at(o) = s1;
9602 *p2.at(o) = s2;
9603 }
9604 }
9605 };
9606 match pool {
9607 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9608 _ => run_range(0, rows),
9609 }
9610 });
9611 });
9612}
9613
9614#[derive(Clone, Copy)]
9615struct SendMut(*mut f32);
9616unsafe impl Send for SendMut {}
9617unsafe impl Sync for SendMut {}
9618
9619impl SendMut {
9620 #[inline]
9621 fn at(self, i: usize) -> *mut f32 {
9622 unsafe { self.0.add(i) }
9623 }
9624}
9625
9626#[cfg(test)]
9627mod tests {
9628 use super::*;
9629
9630 #[test]
9631 fn f32_matvec_matches_matvec_rows_bitexact() {
9632 let (rows, cols) = (300, 40);
9633 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
9634 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
9635 let qt = QTensor::from_f32(w.clone(), rows, cols);
9636
9637 let mut a = vec![0.0f32; rows];
9638 matvec_rows(None, &w, &x, &mut a);
9639 let mut b = vec![0.0f32; rows];
9640 qt.matvec(&x, &mut b, None);
9641 assert_eq!(a, b);
9642 }
9643
9644 #[test]
9645 fn sdot_kernel_exact_on_grid() {
9646 eprintln!("sdot_enabled = {}", sdot_enabled());
9651 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
9653 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
9654 .collect();
9655 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
9656 let x: Vec<f32> = (0..cols)
9657 .map(|i| match i % 3 {
9658 0 => 1.0,
9659 1 => -1.0,
9660 _ => 0.0,
9661 })
9662 .collect();
9663 let mut a = vec![0.0f32; rows];
9664 qmatvec(
9665 &w,
9666 &[],
9667 &scales,
9668 &x,
9669 &[],
9670 TensorDtype::Q8Row,
9671 rows,
9672 cols,
9673 &mut a,
9674 None,
9675 );
9676 for o in 0..rows {
9677 let mut acc = 0.0f32;
9678 for j in 0..cols {
9679 acc += (w[o * cols + j] as i8) as f32 * x[j];
9680 }
9681 let expect = acc * scales[o];
9682 assert!(
9683 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
9684 "row {o}: {} vs {expect}",
9685 a[o]
9686 );
9687 }
9688 }
9689
9690 #[test]
9691 fn q1_tbl_fast_path_matches_reference() {
9692 let (rows, cols) = (5, 256);
9697 let gpr = cols / GROUP_SIZE;
9698 let mut bytes = Vec::new();
9699 for t in 0..rows * gpr {
9700 let s = 0.007 + (t % 11) as f32 * 0.004;
9701 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9702 for j in 0..4 {
9703 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
9704 }
9705 }
9706 let x: Vec<f32> = (0..cols)
9707 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
9708 .collect();
9709 let mut w = vec![0.0f32; rows * cols];
9710 cortiq_core::quant::dequant_q1(&bytes, &mut w);
9711 let mut got = vec![0.0f32; rows];
9712 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
9713 for o in 0..rows {
9714 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
9715 assert!(
9716 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
9717 "row {o}: {} vs {expect}",
9718 got[o]
9719 );
9720 }
9721 let b = 5usize;
9724 let mut xs_all = Vec::new();
9725 for bi in 0..b {
9726 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
9727 }
9728 let mut mm = vec![0.0f32; b * rows];
9729 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
9730 for bi in 0..b {
9731 let mut single = vec![0.0f32; rows];
9732 q1_matvec(
9733 &bytes,
9734 &xs_all[bi * cols..(bi + 1) * cols],
9735 rows,
9736 cols,
9737 &mut single,
9738 None,
9739 );
9740 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
9741 }
9742 }
9743
9744 #[test]
9745 fn q1_kernels_match_exact_reference() {
9746 let (rows, cols) = (7, 96);
9748 let gpr = cols / GROUP_SIZE;
9749 let mut bytes = Vec::new();
9750 for t in 0..rows * gpr {
9751 let s = 0.01 + (t % 13) as f32 * 0.003;
9752 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9753 for j in 0..4 {
9754 bytes.push(((t * 31 + j * 97) % 251) as u8);
9755 }
9756 }
9757 let x: Vec<f32> = (0..cols)
9759 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
9760 .collect();
9761 let mut w = vec![0.0f32; rows * cols];
9763 cortiq_core::quant::dequant_q1(&bytes, &mut w);
9764 let mut expect = vec![0.0f32; rows];
9765 for o in 0..rows {
9766 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
9767 }
9768 let mut got = vec![0.0f32; rows];
9769 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
9770 for o in 0..rows {
9771 assert!(
9772 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
9773 "row {o}: {} vs {}",
9774 got[o],
9775 expect[o]
9776 );
9777 }
9778 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
9780 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
9781 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
9782 assert_eq!(a1, got);
9783 let mut xs = x.clone();
9784 xs.extend_from_slice(&x2);
9785 let mut mm = vec![0.0f32; 2 * rows];
9786 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
9787 assert_eq!(&mm[..rows], got.as_slice());
9788 assert_eq!(&mm[rows..], a2.as_slice());
9789 }
9790
9791 #[test]
9792 fn repack_is_bit_identical() {
9793 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
9799 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
9800 .collect();
9801 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
9802 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
9803 let rep = q8_repack_layout(&w, rows, cols);
9804 for g in 0..rows / 4 {
9806 for c in 0..cols / 16 {
9807 for lane in 0..4 {
9808 assert_eq!(
9809 &rep[g * 4 * cols + c * 64 + lane * 16
9810 ..g * 4 * cols + c * 64 + lane * 16 + 16],
9811 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
9812 );
9813 }
9814 }
9815 }
9816 let mut a = vec![0.0f32; rows];
9817 qmatvec(
9818 &w,
9819 &[],
9820 &scales,
9821 &x,
9822 &[],
9823 TensorDtype::Q8Row,
9824 rows,
9825 cols,
9826 &mut a,
9827 None,
9828 );
9829 let mut b = vec![0.0f32; rows];
9830 qmatvec(
9831 &w,
9832 &rep,
9833 &scales,
9834 &x,
9835 &[],
9836 TensorDtype::Q8Row,
9837 rows,
9838 cols,
9839 &mut b,
9840 None,
9841 );
9842 assert_eq!(a, b, "full-range repack output diverged");
9843
9844 #[cfg(target_arch = "aarch64")]
9845 if sdot_enabled() {
9846 let act = split_act(&x);
9848 let mut c1 = vec![0.0f32; rows];
9849 let mut c2 = vec![0.0f32; rows];
9850 q8_range_sdot(
9851 &w,
9852 &[],
9853 &scales,
9854 &act,
9855 cols,
9856 SendMut(c1.as_mut_ptr()),
9857 3,
9858 rows - 2,
9859 );
9860 q8_range_sdot(
9861 &w,
9862 &rep,
9863 &scales,
9864 &act,
9865 cols,
9866 SendMut(c2.as_mut_ptr()),
9867 3,
9868 rows - 2,
9869 );
9870 assert_eq!(c1, c2, "unaligned-range repack output diverged");
9871 }
9872 }
9873
9874 #[test]
9875 fn sdot_a8w8_noise_is_bounded() {
9876 let (rows, cols) = (16, 512);
9880 let w: Vec<u8> = (0..rows * cols)
9881 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
9882 .collect();
9883 let scales = vec![0.01f32; rows];
9884 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
9885 let mut a = vec![0.0f32; rows];
9886 qmatvec(
9887 &w,
9888 &[],
9889 &scales,
9890 &x,
9891 &[],
9892 TensorDtype::Q8Row,
9893 rows,
9894 cols,
9895 &mut a,
9896 None,
9897 );
9898 let (mut num, mut den) = (0f64, 0f64);
9899 for o in 0..rows {
9900 let mut acc = 0.0f32;
9901 for j in 0..cols {
9902 acc += (w[o * cols + j] as i8) as f32 * x[j];
9903 }
9904 let expect = acc * scales[o];
9905 num += ((a[o] - expect) as f64).powi(2);
9906 den += (expect as f64).powi(2);
9907 }
9908 let rel = (num / den.max(1e-12)).sqrt();
9909 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
9910 }
9911
9912 #[test]
9913 fn i8_dot_neon_matches_scalar() {
9914 let n = 100;
9915 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
9916 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
9917 let mut scalar = 0.0f32;
9918 for j in 0..n {
9919 scalar += (w[j] as i8) as f32 * x[j];
9920 }
9921 let fast = dot_i8_f32(&w, &x);
9922 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
9923 }
9924
9925 #[test]
9927 fn vbitmatvec_matches_full_dequant() {
9928 let (rows, cols) = (6, 64);
9929 let ng = cols / GROUP_SIZE;
9930 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9932 let mut bytes = bits.clone();
9933 for g in 0..rows * ng {
9934 let s = 0.02 + 0.001 * g as f32;
9935 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9936 }
9937 for r in 0..rows {
9938 let b = bits[r] as usize;
9939 let (mut acc, mut nb) = (0u64, 0usize);
9940 let mut rowbytes = Vec::new();
9941 for i in 0..cols {
9942 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9943 acc = (acc << b) | v;
9944 nb += b;
9945 while nb >= 8 {
9946 nb -= 8;
9947 rowbytes.push(((acc >> nb) & 0xFF) as u8);
9948 }
9949 }
9950 if nb > 0 {
9951 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9952 }
9953 bytes.extend_from_slice(&rowbytes);
9954 }
9955 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9956
9957 let mut reference = vec![0f32; rows * cols];
9958 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
9959 let mut expect = vec![0f32; rows];
9960 for r in 0..rows {
9961 expect[r] = reference[r * cols..(r + 1) * cols]
9962 .iter()
9963 .zip(&x)
9964 .map(|(w, xv)| w * xv)
9965 .sum();
9966 }
9967 let mut got = vec![0f32; rows];
9968 let offsets = vbit_row_offsets(&bytes, rows, cols);
9969 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
9970 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9974 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
9975 for r in 0..rows {
9976 assert!(
9977 (got[r] - expect[r]).abs() < tol * scale,
9978 "row {r}: {} vs {}",
9979 got[r],
9980 expect[r]
9981 );
9982 }
9983 }
9984
9985 #[test]
9990 #[cfg(target_arch = "x86_64")]
9991 fn vbit_matmat_blocked_matches_per_row() {
9992 let (rows, cols, b) = (64usize, 128usize, 9usize);
9993 let ng = cols / GROUP_SIZE;
9994 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
9995 let mut bytes = bits.clone();
9996 for g in 0..rows * ng {
9997 let sc = 0.02 + 0.0005 * g as f32;
9998 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9999 }
10000 for r in 0..rows {
10001 let bw = bits[r] as usize;
10002 let (mut acc, mut nb) = (0u64, 0usize);
10003 let mut rowbytes = Vec::new();
10004 for i in 0..cols {
10005 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10006 acc = (acc << bw) | v;
10007 nb += bw;
10008 while nb >= 8 {
10009 nb -= 8;
10010 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10011 }
10012 }
10013 if nb > 0 {
10014 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10015 }
10016 bytes.extend_from_slice(&rowbytes);
10017 }
10018 let x: Vec<f32> = (0..b * cols)
10019 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10020 .collect();
10021 let offsets = vbit_row_offsets(&bytes, rows, cols);
10022 let mut y_a = vec![0f32; b * rows];
10023 let mut y_b = vec![0f32; b * rows];
10024 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10025 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10026 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10027 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10028 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10029 let max_d = y_a
10030 .iter()
10031 .zip(&y_b)
10032 .map(|(p, q)| (p - q).abs())
10033 .fold(0.0f32, f32::max);
10034 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10035 }
10036
10037 #[test]
10045 fn q4t_matmat_blocked_matches_per_row() {
10046 let (rows, cols, b) = (16usize, 64usize, 9usize);
10047 let gpr = cols / GROUP_SIZE;
10048 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10049 for r in 0..rows {
10050 for g in 0..gpr {
10051 let t = (r * gpr + g) * Q4_TILE;
10052 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10053 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10054 for k in 0..16 {
10055 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10056 }
10057 }
10058 }
10059 let x: Vec<f32> = (0..b * cols)
10060 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10061 .collect();
10062 let mut y_blk = vec![0f32; b * rows];
10063 let mut y_row = vec![0f32; b * rows];
10064 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10065 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10066 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10067 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10068 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10069 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10070 }
10071
10072 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10079 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10080 let gpr = cols / GROUP_SIZE;
10081 let stride = q4tp_code_stride(gpr);
10082 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10083 let mut b = vec![0u8; codes_off + rows * stride];
10084 for r in 0..rows {
10085 for g in 0..gpr {
10086 let t = (r * gpr + g) * Q4TP_NIB;
10087 for k in 0..16 {
10088 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10089 }
10090 }
10091 let lo = -6.0 - 0.03 * (r % 17) as f32;
10092 let step = 0.01 + 0.004 * (r % 11) as f32;
10093 let p = params_off + r * 4;
10094 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10095 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10096 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10097 for g in 0..gpr {
10098 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10099 }
10100 }
10101 b
10102 }
10103
10104 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10108 let gpr = cols / GROUP_SIZE;
10109 let v = Q4tpView::new(bytes, rows, cols);
10110 let mut out = vec![0u8; rows * gpr * Q4_TILE];
10111 let mut sc = vec![0f32; gpr];
10112 for r in 0..rows {
10113 v.scales_into(r, gpr, &mut sc);
10114 for g in 0..gpr {
10115 let t = (r * gpr + g) * Q4_TILE;
10116 let s = sc[g];
10117 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10118 let src = (r * gpr + g) * Q4TP_NIB;
10119 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10120 }
10121 }
10122 out
10123 }
10124
10125 #[test]
10131 fn q4tp_exact_path_matches_dequant_reference() {
10132 let (rows, cols) = (256usize, 512usize);
10133 let gpr = cols / GROUP_SIZE;
10134 let bytes = synth_q4tp(rows, cols);
10135 let mut w = vec![0f32; rows * cols];
10136 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10137
10138 let x: Vec<f32> = (0..cols)
10139 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10140 .collect();
10141 let v = Q4tpView::new(&bytes, rows, cols);
10142 let mut sc = vec![0f32; gpr];
10143 for r in 0..rows {
10144 v.scales_into(r, gpr, &mut sc);
10145 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10146 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10147 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10151 assert!(
10152 (got - want).abs() <= 1e-5 * mag,
10153 "row {r}: kernel {got} vs dequant {want}"
10154 );
10155 }
10156 }
10157
10158 #[test]
10164 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10165 let (rows, cols) = (256usize, 512usize);
10166 let bytes = synth_q4tp(rows, cols);
10167 let twin = q4tp_as_q4t(&bytes, rows, cols);
10168 let x: Vec<f32> = (0..cols)
10169 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10170 .collect();
10171
10172 let mut got = vec![0f32; rows];
10173 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10174 let mut want = vec![0f32; rows];
10175 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10176
10177 let mut w = vec![0f32; rows * cols];
10180 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10181 for r in 0..rows {
10182 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10183 assert!(
10184 (got[r] - want[r]).abs() <= 1e-3 * mag,
10185 "row {r}: q4tp {} vs q4t {}",
10186 got[r],
10187 want[r]
10188 );
10189 }
10190 }
10191
10192 #[test]
10197 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10198 let (rows, cols, b) = (256usize, 512usize, 5usize);
10199 let bytes = synth_q4tp(rows, cols);
10200 let twin = q4tp_as_q4t(&bytes, rows, cols);
10201 let xs: Vec<f32> = (0..b * cols)
10202 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10203 .collect();
10204
10205 let mut got = vec![0f32; b * rows];
10206 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10207 let mut want = vec![0f32; b * rows];
10208 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10209
10210 let mut w = vec![0f32; rows * cols];
10211 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10212 for t in 0..b {
10213 for r in 0..rows {
10214 let mag: f32 = (0..cols)
10215 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10216 .sum();
10217 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10218 assert!(
10219 (g - wa).abs() <= 1e-3 * mag,
10220 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10221 );
10222 }
10223 }
10224 }
10225
10226 #[test]
10227 fn q4tp_matvec2_matches_the_single_stream_kernel() {
10228 let (rows, cols) = (128usize, 256usize);
10229 let gpr = cols / GROUP_SIZE;
10230 let bytes = synth_q4tp(rows, cols);
10231 let xs: Vec<f32> = (0..2 * cols)
10232 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10233 .collect();
10234
10235 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10236 q4tp_matvec2(
10237 &bytes,
10238 &xs[..cols],
10239 &xs[cols..],
10240 rows,
10241 cols,
10242 &mut o1,
10243 &mut o2,
10244 None,
10245 );
10246
10247 let v = Q4tpView::new(&bytes, rows, cols);
10250 let mut sc = vec![0f32; gpr];
10251 for r in 0..rows {
10252 v.scales_into(r, gpr, &mut sc);
10253 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10254 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10255 }
10256 }
10257
10258 #[test]
10265 fn q4tp_matvec_keeps_pace_with_q4t() {
10266 let (rows, cols) = (4096usize, 3072usize);
10267 let bytes = synth_q4tp(rows, cols);
10268 let twin = q4tp_as_q4t(&bytes, rows, cols);
10269 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10270 let mut o = vec![0f32; rows];
10271 let n = 12;
10272 let mut best = (f64::MAX, f64::MAX);
10273 for _ in 0..3 {
10276 let t0 = std::time::Instant::now();
10277 for _ in 0..n {
10278 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10279 }
10280 best.0 = best.0.min(t0.elapsed().as_secs_f64());
10281 let t0 = std::time::Instant::now();
10282 for _ in 0..n {
10283 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10284 }
10285 best.1 = best.1.min(t0.elapsed().as_secs_f64());
10286 }
10287 let ratio = best.1 / best.0;
10288 println!(
10289 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10290 best.0 * 1e3 / n as f64,
10291 best.1 * 1e3 / n as f64
10292 );
10293 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10294 }
10295
10296 #[cfg(target_os = "macos")]
10297 #[test]
10298 fn q4t_matmat_accel_matches_dequant_reference() {
10299 if !accel_gemm_enabled() {
10300 return; }
10302 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
10304 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10305 for r in 0..rows {
10306 for g in 0..gpr {
10307 let t = (r * gpr + g) * Q4_TILE;
10308 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10309 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10310 for k in 0..16 {
10311 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10312 }
10313 }
10314 }
10315 let x: Vec<f32> = (0..b * cols)
10316 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10317 .collect();
10318 let mut got = vec![0f32; b * rows];
10319 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10320 let mut w = vec![0f32; rows * cols];
10322 for r in 0..rows {
10323 for g in 0..gpr {
10324 let t = (r * gpr + g) * Q4_TILE;
10325 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10326 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10327 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10328 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10329 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10330 }
10331 }
10332 }
10333 for bi in 0..b {
10334 for r in 0..rows {
10335 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10336 let d = (got[bi * rows + r] - want).abs();
10337 assert!(
10338 d <= want.abs().max(1.0) * 1e-4,
10339 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10340 got[bi * rows + r]
10341 );
10342 }
10343 }
10344 }
10345
10346 #[test]
10347 fn q4matvec_matches_full_dequant() {
10348 let (rows, cols) = (8, 64);
10349 let groups = rows * cols / GROUP_SIZE;
10350 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10352 for i in 0..groups * 16 {
10353 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10354 }
10355 for g in 0..groups {
10356 let s = 0.01 + 0.003 * g as f32;
10357 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10358 }
10359 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10360
10361 let mut reference = vec![0.0f32; rows * cols];
10362 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10363 let mut expect = vec![0.0f32; rows];
10364 for r in 0..rows {
10365 expect[r] = reference[r * cols..(r + 1) * cols]
10366 .iter()
10367 .zip(&x)
10368 .map(|(w, xv)| w * xv)
10369 .sum();
10370 }
10371
10372 let mut got = vec![0.0f32; rows];
10373 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10374 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10378 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10379 for r in 0..rows {
10380 assert!(
10381 (got[r] - expect[r]).abs() < tol * scale,
10382 "row {r}: {} vs {}",
10383 got[r],
10384 expect[r]
10385 );
10386 }
10387 }
10388
10389 #[test]
10392 fn vbitmatvec2_equals_two_singles() {
10393 let (rows, cols) = (6, 64);
10394 let ng = cols / GROUP_SIZE;
10395 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10396 let mut bytes = bits.clone();
10397 for g in 0..rows * ng {
10398 let s = 0.02 + 0.001 * g as f32;
10399 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10400 }
10401 for r in 0..rows {
10402 let b = bits[r] as usize;
10403 let (mut acc, mut nb) = (0u64, 0usize);
10404 let mut rowbytes = Vec::new();
10405 for i in 0..cols {
10406 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10407 acc = (acc << b) | v;
10408 nb += b;
10409 while nb >= 8 {
10410 nb -= 8;
10411 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10412 }
10413 }
10414 if nb > 0 {
10415 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10416 }
10417 bytes.extend_from_slice(&rowbytes);
10418 }
10419 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10420 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10421 let offsets = vbit_row_offsets(&bytes, rows, cols);
10422
10423 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10424 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10425 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10426 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10427 vbitmatvec2(
10428 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10429 );
10430 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10431 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10432 }
10433
10434 #[test]
10436 fn q4matvec2_equals_two_singles() {
10437 let (rows, cols) = (8, 128);
10438 let groups = rows * cols / GROUP_SIZE;
10439 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10440 for i in 0..groups * 16 {
10441 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10442 }
10443 for g in 0..groups {
10444 let s = 0.01 + 0.003 * g as f32;
10445 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10446 }
10447 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10450 x1[9] = 250.0;
10451 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10452
10453 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10454 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10455 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10456 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10457 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10458 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10459 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10460 }
10461
10462 #[test]
10465 fn matvec_many_equals_separate_matvecs() {
10466 use crate::pool::Pool;
10467 let (r1, r2, cols) = (300, 200, 64);
10468 let mk = |salt: usize, rows: usize| {
10469 QTensor::from_f32(
10470 (0..rows * cols)
10471 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10472 .collect(),
10473 rows,
10474 cols,
10475 )
10476 };
10477 let (a, b) = (mk(1, r1), mk(5, r2));
10478 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10479 let pool = Pool::new(3);
10480
10481 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10482 a.matvec(&x, &mut ea, Some(&pool));
10483 b.matvec(&x, &mut eb, Some(&pool));
10484 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10485 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10486 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10487 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10488 }
10489
10490 #[test]
10493 fn batched_matmat_equals_per_position_matvec() {
10494 let (rows, cols, b) = (8, 64, 5);
10495 let groups = rows * cols / GROUP_SIZE;
10497 let mut q4 = Vec::new();
10498 for i in 0..groups * 16 {
10499 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10500 }
10501 for g in 0..groups {
10502 q4.extend_from_slice(
10503 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10504 );
10505 }
10506 let ng = cols / GROUP_SIZE;
10508 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10509 let mut vb = bits.clone();
10510 for g in 0..rows * ng {
10511 vb.extend_from_slice(
10512 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10513 );
10514 }
10515 for r in 0..rows {
10516 let bw = bits[r] as usize;
10517 let (mut acc, mut nb) = (0u64, 0usize);
10518 let mut rowbytes = Vec::new();
10519 for i in 0..cols {
10520 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10521 acc = (acc << bw) | v;
10522 nb += bw;
10523 while nb >= 8 {
10524 nb -= 8;
10525 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10526 }
10527 }
10528 if nb > 0 {
10529 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10530 }
10531 vb.extend_from_slice(&rowbytes);
10532 }
10533 let offsets = vbit_row_offsets(&vb, rows, cols);
10534
10535 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10536
10537 let mut got = vec![0f32; b * rows];
10539 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
10540 for bi in 0..b {
10541 let mut expect = vec![0f32; rows];
10542 q4matvec(
10543 &q4,
10544 &xs[bi * cols..(bi + 1) * cols],
10545 rows,
10546 cols,
10547 &mut expect,
10548 None,
10549 );
10550 assert_eq!(
10551 &got[bi * rows..(bi + 1) * rows],
10552 &expect[..],
10553 "q4 batch pos {bi}"
10554 );
10555 }
10556
10557 let mut got = vec![0f32; b * rows];
10559 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
10560 for bi in 0..b {
10561 let mut expect = vec![0f32; rows];
10562 vbitmatvec(
10563 &vb,
10564 &offsets,
10565 &xs[bi * cols..(bi + 1) * cols],
10566 rows,
10567 cols,
10568 &mut expect,
10569 None,
10570 );
10571 assert_eq!(
10572 &got[bi * rows..(bi + 1) * rows],
10573 &expect[..],
10574 "vbit batch pos {bi}"
10575 );
10576 }
10577 }
10578
10579 #[test]
10583 fn q4_tiled_matches_q4_block_bitexact() {
10584 let (rows, cols, b) = (8usize, 128usize, 3usize);
10585 let groups = rows * cols / GROUP_SIZE;
10586 let mut split = Vec::with_capacity(groups * 18);
10587 for i in 0..groups * 16 {
10588 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10589 }
10590 for g in 0..groups {
10591 split.extend_from_slice(
10592 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10593 );
10594 }
10595 let (packed, scales) = split.split_at(groups * 16);
10597 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
10598 for g in 0..groups {
10599 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
10600 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
10601 }
10602
10603 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10604 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10606
10607 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
10608 q4matvec(&split, &x1, rows, cols, &mut a, None);
10609 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
10610 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
10611
10612 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10613 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
10614 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
10615 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
10616 assert_eq!(a1, t1);
10617 assert_eq!(a2, t2);
10618
10619 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10620 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
10621 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
10622 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
10623 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
10624 }
10625
10626 #[test]
10633 fn q4matvec_sdot_outlier_exact() {
10634 let (rows, cols) = (4, 128);
10635 let groups = rows * cols / GROUP_SIZE;
10636 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10637 for i in 0..groups * 16 {
10638 bytes.push(((i * 11 + 5) % 256) as u8);
10639 }
10640 for g in 0..groups {
10641 let s = 0.02 + 0.002 * g as f32;
10642 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10643 }
10644 let mut x: Vec<f32> = (0..cols)
10645 .map(|i| match i % 3 {
10646 0 => 1.0,
10647 1 => -1.0,
10648 _ => 0.0,
10649 })
10650 .collect();
10651 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
10654 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10655 let mut expect = vec![0.0f32; rows];
10656 for r in 0..rows {
10657 expect[r] = reference[r * cols..(r + 1) * cols]
10658 .iter()
10659 .zip(&x)
10660 .map(|(w, xv)| w * xv)
10661 .sum();
10662 }
10663 let mut got = vec![0.0f32; rows];
10664 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10665 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10666 for r in 0..rows {
10667 assert!(
10668 (got[r] - expect[r]).abs() < 2e-3 * scale,
10669 "row {r}: {} vs {} (outlier term must be exact)",
10670 got[r],
10671 expect[r]
10672 );
10673 }
10674 }
10675
10676 #[test]
10680 fn q1t_matvec_matches_reference() {
10681 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
10682 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
10684 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
10685 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
10687 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
10688 let mut bytes = Vec::new();
10689 for r in 0..rows {
10690 for g in 0..gpr {
10691 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
10692 let mut c = [0u8; 7];
10693 for k in 0..GROUP_SIZE {
10694 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
10696 0
10697 } else {
10698 ((k + r * 3 + g) % 3) as u8 };
10700 cortiq_core::quant::q1t_pack(&mut c, k, code);
10701 }
10702 bytes.extend_from_slice(&c);
10703 }
10704 }
10705 let mut row_ptr = vec![0u32; rows + 1];
10708 for &(idx, _) in &outliers {
10709 row_ptr[idx as usize / cols + 1] += 1;
10710 }
10711 for r in 0..rows {
10712 row_ptr[r + 1] += row_ptr[r];
10713 }
10714 for &p in &row_ptr {
10715 bytes.extend_from_slice(&p.to_le_bytes());
10716 }
10717 for &(idx, v) in &outliers {
10718 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
10719 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
10720 }
10721
10722 let mut refw = vec![0f32; rows * cols];
10723 dequant_q1t(&bytes, rows, cols, &mut refw);
10724 let x: Vec<f32> = (0..cols)
10727 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
10728 .collect();
10729 let mut expect = vec![0f32; rows];
10730 for r in 0..rows {
10731 let mut a = 0.0f32;
10732 for j in 0..cols {
10733 a += refw[r * cols + j] * x[j];
10734 }
10735 expect[r] = a;
10736 }
10737 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
10738 let mut got = vec![0f32; rows];
10739 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
10740 for r in 0..rows {
10741 assert!(
10742 (got[r] - expect[r]).abs() < tol(expect[r]),
10743 "row {r}: {} vs {}",
10744 got[r],
10745 expect[r]
10746 );
10747 }
10748 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
10750 let mut gm = vec![0f32; 2 * rows];
10751 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
10752 for r in 0..rows {
10753 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
10754 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
10755 }
10756 let xb: Vec<f32> = (0..cols)
10760 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
10761 .collect();
10762 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10763 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
10764 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
10765 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10766 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
10767 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
10768 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
10769 }
10770
10771 #[test]
10774 fn q1t_matvec2_odd_gpr_matches_singles() {
10775 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
10776 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
10778 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
10779 for r in 0..rows {
10780 for g in 0..gpr {
10781 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
10782 let mut c = [0u8; 7];
10783 for k in 0..GROUP_SIZE {
10784 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
10785 }
10786 bytes.extend_from_slice(&c);
10787 }
10788 }
10789 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
10790 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
10791 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10792 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10793 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
10794 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10795 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10796 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
10797 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
10798 }
10799
10800 #[test]
10804 #[ignore]
10805 fn q1t_matvec2_speed() {
10806 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
10807 use std::time::Instant;
10808 let (rows, cols) = (8192usize, 4096usize);
10809 let gpr = cols / GROUP_SIZE;
10810 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
10811 for r in 0..rows {
10812 for g in 0..gpr {
10813 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
10814 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
10815 let mut c = [0u8; 7];
10816 for k in 0..GROUP_SIZE {
10817 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
10818 }
10819 bytes.extend_from_slice(&c);
10820 }
10821 }
10822 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
10823 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
10824 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10825 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10826 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10828 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10829 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
10830 for _ in 0..8 {
10831 let t0 = Instant::now();
10832 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10833 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
10834 let t1 = Instant::now();
10835 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10836 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
10837 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
10838 }
10839 assert_eq!(p1, s1);
10840 assert_eq!(p2, s2);
10841 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
10842 }
10843
10844 #[test]
10848 #[ignore]
10849 fn q1t_matvec_speed() {
10850 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
10851 use std::time::Instant;
10852 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
10854 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
10855 for r in 0..rows {
10856 for g in 0..gpr {
10857 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
10858 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
10859 let mut c = [0u8; 7];
10860 for k in 0..GROUP_SIZE {
10861 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
10862 }
10863 bytes.extend_from_slice(&c);
10864 }
10865 }
10866 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
10868 let mut idx = 0usize;
10869 while idx < n {
10870 row_ptr[idx / cols + 1] += 1;
10871 idx += stride;
10872 }
10873 for r in 0..rows {
10874 row_ptr[r + 1] += row_ptr[r];
10875 }
10876 for &p in &row_ptr {
10877 bytes.extend_from_slice(&p.to_le_bytes());
10878 }
10879 let mut idx = 0usize;
10880 while idx < n {
10881 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
10882 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
10883 idx += stride;
10884 }
10885 let x: Vec<f32> = (0..cols)
10888 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
10889 .collect();
10890 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
10891
10892 let slow = |out: &mut [f32]| {
10894 let mut buf = vec![0f32; cols];
10895 for r in 0..rows {
10896 for g in 0..gpr {
10897 let off = (r * gpr + g) * Q1T_TILE;
10898 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10899 let codes = &bytes[off + 2..off + Q1T_TILE];
10900 for k in 0..GROUP_SIZE {
10901 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
10902 1 => s,
10903 2 => -s,
10904 _ => 0.0,
10905 };
10906 }
10907 }
10908 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
10909 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
10910 }
10911 };
10912 let iters = 5;
10913 let mut a = vec![0f32; rows];
10914 slow(&mut a); let t = Instant::now();
10916 for _ in 0..iters {
10917 slow(&mut a);
10918 }
10919 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
10920
10921 let mut b = vec![0f32; rows];
10922 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
10924 for _ in 0..iters {
10925 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
10926 }
10927 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
10928
10929 for r in 0..rows {
10930 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
10931 }
10932 println!(
10933 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
10934 slow_ms / fast_ms
10935 );
10936 }
10937}
10938
10939
10940#[cfg(test)]
10941mod gemm_bench {
10942 #[test]
10953 #[ignore]
10954 fn q4tp_matmat_throughput() {
10955 let b: usize = std::env::var("CMF_BENCH_B")
10959 .ok()
10960 .and_then(|v| v.parse().ok())
10961 .unwrap_or(296);
10962 let (rows, cols) = (9216usize, 2304usize);
10963 let (_, _, _) = (rows, cols, b);
10964 let total = cortiq_core::quant::expected_nbytes(
10965 cortiq_core::TensorDtype::Q4TiledP,
10966 &[rows, cols],
10967 )
10968 .unwrap();
10969 let (params_off, codes_off, _) =
10974 cortiq_core::quant::q4tp_sections(rows, cols);
10975 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
10976 let lo = cortiq_core::quant::f32_to_f16(-4.0);
10977 let step = cortiq_core::quant::f32_to_f16(0.1);
10978 for r in 0..rows {
10979 let o = params_off + r * 4;
10980 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
10981 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
10982 }
10983 let _ = codes_off;
10984 let xs: Vec<f32> = (0..b * cols)
10985 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
10986 .collect();
10987 let mut out = vec![0f32; b * rows];
10988 let pool = crate::pool::Pool::from_env();
10989 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
10995 let reps: usize = std::env::var("CMF_BENCH_REPS")
10996 .ok()
10997 .and_then(|v| v.parse().ok())
10998 .unwrap_or(10);
10999 let mut best = [f64::MAX; 2];
11000 let mut sums = [0f32; 2];
11001 for _ in 0..reps {
11002 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11003 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11004 let t = std::time::Instant::now();
11005 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11006 best[k] = best[k].min(t.elapsed().as_secs_f64());
11007 sums[k] = out.iter().take(64).sum::<f32>();
11008 }
11009 }
11010 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11011 for (k, name) in ["previous", "tuned "].iter().enumerate() {
11012 println!(
11013 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11014 best[k] * 1e3,
11015 flops / best[k] / 1e9,
11016 sums[k]
11017 );
11018 }
11019 assert!(
11020 (sums[0] - sums[1]).abs() < 1e-2,
11021 "the tuned kernel changed the result: {} vs {}",
11022 sums[0],
11023 sums[1]
11024 );
11025 }
11026
11027 #[test]
11033 fn q4tp_matmat_blocked_matches_scalar() {
11034 use std::sync::atomic::Ordering::Relaxed;
11035 for &(rows, cols, b) in &[
11043 (64usize, 128usize, 7usize),
11044 (33, 96, 4),
11045 (16, 256, 9),
11046 (192, 2304, 37),
11047 ] {
11048 let total = cortiq_core::quant::expected_nbytes(
11049 cortiq_core::TensorDtype::Q4TiledP,
11050 &[rows, cols],
11051 )
11052 .unwrap();
11053 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11054 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11055 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11056 let step = cortiq_core::quant::f32_to_f16(0.1);
11057 for r in 0..rows {
11058 let o = params_off + r * 4;
11059 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11060 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11061 }
11062 let xs: Vec<f32> = (0..b * cols)
11063 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11064 .collect();
11065 let mut got = vec![0f32; b * rows];
11066 let mut want = vec![0f32; b * rows];
11067 let gpr = cols / 32;
11068 let view = super::Q4tpView::new(&bytes, rows, cols);
11069 let pool = crate::pool::Pool::from_env();
11070 super::Q4TP_ALT.store(2, Relaxed);
11071 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11072 super::Q4TP_ALT.store(1, Relaxed);
11073 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11074 super::Q4TP_ALT.store(0, Relaxed);
11075 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11082 let (mut worst, mut at) = (0f32, 0usize);
11083 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11084 if (g - w).abs() > worst {
11085 worst = (g - w).abs();
11086 at = i;
11087 }
11088 }
11089 assert!(
11090 worst <= 1e-4 * scale,
11091 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11092 (scale {scale:.3e}) at cell {at}: {} vs {}",
11093 got[at],
11094 want[at]
11095 );
11096
11097 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11104 for bi in 0..b {
11105 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11106 for r in 0..rows {
11107 let mut sc = vec![0f32; gpr];
11108 view.scales_into(r, gpr, &mut sc);
11109 let mut exact = 0f64;
11110 for j in 0..cols {
11111 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11112 exact += w as f64 * sq as f64 * act.xq[j] as f64;
11113 }
11114 exact *= act.sx as f64;
11115 for &(j, xv) in &act.outliers {
11116 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11117 exact += w as f64 * sq as f64 * xv as f64;
11118 }
11119 let i = bi * rows + r;
11120 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11121 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11122 }
11123 }
11124 println!(
11125 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11126 per-column {e_scalar:.3e}"
11127 );
11128 assert!(
11133 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11134 "{rows}x{cols} b={b}: error against f64 too large — blocked \
11135 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11136 );
11137 }
11138 }
11139
11140 #[test]
11144 #[ignore]
11145 fn q4t_matmat_throughput() {
11146 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11147 let total = cortiq_core::quant::expected_nbytes(
11148 cortiq_core::TensorDtype::Q4Tiled,
11149 &[rows, cols],
11150 )
11151 .unwrap();
11152 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11155 let sc = cortiq_core::quant::f32_to_f16(0.02);
11156 for t in bytes.chunks_mut(super::Q4_TILE) {
11157 t[..2].copy_from_slice(&sc.to_le_bytes());
11158 }
11159 let xs: Vec<f32> = (0..b * cols)
11160 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11161 .collect();
11162 let mut out = vec![0f32; b * rows];
11163 let pool = crate::pool::Pool::from_env();
11164 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11165 let reps: usize = std::env::var("CMF_BENCH_REPS")
11166 .ok()
11167 .and_then(|v| v.parse().ok())
11168 .unwrap_or(10);
11169 let mut best = f64::MAX;
11170 for _ in 0..reps {
11171 let t = std::time::Instant::now();
11172 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11173 best = best.min(t.elapsed().as_secs_f64());
11174 }
11175 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11176 println!(
11177 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11178 best * 1e3,
11179 flops / best / 1e9,
11180 out.iter().take(64).sum::<f32>()
11181 );
11182 }
11183
11184}