1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16 GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17 q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23 F32 {
24 data: Vec<f32>,
25 rows: usize,
26 cols: usize,
27 },
28 Mapped {
29 model: Arc<CmfModel>,
30 idx: usize,
32 dtype: TensorDtype,
33 rows: usize,
34 cols: usize,
35 row_scale: Vec<f32>,
37 col_field: Vec<f32>,
39 vbit_offsets: Vec<usize>,
43 repack: Vec<u8>,
51 },
52}
53
54fn repack_enabled() -> bool {
61 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62 *ON.get_or_init(|| {
63 std::env::var("CMF_REPACK")
64 .map(|v| v == "1")
65 .unwrap_or(cfg!(target_os = "android"))
66 })
67}
68
69fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73 #[cfg(target_arch = "aarch64")]
74 let arch_ok = sdot_enabled();
75 #[cfg(not(target_arch = "aarch64"))]
76 let arch_ok = false;
77 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78 return Vec::new();
79 }
80 q8_repack_layout(bytes, rows, cols)
81}
82
83fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86 let groups = rows / 4;
87 let mut rep = vec![0u8; groups * 4 * cols];
88 for g in 0..groups {
89 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90 for c in 0..cols / 16 {
91 for lane in 0..4 {
92 let src = (g * 4 + lane) * cols + c * 16;
93 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94 .copy_from_slice(&bytes[src..src + 16]);
95 }
96 }
97 }
98 rep
99}
100
101fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104 let ng = cols / GROUP_SIZE;
105 let bits = &bytes[..rows];
106 let mut offsets = Vec::with_capacity(rows + 1);
107 let mut off = rows + rows * ng * 2;
108 for r in 0..rows {
109 offsets.push(off);
110 off += (cols * bits[r] as usize).div_ceil(8);
111 }
112 offsets.push(off);
113 offsets
114}
115
116fn blocked_enabled() -> bool {
122 use std::sync::atomic::Ordering::Relaxed;
123 match BLOCKED_OVERRIDE.load(Relaxed) {
124 1 => false,
125 2 => true,
126 _ => {
127 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128 *ON.get_or_init(|| {
129 std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
130 })
131 }
132 }
133}
134
135static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
136
137pub fn set_blocked_override(on: Option<bool>) {
144 let v = match on {
145 None => 0,
146 Some(false) => 1,
147 Some(true) => 2,
148 };
149 BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
150}
151
152fn gpu_lmhead_enabled() -> bool {
153 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
154 *ON.get_or_init(|| std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true))
155}
156
157fn gpu_split_frac() -> f32 {
158 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
159 *F.get_or_init(|| {
160 std::env::var("CMF_GPU_SPLIT")
161 .ok()
162 .and_then(|v| v.parse::<f32>().ok())
163 .unwrap_or(0.5)
164 .clamp(0.0, 1.0)
165 })
166}
167
168impl QTensor {
169 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
170 debug_assert_eq!(data.len(), rows * cols);
171 Self::F32 { data, rows, cols }
172 }
173
174 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
177 let idx = model
180 .tensor_index(name)
181 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
182 let entry = &model.tensors[idx];
183 if entry.shape.len() != 2 {
184 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
185 }
186 let (rows, cols) = (entry.shape[0], entry.shape[1]);
187 let bytes = model.entry_bytes(entry);
188
189 match entry.dtype {
190 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
191 let n = rows * cols;
192 let scales_off = n;
193 let row_scale: Vec<f32> = (0..rows)
194 .map(|o| {
195 f16_to_f32(u16::from_le_bytes([
196 bytes[scales_off + o * 2],
197 bytes[scales_off + o * 2 + 1],
198 ]))
199 })
200 .collect();
201 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
202 let col_off = n + rows * 2;
203 (0..cols)
204 .map(|i| {
205 f16_to_f32(u16::from_le_bytes([
206 bytes[col_off + i * 2],
207 bytes[col_off + i * 2 + 1],
208 ]))
209 })
210 .collect()
211 } else {
212 Vec::new()
213 };
214 Ok(Self::Mapped {
215 model: model.clone(),
216 idx,
217 dtype: entry.dtype,
218 rows,
219 cols,
220 row_scale,
221 col_field,
222 vbit_offsets: Vec::new(),
223 repack: q8_repack(bytes, rows, cols),
224 })
225 }
226 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
228 model: model.clone(),
229 idx,
230 dtype: entry.dtype,
231 rows,
232 cols,
233 row_scale: Vec::new(),
234 col_field: Vec::new(),
235 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
236 repack: Vec::new(),
237 }),
238 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
242 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
243 let offsets: Vec<usize> = (0..=rows)
244 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
245 .collect();
246 Ok(Self::Mapped {
247 model: model.clone(),
248 idx,
249 dtype: entry.dtype,
250 rows,
251 cols,
252 row_scale: Vec::new(),
253 col_field: Vec::new(),
254 vbit_offsets: offsets,
255 repack: Vec::new(),
256 })
257 }
258 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
264 model: model.clone(),
265 idx,
266 dtype: entry.dtype,
267 rows,
268 cols,
269 row_scale: Vec::new(),
270 col_field: Vec::new(),
271 vbit_offsets: Vec::new(),
272 repack: Vec::new(),
273 }),
274 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
277 model: model.clone(),
278 idx,
279 dtype: entry.dtype,
280 rows,
281 cols,
282 row_scale: Vec::new(),
283 col_field: Vec::new(),
284 vbit_offsets: Vec::new(),
285 repack: Vec::new(),
286 }),
287 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
289 model: model.clone(),
290 idx,
291 dtype: entry.dtype,
292 rows,
293 cols,
294 row_scale: Vec::new(),
295 col_field: Vec::new(),
296 vbit_offsets: Vec::new(),
297 repack: Vec::new(),
298 }),
299 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
300 model: model.clone(),
301 idx,
302 dtype: entry.dtype,
303 rows,
304 cols,
305 row_scale: Vec::new(),
306 col_field: Vec::new(),
307 vbit_offsets: Vec::new(),
308 repack: Vec::new(),
309 }),
310 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
312 model: model.clone(),
313 idx,
314 dtype: entry.dtype,
315 rows,
316 cols,
317 row_scale: Vec::new(),
318 col_field: Vec::new(),
319 vbit_offsets: Vec::new(),
320 repack: Vec::new(),
321 }),
322 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
326 model: model.clone(),
327 idx,
328 dtype: entry.dtype,
329 rows,
330 cols,
331 row_scale: Vec::new(),
332 col_field: Vec::new(),
333 vbit_offsets: Vec::new(),
334 repack: Vec::new(),
335 }),
336 _ => {
338 let mut data = vec![0.0f32; rows * cols];
339 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
340 Ok(Self::from_f32(data, rows, cols))
341 }
342 }
343 }
344
345 pub(crate) fn is_q1(&self) -> bool {
348 matches!(
349 self,
350 Self::Mapped {
351 dtype: TensorDtype::Q1,
352 ..
353 }
354 )
355 }
356
357 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
360 match self {
361 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
362 _ => None,
363 }
364 }
365
366 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
373 match self {
374 #[cfg(target_os = "macos")]
375 Self::Mapped {
376 dtype: TensorDtype::Q1T,
377 ..
378 } if !crate::gpu::metal_q1t_enabled() => None,
379 Self::Mapped {
380 idx,
381 dtype:
382 TensorDtype::Q1
383 | TensorDtype::Q1T
384 | TensorDtype::Q4Block
385 | TensorDtype::Q4Tiled
386 | TensorDtype::Q4TiledP
390 | TensorDtype::Q8Row
391 | TensorDtype::Q8_2f,
392 rows,
393 cols,
394 ..
395 } => Some((*idx, *rows, *cols)),
396 _ => None,
397 }
398 }
399
400 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
406 match self {
407 Self::Mapped {
408 idx,
409 dtype: TensorDtype::Q4Tiled,
410 rows,
411 cols,
412 ..
413 } => Some((*idx, *rows, *cols)),
414 _ => None,
415 }
416 }
417
418 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
422 match self {
423 Self::Mapped {
424 idx,
425 dtype: TensorDtype::Q4TiledP,
426 rows,
427 cols,
428 ..
429 } => Some((*idx, *rows, *cols)),
430 _ => None,
431 }
432 }
433
434 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
439 match self {
440 Self::Mapped {
441 idx,
442 dtype: TensorDtype::Q8Row,
443 rows,
444 cols,
445 row_scale,
446 col_field,
447 ..
448 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
449 _ => None,
450 }
451 }
452
453 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
457 match self {
458 Self::Mapped { dtype, .. } => Some(*dtype),
459 _ => None,
460 }
461 }
462
463 pub fn model_idx(&self) -> Option<usize> {
468 match self {
469 Self::Mapped { idx, .. } => Some(*idx),
470 _ => None,
471 }
472 }
473
474 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
479 match self {
480 Self::Mapped { model, .. } => Some(model.clone()),
481 _ => None,
482 }
483 }
484
485 pub fn rows(&self) -> usize {
486 match self {
487 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
488 }
489 }
490
491 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
494 match self {
495 Self::Mapped {
496 model,
497 idx,
498 dtype: TensorDtype::Q4Tiled,
499 ..
500 } => Some((model, *idx)),
501 _ => None,
502 }
503 }
504
505 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
508 match self {
509 Self::Mapped {
510 model,
511 idx,
512 dtype: TensorDtype::Q4TiledP,
513 ..
514 } => Some((model, *idx)),
515 _ => None,
516 }
517 }
518
519 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
522 match self {
523 Self::Mapped {
524 model,
525 idx,
526 dtype: TensorDtype::Q2TiledP,
527 ..
528 } => Some((model, *idx)),
529 _ => None,
530 }
531 }
532
533 pub fn cols(&self) -> usize {
534 match self {
535 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
536 }
537 }
538
539 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
542 match self {
543 Self::Mapped {
544 model,
545 idx,
546 dtype: TensorDtype::Q1,
547 ..
548 } => Some((model, *idx)),
549 _ => None,
550 }
551 }
552
553 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
557 match self {
558 Self::Mapped {
559 model,
560 idx,
561 dtype: TensorDtype::Q8Row,
562 row_scale,
563 ..
564 } => Some((model, *idx, 0, row_scale.as_slice())),
565 Self::Mapped {
566 model,
567 idx,
568 dtype: TensorDtype::Q1,
569 ..
570 } => Some((model, *idx, 1, &[])),
571 Self::Mapped {
576 model,
577 idx,
578 dtype: TensorDtype::Q4Tiled,
579 ..
580 } => Some((model, *idx, 5, &[])),
581 Self::Mapped {
585 model,
586 idx,
587 dtype: TensorDtype::Q4TiledP,
588 ..
589 } => Some((model, *idx, 6, &[])),
590 Self::Mapped {
591 model,
592 idx,
593 dtype: TensorDtype::Q4Block,
594 ..
595 } => Some((model, *idx, 2, &[])),
596 Self::Mapped {
597 model,
598 idx,
599 dtype: TensorDtype::Q1T,
600 ..
601 } => Some((model, *idx, 3, &[])),
602 _ => None,
603 }
604 }
605
606 pub fn as_f32(&self) -> Option<&[f32]> {
609 match self {
610 Self::F32 { data, .. } => Some(data),
611 Self::Mapped { .. } => None,
612 }
613 }
614
615 fn quant_bytes(&self) -> &[u8] {
616 match self {
617 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
618 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
619 }
620 }
621
622 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
624 let cols = self.cols();
625 debug_assert_eq!(dst.len(), cols);
626 match self {
627 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
628 Self::Mapped {
629 dtype,
630 row_scale,
631 col_field,
632 vbit_offsets,
633 ..
634 } => {
635 if *dtype == TensorDtype::Q4Tiled {
636 let bytes = self.quant_bytes();
637 let gpr = cols / GROUP_SIZE;
638 for gi in 0..gpr {
639 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
640 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
641 for (k, &b) in tile[2..].iter().enumerate() {
642 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
643 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
644 }
645 }
646 return;
647 }
648 if *dtype == TensorDtype::Q4TiledP {
649 let bytes = self.quant_bytes();
650 let gpr = cols / GROUP_SIZE;
651 let v = Q4tpView::new(bytes, self.rows(), cols);
652 let mut sc = vec![0f32; gpr];
653 v.scales_into(r, gpr, &mut sc);
654 for gi in 0..gpr {
655 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
656 let s = sc[gi];
657 for (k, &b) in tile.iter().enumerate() {
658 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
659 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
660 }
661 }
662 return;
663 }
664 if *dtype == TensorDtype::Q2TiledP {
665 let bytes = self.quant_bytes();
666 let gpr = cols / GROUP_SIZE;
667 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
668 let mut sc = vec![0f32; gpr];
669 v.scales_into(r, gpr, &mut sc);
670 for gi in 0..gpr {
671 let ch =
672 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
673 let s = sc[gi];
674 for (k, &b) in ch.iter().enumerate() {
675 for j in 0..4 {
676 dst[gi * GROUP_SIZE + k * 4 + j] =
677 (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
678 }
679 }
680 }
681 return;
682 }
683 if *dtype == TensorDtype::Q4Block {
684 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
685 let gpr = cols / GROUP_SIZE;
686 for gi in 0..gpr {
687 let g = r * gpr + gi;
688 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
689 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
690 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
691 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
692 }
693 }
694 return;
695 }
696 if *dtype == TensorDtype::Q1 {
697 let bytes = self.quant_bytes();
698 let gpr = cols / GROUP_SIZE;
699 for gi in 0..gpr {
700 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
701 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
702 for (j, &b) in tile[2..].iter().enumerate() {
703 for k in 0..8 {
704 dst[gi * GROUP_SIZE + j * 8 + k] =
705 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
706 }
707 }
708 }
709 return;
710 }
711 if *dtype == TensorDtype::Q1T {
712 let bytes = self.quant_bytes();
713 let gpr = cols / GROUP_SIZE;
714 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
715 for gi in 0..gpr {
716 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
717 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
718 bytes[off],
719 bytes[off + 1],
720 ]));
721 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
722 for k in 0..GROUP_SIZE {
723 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
724 {
725 1 => s,
726 2 => -s,
727 _ => 0.0,
728 };
729 }
730 }
731 let rows = self.rows();
733 let entries = base_len + (rows + 1) * 4;
734 if entries <= bytes.len() {
735 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
736 let r0 = u32::from_le_bytes([
737 ptrs[r * 4],
738 ptrs[r * 4 + 1],
739 ptrs[r * 4 + 2],
740 ptrs[r * 4 + 3],
741 ]) as usize;
742 let r1 = u32::from_le_bytes([
743 ptrs[(r + 1) * 4],
744 ptrs[(r + 1) * 4 + 1],
745 ptrs[(r + 1) * 4 + 2],
746 ptrs[(r + 1) * 4 + 3],
747 ]) as usize;
748 let off = entries + r0 * 4;
749 for i in 0..r1 - r0 {
750 let item = &bytes[off + i * 4..off + i * 4 + 4];
751 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
752 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
753 item[2], item[3],
754 ]));
755 if c < cols {
756 dst[c] = v;
757 }
758 }
759 }
760 return;
761 }
762 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
763 let bytes = self.quant_bytes();
764 let rows = self.rows();
765 let ng = cols / GROUP_SIZE;
766 let bits = &bytes[..rows];
767 let sc_off = rows;
768 let off = vbit_offsets[r];
771 let b = bits[r] as usize;
772 let l = ((1usize << (b - 1)) - 1) as f32;
773 let data = &bytes[off..];
774 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
775 for (i, d) in dst.iter_mut().enumerate() {
776 while nbits < b {
777 acc = (acc << 8) | data[idx] as u64;
778 idx += 1;
779 nbits += 8;
780 }
781 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
782 nbits -= b;
783 let so = (r * ng + i / GROUP_SIZE) * 2;
784 let sv = f16_to_f32(u16::from_le_bytes([
785 bytes[sc_off + so],
786 bytes[sc_off + so + 1],
787 ]));
788 *d = (u - l) * sv;
789 }
790 return;
791 }
792 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
793 let s = row_scale[r];
794 match dtype {
795 TensorDtype::Q8Row => {
796 for (d, &b) in dst.iter_mut().zip(q) {
797 *d = (b as i8) as f32 * s;
798 }
799 }
800 TensorDtype::Q8_2f => {
801 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
802 *d = (b as i8) as f32 * s * col_field[i];
803 }
804 }
805 _ => unreachable!(),
806 }
807 }
808 }
809 }
810
811 pub fn sparse_col_ok(&self) -> bool {
816 match self {
817 Self::F32 { .. } => true,
818 Self::Mapped { dtype, .. } => {
819 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
820 }
821 }
822 }
823
824 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
828 let inter = self.cols();
829 let hidden = self.rows();
830 debug_assert_eq!(out.len(), hidden);
831 match self {
832 Self::F32 { data, .. } => {
833 for (k, o) in out.iter_mut().enumerate() {
834 *o += w * data[k * inter + c];
835 }
836 }
837 Self::Mapped {
838 dtype,
839 row_scale,
840 col_field,
841 ..
842 } => {
843 let q = self.quant_bytes();
844 let colf = if *dtype == TensorDtype::Q8_2f {
845 col_field[c]
846 } else {
847 1.0
848 };
849 let wc = w * colf;
850 for (k, o) in out.iter_mut().enumerate() {
851 let b = q[k * inter + c] as i8 as f32;
852 *o += wc * b * row_scale[k];
853 }
854 }
855 }
856 }
857
858 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
862 let cols = self.cols();
863 match self {
864 Self::F32 { data, .. } => {
865 let row = &data[r * cols..(r + 1) * cols];
866 row.iter().zip(x).map(|(w, v)| w * v).sum()
867 }
868 Self::Mapped {
869 dtype,
870 row_scale,
871 col_field,
872 ..
873 } => match dtype {
874 TensorDtype::Q8Row => {
875 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
876 dot_i8_f32(q, x) * row_scale[r]
877 }
878 TensorDtype::Q8_2f => {
879 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
880 dot_i8_col_f32(q, x, col_field) * row_scale[r]
881 }
882 _ => {
883 self.row_f32(r, scratch);
884 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
885 }
886 },
887 }
888 }
889
890 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
893 match self {
894 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
898 Self::Mapped {
899 model,
900 idx,
901 dtype,
902 rows,
903 cols,
904 row_scale,
905 col_field,
906 vbit_offsets,
907 repack,
908 } => {
909 let _ = (model, idx);
910 assert!(
920 out.len() >= *rows && x.len() >= *cols,
921 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
922 out.len(),
923 x.len(),
924 );
925 if *dtype == TensorDtype::Q4Block {
926 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
930 let t0 = std::time::Instant::now();
931 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
932 crate::gpu::ProbeArm::Gpu => {
933 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
934 crate::gpu::probe_record(
935 crate::gpu::OpClass::Matvec,
936 true,
937 t0.elapsed(),
938 );
939 return;
940 }
941 }
942 crate::gpu::ProbeArm::CpuTimed => {
943 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
944 crate::gpu::probe_record(
945 crate::gpu::OpClass::Matvec,
946 false,
947 t0.elapsed(),
948 );
949 return;
950 }
951 crate::gpu::ProbeArm::Cpu => {}
952 }
953 }
954 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
955 return;
956 }
957 if *dtype == TensorDtype::Q4Tiled {
958 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
963 let t0 = std::time::Instant::now();
964 let cls = crate::gpu::matvec_class(*rows, *cols);
965 match crate::gpu::probe_arm(cls) {
966 crate::gpu::ProbeArm::Gpu => {
967 if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
968 crate::gpu::probe_record(cls, true, t0.elapsed());
969 return;
970 }
971 }
972 crate::gpu::ProbeArm::CpuTimed => {
973 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
974 crate::gpu::probe_record(cls, false, t0.elapsed());
975 return;
976 }
977 crate::gpu::ProbeArm::Cpu => {}
978 }
979 }
980 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
981 return;
982 }
983 if *dtype == TensorDtype::Q4TiledP {
984 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
990 let t0 = std::time::Instant::now();
991 let cls = crate::gpu::matvec_class(*rows, *cols);
992 match crate::gpu::probe_arm(cls) {
993 crate::gpu::ProbeArm::Gpu => {
994 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
995 crate::gpu::probe_record(cls, true, t0.elapsed());
996 return;
997 }
998 }
999 crate::gpu::ProbeArm::CpuTimed => {
1000 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1001 crate::gpu::probe_record(cls, false, t0.elapsed());
1002 return;
1003 }
1004 crate::gpu::ProbeArm::Cpu => {}
1005 }
1006 }
1007 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1008 return;
1009 }
1010 if *dtype == TensorDtype::Q2TiledP {
1011 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1012 return;
1013 }
1014 if *dtype == TensorDtype::Q1 {
1015 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1020 let t0 = std::time::Instant::now();
1021 let arm = if crate::gpu::q1_force() {
1022 crate::gpu::ProbeArm::Gpu
1023 } else {
1024 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1025 };
1026 match arm {
1027 crate::gpu::ProbeArm::Gpu => {
1028 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1029 crate::gpu::probe_record(
1030 crate::gpu::OpClass::Matvec,
1031 true,
1032 t0.elapsed(),
1033 );
1034 return;
1035 }
1036 }
1037 crate::gpu::ProbeArm::CpuTimed => {
1038 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1039 crate::gpu::probe_record(
1040 crate::gpu::OpClass::Matvec,
1041 false,
1042 t0.elapsed(),
1043 );
1044 return;
1045 }
1046 crate::gpu::ProbeArm::Cpu => {}
1047 }
1048 }
1049 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1050 return;
1051 }
1052 if *dtype == TensorDtype::Q1T {
1053 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1057 let t0 = std::time::Instant::now();
1058 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1059 crate::gpu::ProbeArm::Gpu => {
1060 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1061 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1062 crate::gpu::probe_record(
1063 crate::gpu::OpClass::Matvec,
1064 true,
1065 t0.elapsed(),
1066 );
1067 return;
1068 }
1069 }
1070 crate::gpu::ProbeArm::CpuTimed => {
1071 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1072 crate::gpu::probe_record(
1073 crate::gpu::OpClass::Matvec,
1074 false,
1075 t0.elapsed(),
1076 );
1077 return;
1078 }
1079 crate::gpu::ProbeArm::Cpu => {}
1080 }
1081 }
1082 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1083 return;
1084 }
1085 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1086 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1087 return;
1088 }
1089 let xs = prescale(x, col_field, *dtype);
1090 if *rows >= crate::gpu::min_rows()
1095 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1096 && gpu_lmhead_enabled()
1097 && crate::gpu::enabled_here()
1098 {
1099 let t0 = std::time::Instant::now();
1102 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1103 crate::gpu::ProbeArm::Gpu => {}
1104 crate::gpu::ProbeArm::CpuTimed => {
1105 qmatvec(
1106 self.quant_bytes(),
1107 repack,
1108 row_scale,
1109 x,
1110 col_field,
1111 *dtype,
1112 *rows,
1113 *cols,
1114 out,
1115 pool,
1116 );
1117 crate::gpu::probe_record(
1118 crate::gpu::OpClass::Matvec,
1119 false,
1120 t0.elapsed(),
1121 );
1122 return;
1123 }
1124 crate::gpu::ProbeArm::Cpu => {
1125 qmatvec(
1126 self.quant_bytes(),
1127 repack,
1128 row_scale,
1129 x,
1130 col_field,
1131 *dtype,
1132 *rows,
1133 *cols,
1134 out,
1135 pool,
1136 );
1137 return;
1138 }
1139 }
1140 let frac = gpu_split_frac();
1141 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1142 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1143 let bytes = self.quant_bytes();
1144 let ok = std::thread::scope(|sc| {
1145 let g = sc.spawn(|| {
1146 crate::gpu::q8_matvec_range(
1147 model,
1148 *idx,
1149 cpu_rows,
1150 &row_scale[cpu_rows..],
1151 &xs,
1152 *rows - cpu_rows,
1153 *cols,
1154 out_gpu,
1155 )
1156 });
1157 if cpu_rows > 0 {
1158 let rep_cpu = if repack.is_empty() {
1161 &[][..]
1162 } else {
1163 &repack[..(cpu_rows / 4) * 4 * *cols]
1164 };
1165 qmatvec(
1166 &bytes[..cpu_rows * *cols],
1167 rep_cpu,
1168 &row_scale[..cpu_rows],
1169 x,
1170 col_field,
1171 *dtype,
1172 cpu_rows,
1173 *cols,
1174 out_cpu,
1175 pool,
1176 );
1177 }
1178 g.join().unwrap_or(false)
1179 });
1180 if ok {
1181 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1182 return;
1183 }
1184 qmatvec(
1187 &bytes[cpu_rows * *cols..(*rows) * *cols],
1188 &[],
1189 &row_scale[cpu_rows..],
1190 x,
1191 col_field,
1192 *dtype,
1193 *rows - cpu_rows,
1194 *cols,
1195 out_gpu,
1196 pool,
1197 );
1198 return;
1199 }
1200 qmatvec(
1201 self.quant_bytes(),
1202 repack,
1203 row_scale,
1204 x,
1205 col_field,
1206 *dtype,
1207 *rows,
1208 *cols,
1209 out,
1210 pool,
1211 );
1212 }
1213 }
1214 }
1215
1216 pub fn matvec2(
1218 &self,
1219 x1: &[f32],
1220 x2: &[f32],
1221 o1: &mut [f32],
1222 o2: &mut [f32],
1223 pool: Option<&Pool>,
1224 ) {
1225 match self {
1226 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1227 Self::Mapped {
1228 dtype,
1229 rows,
1230 cols,
1231 row_scale,
1232 col_field,
1233 vbit_offsets,
1234 ..
1235 } => {
1236 if *dtype == TensorDtype::Q4Block {
1237 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1238 return;
1239 }
1240 if *dtype == TensorDtype::Q4Tiled {
1241 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1242 return;
1243 }
1244 if *dtype == TensorDtype::Q4TiledP {
1245 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1246 return;
1247 }
1248 if *dtype == TensorDtype::Q2TiledP {
1249 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1250 return;
1251 }
1252 if *dtype == TensorDtype::Q1 {
1253 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1254 return;
1255 }
1256 if *dtype == TensorDtype::Q1T {
1257 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1263 return;
1264 }
1265 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1266 vbitmatvec2(
1267 self.quant_bytes(),
1268 vbit_offsets,
1269 x1,
1270 x2,
1271 *rows,
1272 *cols,
1273 o1,
1274 o2,
1275 pool,
1276 );
1277 return;
1278 }
1279 qmatvec2(
1280 self.quant_bytes(),
1281 row_scale,
1282 x1,
1283 x2,
1284 col_field,
1285 *dtype,
1286 *rows,
1287 *cols,
1288 o1,
1289 o2,
1290 pool,
1291 );
1292 }
1293 }
1294 }
1295}
1296
1297impl QTensor {
1298 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1303 let cols = self.cols();
1304 let rows = self.rows();
1305 debug_assert_eq!(xs_all.len(), b * cols);
1306 debug_assert_eq!(out.len(), b * rows);
1307 if crate::gptq_capture::capturing() {
1311 if let Self::Mapped { model, idx, .. } = self {
1312 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1313 }
1314 }
1315 match self {
1316 Self::F32 { data, .. } => {
1317 let out_addr = SendMut(out.as_mut_ptr());
1318 let run = |start: usize, end: usize| {
1319 for o in start..end {
1320 let row = &data[o * cols..(o + 1) * cols];
1321 for bi in 0..b {
1322 let x = &xs_all[bi * cols..(bi + 1) * cols];
1323 let mut acc = 0f32;
1324 for j in 0..cols {
1325 acc += row[j] * x[j];
1326 }
1327 unsafe { *out_addr.at(bi * rows + o) = acc };
1328 }
1329 }
1330 };
1331 dispatch_rows(pool, rows, &run);
1332 }
1333 Self::Mapped {
1334 dtype,
1335 row_scale,
1336 col_field,
1337 vbit_offsets,
1338 ..
1339 } => {
1340 if *dtype == TensorDtype::Q4Block {
1341 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1342 return;
1343 }
1344 if *dtype == TensorDtype::Q4TiledP {
1345 if b >= 32
1357 && b * rows * cols >= 128_000_000
1358 && cols % 32 == 0
1359 && !crate::gpu::mm_killed()
1360 && crate::gpu::enabled_here()
1361 {
1362 let class = if b >= 128 {
1363 crate::gpu::OpClass::MatmatWide
1364 } else {
1365 crate::gpu::OpClass::Matmat
1366 };
1367 if let Self::Mapped { model, idx, .. } = self {
1368 let t0 = std::time::Instant::now();
1369 let resident = crate::gpu::weight_is_resident(model, *idx);
1373 match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1374 crate::gpu::ProbeArm::Gpu => {
1375 if crate::gpu::q4tp_matmat(
1376 model, *idx, xs_all, b, rows, cols, out,
1377 ) {
1378 let el = t0.elapsed();
1379 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1389 let budget = std::time::Duration::from_secs_f64(
1390 flops / 1.5e12 * 8.0 + 0.020,
1391 );
1392 if el > budget && !crate::gpu::probe_was_cold() {
1393 tracing::warn!(
1394 "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1395 device contended, CPU for the rest of the process"
1396 );
1397 crate::gpu::mm_kill();
1398 }
1399 crate::gpu::probe_record(class, true, el);
1400 return;
1401 }
1402 }
1403 crate::gpu::ProbeArm::CpuTimed => {
1404 q4tp_matmat(
1405 self.quant_bytes(),
1406 xs_all,
1407 b,
1408 rows,
1409 cols,
1410 out,
1411 pool,
1412 );
1413 crate::gpu::probe_record(class, false, t0.elapsed());
1414 return;
1415 }
1416 crate::gpu::ProbeArm::Cpu => {}
1417 }
1418 }
1419 }
1420 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1421 return;
1422 }
1423 if *dtype == TensorDtype::Q2TiledP {
1424 if b >= 32
1430 && b * rows * cols >= 128_000_000
1431 && cols % 32 == 0
1432 && !crate::gpu::mm_killed()
1433 && crate::gpu::enabled_here()
1434 {
1435 let class = if b >= 128 {
1436 crate::gpu::OpClass::MatmatWide
1437 } else {
1438 crate::gpu::OpClass::Matmat
1439 };
1440 if let Self::Mapped { model, idx, .. } = self {
1441 let t0 = std::time::Instant::now();
1442 match crate::gpu::probe_arm(class) {
1443 crate::gpu::ProbeArm::Gpu => {
1444 if crate::gpu::q2tp_matmat(
1445 model, *idx, xs_all, b, rows, cols, out,
1446 ) {
1447 crate::gpu::probe_record(class, true, t0.elapsed());
1448 return;
1449 }
1450 }
1451 crate::gpu::ProbeArm::CpuTimed => {
1452 q2tp_matmat(
1453 self.quant_bytes(),
1454 xs_all,
1455 b,
1456 rows,
1457 cols,
1458 out,
1459 pool,
1460 );
1461 crate::gpu::probe_record(class, false, t0.elapsed());
1462 return;
1463 }
1464 crate::gpu::ProbeArm::Cpu => {}
1465 }
1466 }
1467 }
1468 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1473 return;
1474 }
1475 if *dtype == TensorDtype::Q4Tiled {
1476 if b >= 32
1488 && b * rows * cols >= 128_000_000
1489 && cols % 32 == 0
1490 && !crate::gpu::mm_killed()
1491 && crate::gpu::enabled_here()
1492 {
1493 let class = if b >= 128 {
1494 crate::gpu::OpClass::MatmatWide
1495 } else {
1496 crate::gpu::OpClass::Matmat
1497 };
1498 if let Self::Mapped { model, idx, .. } = self {
1499 let t0 = std::time::Instant::now();
1500 match crate::gpu::probe_arm(class) {
1501 crate::gpu::ProbeArm::Gpu => {
1502 if crate::gpu::q4t_matmat(
1503 model, *idx, xs_all, b, rows, cols, out,
1504 ) {
1505 let el = t0.elapsed();
1506 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1516 let budget = std::time::Duration::from_secs_f64(
1517 flops / 1.5e12 * 8.0 + 0.020,
1518 );
1519 if el > budget && !crate::gpu::probe_was_cold() {
1520 tracing::warn!(
1521 "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1522 device contended, CPU for the rest of the process"
1523 );
1524 crate::gpu::mm_kill();
1525 }
1526 crate::gpu::probe_record(class, true, el);
1527 return;
1528 }
1529 }
1530 crate::gpu::ProbeArm::CpuTimed => {
1531 q4t_matmat(
1532 self.quant_bytes(),
1533 xs_all,
1534 b,
1535 rows,
1536 cols,
1537 out,
1538 pool,
1539 );
1540 crate::gpu::probe_record(class, false, t0.elapsed());
1541 return;
1542 }
1543 crate::gpu::ProbeArm::Cpu => {}
1544 }
1545 }
1546 }
1547 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1548 return;
1549 }
1550 if *dtype == TensorDtype::Q1 {
1551 if b >= 32
1554 && b * rows * cols >= 128_000_000
1555 && cols % 64 == 0
1556 && crate::gpu::enabled_here()
1557 {
1558 if let Self::Mapped { model, idx, .. } = self {
1559 let t0 = std::time::Instant::now();
1560 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1561 crate::gpu::ProbeArm::Gpu => {
1562 if crate::gpu::q1_matmat(
1563 model, *idx, xs_all, b, rows, cols, out,
1564 ) {
1565 crate::gpu::probe_record(
1566 crate::gpu::OpClass::Matmat,
1567 true,
1568 t0.elapsed(),
1569 );
1570 return;
1571 }
1572 }
1573 crate::gpu::ProbeArm::CpuTimed => {
1574 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1575 crate::gpu::probe_record(
1576 crate::gpu::OpClass::Matmat,
1577 false,
1578 t0.elapsed(),
1579 );
1580 return;
1581 }
1582 crate::gpu::ProbeArm::Cpu => {}
1583 }
1584 }
1585 }
1586 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1587 return;
1588 }
1589 if *dtype == TensorDtype::Q1T {
1590 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1593 if let Self::Mapped { model, idx, .. } = self {
1594 let t0 = std::time::Instant::now();
1595 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1596 crate::gpu::ProbeArm::Gpu => {
1597 if crate::gpu::q1t_matmat(
1598 model, *idx, xs_all, b, rows, cols, out,
1599 ) {
1600 crate::gpu::probe_record(
1601 crate::gpu::OpClass::Matmat,
1602 true,
1603 t0.elapsed(),
1604 );
1605 return;
1606 }
1607 }
1608 crate::gpu::ProbeArm::CpuTimed => {
1609 q1t_matmat(
1610 self.quant_bytes(),
1611 xs_all,
1612 b,
1613 rows,
1614 cols,
1615 out,
1616 pool,
1617 );
1618 crate::gpu::probe_record(
1619 crate::gpu::OpClass::Matmat,
1620 false,
1621 t0.elapsed(),
1622 );
1623 return;
1624 }
1625 crate::gpu::ProbeArm::Cpu => {}
1626 }
1627 }
1628 }
1629 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1630 return;
1631 }
1632 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1633 vbitmatmat(
1634 self.quant_bytes(),
1635 vbit_offsets,
1636 xs_all,
1637 b,
1638 rows,
1639 cols,
1640 out,
1641 pool,
1642 );
1643 return;
1644 }
1645 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1646 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1647 .collect();
1648 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1654 if let Self::Mapped { model, idx, .. } = self {
1655 let t0 = std::time::Instant::now();
1656 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1657 crate::gpu::ProbeArm::Gpu
1658 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1659 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1660 {
1661 let q = self.quant_bytes();
1665 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1666 return;
1667 }
1668 crate::gpu::ProbeArm::Gpu => {
1669 let flat: Vec<f32> =
1670 pre.iter().flat_map(|v| v.iter().copied()).collect();
1671 if crate::gpu::q8_matmat(
1672 model, *idx, row_scale, &flat, b, rows, cols, out,
1673 ) {
1674 crate::gpu::probe_record(
1675 crate::gpu::OpClass::Matmat,
1676 true,
1677 t0.elapsed(),
1678 );
1679 return;
1680 }
1681 }
1682 crate::gpu::ProbeArm::CpuTimed => {
1683 let q = self.quant_bytes();
1684 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1685 crate::gpu::probe_record(
1686 crate::gpu::OpClass::Matmat,
1687 false,
1688 t0.elapsed(),
1689 );
1690 return;
1691 }
1692 crate::gpu::ProbeArm::Cpu => {}
1693 }
1694 }
1695 }
1696 let q = self.quant_bytes();
1697 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1698 }
1699 }
1700 }
1701}
1702
1703impl QTensor {
1704 pub fn matvec_many<const N: usize>(
1711 ts: [&QTensor; N],
1712 x: &[f32],
1713 mut outs: [&mut [f32]; N],
1714 pool: Option<&Pool>,
1715 ) {
1716 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1717 let uniform_q8 = ts.iter().all(|t| {
1718 matches!(
1719 t,
1720 Self::Mapped {
1721 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1722 ..
1723 }
1724 )
1725 });
1726 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1727 let uniform_q4 = ts.iter().all(|t| {
1728 matches!(
1729 t,
1730 Self::Mapped {
1731 dtype: TensorDtype::Q4Block,
1732 ..
1733 }
1734 )
1735 });
1736 let uniform_vbit = ts.iter().all(|t| {
1737 matches!(
1738 t,
1739 Self::Mapped {
1740 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1741 ..
1742 }
1743 )
1744 });
1745 let uniform_q1 = ts.iter().all(|t| {
1746 matches!(
1747 t,
1748 Self::Mapped {
1749 dtype: TensorDtype::Q1,
1750 ..
1751 }
1752 )
1753 });
1754 let uniform_q1t = ts.iter().all(|t| {
1755 matches!(
1756 t,
1757 Self::Mapped {
1758 dtype: TensorDtype::Q1T,
1759 ..
1760 }
1761 )
1762 });
1763 let uniform_q4tp = ts.iter().all(|t| {
1768 matches!(
1769 t,
1770 Self::Mapped {
1771 dtype: TensorDtype::Q4TiledP,
1772 ..
1773 }
1774 )
1775 }) && ts.iter().all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1776 let Some(pool) = pool else {
1777 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1778 t.matvec(x, o, None);
1779 }
1780 return;
1781 };
1782 if total_rows < 256
1783 || !(uniform_q8
1784 || uniform_f32
1785 || uniform_q4
1786 || uniform_vbit
1787 || uniform_q1
1788 || uniform_q1t
1789 || uniform_q4tp)
1790 {
1791 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1792 t.matvec(x, o, Some(pool));
1793 }
1794 return;
1795 }
1796
1797 if uniform_q4tp {
1798 let cols = ts[0].cols();
1804 let gpr = cols / GROUP_SIZE;
1805 let views: Vec<Q4tpView> = ts
1806 .iter()
1807 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1808 .collect();
1809 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1810 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1811 let locate = |flat: usize| -> (usize, usize) {
1813 let mut acc = 0;
1814 for (i, &r) in rows_of.iter().enumerate() {
1815 if flat < acc + r {
1816 return (i, flat - acc);
1817 }
1818 acc += r;
1819 }
1820 (rows_of.len() - 1, 0)
1821 };
1822 let (views, outs_addr) = (&views, &outs_addr);
1823 if a8w8_enabled() {
1824 let act = split_act(x);
1825 let act = &act;
1826 let run = |start: usize, end: usize| {
1827 let mut sc = vec![0f32; gpr];
1828 for flat in start..end {
1829 let (t, r) = locate(flat);
1830 let v = &views[t];
1831 v.scales_into(r, gpr, &mut sc);
1832 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1833 for &(j, xv) in &act.outliers {
1834 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1835 acc += w * s * xv;
1836 }
1837 unsafe { *outs_addr[t].at(r) = acc };
1839 }
1840 };
1841 pool.run_rows(total_rows, &run);
1842 } else {
1843 let run = |start: usize, end: usize| {
1844 let mut sc = vec![0f32; gpr];
1845 for flat in start..end {
1846 let (t, r) = locate(flat);
1847 let v = &views[t];
1848 v.scales_into(r, gpr, &mut sc);
1849 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1851 }
1852 };
1853 pool.run_rows(total_rows, &run);
1854 }
1855 return;
1856 }
1857
1858 if uniform_q1 {
1859 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1862 if a8w8_enabled() {
1863 let act = split_act(x);
1864 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1865 let (act, gsum) = (&act, &gsum);
1866 let closures: [_; N] = std::array::from_fn(|i| {
1867 let (bytes, gpr, out) =
1868 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1869 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1870 });
1871 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1872 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1873 pool.run_many(&parts);
1874 } else {
1875 let closures: [_; N] = std::array::from_fn(|i| {
1876 let (bytes, gpr, out) =
1877 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1878 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1879 });
1880 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1881 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1882 pool.run_many(&parts);
1883 }
1884 return;
1885 }
1886
1887 if uniform_q1t {
1888 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1892 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1893 if a8w8_enabled() {
1894 let act = split_act(x);
1895 let act = &act;
1896 let x_ref = x;
1897 let closures: [_; N] = std::array::from_fn(|i| {
1898 let bytes = ts[i].quant_bytes();
1899 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1900 let gpr = cols / GROUP_SIZE;
1901 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1902 let out = outs_addr[i];
1903 move |s: usize, e: usize| {
1904 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1905 }
1906 });
1907 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1908 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1909 pool.run_many(&parts);
1910 } else {
1911 let x_ref = x;
1912 let closures: [_; N] = std::array::from_fn(|i| {
1913 let bytes = ts[i].quant_bytes();
1914 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1915 let gpr = cols / GROUP_SIZE;
1916 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1917 let out = outs_addr[i];
1918 move |s: usize, e: usize| {
1919 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1920 }
1921 });
1922 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1923 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1924 pool.run_many(&parts);
1925 }
1926 return;
1927 }
1928
1929 if uniform_q4 || uniform_vbit {
1930 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1931 if a8w8_enabled() {
1933 let act = split_act(x);
1934 let act = &act;
1935 if uniform_q4 {
1936 let closures: [_; N] = std::array::from_fn(|i| {
1937 let (packed, scales) =
1938 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1939 let (gpr, cols, out) =
1940 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1941 move |s: usize, e: usize| {
1942 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1943 }
1944 });
1945 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1946 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1947 pool.run_many(&parts);
1948 } else {
1949 let closures: [_; N] = std::array::from_fn(|i| {
1950 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1951 unreachable!()
1952 };
1953 let (bytes, rows, cols, out) = (
1954 ts[i].quant_bytes(),
1955 ts[i].rows(),
1956 ts[i].cols(),
1957 outs_addr[i],
1958 );
1959 move |s: usize, e: usize| {
1960 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1961 }
1962 });
1963 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1964 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1965 pool.run_many(&parts);
1966 }
1967 return;
1968 }
1969 if uniform_q4 {
1970 let closures: [_; N] = std::array::from_fn(|i| {
1971 let (packed, scales) =
1972 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1973 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1974 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1975 });
1976 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1977 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1978 pool.run_many(&parts);
1979 } else {
1980 let closures: [_; N] = std::array::from_fn(|i| {
1981 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1982 unreachable!()
1983 };
1984 let (bytes, rows, cols, out) = (
1985 ts[i].quant_bytes(),
1986 ts[i].rows(),
1987 ts[i].cols(),
1988 outs_addr[i],
1989 );
1990 move |s: usize, e: usize| {
1991 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1992 }
1993 });
1994 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1995 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1996 pool.run_many(&parts);
1997 }
1998 return;
1999 }
2000
2001 if uniform_f32 {
2002 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2003 let closures: [_; N] = std::array::from_fn(|i| {
2004 let Self::F32 { data, cols, .. } = ts[i] else {
2005 unreachable!()
2006 };
2007 let out = outs_addr[i];
2008 move |start: usize, end: usize| {
2009 for o in start..end {
2010 let row = &data[o * cols..(o + 1) * cols];
2011 let mut sum = 0.0f32;
2012 for j in 0..*cols {
2013 sum += row[j] * x[j];
2014 }
2015 unsafe { *out.at(o) = sum };
2017 }
2018 }
2019 });
2020 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2021 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2022 pool.run_many(&parts);
2023 return;
2024 }
2025
2026 struct Ctx<'a> {
2029 bytes: &'a [u8],
2030 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2031 rep: &'a [u8],
2032 row_scale: &'a [f32],
2033 cols: usize,
2034 xs: std::borrow::Cow<'a, [f32]>,
2035 }
2036 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2037 let Self::Mapped {
2038 dtype,
2039 cols,
2040 row_scale,
2041 col_field,
2042 repack,
2043 ..
2044 } = ts[i]
2045 else {
2046 unreachable!()
2047 };
2048 Ctx {
2049 bytes: ts[i].quant_bytes(),
2050 rep: repack,
2051 row_scale,
2052 cols: *cols,
2053 xs: prescale(x, col_field, *dtype),
2054 }
2055 });
2056 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2057 #[cfg(target_arch = "aarch64")]
2058 if sdot_enabled() {
2059 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2060 let closures: [_; N] = std::array::from_fn(|i| {
2061 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2062 move |start: usize, end: usize| {
2063 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2064 }
2065 });
2066 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2067 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2068 pool.run_many(&parts);
2069 return;
2070 }
2071 #[cfg(target_arch = "x86_64")]
2072 if avx2_a8w8_enabled() {
2073 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2074 let closures: [_; N] = std::array::from_fn(|i| {
2075 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2076 move |start: usize, end: usize| {
2077 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2078 }
2079 });
2080 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2081 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2082 pool.run_many(&parts);
2083 return;
2084 }
2085 let closures: [_; N] = std::array::from_fn(|i| {
2086 let (c, out) = (&ctxs[i], outs_addr[i]);
2087 move |start: usize, end: usize| {
2088 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2089 }
2090 });
2091 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2092 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2093 pool.run_many(&parts);
2094 }
2095}
2096
2097impl QTensor {
2098 #[allow(clippy::needless_range_loop)]
2103 pub fn matvec2_many<const N: usize>(
2104 ts: [&QTensor; N],
2105 x1: &[f32],
2106 x2: &[f32],
2107 mut o1s: [&mut [f32]; N],
2108 mut o2s: [&mut [f32]; N],
2109 pool: Option<&Pool>,
2110 ) {
2111 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2112 let uniform_q8 = ts.iter().all(|t| {
2113 matches!(
2114 t,
2115 Self::Mapped {
2116 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2117 ..
2118 }
2119 )
2120 });
2121 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2122 let uniform_q4 = ts.iter().all(|t| {
2123 matches!(
2124 t,
2125 Self::Mapped {
2126 dtype: TensorDtype::Q4Block,
2127 ..
2128 }
2129 )
2130 });
2131 let uniform_vbit = ts.iter().all(|t| {
2132 matches!(
2133 t,
2134 Self::Mapped {
2135 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2136 ..
2137 }
2138 )
2139 });
2140 let fusable = pool.is_some()
2141 && total_rows >= 256
2142 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2143 if !fusable {
2144 for i in 0..N {
2145 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2146 }
2147 return;
2148 }
2149 let pool = pool.unwrap();
2150
2151 if uniform_q4 || uniform_vbit {
2152 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2153 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2154 if a8w8_enabled() {
2156 let a1 = split_act(x1);
2157 let a2 = split_act(x2);
2158 let (a1, a2) = (&a1, &a2);
2159 if uniform_q4 {
2160 let closures: [_; N] = std::array::from_fn(|i| {
2161 let (packed, scales) =
2162 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2163 let (gpr, cols, o1, o2) =
2164 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2165 move |s: usize, e: usize| {
2166 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2167 }
2168 });
2169 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2170 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2171 pool.run_many(&parts);
2172 } else {
2173 let closures: [_; N] = std::array::from_fn(|i| {
2174 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2175 unreachable!()
2176 };
2177 let (bytes, rows, cols, o1, o2) = (
2178 ts[i].quant_bytes(),
2179 ts[i].rows(),
2180 ts[i].cols(),
2181 p1[i],
2182 p2[i],
2183 );
2184 move |s: usize, e: usize| {
2185 vbit_range2_a8w8(
2186 bytes,
2187 vbit_offsets,
2188 x1,
2189 x2,
2190 a1,
2191 a2,
2192 rows,
2193 cols,
2194 o1,
2195 o2,
2196 s,
2197 e,
2198 )
2199 }
2200 });
2201 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2202 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2203 pool.run_many(&parts);
2204 }
2205 return;
2206 }
2207 if uniform_q4 {
2208 let closures: [_; N] = std::array::from_fn(|i| {
2209 let (packed, scales) =
2210 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2211 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2212 move |s: usize, e: usize| {
2213 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2214 }
2215 });
2216 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2217 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2218 pool.run_many(&parts);
2219 } else {
2220 let closures: [_; N] = std::array::from_fn(|i| {
2221 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2222 unreachable!()
2223 };
2224 let (bytes, rows, cols, o1, o2) = (
2225 ts[i].quant_bytes(),
2226 ts[i].rows(),
2227 ts[i].cols(),
2228 p1[i],
2229 p2[i],
2230 );
2231 move |s: usize, e: usize| {
2232 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2233 }
2234 });
2235 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2236 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2237 pool.run_many(&parts);
2238 }
2239 return;
2240 }
2241
2242 if uniform_f32 {
2243 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2244 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2245 let closures: [_; N] = std::array::from_fn(|i| {
2246 let Self::F32 { data, cols, .. } = ts[i] else {
2247 unreachable!()
2248 };
2249 let (o1, o2) = (p1[i], p2[i]);
2250 move |start: usize, end: usize| {
2251 for o in start..end {
2252 let row = &data[o * cols..(o + 1) * cols];
2253 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2254 for j in 0..*cols {
2255 s1 += row[j] * x1[j];
2256 s2 += row[j] * x2[j];
2257 }
2258 unsafe {
2260 *o1.at(o) = s1;
2261 *o2.at(o) = s2;
2262 }
2263 }
2264 }
2265 });
2266 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2267 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2268 pool.run_many(&parts);
2269 return;
2270 }
2271
2272 struct Ctx<'a> {
2273 bytes: &'a [u8],
2274 row_scale: &'a [f32],
2275 cols: usize,
2276 xs1: std::borrow::Cow<'a, [f32]>,
2277 xs2: std::borrow::Cow<'a, [f32]>,
2278 }
2279 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2280 let Self::Mapped {
2281 dtype,
2282 cols,
2283 row_scale,
2284 col_field,
2285 ..
2286 } = ts[i]
2287 else {
2288 unreachable!()
2289 };
2290 Ctx {
2291 bytes: ts[i].quant_bytes(),
2292 row_scale,
2293 cols: *cols,
2294 xs1: prescale(x1, col_field, *dtype),
2295 xs2: prescale(x2, col_field, *dtype),
2296 }
2297 });
2298 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2299 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2300 #[cfg(target_arch = "aarch64")]
2301 if sdot_enabled() {
2302 let acts: [(SplitAct, SplitAct); N] =
2303 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2304 let closures: [_; N] = std::array::from_fn(|i| {
2305 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2306 move |start: usize, end: usize| {
2307 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2308 }
2309 });
2310 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2311 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2312 pool.run_many(&parts);
2313 return;
2314 }
2315 #[cfg(target_arch = "x86_64")]
2316 if avx2_a8w8_enabled() {
2317 let acts: [(SplitAct, SplitAct); N] =
2318 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2319 let closures: [_; N] = std::array::from_fn(|i| {
2320 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2321 move |start: usize, end: usize| {
2322 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2323 }
2324 });
2325 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2326 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2327 pool.run_many(&parts);
2328 return;
2329 }
2330 let closures: [_; N] = std::array::from_fn(|i| {
2331 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2332 move |start: usize, end: usize| {
2333 q8_range2_f32(
2334 c.bytes,
2335 c.row_scale,
2336 &c.xs1,
2337 &c.xs2,
2338 c.cols,
2339 o1,
2340 o2,
2341 start,
2342 end,
2343 )
2344 }
2345 });
2346 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2347 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2348 pool.run_many(&parts);
2349 }
2350
2351 pub fn matvec_silu_mul(
2356 gate: &QTensor,
2357 up: &QTensor,
2358 x: &[f32],
2359 out: &mut [f32],
2360 pool: Option<&Pool>,
2361 ) -> bool {
2362 let inter = gate.rows();
2363 debug_assert_eq!(up.rows(), inter);
2364 debug_assert_eq!(out.len(), inter);
2365 debug_assert_eq!(gate.cols(), up.cols());
2366 if !a8w8_enabled() {
2367 return false;
2368 }
2369 let act = split_act(x);
2370 let act = &act;
2371 let x_ref = x;
2372 let out_addr = SendMut(out.as_mut_ptr());
2373
2374 match (gate, up) {
2375 (
2377 Self::Mapped {
2378 dtype: TensorDtype::Q4Block,
2379 ..
2380 },
2381 Self::Mapped {
2382 dtype: TensorDtype::Q4Block,
2383 ..
2384 },
2385 ) => {
2386 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2387 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2388 let gpr = gate.cols() / GROUP_SIZE;
2389 let cols = gate.cols();
2390 let run = move |start: usize, end: usize| {
2391 for r in start..end {
2392 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2393 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2394 for &(j, xv) in &act.outliers {
2395 let flat = r * cols + j;
2396 let gb = gp[flat / 2];
2397 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2398 let gsc = f16_to_f32(u16::from_le_bytes([
2399 gs[(flat / GROUP_SIZE) * 2],
2400 gs[(flat / GROUP_SIZE) * 2 + 1],
2401 ]));
2402 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2403 let ub = up_p[flat / 2];
2404 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2405 let usc = f16_to_f32(u16::from_le_bytes([
2406 up_s[(flat / GROUP_SIZE) * 2],
2407 up_s[(flat / GROUP_SIZE) * 2 + 1],
2408 ]));
2409 uv += ((un as i32 - 8) as f32) * usc * xv;
2410 }
2411 let silu_g = gv / (1.0 + (-gv).exp());
2412 unsafe { *out_addr.at(r) = silu_g * uv };
2414 }
2415 };
2416 dispatch_rows(pool, inter, &run);
2417 true
2418 }
2419 (
2423 Self::Mapped {
2424 dtype: TensorDtype::Q4Tiled,
2425 ..
2426 },
2427 Self::Mapped {
2428 dtype: TensorDtype::Q4Tiled,
2429 ..
2430 },
2431 ) => {
2432 let g_bytes = gate.quant_bytes();
2433 let u_bytes = up.quant_bytes();
2434 let gpr = gate.cols() / GROUP_SIZE;
2435 let run = move |start: usize, end: usize| {
2436 for r in start..end {
2437 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2438 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2439 for &(j, xv) in &act.outliers {
2440 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2441 gv += w * s * xv;
2442 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2443 uv += w * s * xv;
2444 }
2445 let silu_g = gv / (1.0 + (-gv).exp());
2446 unsafe { *out_addr.at(r) = silu_g * uv };
2448 }
2449 };
2450 dispatch_rows(pool, inter, &run);
2451 true
2452 }
2453 (
2456 Self::Mapped {
2457 dtype: TensorDtype::Q4TiledP,
2458 ..
2459 },
2460 Self::Mapped {
2461 dtype: TensorDtype::Q4TiledP,
2462 ..
2463 },
2464 ) => {
2465 let cols = gate.cols();
2466 let gpr = cols / GROUP_SIZE;
2467 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2468 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2469 let run = |start: usize, end: usize| {
2470 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2471 for r in start..end {
2472 gv_view.scales_into(r, gpr, &mut gsc);
2473 uv_view.scales_into(r, gpr, &mut usc);
2474 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2475 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2476 for &(j, xv) in &act.outliers {
2477 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2478 gv += w * s * xv;
2479 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2480 uv += w * s * xv;
2481 }
2482 let silu_g = gv / (1.0 + (-gv).exp());
2483 unsafe { *out_addr.at(r) = silu_g * uv };
2485 }
2486 };
2487 dispatch_rows(pool, inter, &run);
2488 true
2489 }
2490 (
2496 Self::Mapped {
2497 dtype: TensorDtype::Q1,
2498 ..
2499 },
2500 Self::Mapped {
2501 dtype: TensorDtype::Q1,
2502 ..
2503 },
2504 ) => {
2505 let g_bytes = gate.quant_bytes();
2506 let u_bytes = up.quant_bytes();
2507 let gpr = gate.cols() / GROUP_SIZE;
2508 let gsum = q1_group_sums(&act.xq, gpr);
2509 let gsum = &gsum;
2510 let run = move |start: usize, end: usize| {
2511 for r in start..end {
2512 let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2513 let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2514 for &(j, xv) in &act.outliers {
2515 let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2516 gv += w * s * xv;
2517 let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2518 uv += w * s * xv;
2519 }
2520 let silu_g = gv / (1.0 + (-gv).exp());
2521 unsafe { *out_addr.at(r) = silu_g * uv };
2523 }
2524 };
2525 dispatch_rows(pool, inter, &run);
2526 true
2527 }
2528 (
2532 Self::Mapped {
2533 dtype: TensorDtype::Q2TiledP,
2534 ..
2535 },
2536 Self::Mapped {
2537 dtype: TensorDtype::Q2TiledP,
2538 ..
2539 },
2540 ) => {
2541 let cols = gate.cols();
2542 let gpr = cols / GROUP_SIZE;
2543 let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2544 let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2545 let gsum = q1_group_sums(&act.xq, gpr);
2546 let gsum = &gsum;
2547 let run = move |start: usize, end: usize| {
2548 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2549 for r in start..end {
2550 gv_view.scales_into(r, gpr, &mut gsc);
2551 uv_view.scales_into(r, gpr, &mut usc);
2552 let mut gv =
2553 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2554 let mut uv =
2555 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2556 for &(j, xv) in &act.outliers {
2557 let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2558 gv += w * s * xv;
2559 let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2560 uv += w * s * xv;
2561 }
2562 let silu_g = gv / (1.0 + (-gv).exp());
2563 unsafe { *out_addr.at(r) = silu_g * uv };
2565 }
2566 };
2567 dispatch_rows(pool, inter, &run);
2568 true
2569 }
2570 (
2575 Self::Mapped {
2576 dtype: TensorDtype::Q8Row,
2577 row_scale: g_rs,
2578 ..
2579 },
2580 Self::Mapped {
2581 dtype: TensorDtype::Q8Row,
2582 row_scale: u_rs,
2583 ..
2584 },
2585 ) => {
2586 let g_bytes = gate.quant_bytes();
2587 let u_bytes = up.quant_bytes();
2588 let cols = gate.cols();
2589 let run = move |start: usize, end: usize| {
2590 for r in start..end {
2591 let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2592 let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2593 let silu_g = gv / (1.0 + (-gv).exp());
2594 unsafe { *out_addr.at(r) = silu_g * uv };
2596 }
2597 };
2598 dispatch_rows(pool, inter, &run);
2599 true
2600 }
2601 (
2603 Self::Mapped {
2604 dtype: TensorDtype::Q1T,
2605 ..
2606 },
2607 Self::Mapped {
2608 dtype: TensorDtype::Q1T,
2609 ..
2610 },
2611 ) => {
2612 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2613 let g_bytes = gate.quant_bytes();
2614 let u_bytes = up.quant_bytes();
2615 let gpr = gate.cols() / GROUP_SIZE;
2616 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2617 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2618 let run = move |start: usize, end: usize| {
2619 for r in start..end {
2620 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2621 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2622 for &(j, xv) in &act.outliers {
2623 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2624 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2625 }
2626 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2627 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2628 let silu_g = gv / (1.0 + (-gv).exp());
2629 unsafe { *out_addr.at(r) = silu_g * uv };
2631 }
2632 };
2633 dispatch_rows(pool, inter, &run);
2634 true
2635 }
2636 _ => false,
2637 }
2638 }
2639
2640 pub fn moe_gate_up_many(
2654 pairs: &[(&QTensor, &QTensor)],
2655 x: &[f32],
2656 outs: &mut [Vec<f32>],
2657 pool: Option<&Pool>,
2658 ) -> bool {
2659 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2660 return false;
2661 }
2662 let inter = pairs[0].0.rows();
2663 let cols = pairs[0].0.cols();
2664 if cols % GROUP_SIZE != 0 {
2665 return false;
2666 }
2667 let gpr = cols / GROUP_SIZE;
2668 let q2 = matches!(
2671 pairs[0].0,
2672 Self::Mapped {
2673 dtype: TensorDtype::Q2TiledP,
2674 ..
2675 }
2676 );
2677 let want = if q2 {
2678 TensorDtype::Q2TiledP
2679 } else {
2680 TensorDtype::Q4TiledP
2681 };
2682 let mut views = Vec::with_capacity(pairs.len() * 2);
2683 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2684 let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2685 && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2686 if !both
2687 || g.rows() != inter
2688 || u.rows() != inter
2689 || g.cols() != cols
2690 || u.cols() != cols
2691 || o.len() != inter
2692 {
2693 return false;
2694 }
2695 let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2696 views.push(mk(g.quant_bytes(), inter, cols));
2697 views.push(mk(u.quant_bytes(), inter, cols));
2698 }
2699 let act = split_act(x);
2700 let gsum = if q2 {
2701 q1_group_sums(&act.xq, gpr)
2702 } else {
2703 Vec::new()
2704 };
2705 let (act, gsum) = (&act, &gsum);
2706 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2707 let (views, ptrs) = (&views, &ptrs);
2708 let run = |start: usize, end: usize| {
2709 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2710 for flat in start..end {
2711 let (e, r) = (flat / inter, flat % inter);
2712 let gv_view = &views[e * 2];
2713 let uv_view = &views[e * 2 + 1];
2714 gv_view.scales_into(r, gpr, &mut gsc);
2715 uv_view.scales_into(r, gpr, &mut usc);
2716 let (mut gv, mut uv) = if q2 {
2717 (
2718 dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2719 dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2720 )
2721 } else {
2722 (
2723 dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2724 dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2725 )
2726 };
2727 for &(j, xv) in &act.outliers {
2728 let (og, ou) = if q2 {
2729 (
2730 q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2731 q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2732 )
2733 } else {
2734 (
2735 q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2736 q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2737 )
2738 };
2739 gv += og.0 * og.1 * xv;
2740 uv += ou.0 * ou.1 * xv;
2741 }
2742 let silu_g = gv / (1.0 + (-gv).exp());
2743 unsafe { *ptrs[e].at(r) = silu_g * uv };
2745 }
2746 };
2747 dispatch_rows(pool, pairs.len() * inter, &run);
2748 true
2749 }
2750
2751 pub fn moe_down_many(
2760 downs: &[&QTensor],
2761 gs: &[Vec<f32>],
2762 weights: &[f32],
2763 out: &mut [f32],
2764 pool: Option<&Pool>,
2765 ) -> bool {
2766 if downs.is_empty()
2767 || downs.len() != gs.len()
2768 || downs.len() != weights.len()
2769 || !a8w8_enabled()
2770 {
2771 return false;
2772 }
2773 let rows = out.len();
2774 let cols = downs[0].cols();
2775 if cols % GROUP_SIZE != 0 {
2776 return false;
2777 }
2778 let gpr = cols / GROUP_SIZE;
2779 let mut views = Vec::with_capacity(downs.len());
2780 for (d, g) in downs.iter().zip(gs.iter()) {
2781 if !matches!(
2782 d,
2783 Self::Mapped {
2784 dtype: TensorDtype::Q4TiledP,
2785 ..
2786 }
2787 ) || d.rows() != rows
2788 || d.cols() != cols
2789 || g.len() != cols
2790 {
2791 return false;
2792 }
2793 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2794 }
2795 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2797 let out_addr = SendMut(out.as_mut_ptr());
2805 let (views, acts, weights) = (&views, &acts, &weights);
2806 let run = |start: usize, end: usize| {
2807 let mut sc = vec![0f32; gpr];
2808 for r in start..end {
2809 let mut acc = 0f32;
2810 for (e, v) in views.iter().enumerate() {
2811 v.scales_into(r, gpr, &mut sc);
2812 let a = &acts[e];
2813 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2814 for &(j, xv) in &a.outliers {
2815 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2816 d += w * s * xv;
2817 }
2818 acc += weights[e] * d;
2819 }
2820 unsafe { *out_addr.at(r) = acc };
2822 }
2823 };
2824 dispatch_rows(pool, rows, &run);
2825 true
2826 }
2827}
2828
2829#[cfg(target_os = "macos")]
2834mod accel_blas {
2835 #[link(name = "Accelerate", kind = "framework")]
2836 unsafe extern "C" {
2837 pub fn cblas_sgemm(
2838 order: i32,
2839 trans_a: i32,
2840 trans_b: i32,
2841 m: i32,
2842 n: i32,
2843 k: i32,
2844 alpha: f32,
2845 a: *const f32,
2846 lda: i32,
2847 b: *const f32,
2848 ldb: i32,
2849 beta: f32,
2850 c: *mut f32,
2851 ldc: i32,
2852 );
2853 }
2854}
2855
2856#[cfg(target_os = "macos")]
2857pub(crate) fn accel_gemm_enabled() -> bool {
2858 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2859 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2860}
2861
2862#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2865pub(crate) fn accel_gemm_enabled() -> bool {
2866 true
2867}
2868
2869#[cfg(target_arch = "aarch64")]
2876#[allow(clippy::too_many_arguments)]
2877pub(crate) fn neon_gemm_rm(
2878 m: usize,
2879 n: usize,
2880 k: usize,
2881 alpha: f32,
2882 a: &[f32],
2883 lda: usize,
2884 b_mat: &[f32],
2885 ldb: usize,
2886 b_rows_are_n: bool,
2887 c: &mut [f32],
2888 ldc: usize,
2889) {
2890 debug_assert!(a.len() >= (m - 1) * lda + k);
2891 debug_assert!(c.len() >= (m - 1) * ldc + n);
2892 unsafe {
2894 use core::arch::aarch64::*;
2895 let mut i = 0usize;
2896 while i < m {
2897 let mi = (m - i).min(4);
2898 let mut j = 0usize;
2899 while j < n {
2900 let nj = (n - j).min(8);
2901 if mi == 4 && nj == 8 {
2902 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2903 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2904 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2905 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2906 for p in 0..k {
2907 let (b0, b1) = if b_rows_are_n {
2908 let base = b_mat.as_ptr().add(j * ldb + p);
2911 let g = |o: usize| *base.add(o * ldb);
2912 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2913 } else {
2914 let base = b_mat.as_ptr().add(p * ldb + j);
2915 (
2916 [*base, *base.add(1), *base.add(2), *base.add(3)],
2917 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2918 )
2919 };
2920 let bv0 = vld1q_f32(b0.as_ptr());
2921 let bv1 = vld1q_f32(b1.as_ptr());
2922 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2923 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2924 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2925 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2926 c0a = vfmaq_f32(c0a, a0, bv0);
2927 c0b = vfmaq_f32(c0b, a0, bv1);
2928 c1a = vfmaq_f32(c1a, a1, bv0);
2929 c1b = vfmaq_f32(c1b, a1, bv1);
2930 c2a = vfmaq_f32(c2a, a2, bv0);
2931 c2b = vfmaq_f32(c2b, a2, bv1);
2932 c3a = vfmaq_f32(c3a, a3, bv0);
2933 c3b = vfmaq_f32(c3b, a3, bv1);
2934 }
2935 let al = vdupq_n_f32(alpha);
2936 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2937 .iter()
2938 .enumerate()
2939 {
2940 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2941 vst1q_f32(dst, vmulq_f32(*ca, al));
2942 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2943 }
2944 } else {
2945 for r in 0..mi {
2946 for q in 0..nj {
2947 let mut acc = 0f32;
2948 for p in 0..k {
2949 let bv = if b_rows_are_n {
2950 b_mat[(j + q) * ldb + p]
2951 } else {
2952 b_mat[p * ldb + j + q]
2953 };
2954 acc += a[(i + r) * lda + p] * bv;
2955 }
2956 c[(i + r) * ldc + j + q] = acc * alpha;
2957 }
2958 }
2959 }
2960 j += nj;
2961 }
2962 i += mi;
2963 }
2964 }
2965}
2966
2967#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2969#[allow(clippy::too_many_arguments)]
2970pub(crate) fn sgemm_rm(
2971 m: usize,
2972 n: usize,
2973 k: usize,
2974 alpha: f32,
2975 a: &[f32],
2976 lda: usize,
2977 b_mat: &[f32],
2978 ldb: usize,
2979 b_rows_are_n: bool,
2980 c: &mut [f32],
2981 ldc: usize,
2982) {
2983 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2984}
2985
2986#[allow(clippy::too_many_arguments)]
2990pub fn sgemm_public(
2991 m: usize,
2992 n: usize,
2993 k: usize,
2994 alpha: f32,
2995 a: &[f32],
2996 lda: usize,
2997 b_mat: &[f32],
2998 ldb: usize,
2999 b_rows_are_n: bool,
3000 c: &mut [f32],
3001 ldc: usize,
3002) {
3003 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3004 {
3005 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3006 }
3007 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3012 {
3013 for i in 0..m {
3014 for j in 0..n {
3015 let mut acc = 0f32;
3016 for p in 0..k {
3017 let bv = if b_rows_are_n {
3018 b_mat[j * ldb + p]
3019 } else {
3020 b_mat[p * ldb + j]
3021 };
3022 acc += a[i * lda + p] * bv;
3023 }
3024 c[i * ldc + j] = alpha * acc;
3025 }
3026 }
3027 }
3028}
3029
3030#[cfg(target_os = "macos")]
3033#[allow(clippy::too_many_arguments)]
3034pub(crate) fn sgemm_rm(
3035 m: usize,
3036 n: usize,
3037 k: usize,
3038 alpha: f32,
3039 a: &[f32],
3040 lda: usize,
3041 b_mat: &[f32],
3042 ldb: usize,
3043 b_rows_are_n: bool,
3044 c: &mut [f32],
3045 ldc: usize,
3046) {
3047 debug_assert!(a.len() >= (m - 1) * lda + k);
3048 debug_assert!(c.len() >= (m - 1) * ldc + n);
3049 #[cfg(target_arch = "aarch64")]
3054 if std::env::var("CMF_FORCE_NEON_GEMM")
3055 .map(|v| v == "1")
3056 .unwrap_or(false)
3057 {
3058 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3059 }
3060 unsafe {
3061 accel_blas::cblas_sgemm(
3062 101, 111, if b_rows_are_n { 112 } else { 111 },
3065 m as i32,
3066 n as i32,
3067 k as i32,
3068 alpha,
3069 a.as_ptr(),
3070 lda as i32,
3071 b_mat.as_ptr(),
3072 ldb as i32,
3073 0.0,
3074 c.as_mut_ptr(),
3075 ldc as i32,
3076 );
3077 }
3078}
3079
3080#[cfg(target_os = "macos")]
3087fn qmatmat_accel(
3088 q: &[u8],
3089 row_scale: &[f32],
3090 pre: &[std::borrow::Cow<'_, [f32]>],
3091 rows: usize,
3092 cols: usize,
3093 out: &mut [f32],
3094 pool: Option<&Pool>,
3095) {
3096 const TR: usize = 2048;
3101 let b = pre.len();
3102 thread_local! {
3103 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3104 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3105 }
3106 XPANEL.with(|xp| {
3107 WTILE.with(|wt| {
3108 let mut xpanel = xp.borrow_mut();
3109 xpanel.clear();
3110 for x in pre {
3111 xpanel.extend_from_slice(x);
3112 }
3113 let mut wtile = wt.borrow_mut();
3114 wtile.resize(TR * cols, 0.0);
3115 let mut r0 = 0usize;
3116 while r0 < rows {
3117 let tr = TR.min(rows - r0);
3118 let wt_addr = SendMut(wtile.as_mut_ptr());
3120 let run = |start: usize, end: usize| {
3121 for r in start..end {
3122 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3123 let s = row_scale[r0 + r];
3124 let dst =
3126 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3127 for (d, &v) in dst.iter_mut().zip(row) {
3128 *d = (v as i8) as f32 * s;
3129 }
3130 }
3131 };
3132 dispatch_rows(pool, tr, &run);
3133 unsafe {
3135 accel_blas::cblas_sgemm(
3136 101, 111, 112, b as i32,
3140 tr as i32,
3141 cols as i32,
3142 1.0,
3143 xpanel.as_ptr(),
3144 cols as i32,
3145 wtile.as_ptr(),
3146 cols as i32,
3147 0.0,
3148 out.as_mut_ptr().add(r0),
3149 rows as i32,
3150 );
3151 }
3152 r0 += tr;
3153 }
3154 })
3155 });
3156}
3157
3158fn qmatmat(
3159 q: &[u8],
3160 row_scale: &[f32],
3161 pre: &[std::borrow::Cow<'_, [f32]>],
3162 rows: usize,
3163 cols: usize,
3164 out: &mut [f32],
3165 pool: Option<&Pool>,
3166) {
3167 let b = pre.len();
3168 debug_assert_eq!(out.len(), b * rows);
3169 #[cfg(target_os = "macos")]
3174 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3175 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3176 return;
3177 }
3178 #[cfg(target_arch = "aarch64")]
3179 if sdot_enabled() {
3180 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3181 let out_addr = SendMut(out.as_mut_ptr());
3182 let blocked_ok = blocked_enabled();
3185 let use_i8mm = i8mm_enabled();
3186 if blocked_ok {
3187 let run = |start: usize, end: usize| {
3188 let mut o = start;
3189 while o < end {
3190 if o + 2 <= end {
3191 let r0 = &q[o * cols..(o + 1) * cols];
3192 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3193 let mut bi = 0usize;
3194 while bi + 4 <= acts.len() {
3195 let xs = [
3196 acts[bi].xq.as_slice(),
3197 acts[bi + 1].xq.as_slice(),
3198 acts[bi + 2].xq.as_slice(),
3199 acts[bi + 3].xq.as_slice(),
3200 ];
3201 let d = if use_i8mm {
3202 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3203 } else {
3204 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3205 };
3206 for (r, row) in [r0, r1].into_iter().enumerate() {
3207 for k in 0..4 {
3208 let act = &acts[bi + k];
3209 let mut v = d[r][k] as f32 * act.sx;
3210 for &(j, xv) in &act.outliers {
3211 v += (row[j] as i8) as f32 * xv;
3212 }
3213 unsafe {
3214 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3215 };
3216 }
3217 }
3218 bi += 4;
3219 }
3220 while bi < acts.len() {
3221 for (r, row) in [r0, r1].into_iter().enumerate() {
3222 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3223 unsafe { *out_addr.at(bi * rows + o + r) = v };
3224 }
3225 bi += 1;
3226 }
3227 o += 2;
3228 } else {
3229 let row = &q[o * cols..(o + 1) * cols];
3230 for (bi, act) in acts.iter().enumerate() {
3231 let v = row_dot_sdot(row, act) * row_scale[o];
3232 unsafe { *out_addr.at(bi * rows + o) = v };
3233 }
3234 o += 1;
3235 }
3236 }
3237 };
3238 dispatch_rows(pool, rows, &run);
3239 return;
3240 }
3241 let run = |start: usize, end: usize| {
3242 for o in start..end {
3243 let row = &q[o * cols..(o + 1) * cols];
3244 for (bi, act) in acts.iter().enumerate() {
3245 let v = row_dot_sdot(row, act) * row_scale[o];
3246 unsafe { *out_addr.at(bi * rows + o) = v };
3247 }
3248 }
3249 };
3250 dispatch_rows(pool, rows, &run);
3251 return;
3252 }
3253 #[cfg(target_arch = "x86_64")]
3258 if avx2_a8w8_enabled() {
3259 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3260 let out_addr = SendMut(out.as_mut_ptr());
3261 let blocked_ok = blocked_enabled();
3264 if !avx512vnni_enabled() && blocked_ok {
3265 let run = |start: usize, end: usize| {
3266 let mut o = start;
3267 while o < end {
3268 if o + 2 <= end {
3269 let r0 = &q[o * cols..(o + 1) * cols];
3270 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3271 let mut bi = 0usize;
3272 while bi + 4 <= acts.len() {
3273 let xs = [
3274 acts[bi].xq.as_slice(),
3275 acts[bi + 1].xq.as_slice(),
3276 acts[bi + 2].xq.as_slice(),
3277 acts[bi + 3].xq.as_slice(),
3278 ];
3279 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3280 for (r, row) in [r0, r1].into_iter().enumerate() {
3281 for k in 0..4 {
3282 let act = &acts[bi + k];
3283 let mut v = d[r][k] as f32 * act.sx;
3284 for &(j, xv) in &act.outliers {
3285 v += (row[j] as i8) as f32 * xv;
3286 }
3287 unsafe {
3288 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3289 };
3290 }
3291 }
3292 bi += 4;
3293 }
3294 while bi < acts.len() {
3295 for (r, row) in [r0, r1].into_iter().enumerate() {
3296 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3297 unsafe { *out_addr.at(bi * rows + o + r) = v };
3298 }
3299 bi += 1;
3300 }
3301 o += 2;
3302 } else {
3303 let row = &q[o * cols..(o + 1) * cols];
3304 for (bi, act) in acts.iter().enumerate() {
3305 let v = row_dot_avx2(row, act) * row_scale[o];
3306 unsafe { *out_addr.at(bi * rows + o) = v };
3307 }
3308 o += 1;
3309 }
3310 }
3311 };
3312 dispatch_rows(pool, rows, &run);
3313 return;
3314 }
3315 let run = |start: usize, end: usize| {
3316 for o in start..end {
3317 let row = &q[o * cols..(o + 1) * cols];
3318 for (bi, act) in acts.iter().enumerate() {
3319 let v = row_dot_avx2(row, act) * row_scale[o];
3320 unsafe { *out_addr.at(bi * rows + o) = v };
3321 }
3322 }
3323 };
3324 dispatch_rows(pool, rows, &run);
3325 return;
3326 }
3327 let out_addr = SendMut(out.as_mut_ptr());
3328 let run = |start: usize, end: usize| {
3329 for o in start..end {
3330 let row = &q[o * cols..(o + 1) * cols];
3331 for (bi, x) in pre.iter().enumerate() {
3332 let mut acc = 0f32;
3333 for j in 0..cols {
3334 acc += (row[j] as i8) as f32 * x[j];
3335 }
3336 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3337 }
3338 }
3339 };
3340 dispatch_rows(pool, rows, &run);
3341}
3342
3343fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3346 match pool {
3347 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3348 _ => run(0, rows),
3349 }
3350}
3351
3352fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3354 let groups = rows * cols / GROUP_SIZE;
3355 bytes.split_at(groups * 16)
3356}
3357
3358#[inline]
3363fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3364 #[cfg(target_arch = "aarch64")]
3365 unsafe {
3366 return vbit_fill4_neon(data, buf);
3367 }
3368 #[cfg(target_arch = "x86_64")]
3369 if avx2_enabled() {
3370 return unsafe { vbit_fill4_avx2(data, buf) };
3371 }
3372 #[allow(unreachable_code)]
3373 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3374 let u = unpack8::<4>(&data[blk * 4..]);
3375 for k in 0..8 {
3376 chunk[k] = (u[k] - 7) as i8 as u8;
3377 }
3378 }
3379}
3380
3381#[cfg(target_arch = "aarch64")]
3382#[target_feature(enable = "neon")]
3383unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3384 unsafe {
3387 use core::arch::aarch64::*;
3388 let n = buf.len();
3389 let mask = vdupq_n_u8(0x0F);
3390 let seven = vdupq_n_s8(7);
3391 let mut g = 0usize;
3392 while g * 32 + 32 <= n {
3393 let b = vld1q_u8(data.as_ptr().add(g * 16));
3394 let hi = vshrq_n_u8::<4>(b);
3395 let lo = vandq_u8(b, mask);
3396 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3397 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3398 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3399 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3400 g += 1;
3401 }
3402 }
3403}
3404
3405#[cfg(target_arch = "x86_64")]
3406#[target_feature(enable = "avx2")]
3407unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3408 unsafe {
3410 use core::arch::x86_64::*;
3411 let n = buf.len();
3412 let mask = _mm_set1_epi8(0x0F);
3413 let seven = _mm256_set1_epi8(7);
3414 let mut g = 0usize;
3415 while g * 32 + 32 <= n {
3416 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3417 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3418 let lo = _mm_and_si128(b, mask);
3419 let z = _mm256_sub_epi8(
3420 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3421 seven,
3422 );
3423 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3424 g += 1;
3425 }
3426 }
3427}
3428
3429#[inline(always)]
3434fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3435 let mut acc = 0u64;
3436 for i in 0..B {
3437 acc = (acc << 8) | data[i] as u64;
3438 }
3439 let mask = (1u64 << B) - 1;
3440 let mut out = [0i32; 8];
3441 for (k, o) in out.iter_mut().enumerate() {
3442 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3443 }
3444 out
3445}
3446
3447#[allow(clippy::too_many_arguments)]
3453fn vbitmatvec(
3454 bytes: &[u8],
3455 offsets: &[usize],
3456 x: &[f32],
3457 rows: usize,
3458 cols: usize,
3459 out: &mut [f32],
3460 pool: Option<&Pool>,
3461) {
3462 debug_assert_eq!(out.len(), rows);
3463 debug_assert_eq!(offsets.len(), rows + 1);
3464
3465 if a8w8_enabled() {
3469 let act = split_act(x);
3470 let out_addr = SendMut(out.as_mut_ptr());
3471 let run = move |start: usize, end: usize| {
3472 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3473 };
3474 dispatch_rows(pool, rows, &run);
3475 return;
3476 }
3477
3478 let out_addr = SendMut(out.as_mut_ptr());
3479 let run = move |start: usize, end: usize| {
3480 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3481 };
3482 dispatch_rows(pool, rows, &run);
3483}
3484
3485#[allow(clippy::too_many_arguments)]
3489fn vbit_range_a8w8(
3490 bytes: &[u8],
3491 offsets: &[usize],
3492 x: &[f32],
3493 act: &SplitAct,
3494 rows: usize,
3495 cols: usize,
3496 out: SendMut,
3497 start: usize,
3498 end: usize,
3499) {
3500 let ng = cols / GROUP_SIZE;
3501 let bits = &bytes[..rows];
3502 let sc_off = rows;
3503 let row_dot = |r: usize| -> f32 {
3504 let b = bits[r] as usize;
3505 let l = (1i32 << (b - 1)) - 1;
3506 let mask = (1u64 << b) - 1;
3507 let data = &bytes[offsets[r]..offsets[r + 1]];
3508 if b == 8 {
3509 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3511 let mut dot = 0f32;
3512 for g in 0..ng {
3513 let so = (r * ng + g) * 2;
3514 let sgf = f16_to_f32(u16::from_le_bytes([
3515 bytes[sc_off + so],
3516 bytes[sc_off + so + 1],
3517 ]));
3518 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3519 let mut gd = 0f32;
3520 for &xv in xg.iter() {
3521 if nbits < 8 {
3522 acc = (acc << 8) | data[idx] as u64;
3523 idx += 1;
3524 nbits += 8;
3525 }
3526 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3527 nbits -= 8;
3528 gd += (u - l) as f32 * xv;
3529 }
3530 dot += gd * sgf;
3531 }
3532 return dot;
3533 }
3534 thread_local! {
3538 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3539 const { std::cell::RefCell::new(Vec::new()) };
3540 }
3541 #[inline(always)]
3542 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3543 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3544 let u = unpack8::<B>(&data[blk * B..]);
3545 for k in 0..8 {
3546 chunk[k] = (u[k] - l) as i8 as u8;
3547 }
3548 }
3549 }
3550 let _ = mask;
3551 VBIT_SCRATCH.with(|scratch| {
3552 let mut buf = scratch.borrow_mut();
3553 buf.resize(cols, 0);
3554 match b {
3555 3 => fill::<3>(data, l, &mut buf),
3556 4 => vbit_fill4(data, &mut buf),
3557 5 => fill::<5>(data, l, &mut buf),
3558 6 => fill::<6>(data, l, &mut buf),
3559 _ => unreachable!(),
3560 }
3561 let mut dot = 0f32;
3562 for g in 0..ng {
3563 let so = (r * ng + g) * 2;
3564 let s = f16_to_f32(u16::from_le_bytes([
3565 bytes[sc_off + so],
3566 bytes[sc_off + so + 1],
3567 ]));
3568 let d = dot_i8_i8(
3569 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3570 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3571 ) as f32
3572 * act.sx;
3573 dot += d * s;
3574 }
3575 for &(j, xv) in &act.outliers {
3576 let so = (r * ng + j / GROUP_SIZE) * 2;
3577 let s = f16_to_f32(u16::from_le_bytes([
3578 bytes[sc_off + so],
3579 bytes[sc_off + so + 1],
3580 ]));
3581 dot += (buf[j] as i8) as f32 * s * xv;
3583 }
3584 dot
3585 })
3586 };
3587 for r in start..end {
3588 unsafe { *out.at(r) = row_dot(r) };
3590 }
3591}
3592
3593#[allow(clippy::too_many_arguments)]
3595fn vbit_range_f32(
3596 bytes: &[u8],
3597 offsets: &[usize],
3598 x: &[f32],
3599 rows: usize,
3600 cols: usize,
3601 out: SendMut,
3602 start: usize,
3603 end: usize,
3604) {
3605 let ng = cols / GROUP_SIZE;
3606 let bits = &bytes[..rows];
3607 let sc_off = rows;
3608 #[inline(always)]
3612 fn dot_row<const B: usize>(
3613 data: &[u8],
3614 bytes: &[u8],
3615 sc_off: usize,
3616 r: usize,
3617 ng: usize,
3618 x: &[f32],
3619 ) -> f32 {
3620 let l = ((1i32 << (B - 1)) - 1) as f32;
3621 let gbytes = GROUP_SIZE * B / 8;
3622 let mut dot = 0f32;
3623 for g in 0..ng {
3624 let so = (r * ng + g) * 2;
3625 let s = f16_to_f32(u16::from_le_bytes([
3626 bytes[sc_off + so],
3627 bytes[sc_off + so + 1],
3628 ]));
3629 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3630 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3631 let mut gd = 0f32;
3632 for blk in 0..GROUP_SIZE / 8 {
3633 let u = unpack8::<B>(&gd0[blk * B..]);
3634 let xb = &xg[blk * 8..blk * 8 + 8];
3635 for k in 0..8 {
3636 gd += (u[k] as f32 - l) * xb[k];
3637 }
3638 }
3639 dot += gd * s;
3640 }
3641 dot
3642 }
3643 for r in start..end {
3644 let data = &bytes[offsets[r]..offsets[r + 1]];
3645 let v = match bits[r] {
3646 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3647 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3648 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3649 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3650 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3651 b => unreachable!("vbit bit-width {b} (validated at load)"),
3652 };
3653 unsafe { *out.at(r) = v };
3655 }
3656}
3657
3658#[allow(clippy::too_many_arguments)]
3663fn vbitmatvec2(
3664 bytes: &[u8],
3665 offsets: &[usize],
3666 x1: &[f32],
3667 x2: &[f32],
3668 rows: usize,
3669 cols: usize,
3670 o1: &mut [f32],
3671 o2: &mut [f32],
3672 pool: Option<&Pool>,
3673) {
3674 debug_assert_eq!(o1.len(), rows);
3675 debug_assert_eq!(o2.len(), rows);
3676
3677 if a8w8_enabled() {
3678 let a1 = split_act(x1);
3679 let a2 = split_act(x2);
3680 let p1 = SendMut(o1.as_mut_ptr());
3681 let p2 = SendMut(o2.as_mut_ptr());
3682 let run = move |start: usize, end: usize| {
3683 vbit_range2_a8w8(
3684 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3685 )
3686 };
3687 dispatch_rows(pool, rows, &run);
3688 return;
3689 }
3690
3691 let p1 = SendMut(o1.as_mut_ptr());
3692 let p2 = SendMut(o2.as_mut_ptr());
3693 let run = move |start: usize, end: usize| {
3694 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3695 };
3696 dispatch_rows(pool, rows, &run);
3697}
3698
3699#[allow(clippy::too_many_arguments)]
3703fn vbit_range2_a8w8(
3704 bytes: &[u8],
3705 offsets: &[usize],
3706 x1: &[f32],
3707 x2: &[f32],
3708 a1: &SplitAct,
3709 a2: &SplitAct,
3710 rows: usize,
3711 cols: usize,
3712 p1: SendMut,
3713 p2: SendMut,
3714 start: usize,
3715 end: usize,
3716) {
3717 let ng = cols / GROUP_SIZE;
3718 let bits = &bytes[..rows];
3719 let sc_off = rows;
3720 let row_dots = |r: usize| -> (f32, f32) {
3721 let b = bits[r] as usize;
3722 let l = (1i32 << (b - 1)) - 1;
3723 let data = &bytes[offsets[r]..offsets[r + 1]];
3724 if b == 8 {
3725 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3728 let (mut d1, mut d2) = (0f32, 0f32);
3729 for g in 0..ng {
3730 let so = (r * ng + g) * 2;
3731 let sgf = f16_to_f32(u16::from_le_bytes([
3732 bytes[sc_off + so],
3733 bytes[sc_off + so + 1],
3734 ]));
3735 let (mut g1, mut g2) = (0f32, 0f32);
3736 for k in 0..GROUP_SIZE {
3737 if nbits < 8 {
3738 acc = (acc << 8) | data[idx] as u64;
3739 idx += 1;
3740 nbits += 8;
3741 }
3742 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3743 nbits -= 8;
3744 let w = (u - l) as f32;
3745 g1 += w * x1[g * GROUP_SIZE + k];
3746 g2 += w * x2[g * GROUP_SIZE + k];
3747 }
3748 d1 += g1 * sgf;
3749 d2 += g2 * sgf;
3750 }
3751 return (d1, d2);
3752 }
3753 thread_local! {
3754 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3755 const { std::cell::RefCell::new(Vec::new()) };
3756 }
3757 #[inline(always)]
3758 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3759 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3760 let u = unpack8::<B>(&data[blk * B..]);
3761 for k in 0..8 {
3762 chunk[k] = (u[k] - l) as i8 as u8;
3763 }
3764 }
3765 }
3766 VBIT_SCRATCH2.with(|scratch| {
3767 let mut buf = scratch.borrow_mut();
3768 buf.resize(cols, 0);
3769 match b {
3770 3 => fill::<3>(data, l, &mut buf),
3771 4 => vbit_fill4(data, &mut buf),
3772 5 => fill::<5>(data, l, &mut buf),
3773 6 => fill::<6>(data, l, &mut buf),
3774 _ => unreachable!(),
3775 }
3776 let (mut d1, mut d2) = (0f32, 0f32);
3777 for g in 0..ng {
3778 let so = (r * ng + g) * 2;
3779 let s = f16_to_f32(u16::from_le_bytes([
3780 bytes[sc_off + so],
3781 bytes[sc_off + so + 1],
3782 ]));
3783 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3784 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3785 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3786 d1 += v1 * s;
3787 d2 += v2 * s;
3788 }
3789 for &(j, xv) in &a1.outliers {
3790 let so = (r * ng + j / GROUP_SIZE) * 2;
3791 let s = f16_to_f32(u16::from_le_bytes([
3792 bytes[sc_off + so],
3793 bytes[sc_off + so + 1],
3794 ]));
3795 d1 += (buf[j] as i8) as f32 * s * xv;
3796 }
3797 for &(j, xv) in &a2.outliers {
3798 let so = (r * ng + j / GROUP_SIZE) * 2;
3799 let s = f16_to_f32(u16::from_le_bytes([
3800 bytes[sc_off + so],
3801 bytes[sc_off + so + 1],
3802 ]));
3803 d2 += (buf[j] as i8) as f32 * s * xv;
3804 }
3805 (d1, d2)
3806 })
3807 };
3808 for r in start..end {
3809 let (v1, v2) = row_dots(r);
3810 unsafe {
3812 *p1.at(r) = v1;
3813 *p2.at(r) = v2;
3814 }
3815 }
3816}
3817
3818#[allow(clippy::too_many_arguments)]
3822fn vbit_range2_f32(
3823 bytes: &[u8],
3824 offsets: &[usize],
3825 x1: &[f32],
3826 x2: &[f32],
3827 rows: usize,
3828 cols: usize,
3829 p1: SendMut,
3830 p2: SendMut,
3831 start: usize,
3832 end: usize,
3833) {
3834 let ng = cols / GROUP_SIZE;
3835 let bits = &bytes[..rows];
3836 let sc_off = rows;
3837 #[inline(always)]
3838 #[allow(clippy::too_many_arguments)]
3839 fn dot_row2<const B: usize>(
3840 data: &[u8],
3841 bytes: &[u8],
3842 sc_off: usize,
3843 r: usize,
3844 ng: usize,
3845 x1: &[f32],
3846 x2: &[f32],
3847 ) -> (f32, f32) {
3848 let l = ((1i32 << (B - 1)) - 1) as f32;
3849 let gbytes = GROUP_SIZE * B / 8;
3850 let (mut d1, mut d2) = (0f32, 0f32);
3851 for g in 0..ng {
3852 let so = (r * ng + g) * 2;
3853 let s = f16_to_f32(u16::from_le_bytes([
3854 bytes[sc_off + so],
3855 bytes[sc_off + so + 1],
3856 ]));
3857 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3858 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3859 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3860 let (mut g1, mut g2) = (0f32, 0f32);
3861 for blk in 0..GROUP_SIZE / 8 {
3862 let u = unpack8::<B>(&gd0[blk * B..]);
3863 for k in 0..8 {
3864 let w = u[k] as f32 - l;
3865 g1 += w * x1g[blk * 8 + k];
3866 g2 += w * x2g[blk * 8 + k];
3867 }
3868 }
3869 d1 += g1 * s;
3870 d2 += g2 * s;
3871 }
3872 (d1, d2)
3873 }
3874 for r in start..end {
3875 let data = &bytes[offsets[r]..offsets[r + 1]];
3876 let (v1, v2) = match bits[r] {
3877 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3878 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3879 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3880 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3881 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3882 b => unreachable!("vbit bit-width {b} (validated at load)"),
3883 };
3884 unsafe {
3886 *p1.at(r) = v1;
3887 *p2.at(r) = v2;
3888 }
3889 }
3890}
3891
3892#[inline]
3899#[allow(unreachable_code)]
3900fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3901 #[cfg(target_arch = "aarch64")]
3902 unsafe {
3903 return dot_q4t_row_sdot(bytes, r, gpr, xq);
3904 }
3905 #[cfg(target_arch = "x86_64")]
3906 unsafe {
3907 if vnni_tiles_enabled() {
3908 return dot_q4t_row_vnni(bytes, r, gpr, xq);
3909 }
3910 return dot_q4t_row_avx2(bytes, r, gpr, xq);
3911 }
3912 let mut acc = 0f32;
3913 for gi in 0..gpr {
3914 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3915 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3916 let mut d = 0i32;
3917 for (k, &b) in tile[2..].iter().enumerate() {
3918 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3919 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3920 }
3921 acc += d as f32 * s;
3922 }
3923 acc
3924}
3925
3926#[cfg(target_arch = "aarch64")]
3927#[target_feature(enable = "neon,dotprod")]
3928unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3929 unsafe {
3932 use core::arch::aarch64::*;
3933 use core::arch::asm;
3934 let lomask = vdupq_n_u8(0x0F);
3935 let eight = vdupq_n_s8(8);
3936 let mut acc = 0f32;
3937 for gi in 0..gpr {
3938 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3939 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3940 let b = vld1q_u8(t.add(2));
3941 let lo = vandq_u8(b, lomask);
3942 let hi = vshrq_n_u8::<4>(b);
3943 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3944 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3945 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3946 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3947 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3948 asm!(
3949 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3950 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3951 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3952 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3953 options(pure, nomem, nostack),
3954 );
3955 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3956 }
3957 acc
3958 }
3959}
3960
3961#[cfg(target_arch = "x86_64")]
3962#[target_feature(enable = "avx2")]
3963unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3964 unsafe {
3966 use core::arch::x86_64::*;
3967 let lomask = _mm_set1_epi8(0x0F);
3968 let eight = _mm256_set1_epi8(8);
3969 let ones = _mm256_set1_epi16(1);
3970 let mut acc = 0f32;
3971 for gi in 0..gpr {
3972 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3973 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3974 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3975 let lo = _mm_and_si128(b, lomask);
3976 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3977 let w = _mm256_sub_epi8(
3978 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3979 eight,
3980 );
3981 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3982 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3983 let d = _mm256_madd_epi16(p16, ones);
3984 let hi128 = _mm256_extracti128_si256::<1>(d);
3985 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3986 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3987 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3988 acc += _mm_cvtsi128_si32(s32) as f32 * s;
3989 }
3990 acc
3991 }
3992}
3993
3994#[cfg(target_arch = "x86_64")]
3998#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3999unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4000 unsafe {
4002 use core::arch::x86_64::*;
4003 let lomask = _mm_set1_epi8(0x0F);
4004 let eight = _mm256_set1_epi8(8);
4005 let mut acc = 0f32;
4006 for gi in 0..gpr {
4007 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4008 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4009 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4010 let lo = _mm_and_si128(b, lomask);
4011 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4012 let w = _mm256_sub_epi8(
4013 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4014 eight,
4015 );
4016 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4017 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4018 acc += d as f32 * s;
4019 }
4020 acc
4021 }
4022}
4023
4024#[cfg(target_arch = "x86_64")]
4029#[target_feature(enable = "avx2,fma")]
4034unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4035 unsafe {
4037 use core::arch::x86_64::*;
4038 let lomask = _mm_set1_epi8(0x0F);
4039 let eight = _mm256_set1_epi8(8);
4040 let ones = _mm256_set1_epi16(1);
4041 let mut f0 = _mm256_setzero_ps();
4054 let mut f1 = _mm256_setzero_ps();
4055 let mut f2 = _mm256_setzero_ps();
4056 let mut f3 = _mm256_setzero_ps();
4057 for gi in 0..gpr {
4058 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4059 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4060 let sv = _mm256_set1_ps(s);
4061 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4062 let lo = _mm_and_si128(bb, lomask);
4063 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4064 let w = _mm256_sub_epi8(
4065 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4066 eight,
4067 );
4068 let aw = _mm256_abs_epi8(w);
4069 let off = gi * GROUP_SIZE;
4070 let dot = |xq: &[i8]| {
4071 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4072 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4073 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4074 };
4075 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4076 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4077 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4078 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4079 }
4080 [
4081 hsum256_ps(f0),
4082 hsum256_ps(f1),
4083 hsum256_ps(f2),
4084 hsum256_ps(f3),
4085 ]
4086 }
4087}
4088
4089#[cfg(target_arch = "x86_64")]
4092#[target_feature(enable = "avx2")]
4093#[inline]
4094unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4095 unsafe {
4097 use core::arch::x86_64::*;
4098 let hi = _mm256_extractf128_ps::<1>(v);
4099 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4100 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4101 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4102 _mm_cvtss_f32(s)
4103 }
4104}
4105
4106#[cfg(target_arch = "x86_64")]
4108#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4109unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4110 unsafe {
4112 use core::arch::x86_64::*;
4113 let lomask = _mm_set1_epi8(0x0F);
4114 let eight = _mm256_set1_epi8(8);
4115 let mut f0 = _mm256_setzero_ps();
4118 let mut f1 = _mm256_setzero_ps();
4119 let mut f2 = _mm256_setzero_ps();
4120 let mut f3 = _mm256_setzero_ps();
4121 for gi in 0..gpr {
4122 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4123 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4124 let sv = _mm256_set1_ps(s);
4125 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4126 let lo = _mm_and_si128(bb, lomask);
4127 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4128 let w = _mm256_sub_epi8(
4129 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4130 eight,
4131 );
4132 let aw = _mm256_abs_epi8(w);
4133 let off = gi * GROUP_SIZE;
4134 let dot = |xq: &[i8]| {
4135 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4136 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4137 _mm256_setzero_si256(),
4138 aw,
4139 _mm256_sign_epi8(x, w),
4140 ))
4141 };
4142 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4143 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4144 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4145 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4146 }
4147 let acc = [
4148 hsum256_ps(f0),
4149 hsum256_ps(f1),
4150 hsum256_ps(f2),
4151 hsum256_ps(f3),
4152 ];
4153 acc
4154 }
4155}
4156
4157#[cfg(target_arch = "aarch64")]
4162#[target_feature(enable = "neon,dotprod")]
4163unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4164 unsafe {
4166 use core::arch::aarch64::*;
4167 use core::arch::asm;
4168 let lomask = vdupq_n_u8(0x0F);
4169 let eight = vdupq_n_s8(8);
4170 let mut acc = [0f32; 4];
4171 for gi in 0..gpr {
4172 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4173 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4174 let b = vld1q_u8(t.add(2));
4175 let lo = vandq_u8(b, lomask);
4176 let hi = vshrq_n_u8::<4>(b);
4177 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4178 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4179 for (k, xq) in xs.iter().enumerate() {
4180 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4181 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4182 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4183 asm!(
4184 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4185 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4186 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4187 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4188 options(pure, nomem, nostack),
4189 );
4190 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4191 }
4192 }
4193 acc
4194 }
4195}
4196
4197#[inline]
4199fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4200 let gi = j / GROUP_SIZE;
4201 let k = j % GROUP_SIZE;
4202 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4203 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4204 let byte = tile[2 + k / 2];
4205 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4206 ((nib as i32 - 8) as f32, s)
4207}
4208
4209#[inline]
4212fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4213 let mut acc = 0f32;
4214 for gi in 0..gpr {
4215 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4216 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4217 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4218 let mut ga = 0f32;
4219 for (k, &b) in tile[2..].iter().enumerate() {
4220 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4221 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4222 }
4223 acc += ga * s;
4224 }
4225 acc
4226}
4227
4228struct Q4tpView<'a> {
4232 nib: &'a [u8],
4233 params: &'a [u8],
4234 codes: &'a [u8],
4235 stride: usize,
4236 zero_rung: bool,
4238}
4239
4240impl<'a> Q4tpView<'a> {
4241 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4242 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4243 Self {
4244 nib: &bytes[..params_off],
4245 params: &bytes[params_off..codes_off],
4246 codes: &bytes[codes_off..],
4247 stride,
4248 zero_rung: false,
4249 }
4250 }
4251
4252 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4254 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4255 Self {
4256 nib: &bytes[..params_off],
4257 params: &bytes[params_off..codes_off],
4258 codes: &bytes[codes_off..],
4259 stride,
4260 zero_rung: true,
4261 }
4262 }
4263
4264 #[inline]
4280 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4281 let tab = if self.zero_rung {
4282 q2tp_ladder(self.params, r)
4283 } else {
4284 q4tp_ladder(self.params, r)
4285 };
4286 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4287 let out = &mut out[..gpr];
4288 let mut chunks = out.chunks_exact_mut(8);
4289 let mut ci = 0usize;
4290 for c in &mut chunks {
4291 let w = u64::from(codes[ci])
4292 | u64::from(codes[ci + 1]) << 8
4293 | u64::from(codes[ci + 2]) << 16
4294 | u64::from(codes[ci + 3]) << 24
4295 | u64::from(codes[ci + 4]) << 32;
4296 for (k, o) in c.iter_mut().enumerate() {
4297 *o = tab[((w >> (5 * k)) & 31) as usize];
4298 }
4299 ci += 5;
4300 }
4301 let tail = &codes[ci..];
4304 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4305 *o = tab[q4tp_code(tail, k)];
4306 }
4307 }
4308}
4309
4310#[inline]
4311fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4312 #[cfg(target_arch = "aarch64")]
4313 unsafe {
4314 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4315 }
4316 #[cfg(target_arch = "x86_64")]
4317 unsafe {
4318 if vnni_tiles_enabled() {
4319 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4320 }
4321 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4322 }
4323 #[allow(unreachable_code)]
4324 {
4325 let mut acc = 0f32;
4326 for gi in 0..gpr {
4327 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4328 let s = scales[gi];
4329 let mut d = 0i32;
4330 for (k, &b) in tile.iter().enumerate() {
4331 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4332 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4333 }
4334 acc += d as f32 * s;
4335 }
4336 acc
4337 }
4338}
4339
4340#[cfg(target_arch = "aarch64")]
4343#[target_feature(enable = "neon,dotprod")]
4344unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4345 unsafe {
4348 use core::arch::aarch64::*;
4349 use core::arch::asm;
4350 let lomask = vdupq_n_u8(0x0F);
4351 let eight = vdupq_n_s8(8);
4352 let mut acc = 0f32;
4353 for gi in 0..gpr {
4354 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4355 let s = *scales.get_unchecked(gi);
4356 let b = vld1q_u8(t);
4357 let lo = vandq_u8(b, lomask);
4358 let hi = vshrq_n_u8::<4>(b);
4359 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4360 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4361 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4362 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4363 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4364 asm!(
4365 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4366 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4367 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4368 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4369 options(pure, nomem, nostack),
4370 );
4371 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4372 }
4373 acc
4374 }
4375}
4376
4377#[cfg(target_arch = "x86_64")]
4378#[target_feature(enable = "avx2")]
4379unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4380 unsafe {
4382 use core::arch::x86_64::*;
4383 let lomask = _mm_set1_epi8(0x0F);
4384 let eight = _mm256_set1_epi8(8);
4385 let ones = _mm256_set1_epi16(1);
4386 let mut acc = 0f32;
4387 for gi in 0..gpr {
4388 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4389 let s = *scales.get_unchecked(gi);
4390 let b = _mm_loadu_si128(t as *const __m128i);
4391 let lo = _mm_and_si128(b, lomask);
4392 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4393 let w = _mm256_sub_epi8(
4394 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4395 eight,
4396 );
4397 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4398 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4399 let d = _mm256_madd_epi16(p16, ones);
4400 let hi128 = _mm256_extracti128_si256::<1>(d);
4401 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4402 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4403 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4404 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4405 }
4406 acc
4407 }
4408}
4409
4410
4411#[cfg(target_arch = "x86_64")]
4414#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4415unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4416 unsafe {
4418 use core::arch::x86_64::*;
4419 let lomask = _mm_set1_epi8(0x0F);
4420 let eight = _mm256_set1_epi8(8);
4421 let mut acc = 0f32;
4422 for gi in 0..gpr {
4423 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4424 let s = *scales.get_unchecked(gi);
4425 let b = _mm_loadu_si128(t as *const __m128i);
4426 let lo = _mm_and_si128(b, lomask);
4427 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4428 let w = _mm256_sub_epi8(
4429 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4430 eight,
4431 );
4432 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4433 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4434 }
4435 acc
4436 }
4437}
4438
4439#[inline]
4442fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4443 let mut acc = 0f32;
4444 for gi in 0..gpr {
4445 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4446 let s = scales[gi];
4447 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4448 let mut ga = 0f32;
4449 for (k, &b) in tile.iter().enumerate() {
4450 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4451 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4452 }
4453 acc += ga * s;
4454 }
4455 acc
4456}
4457
4458#[inline]
4461fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4462 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4463 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4464 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4465 ((n as i32 - 8) as f32, scales[gi])
4466}
4467
4468fn q4tp_matvec(
4470 bytes: &[u8],
4471 x: &[f32],
4472 rows: usize,
4473 cols: usize,
4474 out: &mut [f32],
4475 pool: Option<&Pool>,
4476) {
4477 debug_assert_eq!(out.len(), rows);
4478 let gpr = cols / GROUP_SIZE;
4479 let v = Q4tpView::new(bytes, rows, cols);
4480 let out_addr = SendMut(out.as_mut_ptr());
4481 if a8w8_enabled() {
4482 let act = split_act(x);
4483 let run = |start: usize, end: usize| {
4484 with_krow(gpr, |sc| {
4486 for r in start..end {
4487 v.scales_into(r, gpr, sc);
4488 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4489 for &(j, xv) in &act.outliers {
4490 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4491 acc += w * s * xv;
4492 }
4493 unsafe { *out_addr.at(r) = acc };
4495 }
4496 })
4497 };
4498 dispatch_rows(pool, rows, &run);
4499 return;
4500 }
4501 let run = |start: usize, end: usize| {
4502 with_krow(gpr, |sc| {
4503 for r in start..end {
4504 v.scales_into(r, gpr, sc);
4505 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4507 }
4508 })
4509 };
4510 dispatch_rows(pool, rows, &run);
4511}
4512
4513#[allow(clippy::too_many_arguments)]
4516fn q4tp_matvec2(
4517 bytes: &[u8],
4518 x1: &[f32],
4519 x2: &[f32],
4520 rows: usize,
4521 cols: usize,
4522 o1: &mut [f32],
4523 o2: &mut [f32],
4524 pool: Option<&Pool>,
4525) {
4526 let gpr = cols / GROUP_SIZE;
4527 let v = Q4tpView::new(bytes, rows, cols);
4528 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4529 let run = |start: usize, end: usize| {
4530 let mut sc = vec![0f32; gpr];
4531 for r in start..end {
4532 v.scales_into(r, gpr, &mut sc);
4533 unsafe {
4535 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4536 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4537 }
4538 }
4539 };
4540 dispatch_rows(pool, rows, &run);
4541}
4542
4543#[inline]
4546fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4547 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4548 let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4549 let c = (byte >> (2 * (k % 4))) & 3;
4550 (c as f32 - 1.5, scales[gi])
4551}
4552
4553#[inline]
4560fn dot_q2tp_row_i8(
4561 chunks: &[u8],
4562 r: usize,
4563 gpr: usize,
4564 xq: &[i8],
4565 gsum: &[i32],
4566 scales: &[f32],
4567) -> f32 {
4568 let mut acc = 0f32;
4569 let base = r * gpr * Q2TP_CHUNK;
4570 #[cfg(not(target_arch = "aarch64"))]
4571 let mut codes = [0i8; GROUP_SIZE];
4572 for gi in 0..gpr {
4573 let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4574 let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4575 #[cfg(target_arch = "aarch64")]
4576 let dot = unsafe {
4582 use core::arch::aarch64::*;
4583 let b = vld1_u8(ch.as_ptr());
4584 let three = vdup_n_u8(3);
4585 let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4586 let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4587 let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4588 let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4589 let x4 = vld4_s8(xg.as_ptr());
4590 let mut acc4 = vdupq_n_s32(0);
4591 acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4592 acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4593 acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4594 acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4595 vaddvq_s32(acc4)
4596 };
4597 #[cfg(not(target_arch = "aarch64"))]
4598 let dot: i32 = {
4599 for (k, &b) in ch.iter().enumerate() {
4600 codes[k * 4] = (b & 3) as i8;
4601 codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4602 codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4603 codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4604 }
4605 codes
4606 .iter()
4607 .zip(xg)
4608 .map(|(&c, &x)| c as i32 * x as i32)
4609 .sum()
4610 };
4611 acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4612 }
4613 acc
4614}
4615
4616fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4620 let mut acc = 0f32;
4621 for gi in 0..gpr {
4622 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4623 let s = scales[gi];
4624 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4625 let mut g = 0f32;
4626 for (k, &b) in ch.iter().enumerate() {
4627 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4628 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4629 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4630 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4631 }
4632 acc += s * g;
4633 }
4634 acc
4635}
4636
4637fn q2tp_matvec(
4638 bytes: &[u8],
4639 x: &[f32],
4640 rows: usize,
4641 cols: usize,
4642 out: &mut [f32],
4643 pool: Option<&Pool>,
4644) {
4645 debug_assert_eq!(out.len(), rows);
4646 let gpr = cols / GROUP_SIZE;
4647 let v = Q4tpView::new_q2(bytes, rows, cols);
4648 let out_addr = SendMut(out.as_mut_ptr());
4649 if a8w8_enabled() {
4654 let act = split_act(x);
4655 let gsum = q1_group_sums(&act.xq, gpr);
4656 let (act, gsum) = (&act, &gsum);
4657 let run = move |start: usize, end: usize| {
4658 with_krow(gpr, |sc| {
4659 for r in start..end {
4660 v.scales_into(r, gpr, sc);
4661 let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4662 for &(j, xv) in &act.outliers {
4663 let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4664 acc += w * s * xv;
4665 }
4666 unsafe { *out_addr.at(r) = acc };
4668 }
4669 })
4670 };
4671 dispatch_rows(pool, rows, &run);
4672 return;
4673 }
4674 let run = |start: usize, end: usize| {
4675 with_krow(gpr, |sc| {
4676 for r in start..end {
4677 v.scales_into(r, gpr, sc);
4678 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4680 }
4681 })
4682 };
4683 dispatch_rows(pool, rows, &run);
4684}
4685
4686#[allow(clippy::too_many_arguments)]
4688fn q2tp_matvec2(
4689 bytes: &[u8],
4690 x1: &[f32],
4691 x2: &[f32],
4692 rows: usize,
4693 cols: usize,
4694 o1: &mut [f32],
4695 o2: &mut [f32],
4696 pool: Option<&Pool>,
4697) {
4698 let gpr = cols / GROUP_SIZE;
4699 let v = Q4tpView::new_q2(bytes, rows, cols);
4700 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4701 let run = |start: usize, end: usize| {
4702 let mut sc = vec![0f32; gpr];
4703 for r in start..end {
4704 v.scales_into(r, gpr, &mut sc);
4705 unsafe {
4707 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4708 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4709 }
4710 }
4711 };
4712 dispatch_rows(pool, rows, &run);
4713}
4714
4715pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4722 let gpr = cols / GROUP_SIZE;
4727 let v = Q4tpView::new_q2(bytes, rows, cols);
4728 with_krow(gpr, |sc| {
4729 for r in 0..rows {
4730 v.scales_into(r, gpr, sc);
4731 out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4732 }
4733 });
4734}
4735
4736pub fn q2tp_matmat_for_test(
4737 bytes: &[u8],
4738 xs_all: &[f32],
4739 b: usize,
4740 rows: usize,
4741 cols: usize,
4742 out: &mut [f32],
4743) {
4744 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4745}
4746
4747fn q2tp_matmat(
4748 bytes: &[u8],
4749 xs_all: &[f32],
4750 b: usize,
4751 rows: usize,
4752 cols: usize,
4753 out: &mut [f32],
4754 pool: Option<&Pool>,
4755) {
4756 debug_assert_eq!(out.len(), b * rows);
4757 let gpr = cols / GROUP_SIZE;
4758 let v = Q4tpView::new_q2(bytes, rows, cols);
4759 let out_addr = SendMut(out.as_mut_ptr());
4760 let run = |start: usize, end: usize| {
4761 let mut sc = vec![0f32; gpr];
4762 for r in start..end {
4763 v.scales_into(r, gpr, &mut sc);
4764 for bi in 0..b {
4765 let x = &xs_all[bi * cols..(bi + 1) * cols];
4766 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4768 }
4769 }
4770 };
4771 dispatch_rows(pool, rows, &run);
4772}
4773
4774#[cfg(target_arch = "aarch64")]
4778#[target_feature(enable = "neon,dotprod")]
4779unsafe fn dot_q4tp_row_1x4_sdot_v1(
4780 nib: &[u8],
4781 r: usize,
4782 gpr: usize,
4783 xs: [&[i8]; 4],
4784 scales: &[f32],
4785) -> [f32; 4] {
4786 unsafe {
4787 use core::arch::aarch64::*;
4788 use core::arch::asm;
4789 let lomask = vdupq_n_u8(0x0F);
4790 let eight = vdupq_n_s8(8);
4791 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4792 for gi in 0..gpr {
4793 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4794 let s = *scales.get_unchecked(gi);
4795 let bb = vld1q_u8(t);
4796 let lo = vandq_u8(bb, lomask);
4797 let hi = vshrq_n_u8::<4>(bb);
4798 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4799 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4800 let mut d = [0f32; 4];
4801 for (k, dk) in d.iter_mut().enumerate() {
4802 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4803 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4804 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4805 asm!(
4806 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4807 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4808 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4809 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4810 options(pure, nomem, nostack),
4811 );
4812 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4813 }
4814 f0 += d[0];
4815 f1 += d[1];
4816 f2 += d[2];
4817 f3 += d[3];
4818 }
4819 [f0, f1, f2, f3]
4820 }
4821}
4822
4823#[allow(dead_code)]
4830static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4831
4832#[cfg(target_arch = "x86_64")]
4838fn q4tp_blocked_x86() -> bool {
4839 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4840 1 => false,
4841 2 => avx512vnni_enabled(),
4846 _ => blocked_enabled() && avx512vnni_enabled(),
4850 }
4851}
4852
4853#[cfg(target_arch = "aarch64")]
4855#[allow(dead_code)]
4856fn q4tp_v1() -> bool {
4857 match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4858 1 => true,
4859 2 => false,
4860 _ => {
4861 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4862 *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4863 }
4864 }
4865}
4866
4867#[cfg(target_arch = "x86_64")]
4879#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4880unsafe fn dot_q4tp_2x8_avx512(
4881 nib: &[u8],
4882 r0: usize,
4883 gpr: usize,
4884 xs: [&[i8]; 8],
4885 sc0: &[f32],
4886 sc1: &[f32],
4887) -> [[f32; 8]; 2] {
4888 unsafe {
4891 use core::arch::x86_64::*;
4892 let lomask = _mm256_set1_epi8(0x0F);
4893 let eight = _mm256_set1_epi8(8);
4894 let zero = _mm512_setzero_si512();
4895 let mut v0 = [_mm512_setzero_ps(); 8];
4896 let mut v1 = [_mm512_setzero_ps(); 8];
4897 let pairs = gpr / 2;
4898 let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4899 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4900 let bb = _mm256_loadu_si256(t as *const __m256i);
4901 let lo = _mm256_and_si256(bb, lomask);
4902 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4903 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4904 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4905 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4906 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4907 (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
4908 };
4909 for gp in 0..pairs {
4910 let gi = gp * 2;
4911 let (wa0, neg0) = unpack(r0, gi);
4912 let (wa1, neg1) = unpack(r0 + 1, gi);
4913 let off = gi * GROUP_SIZE;
4914 let sv = |sc: &[f32]| {
4915 _mm512_insertf32x8::<1>(
4916 _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
4917 _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
4918 )
4919 };
4920 let s0 = sv(sc0);
4921 let s1 = sv(sc1);
4922 for k in 0..8 {
4923 let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
4924 let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4925 zero,
4926 wa0,
4927 _mm512_mask_sub_epi8(xv, neg0, zero, xv),
4928 ));
4929 let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4930 zero,
4931 wa1,
4932 _mm512_mask_sub_epi8(xv, neg1, zero, xv),
4933 ));
4934 v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
4935 v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
4936 }
4937 }
4938 let mut acc = [[0f32; 8]; 2];
4939 for k in 0..8 {
4940 acc[0][k] = _mm512_reduce_add_ps(v0[k]);
4941 acc[1][k] = _mm512_reduce_add_ps(v1[k]);
4942 }
4943 if gpr % 2 == 1 {
4944 let off = (gpr - 1) * GROUP_SIZE;
4945 for j in off..off + GROUP_SIZE {
4946 let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
4947 let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
4948 for k in 0..8 {
4949 let x = *xs[k].get_unchecked(j) as f32;
4950 acc[0][k] += w0 * sa * x;
4951 acc[1][k] += w1 * sb * x;
4952 }
4953 }
4954 }
4955 acc
4956 }
4957}
4958
4959#[cfg(target_arch = "x86_64")]
4964#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4965unsafe fn dot_q4tp_row_1x8_avx512(
4966 nib: &[u8],
4967 r: usize,
4968 gpr: usize,
4969 xs: [&[i8]; 8],
4970 scales: &[f32],
4971) -> [f32; 8] {
4972 unsafe {
4974 use core::arch::x86_64::*;
4975 let lomask = _mm256_set1_epi8(0x0F);
4976 let eight = _mm256_set1_epi8(8);
4977 let zero = _mm512_setzero_si512();
4978 let (mut v0, mut v1, mut v2, mut v3) = (
4979 _mm512_setzero_ps(),
4980 _mm512_setzero_ps(),
4981 _mm512_setzero_ps(),
4982 _mm512_setzero_ps(),
4983 );
4984 let (mut v4, mut v5, mut v6, mut v7) = (
4985 _mm512_setzero_ps(),
4986 _mm512_setzero_ps(),
4987 _mm512_setzero_ps(),
4988 _mm512_setzero_ps(),
4989 );
4990 let pairs = gpr / 2;
4991 for gp in 0..pairs {
4992 let gi = gp * 2;
4993 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4994 let bb = _mm256_loadu_si256(t as *const __m256i);
4995 let lo = _mm256_and_si256(bb, lomask);
4996 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4997 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5002 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5003 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5004 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5005 let wabs = _mm512_abs_epi8(w);
5006 let neg = _mm512_movepi8_mask(w);
5007 let off = gi * GROUP_SIZE;
5008 let sv = _mm512_insertf32x8::<1>(
5009 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5010 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5011 );
5012 let dot = |x: &[i8]| -> __m512 {
5013 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5014 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5015 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5016 };
5017 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5018 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5019 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5020 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5021 v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5022 v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5023 v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5024 v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5025 }
5026 let mut acc = [
5027 _mm512_reduce_add_ps(v0),
5028 _mm512_reduce_add_ps(v1),
5029 _mm512_reduce_add_ps(v2),
5030 _mm512_reduce_add_ps(v3),
5031 _mm512_reduce_add_ps(v4),
5032 _mm512_reduce_add_ps(v5),
5033 _mm512_reduce_add_ps(v6),
5034 _mm512_reduce_add_ps(v7),
5035 ];
5036 if gpr % 2 == 1 {
5039 let off = (gpr - 1) * GROUP_SIZE;
5040 for j in off..off + GROUP_SIZE {
5041 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5042 let ws = w * s;
5043 for k in 0..8 {
5044 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5045 }
5046 }
5047 }
5048 acc
5049 }
5050}
5051
5052#[cfg(target_arch = "x86_64")]
5065#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5066unsafe fn dot_q4tp_row_1x4_avx512(
5067 nib: &[u8],
5068 r: usize,
5069 gpr: usize,
5070 xs: [&[i8]; 4],
5071 scales: &[f32],
5072) -> [f32; 4] {
5073 unsafe {
5075 use core::arch::x86_64::*;
5076 let lomask = _mm256_set1_epi8(0x0F);
5077 let eight = _mm256_set1_epi8(8);
5078 let zero = _mm512_setzero_si512();
5079 let (mut v0, mut v1, mut v2, mut v3) = (
5080 _mm512_setzero_ps(),
5081 _mm512_setzero_ps(),
5082 _mm512_setzero_ps(),
5083 _mm512_setzero_ps(),
5084 );
5085 let pairs = gpr / 2;
5086 for gp in 0..pairs {
5087 let gi = gp * 2;
5088 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5089 let bb = _mm256_loadu_si256(t as *const __m256i);
5090 let lo = _mm256_and_si256(bb, lomask);
5091 let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5092 let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5097 let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5098 let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5099 let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5100 let wabs = _mm512_abs_epi8(w);
5101 let neg = _mm512_movepi8_mask(w);
5102 let off = gi * GROUP_SIZE;
5103 let sv = _mm512_insertf32x8::<1>(
5104 _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5105 _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5106 );
5107 let dot = |x: &[i8]| -> __m512 {
5108 let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5109 let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5110 _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5111 };
5112 v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5113 v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5114 v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5115 v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5116 }
5117 let mut acc = [
5118 _mm512_reduce_add_ps(v0),
5119 _mm512_reduce_add_ps(v1),
5120 _mm512_reduce_add_ps(v2),
5121 _mm512_reduce_add_ps(v3),
5122 ];
5123 if gpr % 2 == 1 {
5126 let off = (gpr - 1) * GROUP_SIZE;
5127 for j in off..off + GROUP_SIZE {
5128 let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5129 let ws = w * s;
5130 for k in 0..4 {
5131 acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5132 }
5133 }
5134 }
5135 acc
5136 }
5137}
5138
5139#[cfg(target_arch = "aarch64")]
5143#[target_feature(enable = "neon,dotprod")]
5144unsafe fn dot_q4tp_row_1x4_sdot(
5145 nib: &[u8],
5146 r: usize,
5147 gpr: usize,
5148 xs: [&[i8]; 4],
5149 scales: &[f32],
5150) -> [f32; 4] {
5151 unsafe {
5153 use core::arch::aarch64::*;
5154 use core::arch::asm;
5155 let lomask = vdupq_n_u8(0x0F);
5156 let eight = vdupq_n_s8(8);
5157 let (mut v0, mut v1, mut v2, mut v3) = (
5173 vdupq_n_f32(0.0),
5174 vdupq_n_f32(0.0),
5175 vdupq_n_f32(0.0),
5176 vdupq_n_f32(0.0),
5177 );
5178 for gi in 0..gpr {
5179 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5180 let s = *scales.get_unchecked(gi);
5181 let bb = vld1q_u8(t);
5182 let lo = vandq_u8(bb, lomask);
5183 let hi = vshrq_n_u8::<4>(bb);
5184 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5185 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5186 let off = gi * GROUP_SIZE;
5187 let dot4 = |x: &[i8]| -> int32x4_t {
5188 let x0 = vld1q_s8(x.as_ptr().add(off));
5189 let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5190 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5191 asm!(
5192 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5193 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5194 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5195 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5196 options(pure, nomem, nostack),
5197 );
5198 vaddq_s32(a0, a1)
5199 };
5200 v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5201 v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5202 v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5203 v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5204 }
5205 [
5206 vaddvq_f32(v0),
5207 vaddvq_f32(v1),
5208 vaddvq_f32(v2),
5209 vaddvq_f32(v3),
5210 ]
5211 }
5212}
5213
5214fn q4tp_matmat(
5218 bytes: &[u8],
5219 xs_all: &[f32],
5220 b: usize,
5221 rows: usize,
5222 cols: usize,
5223 out: &mut [f32],
5224 pool: Option<&Pool>,
5225) {
5226 debug_assert_eq!(out.len(), b * rows);
5227 let gpr = cols / GROUP_SIZE;
5228 let v = Q4tpView::new(bytes, rows, cols);
5229
5230 #[cfg(target_os = "macos")]
5232 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5233 dequant_matmat_accel(
5234 &|r, dst| {
5235 let mut sc = [0f32; 32];
5236 let mut scv;
5237 let s: &[f32] = if gpr <= 32 {
5238 v.scales_into(r, gpr, &mut sc);
5239 &sc[..gpr]
5240 } else {
5241 scv = vec![0f32; gpr];
5242 v.scales_into(r, gpr, &mut scv);
5243 &scv
5244 };
5245 for gi in 0..gpr {
5246 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5247 for (k, &bb) in tile.iter().enumerate() {
5248 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5249 dst[gi * GROUP_SIZE + k * 2 + 1] =
5250 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5251 }
5252 }
5253 },
5254 xs_all,
5255 b,
5256 rows,
5257 cols,
5258 out,
5259 pool,
5260 );
5261 return;
5262 }
5263
5264 let out_addr = SendMut(out.as_mut_ptr());
5265 if a8w8_enabled() {
5266 let acts: Vec<SplitAct> = (0..b)
5267 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5268 .collect();
5269 let acts = &acts;
5270 #[cfg(target_arch = "aarch64")]
5271 let blocked_ok = sdot_enabled() && blocked_enabled();
5272 #[cfg(target_arch = "x86_64")]
5279 let blocked_ok = q4tp_blocked_x86();
5280 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5281 let blocked_ok = false;
5282 let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5292 .ok()
5293 .and_then(|v| v.parse().ok())
5294 .filter(|v| *v > 0)
5295 .unwrap_or(256);
5296 let run = |start: usize, end: usize| {
5297 for abase in (0..acts.len()).step_by(panel_cols) {
5298 let alen = (acts.len() - abase).min(panel_cols);
5299 let mut sc = vec![0f32; gpr];
5300 #[cfg(target_arch = "x86_64")]
5301 let mut r_lo = start;
5302 #[cfg(target_arch = "x86_64")]
5303 if blocked_ok && alen >= 8 {
5304 let mut sc1 = vec![0f32; gpr];
5305 while r_lo + 2 <= end {
5306 v.scales_into(r_lo, gpr, &mut sc);
5307 v.scales_into(r_lo + 1, gpr, &mut sc1);
5308 let mut bi = 0usize;
5309 while bi + 8 <= alen {
5310 let xs = [
5311 acts[abase + bi].xq.as_slice(),
5312 acts[abase + bi + 1].xq.as_slice(),
5313 acts[abase + bi + 2].xq.as_slice(),
5314 acts[abase + bi + 3].xq.as_slice(),
5315 acts[abase + bi + 4].xq.as_slice(),
5316 acts[abase + bi + 5].xq.as_slice(),
5317 acts[abase + bi + 6].xq.as_slice(),
5318 acts[abase + bi + 7].xq.as_slice(),
5319 ];
5320 let d =
5321 unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5322 for (row, dr, scr) in
5323 [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)]
5324 {
5325 for k in 0..8 {
5326 let act = &acts[abase + bi + k];
5327 let mut acc = dr[k] * act.sx;
5328 for &(j, xv) in &act.outliers {
5329 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5330 acc += w * s * xv;
5331 }
5332 unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5334 }
5335 }
5336 bi += 8;
5337 }
5338 for row in [r_lo, r_lo + 1] {
5341 let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5342 for b2 in bi..alen {
5343 let act = &acts[abase + b2];
5344 let xs4 = [
5345 act.xq.as_slice(),
5346 act.xq.as_slice(),
5347 act.xq.as_slice(),
5348 act.xq.as_slice(),
5349 ];
5350 let d =
5351 unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5352 let mut acc = d[0] * act.sx;
5353 for &(j, xv) in &act.outliers {
5354 let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5355 acc += w * s * xv;
5356 }
5357 unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5359 }
5360 }
5361 r_lo += 2;
5362 }
5363 }
5364 #[cfg(target_arch = "x86_64")]
5365 let row_start = r_lo;
5366 #[cfg(not(target_arch = "x86_64"))]
5367 let row_start = start;
5368 for r in row_start..end {
5369 v.scales_into(r, gpr, &mut sc);
5370 let mut bi = 0usize;
5371 #[cfg(target_arch = "x86_64")]
5372 if blocked_ok {
5373 while bi + 8 <= alen {
5374 let xs = [
5375 acts[abase + bi].xq.as_slice(),
5376 acts[abase + bi + 1].xq.as_slice(),
5377 acts[abase + bi + 2].xq.as_slice(),
5378 acts[abase + bi + 3].xq.as_slice(),
5379 acts[abase + bi + 4].xq.as_slice(),
5380 acts[abase + bi + 5].xq.as_slice(),
5381 acts[abase + bi + 6].xq.as_slice(),
5382 acts[abase + bi + 7].xq.as_slice(),
5383 ];
5384 let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5385 for k in 0..8 {
5386 let act = &acts[abase + bi + k];
5387 let mut acc = d[k] * act.sx;
5388 for &(j, xv) in &act.outliers {
5389 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5390 acc += w * s * xv;
5391 }
5392 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5394 }
5395 bi += 8;
5396 }
5397 while bi + 4 <= alen {
5398 let xs = [
5399 acts[abase + bi].xq.as_slice(),
5400 acts[abase + bi + 1].xq.as_slice(),
5401 acts[abase + bi + 2].xq.as_slice(),
5402 acts[abase + bi + 3].xq.as_slice(),
5403 ];
5404 let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5405 for k in 0..4 {
5406 let act = &acts[abase + bi + k];
5407 let mut acc = d[k] * act.sx;
5408 for &(j, xv) in &act.outliers {
5409 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5410 acc += w * s * xv;
5411 }
5412 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5414 }
5415 bi += 4;
5416 }
5417 }
5418 #[cfg(target_arch = "aarch64")]
5419 if blocked_ok {
5420 while bi + 4 <= alen {
5421 let xs = [
5422 acts[abase + bi].xq.as_slice(),
5423 acts[abase + bi + 1].xq.as_slice(),
5424 acts[abase + bi + 2].xq.as_slice(),
5425 acts[abase + bi + 3].xq.as_slice(),
5426 ];
5427 let d = unsafe {
5428 if q4tp_v1() {
5429 dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5430 } else {
5431 dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5432 }
5433 };
5434 for k in 0..4 {
5435 let act = &acts[abase + bi + k];
5436 let mut acc = d[k] * act.sx;
5437 for &(j, xv) in &act.outliers {
5438 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5439 acc += w * s * xv;
5440 }
5441 unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5443 }
5444 bi += 4;
5445 }
5446 }
5447 let _ = blocked_ok;
5448 while bi < alen {
5449 let act = &acts[abase + bi];
5450 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5451 for &(j, xv) in &act.outliers {
5452 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5453 acc += w * s * xv;
5454 }
5455 unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5457 bi += 1;
5458 }
5459 }
5460
5461 }
5462 };
5463 dispatch_rows(pool, rows, &run);
5464 return;
5465 }
5466
5467 let run = |start: usize, end: usize| {
5468 let mut sc = vec![0f32; gpr];
5469 for r in start..end {
5470 v.scales_into(r, gpr, &mut sc);
5471 for bi in 0..b {
5472 let x = &xs_all[bi * cols..(bi + 1) * cols];
5473 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5475 }
5476 }
5477 };
5478 dispatch_rows(pool, rows, &run);
5479}
5480
5481fn q4t_matvec(
5483 bytes: &[u8],
5484 x: &[f32],
5485 rows: usize,
5486 cols: usize,
5487 out: &mut [f32],
5488 pool: Option<&Pool>,
5489) {
5490 debug_assert_eq!(out.len(), rows);
5491 let gpr = cols / GROUP_SIZE;
5492 let out_addr = SendMut(out.as_mut_ptr());
5493 if a8w8_enabled() {
5494 let act = split_act(x);
5495 let run = move |start: usize, end: usize| {
5496 for r in start..end {
5497 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5498 for &(j, xv) in &act.outliers {
5499 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5500 acc += w * s * xv;
5501 }
5502 unsafe { *out_addr.at(r) = acc };
5504 }
5505 };
5506 dispatch_rows(pool, rows, &run);
5507 return;
5508 }
5509 let run = move |start: usize, end: usize| {
5510 for r in start..end {
5511 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5513 }
5514 };
5515 dispatch_rows(pool, rows, &run);
5516}
5517
5518#[allow(clippy::too_many_arguments)]
5520fn q4t_matvec2(
5521 bytes: &[u8],
5522 x1: &[f32],
5523 x2: &[f32],
5524 rows: usize,
5525 cols: usize,
5526 o1: &mut [f32],
5527 o2: &mut [f32],
5528 pool: Option<&Pool>,
5529) {
5530 let gpr = cols / GROUP_SIZE;
5531 let p1 = SendMut(o1.as_mut_ptr());
5532 let p2 = SendMut(o2.as_mut_ptr());
5533 if a8w8_enabled() {
5534 let a1 = split_act(x1);
5535 let a2 = split_act(x2);
5536 let run = move |start: usize, end: usize| {
5537 for r in start..end {
5538 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5539 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5540 for &(j, xv) in &a1.outliers {
5541 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5542 v1 += w * s * xv;
5543 }
5544 for &(j, xv) in &a2.outliers {
5545 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5546 v2 += w * s * xv;
5547 }
5548 unsafe {
5550 *p1.at(r) = v1;
5551 *p2.at(r) = v2;
5552 }
5553 }
5554 };
5555 dispatch_rows(pool, rows, &run);
5556 return;
5557 }
5558 let run = move |start: usize, end: usize| {
5559 for r in start..end {
5560 unsafe {
5562 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5563 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5564 }
5565 }
5566 };
5567 dispatch_rows(pool, rows, &run);
5568}
5569
5570#[allow(clippy::too_many_arguments)]
5572#[cfg(target_os = "macos")]
5578fn dequant_matmat_accel(
5579 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5580 xs_all: &[f32],
5581 b: usize,
5582 rows: usize,
5583 cols: usize,
5584 out: &mut [f32],
5585 pool: Option<&Pool>,
5586) {
5587 const TR: usize = 2048;
5588 thread_local! {
5589 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5590 }
5591 WTILE.with(|wt| {
5592 let mut wtile = wt.borrow_mut();
5593 wtile.resize(TR * cols, 0.0);
5594 let mut r0 = 0usize;
5595 while r0 < rows {
5596 let tr = TR.min(rows - r0);
5597 let wt_addr = SendMut(wtile.as_mut_ptr());
5598 let run = |start: usize, end: usize| {
5599 for r in start..end {
5600 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5602 dequant_row(r0 + r, dst);
5603 }
5604 };
5605 dispatch_rows(pool, tr, &run);
5606 unsafe {
5607 accel_blas::cblas_sgemm(
5608 101, 111, 112, b as i32,
5612 tr as i32,
5613 cols as i32,
5614 1.0,
5615 xs_all.as_ptr(),
5616 cols as i32,
5617 wtile.as_ptr(),
5618 cols as i32,
5619 0.0,
5620 out.as_mut_ptr().add(r0),
5621 rows as i32,
5622 );
5623 }
5624 r0 += tr;
5625 }
5626 });
5627}
5628
5629fn q4t_matmat(
5630 bytes: &[u8],
5631 xs_all: &[f32],
5632 b: usize,
5633 rows: usize,
5634 cols: usize,
5635 out: &mut [f32],
5636 pool: Option<&Pool>,
5637) {
5638 debug_assert_eq!(out.len(), b * rows);
5639 let gpr = cols / GROUP_SIZE;
5640 #[cfg(target_os = "macos")]
5644 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5645 dequant_matmat_accel(
5646 &|r, dst| {
5647 for gi in 0..gpr {
5648 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5649 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5650 for (k, &bb) in tile[2..].iter().enumerate() {
5651 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5652 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5653 }
5654 }
5655 },
5656 xs_all,
5657 b,
5658 rows,
5659 cols,
5660 out,
5661 pool,
5662 );
5663 return;
5664 }
5665 let out_addr = SendMut(out.as_mut_ptr());
5666 if a8w8_enabled() {
5667 let acts: Vec<SplitAct> = (0..b)
5668 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5669 .collect();
5670 let acts = &acts;
5671 #[cfg(target_arch = "x86_64")]
5672 let blocked_ok = avx2_enabled()
5673 && blocked_enabled();
5674 #[cfg(target_arch = "aarch64")]
5675 let blocked_ok = sdot_enabled()
5676 && blocked_enabled();
5677 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5678 let blocked_ok = false;
5679 let run = move |start: usize, end: usize| {
5680 for r in start..end {
5681 let mut bi = 0usize;
5682 #[cfg(target_arch = "aarch64")]
5683 if blocked_ok {
5684 while bi + 4 <= acts.len() {
5685 let xs = [
5686 acts[bi].xq.as_slice(),
5687 acts[bi + 1].xq.as_slice(),
5688 acts[bi + 2].xq.as_slice(),
5689 acts[bi + 3].xq.as_slice(),
5690 ];
5691 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5692 for k in 0..4 {
5693 let act = &acts[bi + k];
5694 let mut acc = d[k] * act.sx;
5695 for &(j, xv) in &act.outliers {
5696 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5697 acc += w * sc * xv;
5698 }
5699 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5701 }
5702 bi += 4;
5703 }
5704 }
5705 #[cfg(target_arch = "x86_64")]
5706 if blocked_ok {
5707 while bi + 4 <= acts.len() {
5708 let xs = [
5709 acts[bi].xq.as_slice(),
5710 acts[bi + 1].xq.as_slice(),
5711 acts[bi + 2].xq.as_slice(),
5712 acts[bi + 3].xq.as_slice(),
5713 ];
5714 let d = unsafe {
5715 if vnni_tiles_enabled() {
5716 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5717 } else {
5718 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5719 }
5720 };
5721 for k in 0..4 {
5722 let act = &acts[bi + k];
5723 let mut acc = d[k] * act.sx;
5724 for &(j, xv) in &act.outliers {
5725 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5726 acc += w * sc * xv;
5727 }
5728 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5730 }
5731 bi += 4;
5732 }
5733 }
5734 let _ = blocked_ok;
5735 while bi < acts.len() {
5736 let act = &acts[bi];
5737 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5738 for &(j, xv) in &act.outliers {
5739 let (w, s) = q4t_outlier(bytes, r, gpr, j);
5740 acc += w * s * xv;
5741 }
5742 unsafe { *out_addr.at(bi * rows + r) = acc };
5744 bi += 1;
5745 }
5746 }
5747 };
5748 dispatch_rows(pool, rows, &run);
5749 return;
5750 }
5751 let run = move |start: usize, end: usize| {
5752 for r in start..end {
5753 for bi in 0..b {
5754 let x = &xs_all[bi * cols..(bi + 1) * cols];
5755 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5757 }
5758 }
5759 };
5760 dispatch_rows(pool, rows, &run);
5761}
5762
5763fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5772 (0..gpr)
5773 .map(|gi| {
5774 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5775 .iter()
5776 .map(|&v| v as i32)
5777 .sum()
5778 })
5779 .collect()
5780}
5781
5782#[inline]
5786#[allow(unreachable_code)]
5787#[cfg(target_arch = "x86_64")]
5792#[target_feature(enable = "avx2")]
5793unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5794 unsafe {
5796 use core::arch::x86_64::*;
5797 let expand = _mm256_setr_epi8(
5799 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,
5800 3, 3, 3,
5801 );
5802 let bitsel = _mm256_setr_epi8(
5803 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5804 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5805 );
5806 let ones8 = _mm256_set1_epi8(1);
5807 let ones16 = _mm256_set1_epi16(1);
5808 let mut acc = 0f32;
5809 for gi in 0..gpr {
5810 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5811 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5812 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5813 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5814 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5815 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5816 let sel = _mm256_and_si256(x, mask);
5817 let p16 = _mm256_maddubs_epi16(ones8, sel);
5819 let d32 = _mm256_madd_epi16(p16, ones16);
5820 let hi128 = _mm256_extracti128_si256::<1>(d32);
5821 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5822 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5823 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5824 let msum = _mm_cvtsi128_si32(s32);
5825 let d = 2 * msum - gsum[gi];
5828 acc += d as f32 * s;
5829 }
5830 acc
5831 }
5832}
5833
5834#[cfg(target_arch = "x86_64")]
5837#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5838unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5839 unsafe {
5841 use core::arch::x86_64::*;
5842 let expand = _mm256_setr_epi8(
5843 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,
5844 3, 3, 3,
5845 );
5846 let bitsel = _mm256_setr_epi8(
5847 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5848 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5849 );
5850 let ones8 = _mm256_set1_epi8(1);
5851 let mut acc = 0f32;
5852 for gi in 0..gpr {
5853 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5854 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5855 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5856 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5857 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5858 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5859 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5860 let d = 2 * msum - gsum[gi];
5861 acc += d as f32 * s;
5862 }
5863 acc
5864 }
5865}
5866
5867#[cfg(target_arch = "x86_64")]
5869#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5870unsafe fn dot_q1_row_1x4_vnni(
5871 bytes: &[u8],
5872 r: usize,
5873 gpr: usize,
5874 xs: [&[i8]; 4],
5875 gsums: [&[i32]; 4],
5876) -> [f32; 4] {
5877 unsafe {
5879 use core::arch::x86_64::*;
5880 let expand = _mm256_setr_epi8(
5881 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,
5882 3, 3, 3,
5883 );
5884 let bitsel = _mm256_setr_epi8(
5885 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5886 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5887 );
5888 let ones8 = _mm256_set1_epi8(1);
5889 let mut acc = [0f32; 4];
5890 for gi in 0..gpr {
5891 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5892 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5893 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5894 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5895 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5896 for (k, xq) in xs.iter().enumerate() {
5897 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5898 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5899 let d = 2 * msum - gsums[k][gi];
5900 acc[k] += d as f32 * s;
5901 }
5902 }
5903 acc
5904 }
5905}
5906
5907#[cfg(target_arch = "x86_64")]
5910#[target_feature(enable = "avx2")]
5911unsafe fn dot_q1_row_1x4_avx2(
5912 bytes: &[u8],
5913 r: usize,
5914 gpr: usize,
5915 xs: [&[i8]; 4],
5916 gsums: [&[i32]; 4],
5917) -> [f32; 4] {
5918 unsafe {
5920 use core::arch::x86_64::*;
5921 let expand = _mm256_setr_epi8(
5922 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,
5923 3, 3, 3,
5924 );
5925 let bitsel = _mm256_setr_epi8(
5926 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5927 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5928 );
5929 let ones8 = _mm256_set1_epi8(1);
5930 let ones16 = _mm256_set1_epi16(1);
5931 let mut acc = [0f32; 4];
5932 for gi in 0..gpr {
5933 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5934 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5935 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5936 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5937 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5938 for (k, xq) in xs.iter().enumerate() {
5939 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5940 let sel = _mm256_and_si256(x, mask);
5941 let p16 = _mm256_maddubs_epi16(ones8, sel);
5942 let d32 = _mm256_madd_epi16(p16, ones16);
5943 let hi128 = _mm256_extracti128_si256::<1>(d32);
5944 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5945 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5946 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5947 let msum = _mm_cvtsi128_si32(s32);
5948 let d = 2 * msum - gsums[k][gi];
5949 acc[k] += d as f32 * s;
5950 }
5951 }
5952 acc
5953 }
5954}
5955
5956#[allow(unreachable_code)]
5957fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5958 #[cfg(target_arch = "aarch64")]
5959 unsafe {
5960 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
5961 }
5962 #[cfg(target_arch = "x86_64")]
5963 if avx2_enabled() {
5964 unsafe {
5965 if vnni_tiles_enabled() {
5966 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
5967 }
5968 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
5969 }
5970 }
5971 let _ = gsum;
5972 let mut acc = 0f32;
5973 for gi in 0..gpr {
5974 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5975 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5976 let mut d = 0i32;
5977 for (j, &b) in tile[2..].iter().enumerate() {
5978 for k in 0..8 {
5979 let w = ((b >> k) & 1) as i32 * 2 - 1;
5980 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
5981 }
5982 }
5983 acc += d as f32 * s;
5984 }
5985 acc
5986}
5987
5988#[cfg(target_arch = "aarch64")]
5997#[target_feature(enable = "neon,dotprod")]
5998unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5999 unsafe {
6002 use core::arch::aarch64::*;
6003 use core::arch::asm;
6004 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6005 let m = vld1q_u8(MASKS.as_ptr());
6006 macro_rules! tile_dot {
6008 ($t:expr, $x:expr) => {{
6009 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6010 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6011 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6012 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6013 let x0 = vld1q_s8($x);
6014 let x1 = vld1q_s8($x.add(16));
6015 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6016 asm!(
6017 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6018 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6019 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6020 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6021 options(pure, nomem, nostack),
6022 );
6023 vaddq_s32(a0, a1)
6024 }};
6025 }
6026 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6035 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6036 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6037 const IW11: [u8; 16] = [
6038 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6039 ];
6040 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6041 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6042 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6043 let isc = vld1_u8(ISC.as_ptr());
6044 macro_rules! tile_dot_tbl {
6046 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6047 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6048 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6049 let x0 = vld1q_s8($x);
6050 let x1 = vld1q_s8($x.add(16));
6051 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6052 asm!(
6053 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6054 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6055 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6056 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6057 options(pure, nomem, nostack),
6058 );
6059 vaddq_s32(a0, a1)
6060 }};
6061 }
6062 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6063 let row_base = r * gpr * Q1_TILE;
6064 let abs_end = bytes.len();
6065 let xp = xq.as_ptr();
6066 let gp = gsum.as_ptr();
6067 let mut accv = vdupq_n_f32(0.0);
6068 let mut gi = 0;
6069 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6072 let t0 = base.add(gi * Q1_TILE);
6073 let ld_a = vld1q_u8(t0);
6074 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6075 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6076 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6077 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6078 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6079 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6081 let g = vld1q_s32(gp.add(gi));
6082 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6083 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6084 let scf: float32x4_t;
6085 asm!(
6086 "fcvtl {o:v}.4s, {i:v}.4h",
6087 o = out(vreg) scf, i = in(vreg) sc16,
6088 options(pure, nomem, nostack),
6089 );
6090 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6091 gi += 4;
6092 }
6093 let mut acc = vaddvq_f32(accv);
6094 while gi < gpr {
6095 let t = base.add(gi * Q1_TILE);
6096 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6097 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6098 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6099 gi += 1;
6100 }
6101 acc
6102 }
6103}
6104
6105#[cfg(target_arch = "aarch64")]
6110#[target_feature(enable = "neon,dotprod")]
6111unsafe fn dot_q1_row_1x4_sdot(
6112 bytes: &[u8],
6113 r: usize,
6114 gpr: usize,
6115 xs: [&[i8]; 4],
6116 gs: [&[i32]; 4],
6117) -> [f32; 4] {
6118 unsafe {
6120 use core::arch::aarch64::*;
6121 use core::arch::asm;
6122 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6123 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6124 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6125 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6126 const IW11: [u8; 16] = [
6127 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6128 ];
6129 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6130 let m = vld1q_u8(MASKS.as_ptr());
6131 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6132 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6133 let isc = vld1_u8(ISC.as_ptr());
6134 macro_rules! sdot2 {
6135 ($w0:expr, $w1:expr, $x:expr) => {{
6136 let x0 = vld1q_s8($x);
6137 let x1 = vld1q_s8($x.add(16));
6138 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6139 asm!(
6140 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6141 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6142 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6143 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6144 options(pure, nomem, nostack),
6145 );
6146 vaddq_s32(a0, a1)
6147 }};
6148 }
6149 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6150 let row_base = r * gpr * Q1_TILE;
6151 let abs_end = bytes.len();
6152 let mut accv = [vdupq_n_f32(0.0); 4];
6153 let mut gi = 0;
6154 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6155 let t0 = base.add(gi * Q1_TILE);
6156 let ld_a = vld1q_u8(t0);
6157 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6158 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6160 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6161 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6162 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6163 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6164 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6165 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6166 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6167 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6168 let scf: float32x4_t;
6169 asm!(
6170 "fcvtl {o:v}.4s, {i:v}.4h",
6171 o = out(vreg) scf, i = in(vreg) sc16,
6172 options(pure, nomem, nostack),
6173 );
6174 for k in 0..4 {
6175 let xp = xs[k].as_ptr();
6176 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6177 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6178 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6179 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6180 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6181 let g = vld1q_s32(gs[k].as_ptr().add(gi));
6182 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6183 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6184 }
6185 gi += 4;
6186 }
6187 let mut acc = [
6188 vaddvq_f32(accv[0]),
6189 vaddvq_f32(accv[1]),
6190 vaddvq_f32(accv[2]),
6191 vaddvq_f32(accv[3]),
6192 ];
6193 while gi < gpr {
6194 let t = base.add(gi * Q1_TILE);
6195 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6196 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6197 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6198 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6199 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6200 for k in 0..4 {
6201 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6202 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6203 }
6204 gi += 1;
6205 }
6206 acc
6207 }
6208}
6209
6210#[inline]
6212fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6213 let gi = j / GROUP_SIZE;
6214 let k = j % GROUP_SIZE;
6215 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6216 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6217 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6218 ((bit as i32 * 2 - 1) as f32, s)
6219}
6220
6221#[inline]
6223fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6224 let mut acc = 0f32;
6225 for gi in 0..gpr {
6226 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6227 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6228 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6229 let mut ga = 0f32;
6230 for (j, &b) in tile[2..].iter().enumerate() {
6231 for k in 0..8 {
6232 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6233 }
6234 }
6235 acc += ga * s;
6236 }
6237 acc
6238}
6239
6240#[allow(clippy::too_many_arguments)]
6243fn q1_range_a8w8(
6244 bytes: &[u8],
6245 gpr: usize,
6246 act: &SplitAct,
6247 gsum: &[i32],
6248 out: SendMut,
6249 start: usize,
6250 end: usize,
6251) {
6252 for r in start..end {
6253 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6254 for &(j, xv) in &act.outliers {
6255 let (w, s) = q1_outlier(bytes, r, gpr, j);
6256 acc += w * s * xv;
6257 }
6258 unsafe { *out.at(r) = acc };
6260 }
6261}
6262
6263fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6265 for r in start..end {
6266 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6268 }
6269}
6270
6271fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6276 let entries = base_len + (rows + 1) * 4;
6277 (base_len, entries, entries <= bytes.len())
6278}
6279
6280#[inline]
6282fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6283 let o = rp_off + r * 4;
6284 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6285}
6286
6287const SIGN5: [[f32; 5]; 256] = {
6291 let mut lut = [[0.0f32; 5]; 256];
6292 let pow3 = [1u16, 3, 9, 27, 81];
6293 let mut byte = 0usize;
6294 while byte < 256 {
6295 let mut i = 0usize;
6296 while i < 5 {
6297 let code = (byte as u16 / pow3[i]) % 3;
6298 lut[byte][i] = if code == 1 {
6299 1.0
6300 } else if code == 2 {
6301 -1.0
6302 } else {
6303 0.0
6304 };
6305 i += 1;
6306 }
6307 byte += 1;
6308 }
6309 lut
6310};
6311
6312const SIGN5_I8: [[i8; 5]; 256] = {
6314 let mut lut = [[0i8; 5]; 256];
6315 let pow3 = [1u16, 3, 9, 27, 81];
6316 let mut byte = 0usize;
6317 while byte < 256 {
6318 let mut i = 0usize;
6319 while i < 5 {
6320 let code = (byte as u16 / pow3[i]) % 3;
6321 lut[byte][i] = if code == 1 {
6322 1
6323 } else if code == 2 {
6324 -1
6325 } else {
6326 0
6327 };
6328 i += 1;
6329 }
6330 byte += 1;
6331 }
6332 lut
6333};
6334
6335const SIGN5_U64: [u64; 256] = {
6341 let mut lut = [0u64; 256];
6342 let pow3 = [1u16, 3, 9, 27, 81];
6343 let mut byte = 0usize;
6344 while byte < 256 {
6345 let mut v = 0u64;
6346 let mut i = 0usize;
6347 while i < 5 {
6348 let code = (byte as u16 / pow3[i]) % 3;
6349 let s: u8 = if code == 1 {
6350 1
6351 } else if code == 2 {
6352 0xFF
6353 } else {
6354 0
6355 };
6356 v |= (s as u64) << (i * 8);
6357 i += 1;
6358 }
6359 lut[byte] = v;
6360 byte += 1;
6361 }
6362 lut
6363};
6364
6365#[inline]
6370fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6371 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6372 let off = (r * gpr + j / GROUP_SIZE) * TILE;
6373 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6374 let within = j % GROUP_SIZE;
6375 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6376}
6377
6378#[cfg(target_arch = "aarch64")]
6381#[target_feature(enable = "neon,dotprod")]
6382#[inline]
6383unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6384 unsafe {
6386 use core::arch::aarch64::*;
6387 use core::arch::asm;
6388 let w0 = vld1q_s8(w);
6389 let w1 = vld1q_s8(w.add(16));
6390 let x0 = vld1q_s8(x);
6391 let x1 = vld1q_s8(x.add(16));
6392 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6393 asm!(
6394 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6395 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6396 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6397 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6398 options(pure, nomem, nostack),
6399 );
6400 vaddvq_s32(vaddq_s32(a0, a1))
6401 }
6402}
6403
6404#[cfg(target_arch = "x86_64")]
6407#[target_feature(enable = "avx2")]
6408#[inline]
6409unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6410 unsafe {
6412 use core::arch::x86_64::*;
6413 let wv = _mm256_loadu_si256(w as *const __m256i);
6414 let xv = _mm256_loadu_si256(x as *const __m256i);
6415 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6416 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6417 let hi128 = _mm256_extracti128_si256::<1>(d);
6418 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6419 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6420 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6421 _mm_cvtsi128_si32(s32)
6422 }
6423}
6424
6425#[inline]
6430fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6431 debug_assert!(dst.len() >= 40);
6432 unsafe {
6435 let p = dst.as_mut_ptr();
6436 for bi in 0..7 {
6437 core::ptr::write_unaligned(
6438 p.add(bi * 5) as *mut u64,
6439 SIGN5_U64[*codes.add(bi) as usize],
6440 );
6441 }
6442 }
6443}
6444
6445#[inline]
6450fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6451 #[cfg(target_arch = "aarch64")]
6452 unsafe {
6453 return sdot32_i8(w, x);
6454 }
6455 #[cfg(target_arch = "x86_64")]
6456 unsafe {
6457 return i8dot32_avx2(w, x);
6458 }
6459 #[allow(unreachable_code)]
6460 unsafe {
6461 let mut s = 0i32;
6462 for k in 0..GROUP_SIZE {
6463 s += *w.add(k) as i32 * *x.add(k) as i32;
6464 }
6465 s
6466 }
6467}
6468
6469#[inline]
6470unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6471 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6472 (
6473 SIGN5_U64[*codes as usize],
6474 SIGN5_U64[*codes.add(1) as usize],
6475 SIGN5_U64[*codes.add(2) as usize],
6476 SIGN5_U64[*codes.add(3) as usize],
6477 SIGN5_U64[*codes.add(4) as usize],
6478 SIGN5_U64[*codes.add(5) as usize],
6479 SIGN5_U64[*codes.add(6) as usize],
6480 )
6481 };
6482
6483 let u0 = s0 | (s1 << 40);
6484 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6485 let u2 = (s3 >> 8) | (s4 << 32);
6486 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6487
6488 (u0, u1, u2, u3)
6489}
6490
6491#[cfg(target_arch = "aarch64")]
6495#[target_feature(enable = "neon,dotprod")]
6496unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6497 use core::arch::aarch64::*;
6498 use core::arch::asm;
6499 unsafe {
6500 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6501 let mut acc = 0f32;
6502 let bytes_ptr = bytes.as_ptr();
6503 let xq_ptr = xq.as_ptr();
6504 let row_off = r * gpr * TILE;
6505
6506 let gpr2 = gpr & !1;
6507 let mut gi = 0;
6508 while gi < gpr2 {
6509 let off0 = row_off + gi * TILE;
6510 let off1 = off0 + TILE;
6511 let s0 = f16_to_f32(u16::from_le_bytes([
6512 *bytes_ptr.add(off0),
6513 *bytes_ptr.add(off0 + 1),
6514 ]));
6515 let s1 = f16_to_f32(u16::from_le_bytes([
6516 *bytes_ptr.add(off1),
6517 *bytes_ptr.add(off1 + 1),
6518 ]));
6519
6520 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6521 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6522
6523 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6524 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6525 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6526 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6527
6528 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6529 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6530 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6531 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6532
6533 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6534 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6535 asm!(
6536 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6537 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6538 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6539 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6540 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6541 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6542 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6543 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6544 options(pure, nomem, nostack),
6545 );
6546 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6547 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6548 acc += d0 as f32 * s0 + d1 as f32 * s1;
6549 gi += 2;
6550 }
6551
6552 if gi < gpr {
6553 let off = row_off + gi * TILE;
6554 let s = f16_to_f32(u16::from_le_bytes([
6555 *bytes_ptr.add(off),
6556 *bytes_ptr.add(off + 1),
6557 ]));
6558 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6559 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6560 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6561 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6562 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6563 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6564 asm!(
6565 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6566 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6567 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6568 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6569 options(pure, nomem, nostack),
6570 );
6571 let d = vaddvq_s32(vaddq_s32(a0, a1));
6572 acc += d as f32 * s;
6573 }
6574 acc
6575 }
6576}
6577
6578#[cfg(target_arch = "x86_64")]
6580#[target_feature(enable = "avx2")]
6581unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6582 use core::arch::x86_64::*;
6583 unsafe {
6584 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6585 let mut acc = 0f32;
6586 let bytes_ptr = bytes.as_ptr();
6587 let xq_ptr = xq.as_ptr();
6588 let row_off = r * gpr * TILE;
6589
6590 let ones = _mm256_set1_epi16(1);
6591 for gi in 0..gpr {
6592 let off = row_off + gi * TILE;
6593 let s = f16_to_f32(u16::from_le_bytes([
6594 *bytes_ptr.add(off),
6595 *bytes_ptr.add(off + 1),
6596 ]));
6597 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6598 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6599 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6600 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6601 let d256 = _mm256_madd_epi16(p16, ones);
6602 let d128 = _mm_add_epi32(
6603 _mm256_castsi256_si128(d256),
6604 _mm256_extracti128_si256(d256, 1),
6605 );
6606 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6607 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6608 acc += d32 as f32 * s;
6609 }
6610 acc
6611 }
6612}
6613
6614#[cfg(target_arch = "x86_64")]
6616#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6617unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6618 use core::arch::x86_64::*;
6619 unsafe {
6621 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6622 let mut acc = 0f32;
6623 let bytes_ptr = bytes.as_ptr();
6624 let xq_ptr = xq.as_ptr();
6625 let row_off = r * gpr * TILE;
6626 for gi in 0..gpr {
6627 let off = row_off + gi * TILE;
6628 let s = f16_to_f32(u16::from_le_bytes([
6629 *bytes_ptr.add(off),
6630 *bytes_ptr.add(off + 1),
6631 ]));
6632 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6633 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6634 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6635 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6636 acc += d as f32 * s;
6637 }
6638 acc
6639 }
6640}
6641
6642#[inline]
6646fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6647 #[cfg(target_arch = "aarch64")]
6648 unsafe {
6649 return q1t_dot_row_sdot(bytes, r, gpr, xq);
6650 }
6651 #[cfg(target_arch = "x86_64")]
6652 unsafe {
6653 if vnni_tiles_enabled() {
6654 return q1t_dot_row_vnni(bytes, r, gpr, xq);
6655 }
6656 return q1t_dot_row_avx2(bytes, r, gpr, xq);
6657 }
6658 #[allow(unreachable_code)]
6659 {
6660 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6661 let mut acc = 0f32;
6662 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
6664 let off = (r * gpr + gi) * TILE;
6665 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6666 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6667 let mut d = 0i32;
6668 for k in 0..GROUP_SIZE {
6669 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6670 }
6671 acc += d as f32 * s;
6672 }
6673 acc
6674 }
6675}
6676
6677fn q1t_row_outlier_correction(
6684 bytes: &[u8],
6685 r: usize,
6686 rp_off: usize,
6687 entries_off: usize,
6688 has_ov: bool,
6689 x: &[f32],
6690) -> f32 {
6691 if !has_ov {
6692 return 0.0;
6693 }
6694 let (c0, c1) = (
6695 q1t_rowptr(bytes, rp_off, r),
6696 q1t_rowptr(bytes, rp_off, r + 1),
6697 );
6698 let mut corr = 0f32;
6699 for p in c0..c1 {
6700 let e = entries_off + p * 4;
6701 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6702 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6703 corr += val * x[col];
6704 }
6705 corr
6706}
6707
6708fn q1t_dequant_row(
6712 bytes: &[u8],
6713 r: usize,
6714 gpr: usize,
6715 rp_off: usize,
6716 entries_off: usize,
6717 has_ov: bool,
6718 buf: &mut [f32],
6719) {
6720 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6721 for g in 0..gpr {
6722 let off = (r * gpr + g) * TILE;
6723 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6724 let codes = &bytes[off + 2..off + TILE];
6725 let bc = g * GROUP_SIZE;
6726 for bi in 0..6 {
6728 let lut = &SIGN5[codes[bi] as usize];
6729 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6730 for i in 0..5 {
6731 d[i] = lut[i] * s;
6732 }
6733 }
6734 let lut = &SIGN5[codes[6] as usize];
6735 buf[bc + 30] = lut[0] * s;
6736 buf[bc + 31] = lut[1] * s;
6737 }
6738 if !has_ov {
6739 return;
6740 }
6741 let (c0, c1) = (
6742 q1t_rowptr(bytes, rp_off, r),
6743 q1t_rowptr(bytes, rp_off, r + 1),
6744 );
6745 for p in c0..c1 {
6746 let e = entries_off + p * 4;
6747 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6748 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6749 }
6750}
6751
6752fn q1t_add_overlay(
6756 bytes: &[u8],
6757 x: &[f32],
6758 rows: usize,
6759 cols: usize,
6760 out: &mut [f32],
6761 pool: Option<&Pool>,
6762) {
6763 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6764 let gpr = cols / GROUP_SIZE;
6765 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6766 if !has_ov {
6767 return;
6768 }
6769 let out_addr = SendMut(out.as_mut_ptr());
6770 let run = move |start: usize, end: usize| {
6771 for r in start..end {
6772 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6773 unsafe { *out_addr.at(r) += corr };
6775 }
6776 };
6777 dispatch_rows(pool, rows, &run);
6778}
6779
6780#[allow(clippy::too_many_arguments)]
6783fn q1t_range_a8w8(
6784 bytes: &[u8],
6785 gpr: usize,
6786 rp_off: usize,
6787 ent_off: usize,
6788 has_ov: bool,
6789 act: &SplitAct,
6790 x: &[f32],
6791 out: SendMut,
6792 start: usize,
6793 end: usize,
6794) {
6795 for r in start..end {
6796 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6797 for &(j, xv) in &act.outliers {
6798 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6799 }
6800 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6801 unsafe { *out.at(r) = acc };
6803 }
6804}
6805
6806#[allow(clippy::too_many_arguments)]
6809fn q1t_range_f32_batch(
6810 bytes: &[u8],
6811 gpr: usize,
6812 rp_off: usize,
6813 ent_off: usize,
6814 has_ov: bool,
6815 x: &[f32],
6816 out: SendMut,
6817 start: usize,
6818 end: usize,
6819) {
6820 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6821 let mut sg = [0f32; GROUP_SIZE];
6822 for r in start..end {
6823 let mut acc = 0f32;
6824 for g in 0..gpr {
6825 let off = (r * gpr + g) * TILE;
6826 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6827 let codes = &bytes[off + 2..off + TILE];
6828 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6829 for bi in 0..6 {
6830 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6831 }
6832 let lut = &SIGN5[codes[6] as usize];
6833 sg[30] = lut[0];
6834 sg[31] = lut[1];
6835 let mut gsum = 0f32;
6836 for k in 0..GROUP_SIZE {
6837 gsum += sg[k] * xg[k];
6838 }
6839 acc += s * gsum;
6840 }
6841 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6842 unsafe { *out.at(r) = acc };
6844 }
6845}
6846
6847fn q1t_matvec(
6851 bytes: &[u8],
6852 x: &[f32],
6853 rows: usize,
6854 cols: usize,
6855 out: &mut [f32],
6856 pool: Option<&Pool>,
6857) {
6858 debug_assert_eq!(out.len(), rows);
6859 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6860 let gpr = cols / GROUP_SIZE;
6861 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6862 let out_addr = SendMut(out.as_mut_ptr());
6863 if a8w8_enabled() {
6867 let act = split_act(x);
6868 let act = &act;
6869 let run = move |start: usize, end: usize| {
6870 for r in start..end {
6871 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6872 for &(j, xv) in &act.outliers {
6873 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6874 }
6875 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6876 unsafe { *out_addr.at(r) = acc };
6878 }
6879 };
6880 dispatch_rows(pool, rows, &run);
6881 return;
6882 }
6883 let run = move |start: usize, end: usize| {
6884 let mut sg = [0f32; GROUP_SIZE];
6888 for r in start..end {
6889 let mut acc = 0f32;
6890 for g in 0..gpr {
6891 let off = (r * gpr + g) * TILE;
6892 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6893 let codes = &bytes[off + 2..off + TILE];
6894 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6895 for bi in 0..6 {
6896 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6897 }
6898 let lut = &SIGN5[codes[6] as usize];
6899 sg[30] = lut[0];
6900 sg[31] = lut[1];
6901 let mut gsum = 0f32;
6902 for k in 0..GROUP_SIZE {
6903 gsum += sg[k] * xg[k];
6904 }
6905 acc += s * gsum;
6906 }
6907 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6908 unsafe { *out_addr.at(r) = acc };
6909 }
6910 };
6911 dispatch_rows(pool, rows, &run);
6912}
6913
6914#[cfg(target_arch = "aarch64")]
6920#[target_feature(enable = "neon,dotprod")]
6921unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6922 use core::arch::aarch64::*;
6923 use core::arch::asm;
6924 unsafe {
6926 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6927 let bytes_ptr = bytes.as_ptr();
6928 let row_off = r * gpr * TILE;
6929 let xp = [xa.as_ptr(), xb.as_ptr()];
6930 let mut acc = [0f32; 2];
6931 macro_rules! sdot2 {
6932 ($w0:expr, $w1:expr, $x:expr) => {{
6933 let x0 = vld1q_s8($x);
6934 let x1 = vld1q_s8($x.add(16));
6935 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6936 asm!(
6937 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6938 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6939 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6940 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6941 options(pure, nomem, nostack),
6942 );
6943 vaddvq_s32(vaddq_s32(a0, a1))
6944 }};
6945 }
6946 let gpr2 = gpr & !1;
6947 let mut gi = 0;
6948 while gi < gpr2 {
6949 let off0 = row_off + gi * TILE;
6950 let off1 = off0 + TILE;
6951 let s0 = f16_to_f32(u16::from_le_bytes([
6952 *bytes_ptr.add(off0),
6953 *bytes_ptr.add(off0 + 1),
6954 ]));
6955 let s1 = f16_to_f32(u16::from_le_bytes([
6956 *bytes_ptr.add(off1),
6957 *bytes_ptr.add(off1 + 1),
6958 ]));
6959 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6960 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6961 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6962 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6963 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6964 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6965 for k in 0..2 {
6966 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
6967 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
6968 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
6969 }
6970 gi += 2;
6971 }
6972 if gi < gpr {
6973 let off = row_off + gi * TILE;
6974 let s = f16_to_f32(u16::from_le_bytes([
6975 *bytes_ptr.add(off),
6976 *bytes_ptr.add(off + 1),
6977 ]));
6978 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6979 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6980 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6981 for k in 0..2 {
6982 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
6983 acc[k] += d as f32 * s;
6984 }
6985 }
6986 acc
6987 }
6988}
6989
6990fn q1t_matvec2(
6996 bytes: &[u8],
6997 x1: &[f32],
6998 x2: &[f32],
6999 rows: usize,
7000 cols: usize,
7001 o1: &mut [f32],
7002 o2: &mut [f32],
7003 pool: Option<&Pool>,
7004) {
7005 debug_assert_eq!(o1.len(), rows);
7006 debug_assert_eq!(o2.len(), rows);
7007 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7008 let gpr = cols / GROUP_SIZE;
7009 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7010 let out1 = SendMut(o1.as_mut_ptr());
7011 let out2 = SendMut(o2.as_mut_ptr());
7012 if a8w8_enabled() {
7013 let a1 = split_act(x1);
7014 let a2 = split_act(x2);
7015 let (a1, a2) = (&a1, &a2);
7016 let run = move |start: usize, end: usize| {
7017 for r in start..end {
7018 #[cfg(target_arch = "aarch64")]
7019 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7022 #[cfg(not(target_arch = "aarch64"))]
7023 let ds = [
7024 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7025 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7026 ];
7027 let mut acc1 = ds[0] * a1.sx;
7028 for &(j, xv) in &a1.outliers {
7029 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7030 }
7031 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7032 let mut acc2 = ds[1] * a2.sx;
7033 for &(j, xv) in &a2.outliers {
7034 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7035 }
7036 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7037 unsafe {
7039 *out1.at(r) = acc1;
7040 *out2.at(r) = acc2;
7041 }
7042 }
7043 };
7044 dispatch_rows(pool, rows, &run);
7045 return;
7046 }
7047 let run = move |start: usize, end: usize| {
7048 let mut sg = [0f32; GROUP_SIZE];
7051 for r in start..end {
7052 let mut acc1 = 0f32;
7053 let mut acc2 = 0f32;
7054 for g in 0..gpr {
7055 let off = (r * gpr + g) * TILE;
7056 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7057 let codes = &bytes[off + 2..off + TILE];
7058 for bi in 0..6 {
7059 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7060 }
7061 let lut = &SIGN5[codes[6] as usize];
7062 sg[30] = lut[0];
7063 sg[31] = lut[1];
7064 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7065 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7066 let mut gsum1 = 0f32;
7067 for k in 0..GROUP_SIZE {
7068 gsum1 += sg[k] * xg1[k];
7069 }
7070 acc1 += s * gsum1;
7071 let mut gsum2 = 0f32;
7072 for k in 0..GROUP_SIZE {
7073 gsum2 += sg[k] * xg2[k];
7074 }
7075 acc2 += s * gsum2;
7076 }
7077 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7078 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7079 unsafe {
7081 *out1.at(r) = acc1;
7082 *out2.at(r) = acc2;
7083 }
7084 }
7085 };
7086 dispatch_rows(pool, rows, &run);
7087}
7088
7089fn q1t_matmat(
7092 bytes: &[u8],
7093 xs: &[f32],
7094 b: usize,
7095 rows: usize,
7096 cols: usize,
7097 out: &mut [f32],
7098 pool: Option<&Pool>,
7099) {
7100 debug_assert_eq!(out.len(), b * rows);
7101 const TILE: usize = cortiq_core::quant::Q1T_TILE;
7102 let gpr = cols / GROUP_SIZE;
7103 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7104 let out_addr = SendMut(out.as_mut_ptr());
7105 if a8w8_enabled() {
7109 let acts: Vec<SplitAct> = (0..b)
7110 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7111 .collect();
7112 let acts = &acts;
7113 let run = move |start: usize, end: usize| {
7114 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
7118 for g in 0..gpr {
7119 let off = (r * gpr + g) * TILE;
7120 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7121 q1t_unpack_group_i8(
7122 bytes.as_ptr().wrapping_add(off + 2),
7123 &mut sg[g * GROUP_SIZE..],
7124 );
7125 }
7126 for bi in 0..b {
7127 let act = &acts[bi];
7128 let mut isum = 0f32;
7129 for g in 0..gpr {
7130 let d = q1t_i8dot32(
7131 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7132 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7133 );
7134 isum += d as f32 * sc[g];
7135 }
7136 let mut acc = isum * act.sx;
7137 for &(j, xv) in &act.outliers {
7138 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7139 }
7140 accs[bi] = acc;
7141 }
7142 if has_ov {
7146 let (c0, c1) = (
7147 q1t_rowptr(bytes, rp_off, r),
7148 q1t_rowptr(bytes, rp_off, r + 1),
7149 );
7150 for p in c0..c1 {
7151 let e = ent_off + p * 4;
7152 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7153 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7154 for bi in 0..b {
7155 accs[bi] += val * xs[bi * cols + col];
7156 }
7157 }
7158 }
7159 for bi in 0..b {
7160 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7161 }
7162 }
7163 };
7164 dispatch_rows(pool, rows, &run);
7165 return;
7166 }
7167 let run = move |start: usize, end: usize| {
7168 let mut buf = vec![0f32; cols];
7169 for r in start..end {
7170 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7171 for bi in 0..b {
7172 let xr = &xs[bi * cols..(bi + 1) * cols];
7173 let mut acc = 0f32;
7174 for j in 0..cols {
7175 acc += buf[j] * xr[j];
7176 }
7177 unsafe { *out_addr.at(bi * rows + r) = acc };
7178 }
7179 }
7180 };
7181 dispatch_rows(pool, rows, &run);
7182}
7183
7184fn q1_matvec(
7185 bytes: &[u8],
7186 x: &[f32],
7187 rows: usize,
7188 cols: usize,
7189 out: &mut [f32],
7190 pool: Option<&Pool>,
7191) {
7192 debug_assert_eq!(out.len(), rows);
7193 let gpr = cols / GROUP_SIZE;
7194 let out_addr = SendMut(out.as_mut_ptr());
7195 if a8w8_enabled() {
7196 let act = split_act(x);
7197 let gsum = q1_group_sums(&act.xq, gpr);
7198 let (act, gsum) = (&act, &gsum);
7199 let run = move |start: usize, end: usize| {
7200 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7201 };
7202 dispatch_rows(pool, rows, &run);
7203 return;
7204 }
7205 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7206 dispatch_rows(pool, rows, &run);
7207}
7208
7209#[allow(clippy::too_many_arguments)]
7211fn q1_matvec2(
7212 bytes: &[u8],
7213 x1: &[f32],
7214 x2: &[f32],
7215 rows: usize,
7216 cols: usize,
7217 o1: &mut [f32],
7218 o2: &mut [f32],
7219 pool: Option<&Pool>,
7220) {
7221 let gpr = cols / GROUP_SIZE;
7222 let p1 = SendMut(o1.as_mut_ptr());
7223 let p2 = SendMut(o2.as_mut_ptr());
7224 if a8w8_enabled() {
7225 let a1 = split_act(x1);
7226 let a2 = split_act(x2);
7227 let g1 = q1_group_sums(&a1.xq, gpr);
7228 let g2 = q1_group_sums(&a2.xq, gpr);
7229 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7230 let run = move |start: usize, end: usize| {
7231 for r in start..end {
7232 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7233 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7234 for &(j, xv) in &a1.outliers {
7235 let (w, s) = q1_outlier(bytes, r, gpr, j);
7236 v1 += w * s * xv;
7237 }
7238 for &(j, xv) in &a2.outliers {
7239 let (w, s) = q1_outlier(bytes, r, gpr, j);
7240 v2 += w * s * xv;
7241 }
7242 unsafe {
7244 *p1.at(r) = v1;
7245 *p2.at(r) = v2;
7246 }
7247 }
7248 };
7249 dispatch_rows(pool, rows, &run);
7250 return;
7251 }
7252 let run = move |start: usize, end: usize| {
7253 for r in start..end {
7254 unsafe {
7256 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7257 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7258 }
7259 }
7260 };
7261 dispatch_rows(pool, rows, &run);
7262}
7263
7264#[allow(clippy::too_many_arguments)]
7266fn q1_matmat(
7267 bytes: &[u8],
7268 xs_all: &[f32],
7269 b: usize,
7270 rows: usize,
7271 cols: usize,
7272 out: &mut [f32],
7273 pool: Option<&Pool>,
7274) {
7275 debug_assert_eq!(out.len(), b * rows);
7276 let gpr = cols / GROUP_SIZE;
7277 let out_addr = SendMut(out.as_mut_ptr());
7278 if a8w8_enabled() {
7279 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7280 .map(|bi| {
7281 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7282 let gsum = q1_group_sums(&act.xq, gpr);
7283 (act, gsum)
7284 })
7285 .collect();
7286 let acts = &acts;
7287 #[cfg(target_arch = "x86_64")]
7288 let blocked_ok = avx2_enabled()
7289 && blocked_enabled();
7290 #[cfg(target_arch = "aarch64")]
7291 let blocked_ok = sdot_enabled()
7292 && blocked_enabled();
7293 let run = move |start: usize, end: usize| {
7294 for r in start..end {
7295 let mut bi = 0usize;
7296 #[cfg(target_arch = "aarch64")]
7299 if blocked_ok {
7300 while bi + 4 <= acts.len() {
7301 let xs = [
7302 acts[bi].0.xq.as_slice(),
7303 acts[bi + 1].0.xq.as_slice(),
7304 acts[bi + 2].0.xq.as_slice(),
7305 acts[bi + 3].0.xq.as_slice(),
7306 ];
7307 let gs = [
7308 acts[bi].1.as_slice(),
7309 acts[bi + 1].1.as_slice(),
7310 acts[bi + 2].1.as_slice(),
7311 acts[bi + 3].1.as_slice(),
7312 ];
7313 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7314 for k in 0..4 {
7315 let (act, _) = &acts[bi + k];
7316 let mut acc = d[k] * act.sx;
7317 for &(j, xv) in &act.outliers {
7318 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7319 acc += w * sc * xv;
7320 }
7321 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7323 }
7324 bi += 4;
7325 }
7326 }
7327 #[cfg(target_arch = "x86_64")]
7328 if blocked_ok {
7329 while bi + 4 <= acts.len() {
7330 let xs = [
7331 acts[bi].0.xq.as_slice(),
7332 acts[bi + 1].0.xq.as_slice(),
7333 acts[bi + 2].0.xq.as_slice(),
7334 acts[bi + 3].0.xq.as_slice(),
7335 ];
7336 let gs = [
7337 acts[bi].1.as_slice(),
7338 acts[bi + 1].1.as_slice(),
7339 acts[bi + 2].1.as_slice(),
7340 acts[bi + 3].1.as_slice(),
7341 ];
7342 let d = unsafe {
7343 if vnni_tiles_enabled() {
7344 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7345 } else {
7346 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7347 }
7348 };
7349 for k in 0..4 {
7350 let (act, _) = &acts[bi + k];
7351 let mut acc = d[k] * act.sx;
7352 for &(j, xv) in &act.outliers {
7353 let (w, sc) = q1_outlier(bytes, r, gpr, j);
7354 acc += w * sc * xv;
7355 }
7356 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7358 }
7359 bi += 4;
7360 }
7361 }
7362 while bi < acts.len() {
7363 let (act, gsum) = &acts[bi];
7364 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7365 for &(j, xv) in &act.outliers {
7366 let (w, s) = q1_outlier(bytes, r, gpr, j);
7367 acc += w * s * xv;
7368 }
7369 unsafe { *out_addr.at(bi * rows + r) = acc };
7371 bi += 1;
7372 }
7373 }
7374 };
7375 dispatch_rows(pool, rows, &run);
7376 return;
7377 }
7378 let run = move |start: usize, end: usize| {
7379 for r in start..end {
7380 for bi in 0..b {
7381 let x = &xs_all[bi * cols..(bi + 1) * cols];
7382 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7384 }
7385 }
7386 };
7387 dispatch_rows(pool, rows, &run);
7388}
7389
7390fn q4matvec(
7396 bytes: &[u8],
7397 x: &[f32],
7398 rows: usize,
7399 cols: usize,
7400 out: &mut [f32],
7401 pool: Option<&Pool>,
7402) {
7403 debug_assert_eq!(out.len(), rows);
7404 let (packed, scales) = q4_split(bytes, rows, cols);
7405 let gpr = cols / GROUP_SIZE;
7406 let out_addr = SendMut(out.as_mut_ptr());
7407
7408 if a8w8_enabled() {
7409 let act = split_act(x);
7410 let run = move |start: usize, end: usize| {
7411 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7412 };
7413 dispatch_rows(pool, rows, &run);
7414 return;
7415 }
7416
7417 let run =
7418 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7419 dispatch_rows(pool, rows, &run);
7420}
7421
7422#[inline]
7425#[allow(unreachable_code)]
7426#[cfg(target_arch = "x86_64")]
7431#[target_feature(enable = "avx2")]
7432unsafe fn dot_q4b_row_1x4_avx2(
7433 buf: &[u8],
7434 scales: &[u8],
7435 g0: usize,
7436 gpr: usize,
7437 xs: [&[i8]; 4],
7438) -> [f32; 4] {
7439 unsafe {
7441 use core::arch::x86_64::*;
7442 let ones = _mm256_set1_epi16(1);
7443 let mut acc = [0f32; 4];
7444 for gi in 0..gpr {
7445 let s = f16_to_f32(u16::from_le_bytes([
7446 scales[(g0 + gi) * 2],
7447 scales[(g0 + gi) * 2 + 1],
7448 ]));
7449 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7450 let aw = _mm256_abs_epi8(w);
7451 for (k, xq) in xs.iter().enumerate() {
7452 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7453 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7454 let d = _mm256_madd_epi16(p16, ones);
7455 let hi128 = _mm256_extracti128_si256::<1>(d);
7456 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7457 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7458 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7459 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7460 }
7461 }
7462 acc
7463 }
7464}
7465
7466#[cfg(target_arch = "x86_64")]
7468#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7469unsafe fn dot_q4b_row_1x4_vnni(
7470 buf: &[u8],
7471 scales: &[u8],
7472 g0: usize,
7473 gpr: usize,
7474 xs: [&[i8]; 4],
7475) -> [f32; 4] {
7476 unsafe {
7478 use core::arch::x86_64::*;
7479 let mut acc = [0f32; 4];
7480 for gi in 0..gpr {
7481 let s = f16_to_f32(u16::from_le_bytes([
7482 scales[(g0 + gi) * 2],
7483 scales[(g0 + gi) * 2 + 1],
7484 ]));
7485 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7486 let aw = _mm256_abs_epi8(w);
7487 for (k, xq) in xs.iter().enumerate() {
7488 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7489 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7490 acc[k] += d as f32 * s;
7491 }
7492 }
7493 acc
7494 }
7495}
7496
7497#[cfg(target_arch = "x86_64")]
7503#[target_feature(enable = "avx2")]
7504unsafe fn dot_q4b_row_1x4_sx_avx2(
7505 buf: &[u8],
7506 scales: &[u8],
7507 g0: usize,
7508 gpr: usize,
7509 xs: [&[i8]; 4],
7510 sxs: [f32; 4],
7511) -> [f32; 4] {
7512 unsafe {
7514 use core::arch::x86_64::*;
7515 let ones = _mm256_set1_epi16(1);
7516 let mut acc = [0f32; 4];
7517 for gi in 0..gpr {
7518 let s = f16_to_f32(u16::from_le_bytes([
7519 scales[(g0 + gi) * 2],
7520 scales[(g0 + gi) * 2 + 1],
7521 ]));
7522 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7523 let aw = _mm256_abs_epi8(w);
7524 for (k, xq) in xs.iter().enumerate() {
7525 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7526 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7527 let d = _mm256_madd_epi16(p16, ones);
7528 let hi128 = _mm256_extracti128_si256::<1>(d);
7529 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7530 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7531 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7532 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7533 }
7534 }
7535 acc
7536 }
7537}
7538
7539#[cfg(target_arch = "x86_64")]
7542#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7543unsafe fn dot_q4b_row_1x4_sx_vnni(
7544 buf: &[u8],
7545 scales: &[u8],
7546 g0: usize,
7547 gpr: usize,
7548 xs: [&[i8]; 4],
7549 sxs: [f32; 4],
7550) -> [f32; 4] {
7551 unsafe {
7553 use core::arch::x86_64::*;
7554 let mut acc = [0f32; 4];
7555 for gi in 0..gpr {
7556 let s = f16_to_f32(u16::from_le_bytes([
7557 scales[(g0 + gi) * 2],
7558 scales[(g0 + gi) * 2 + 1],
7559 ]));
7560 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7561 let aw = _mm256_abs_epi8(w);
7562 for (k, xq) in xs.iter().enumerate() {
7563 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7564 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7565 acc[k] += (d as f32 * sxs[k]) * s;
7566 }
7567 }
7568 acc
7569 }
7570}
7571
7572#[allow(unreachable_code)]
7573fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7574 #[cfg(target_arch = "aarch64")]
7575 unsafe {
7576 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7577 }
7578 #[cfg(target_arch = "x86_64")]
7579 unsafe {
7580 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7581 }
7582 let mut acc = 0f32;
7583 for gi in 0..gpr {
7584 let g = g0 + gi;
7585 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7586 let mut d = 0i32;
7587 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7588 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7589 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7590 }
7591 acc += d as f32 * s;
7592 }
7593 acc
7594}
7595
7596#[inline]
7598#[allow(unreachable_code)]
7599fn dot_q4_row_i8_2(
7600 packed: &[u8],
7601 scales: &[u8],
7602 g0: usize,
7603 gpr: usize,
7604 xq1: &[i8],
7605 xq2: &[i8],
7606) -> (f32, f32) {
7607 #[cfg(target_arch = "aarch64")]
7608 unsafe {
7609 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7610 }
7611 #[cfg(target_arch = "x86_64")]
7612 unsafe {
7613 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7614 }
7615 (
7616 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7617 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7618 )
7619}
7620
7621#[allow(clippy::too_many_arguments)]
7624fn q4_range_a8w8(
7625 packed: &[u8],
7626 scales: &[u8],
7627 gpr: usize,
7628 cols: usize,
7629 act: &SplitAct,
7630 out: SendMut,
7631 start: usize,
7632 end: usize,
7633) {
7634 for r in start..end {
7635 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7636 for &(j, xv) in &act.outliers {
7638 let flat = r * cols + j;
7639 let byte = packed[flat / 2];
7640 let nib = if flat & 1 == 0 {
7641 byte & 0x0F
7642 } else {
7643 byte >> 4
7644 };
7645 let s = f16_to_f32(u16::from_le_bytes([
7646 scales[(flat / GROUP_SIZE) * 2],
7647 scales[(flat / GROUP_SIZE) * 2 + 1],
7648 ]));
7649 acc += ((nib as i32 - 8) as f32) * s * xv;
7650 }
7651 unsafe { *out.at(r) = acc };
7653 }
7654}
7655
7656#[allow(clippy::too_many_arguments)]
7659fn q4_range2_a8w8(
7660 packed: &[u8],
7661 scales: &[u8],
7662 gpr: usize,
7663 cols: usize,
7664 a1: &SplitAct,
7665 a2: &SplitAct,
7666 p1: SendMut,
7667 p2: SendMut,
7668 start: usize,
7669 end: usize,
7670) {
7671 for r in start..end {
7672 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7673 let mut acc1 = s1 * a1.sx;
7674 let mut acc2 = s2 * a2.sx;
7675 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7677 for &(j, xv) in outliers {
7678 let flat = r * cols + j;
7679 let byte = packed[flat / 2];
7680 let nib = if flat & 1 == 0 {
7681 byte & 0x0F
7682 } else {
7683 byte >> 4
7684 };
7685 let s = f16_to_f32(u16::from_le_bytes([
7686 scales[(flat / GROUP_SIZE) * 2],
7687 scales[(flat / GROUP_SIZE) * 2 + 1],
7688 ]));
7689 *acc += ((nib as i32 - 8) as f32) * s * xv;
7690 }
7691 };
7692 fix(&a1.outliers, &mut acc1);
7693 fix(&a2.outliers, &mut acc2);
7694 unsafe {
7696 *p1.at(r) = acc1;
7697 *p2.at(r) = acc2;
7698 }
7699 }
7700}
7701
7702fn q4_range_f32(
7704 packed: &[u8],
7705 scales: &[u8],
7706 gpr: usize,
7707 x: &[f32],
7708 out: SendMut,
7709 start: usize,
7710 end: usize,
7711) {
7712 for r in start..end {
7713 let mut acc = 0f32;
7714 for gi in 0..gpr {
7715 let g = r * gpr + gi;
7716 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7717 let pk = &packed[g * 16..(g + 1) * 16];
7718 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7719 let mut ga = 0f32;
7720 for (k, &b) in pk.iter().enumerate() {
7721 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7722 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7723 }
7724 acc += ga * s;
7725 }
7726 unsafe { *out.at(r) = acc };
7728 }
7729}
7730
7731#[allow(clippy::too_many_arguments)]
7735fn q4matvec2(
7736 bytes: &[u8],
7737 x1: &[f32],
7738 x2: &[f32],
7739 rows: usize,
7740 cols: usize,
7741 o1: &mut [f32],
7742 o2: &mut [f32],
7743 pool: Option<&Pool>,
7744) {
7745 debug_assert_eq!(o1.len(), rows);
7746 debug_assert_eq!(o2.len(), rows);
7747 let (packed, scales) = q4_split(bytes, rows, cols);
7748 let gpr = cols / GROUP_SIZE;
7749
7750 if a8w8_enabled() {
7751 let a1 = split_act(x1);
7752 let a2 = split_act(x2);
7753 let p1 = SendMut(o1.as_mut_ptr());
7754 let p2 = SendMut(o2.as_mut_ptr());
7755 let run = move |start: usize, end: usize| {
7756 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7757 };
7758 dispatch_rows(pool, rows, &run);
7759 return;
7760 }
7761
7762 let p1 = SendMut(o1.as_mut_ptr());
7763 let p2 = SendMut(o2.as_mut_ptr());
7764 let run = move |start: usize, end: usize| {
7765 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7766 };
7767 dispatch_rows(pool, rows, &run);
7768}
7769
7770#[allow(clippy::too_many_arguments)]
7772fn q4_range2_f32(
7773 packed: &[u8],
7774 scales: &[u8],
7775 gpr: usize,
7776 x1: &[f32],
7777 x2: &[f32],
7778 p1: SendMut,
7779 p2: SendMut,
7780 start: usize,
7781 end: usize,
7782) {
7783 for r in start..end {
7784 let (mut acc1, mut acc2) = (0f32, 0f32);
7785 for gi in 0..gpr {
7786 let g = r * gpr + gi;
7787 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7788 let pk = &packed[g * 16..(g + 1) * 16];
7789 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7790 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7791 let (mut g1, mut g2) = (0f32, 0f32);
7792 for (k, &b) in pk.iter().enumerate() {
7793 let wl = (b & 0x0F) as f32 - 8.0;
7794 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7795 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7796 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7797 }
7798 acc1 += g1 * s;
7799 acc2 += g2 * s;
7800 }
7801 unsafe {
7803 *p1.at(r) = acc1;
7804 *p2.at(r) = acc2;
7805 }
7806 }
7807}
7808
7809thread_local! {
7810 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7813 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7814}
7815
7816#[allow(clippy::too_many_arguments)]
7822fn q4matmat(
7823 bytes: &[u8],
7824 xs_all: &[f32],
7825 b: usize,
7826 rows: usize,
7827 cols: usize,
7828 out: &mut [f32],
7829 pool: Option<&Pool>,
7830) {
7831 debug_assert_eq!(xs_all.len(), b * cols);
7832 debug_assert_eq!(out.len(), b * rows);
7833 let (packed, scales) = q4_split(bytes, rows, cols);
7834 let gpr = cols / GROUP_SIZE;
7835 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7836
7837 if a8w8_enabled() {
7838 let acts: Vec<SplitAct> = (0..b)
7839 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7840 .collect();
7841 let acts = &acts;
7842 let out_addr = SendMut(out.as_mut_ptr());
7843 let run = move |start: usize, end: usize| {
7844 ROW_I8.with(|rb| {
7845 let mut buf = rb.borrow_mut();
7846 buf.resize(cols, 0);
7847 for r in start..end {
7848 for gi in 0..gpr {
7852 let g = r * gpr + gi;
7853 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7854 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7855 buf[gi * GROUP_SIZE + k * 2 + 1] =
7856 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7857 }
7858 }
7859 let mut bi = 0usize;
7860 #[cfg(target_arch = "x86_64")]
7861 if avx2_enabled()
7862 && blocked_enabled()
7863 {
7864 while bi + 4 <= acts.len() {
7865 let xs = [
7866 acts[bi].xq.as_slice(),
7867 acts[bi + 1].xq.as_slice(),
7868 acts[bi + 2].xq.as_slice(),
7869 acts[bi + 3].xq.as_slice(),
7870 ];
7871 let d = unsafe {
7872 if vnni_tiles_enabled() {
7873 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7874 } else {
7875 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7876 }
7877 };
7878 for k in 0..4 {
7879 let act = &acts[bi + k];
7880 let mut acc = d[k] * act.sx;
7881 for &(j, xv) in &act.outliers {
7882 acc += (buf[j] as i8) as f32
7883 * gscale((r * cols + j) / GROUP_SIZE)
7884 * xv;
7885 }
7886 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7888 }
7889 bi += 4;
7890 }
7891 }
7892 while bi < acts.len() {
7893 let act = &acts[bi];
7894 let mut acc = 0f32;
7895 for gi in 0..gpr {
7896 let d = dot_i8_i8(
7897 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7898 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7899 );
7900 acc += d as f32 * gscale(r * gpr + gi);
7901 }
7902 acc *= act.sx;
7903 for &(j, xv) in &act.outliers {
7905 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7906 }
7907 unsafe { *out_addr.at(bi * rows + r) = acc };
7909 bi += 1;
7910 }
7911 }
7912 })
7913 };
7914 dispatch_rows(pool, rows, &run);
7915 return;
7916 }
7917
7918 let out_addr = SendMut(out.as_mut_ptr());
7919 let run = move |start: usize, end: usize| {
7920 ROW_F32.with(|rb| {
7921 let mut buf = rb.borrow_mut();
7922 buf.resize(cols, 0.0);
7923 for r in start..end {
7924 for gi in 0..gpr {
7927 let g = r * gpr + gi;
7928 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7929 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7930 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7931 }
7932 }
7933 for bi in 0..b {
7934 let x = &xs_all[bi * cols..(bi + 1) * cols];
7935 let mut acc = 0f32;
7936 for gi in 0..gpr {
7937 let mut ga = 0f32;
7938 for k in 0..GROUP_SIZE / 2 {
7943 let e = gi * GROUP_SIZE + k * 2;
7944 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
7945 }
7946 acc += ga * gscale(r * gpr + gi);
7947 }
7948 unsafe { *out_addr.at(bi * rows + r) = acc };
7950 }
7951 }
7952 })
7953 };
7954 dispatch_rows(pool, rows, &run);
7955}
7956
7957#[allow(clippy::too_many_arguments)]
7962fn vbitmatmat(
7963 bytes: &[u8],
7964 offsets: &[usize],
7965 xs_all: &[f32],
7966 b: usize,
7967 rows: usize,
7968 cols: usize,
7969 out: &mut [f32],
7970 pool: Option<&Pool>,
7971) {
7972 debug_assert_eq!(xs_all.len(), b * cols);
7973 debug_assert_eq!(out.len(), b * rows);
7974 debug_assert_eq!(offsets.len(), rows + 1);
7975 let ng = cols / GROUP_SIZE;
7976 let bits = &bytes[..rows];
7977 let sc_off = rows;
7978 let gscale = |r: usize, g: usize| {
7979 let so = (r * ng + g) * 2;
7980 f16_to_f32(u16::from_le_bytes([
7981 bytes[sc_off + so],
7982 bytes[sc_off + so + 1],
7983 ]))
7984 };
7985
7986 let decode_f32 = |r: usize, dst: &mut [f32]| {
7988 let bw = bits[r] as usize;
7989 let l = ((1i32 << (bw - 1)) - 1) as f32;
7990 let data = &bytes[offsets[r]..offsets[r + 1]];
7991 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
7992 for d in dst.iter_mut() {
7993 while nbits < bw {
7994 acc = (acc << 8) | data[idx] as u64;
7995 idx += 1;
7996 nbits += 8;
7997 }
7998 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
7999 nbits -= bw;
8000 *d = u - l;
8001 }
8002 };
8003
8004 if a8w8_enabled() {
8005 let acts: Vec<SplitAct> = (0..b)
8006 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8007 .collect();
8008 let acts = &acts;
8009 let out_addr = SendMut(out.as_mut_ptr());
8010 let run = move |start: usize, end: usize| {
8011 for r in start..end {
8012 let bw = bits[r] as usize;
8013 if bw == 8 {
8014 ROW_F32.with(|rb| {
8017 let mut buf = rb.borrow_mut();
8018 buf.resize(cols, 0.0);
8019 decode_f32(r, &mut buf);
8020 for bi in 0..b {
8021 let x = &xs_all[bi * cols..(bi + 1) * cols];
8022 let mut dot = 0f32;
8023 for g in 0..ng {
8024 let mut gd = 0f32;
8025 for k in 0..GROUP_SIZE {
8026 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8027 }
8028 dot += gd * gscale(r, g);
8029 }
8030 unsafe { *out_addr.at(bi * rows + r) = dot };
8032 }
8033 });
8034 continue;
8035 }
8036 let l = (1i32 << (bw - 1)) - 1;
8037 let data = &bytes[offsets[r]..offsets[r + 1]];
8038 ROW_I8.with(|rb| {
8039 let mut buf = rb.borrow_mut();
8040 buf.resize(cols, 0);
8041 #[inline(always)]
8042 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8043 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8044 let u = unpack8::<B>(&data[blk * B..]);
8045 for k in 0..8 {
8046 chunk[k] = (u[k] - l) as i8 as u8;
8047 }
8048 }
8049 }
8050 match bw {
8051 3 => fill::<3>(data, l, &mut buf),
8052 4 => vbit_fill4(data, &mut buf),
8053 5 => fill::<5>(data, l, &mut buf),
8054 6 => fill::<6>(data, l, &mut buf),
8055 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8056 }
8057 let mut bi = 0usize;
8058 #[cfg(target_arch = "x86_64")]
8062 if avx2_enabled()
8063 && blocked_enabled()
8064 {
8065 while bi + 4 <= acts.len() {
8066 let xs = [
8067 acts[bi].xq.as_slice(),
8068 acts[bi + 1].xq.as_slice(),
8069 acts[bi + 2].xq.as_slice(),
8070 acts[bi + 3].xq.as_slice(),
8071 ];
8072 let sxs = [
8073 acts[bi].sx,
8074 acts[bi + 1].sx,
8075 acts[bi + 2].sx,
8076 acts[bi + 3].sx,
8077 ];
8078 let d = unsafe {
8079 if vnni_tiles_enabled() {
8080 dot_q4b_row_1x4_sx_vnni(
8081 &buf,
8082 &bytes[sc_off..],
8083 r * ng,
8084 ng,
8085 xs,
8086 sxs,
8087 )
8088 } else {
8089 dot_q4b_row_1x4_sx_avx2(
8090 &buf,
8091 &bytes[sc_off..],
8092 r * ng,
8093 ng,
8094 xs,
8095 sxs,
8096 )
8097 }
8098 };
8099 for k in 0..4 {
8100 let act = &acts[bi + k];
8101 let mut dot = d[k];
8102 for &(j, xv) in &act.outliers {
8103 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8104 }
8105 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8107 }
8108 bi += 4;
8109 }
8110 }
8111 while bi < acts.len() {
8112 let act = &acts[bi];
8113 let mut dot = 0f32;
8114 for g in 0..ng {
8115 let d = dot_i8_i8(
8116 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8117 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8118 ) as f32
8119 * act.sx;
8120 dot += d * gscale(r, g);
8121 }
8122 for &(j, xv) in &act.outliers {
8123 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8124 }
8125 unsafe { *out_addr.at(bi * rows + r) = dot };
8127 bi += 1;
8128 }
8129 });
8130 }
8131 };
8132 dispatch_rows(pool, rows, &run);
8133 return;
8134 }
8135
8136 let out_addr = SendMut(out.as_mut_ptr());
8137 let run = move |start: usize, end: usize| {
8138 ROW_F32.with(|rb| {
8139 let mut buf = rb.borrow_mut();
8140 buf.resize(cols, 0.0);
8141 for r in start..end {
8142 decode_f32(r, &mut buf);
8143 for bi in 0..b {
8144 let x = &xs_all[bi * cols..(bi + 1) * cols];
8145 let mut dot = 0f32;
8146 for g in 0..ng {
8147 let mut gd = 0f32;
8148 for k in 0..GROUP_SIZE {
8149 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8150 }
8151 dot += gd * gscale(r, g);
8152 }
8153 unsafe { *out_addr.at(bi * rows + r) = dot };
8155 }
8156 }
8157 })
8158 };
8159 dispatch_rows(pool, rows, &run);
8160}
8161
8162pub(crate) fn gpu_batch_job<'a>(
8166 t: &'a QTensor,
8167 x: &[f32],
8168) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8169 match t {
8170 QTensor::Mapped {
8171 model,
8172 idx,
8173 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8174 rows,
8175 cols,
8176 row_scale,
8177 col_field,
8178 ..
8179 } => Some((
8180 model.clone(),
8181 crate::gpu::BatchJob {
8182 idx: *idx,
8183 rows: *rows,
8184 cols: *cols,
8185 row_scale,
8186 xs: prescale(x, col_field, *dt).into_owned(),
8187 layout: crate::gpu::BatchLayout::Q8,
8188 },
8189 )),
8190 QTensor::Mapped {
8192 model,
8193 idx,
8194 dtype: TensorDtype::Q1,
8195 rows,
8196 cols,
8197 ..
8198 } => Some((
8199 model.clone(),
8200 crate::gpu::BatchJob {
8201 idx: *idx,
8202 rows: *rows,
8203 cols: *cols,
8204 row_scale: &[],
8205 xs: x.to_vec(),
8206 layout: crate::gpu::BatchLayout::Q1,
8207 },
8208 )),
8209 QTensor::Mapped {
8214 model,
8215 idx,
8216 dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8217 rows,
8218 cols,
8219 ..
8220 } => Some((
8221 model.clone(),
8222 crate::gpu::BatchJob {
8223 idx: *idx,
8224 rows: *rows,
8225 cols: *cols,
8226 row_scale: &[],
8227 xs: x.to_vec(),
8228 layout: if *dt == TensorDtype::Q4Tiled {
8229 crate::gpu::BatchLayout::Q4t
8230 } else {
8231 crate::gpu::BatchLayout::Q4tp
8232 },
8233 },
8234 )),
8235 _ => None,
8236 }
8237}
8238
8239thread_local! {
8240 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8241 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8242}
8243
8244pub(crate) fn prescale<'a>(
8245 x: &'a [f32],
8246 col_field: &[f32],
8247 dtype: TensorDtype,
8248) -> std::borrow::Cow<'a, [f32]> {
8249 if dtype == TensorDtype::Q8_2f {
8250 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8251 } else {
8252 std::borrow::Cow::Borrowed(x)
8253 }
8254}
8255
8256pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8259 x: &[f32],
8260 col_field: &[f32],
8261 dtype: TensorDtype,
8262 buf_id: u8,
8263 f: F,
8264) -> R {
8265 if dtype == TensorDtype::Q8_2f {
8266 if buf_id == 1 {
8267 PRESCALE_BUF1.with(|b| {
8268 let mut buf = b.borrow_mut();
8269 buf.clear();
8270 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8271 f(&buf)
8272 })
8273 } else {
8274 PRESCALE_BUF2.with(|b| {
8275 let mut buf = b.borrow_mut();
8276 buf.clear();
8277 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8278 f(&buf)
8279 })
8280 }
8281 } else {
8282 f(x)
8283 }
8284}
8285
8286#[cfg(target_arch = "x86_64")]
8291pub(crate) fn avx2_enabled() -> bool {
8292 use std::sync::OnceLock;
8293 static ON: OnceLock<bool> = OnceLock::new();
8294 *ON.get_or_init(|| {
8295 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8296 && std::arch::is_x86_feature_detected!("avx2")
8297 && std::arch::is_x86_feature_detected!("fma")
8298 })
8299}
8300
8301#[cfg(target_arch = "x86_64")]
8306fn avx2_a8w8_enabled() -> bool {
8307 use std::sync::OnceLock;
8308 static ON: OnceLock<bool> = OnceLock::new();
8309 *ON.get_or_init(|| {
8310 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8311 })
8312}
8313
8314#[inline]
8318pub(crate) fn a8w8_enabled() -> bool {
8319 #[cfg(target_arch = "aarch64")]
8320 {
8321 sdot_enabled()
8322 }
8323 #[cfg(target_arch = "x86_64")]
8324 {
8325 avx2_a8w8_enabled()
8326 }
8327 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8328 {
8329 false
8330 }
8331}
8332
8333#[inline]
8336#[allow(unreachable_code)]
8337fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8338 #[cfg(target_arch = "aarch64")]
8339 unsafe {
8340 return dot_i8_sdot(w, xq);
8341 }
8342 #[cfg(target_arch = "x86_64")]
8343 unsafe {
8344 if avx512vnni_enabled() {
8345 return dot_i8_i8_vnni(w, xq);
8346 }
8347 return dot_i8_i8_avx2(w, xq);
8348 }
8349 w.iter()
8350 .zip(xq)
8351 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8352 .sum()
8353}
8354
8355#[cfg(target_arch = "x86_64")]
8359fn avx512vnni_enabled() -> bool {
8360 use std::sync::OnceLock;
8361 static ON: OnceLock<bool> = OnceLock::new();
8362 *ON.get_or_init(|| {
8363 std::env::var("CMF_AVX512")
8364 .map(|v| v != "0")
8365 .unwrap_or(true)
8366 && std::arch::is_x86_feature_detected!("avx512f")
8367 && std::arch::is_x86_feature_detected!("avx512bw")
8368 && std::arch::is_x86_feature_detected!("avx512vl")
8369 && std::arch::is_x86_feature_detected!("avx512vnni")
8370 })
8371}
8372
8373#[cfg(target_arch = "x86_64")]
8381fn vnni_tiles_enabled() -> bool {
8382 use std::sync::OnceLock;
8383 static ON: OnceLock<bool> = OnceLock::new();
8384 *ON.get_or_init(|| {
8385 std::env::var("CMF_VNNI_TILES")
8386 .map(|v| v != "0")
8387 .unwrap_or(true)
8388 && avx512vnni_enabled()
8389 })
8390}
8391
8392#[cfg(target_arch = "x86_64")]
8397#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8398#[inline]
8399unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8400 unsafe {
8402 use core::arch::x86_64::*;
8403 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8404 let hi128 = _mm256_extracti128_si256::<1>(d);
8405 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8406 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8407 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8408 _mm_cvtsi128_si32(s32)
8409 }
8410}
8411
8412#[cfg(target_arch = "x86_64")]
8417#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8418unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8419 unsafe {
8421 use core::arch::x86_64::*;
8422 let n = w.len();
8423 let mut j = 0usize;
8424 let mut total: i32;
8425 {
8430 #[inline(always)]
8431 unsafe fn step(
8432 w: *const u8,
8433 x: *const i8,
8434 acc: core::arch::x86_64::__m512i,
8435 ) -> core::arch::x86_64::__m512i {
8436 unsafe {
8437 use core::arch::x86_64::*;
8438 let wv = _mm512_loadu_si512(w as *const _);
8439 let xv = _mm512_loadu_si512(x as *const _);
8440 let aw = _mm512_abs_epi8(wv);
8441 let neg = _mm512_movepi8_mask(wv);
8442 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8443 _mm512_dpbusd_epi32(acc, aw, sx)
8444 }
8445 }
8446 let (mut a0, mut a1, mut a2, mut a3) = (
8447 _mm512_setzero_si512(),
8448 _mm512_setzero_si512(),
8449 _mm512_setzero_si512(),
8450 _mm512_setzero_si512(),
8451 );
8452 while j + 256 <= n {
8453 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8454 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8455 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8456 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8457 j += 256;
8458 }
8459 while j + 64 <= n {
8460 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8461 j += 64;
8462 }
8463 let s01 = _mm512_add_epi32(a0, a1);
8464 let s23 = _mm512_add_epi32(a2, a3);
8465 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8466 }
8467 if j + 32 <= n {
8469 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8470 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8471 let d = _mm256_dpbusd_epi32(
8472 _mm256_setzero_si256(),
8473 _mm256_abs_epi8(wv),
8474 _mm256_sign_epi8(xv, wv),
8475 );
8476 let hi128 = _mm256_extracti128_si256::<1>(d);
8477 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8478 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8479 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8480 total += _mm_cvtsi128_si32(s32);
8481 j += 32;
8482 }
8483 while j < n {
8484 total += (w[j] as i8) as i32 * xq[j] as i32;
8485 j += 1;
8486 }
8487 total
8488 }
8489}
8490
8491#[cfg(target_arch = "x86_64")]
8493#[target_feature(enable = "avx2,fma")]
8494unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8495 unsafe {
8497 use core::arch::x86_64::*;
8498 let n = x.len();
8499 let wp = w.as_ptr();
8500 let xp = x.as_ptr();
8501 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8502 let mut j = 0usize;
8503 while j + 16 <= n {
8504 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8505 let lo = _mm256_cvtepi8_epi32(wb);
8506 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8507 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8508 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8509 j += 16;
8510 }
8511 let acc = _mm256_add_ps(a0, a1);
8512 let hi128 = _mm256_extractf128_ps::<1>(acc);
8513 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8514 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8515 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8516 let mut sum = _mm_cvtss_f32(s32);
8517 while j < n {
8518 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8519 j += 1;
8520 }
8521 sum
8522 }
8523}
8524
8525#[cfg(target_arch = "x86_64")]
8530#[target_feature(enable = "avx2")]
8531unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8532 unsafe {
8534 use core::arch::x86_64::*;
8535 let n = w.len();
8536 let ones = _mm256_set1_epi16(1);
8537 let mut acc = _mm256_setzero_si256();
8538 let mut j = 0usize;
8539 while j + 32 <= n {
8540 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8541 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8542 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8543 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8544 j += 32;
8545 }
8546 let hi128 = _mm256_extracti128_si256::<1>(acc);
8547 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8548 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8549 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8550 let mut s = _mm_cvtsi128_si32(s32);
8551 while j < n {
8552 s += (w[j] as i8) as i32 * xq[j] as i32;
8553 j += 1;
8554 }
8555 s
8556 }
8557}
8558
8559#[cfg(target_arch = "aarch64")]
8563#[target_feature(enable = "neon,i8mm")]
8564unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8565 unsafe {
8567 use core::arch::aarch64::*;
8568 use core::arch::asm;
8569 let n = w0.len();
8570 let w0p = w0.as_ptr() as *const i8;
8571 let w1p = w1.as_ptr() as *const i8;
8572 let mut acc01 = vdupq_n_s32(0);
8575 let mut acc23 = vdupq_n_s32(0);
8576 let mut i = 0usize;
8577 while i + 8 <= n {
8578 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8579 let xb01 = vcombine_s8(
8580 vld1_s8(xs[0].as_ptr().add(i)),
8581 vld1_s8(xs[1].as_ptr().add(i)),
8582 );
8583 let xb23 = vcombine_s8(
8584 vld1_s8(xs[2].as_ptr().add(i)),
8585 vld1_s8(xs[3].as_ptr().add(i)),
8586 );
8587 asm!(
8588 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8589 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8590 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8591 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8592 options(pure, nomem, nostack),
8593 );
8594 i += 8;
8595 }
8596 let mut out = [[0i32; 4]; 2];
8597 let a01: [i32; 4] = core::mem::transmute(acc01);
8598 let a23: [i32; 4] = core::mem::transmute(acc23);
8599 out[0][0] = a01[0];
8600 out[0][1] = a01[1];
8601 out[1][0] = a01[2];
8602 out[1][1] = a01[3];
8603 out[0][2] = a23[0];
8604 out[0][3] = a23[1];
8605 out[1][2] = a23[2];
8606 out[1][3] = a23[3];
8607 if i < n {
8608 for (k, x) in xs.iter().enumerate() {
8609 for j in i..n {
8610 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8611 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8612 }
8613 }
8614 }
8615 out
8616 }
8617}
8618
8619#[cfg(target_arch = "aarch64")]
8623#[target_feature(enable = "neon,dotprod")]
8624unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8625 unsafe {
8627 use core::arch::aarch64::*;
8628 use core::arch::asm;
8629 let n = w0.len();
8630 let w0p = w0.as_ptr() as *const i8;
8631 let w1p = w1.as_ptr() as *const i8;
8632 let mut acc = [[vdupq_n_s32(0); 4]; 2];
8633 let mut i = 0usize;
8634 while i + 16 <= n {
8635 let wv0 = vld1q_s8(w0p.add(i));
8636 let wv1 = vld1q_s8(w1p.add(i));
8637 for (k, x) in xs.iter().enumerate() {
8638 let xv = vld1q_s8(x.as_ptr().add(i));
8639 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8640 asm!(
8641 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8642 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8643 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8644 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8645 options(pure, nomem, nostack),
8646 );
8647 acc[0][k] = a0;
8648 acc[1][k] = a1;
8649 }
8650 i += 16;
8651 }
8652 let mut out = [[0i32; 4]; 2];
8653 for r in 0..2 {
8654 for k in 0..4 {
8655 out[r][k] = vaddvq_s32(acc[r][k]);
8656 }
8657 }
8658 if i < n {
8659 for (k, x) in xs.iter().enumerate() {
8660 for j in i..n {
8661 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8662 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8663 }
8664 }
8665 }
8666 out
8667 }
8668}
8669
8670#[cfg(target_arch = "x86_64")]
8676#[target_feature(enable = "avx2")]
8677unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8678 unsafe {
8680 use core::arch::x86_64::*;
8681 let n = w0.len();
8682 let ones = _mm256_set1_epi16(1);
8683 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8684 let mut j = 0usize;
8685 while j + 32 <= n {
8686 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8687 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8688 let aw0 = _mm256_abs_epi8(wv0);
8689 let aw1 = _mm256_abs_epi8(wv1);
8690 for (k, x) in xs.iter().enumerate() {
8691 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8692 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8693 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8694 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8695 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8696 }
8697 j += 32;
8698 }
8699 let mut out = [[0i32; 4]; 2];
8700 for r in 0..2 {
8701 for k in 0..4 {
8702 let a = acc[r][k];
8703 let hi128 = _mm256_extracti128_si256::<1>(a);
8704 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8705 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8706 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8707 out[r][k] = _mm_cvtsi128_si32(s32);
8708 }
8709 }
8710 if j < n {
8711 for (k, x) in xs.iter().enumerate() {
8712 for i in j..n {
8713 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8714 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8715 }
8716 }
8717 }
8718 out
8719 }
8720}
8721
8722#[cfg(target_arch = "x86_64")]
8727#[inline]
8728fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8729 let dot = if avx512vnni_enabled() && row.len() >= 64 {
8730 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8731 } else {
8732 unsafe { dot_i8_i8_avx2(row, &act.xq) }
8733 };
8734 let mut acc = dot as f32 * act.sx;
8735 for &(j, xv) in &act.outliers {
8736 acc += (row[j] as i8) as f32 * xv;
8737 }
8738 acc
8739}
8740
8741#[cfg(target_arch = "x86_64")]
8745#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8746unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8747 unsafe {
8749 use core::arch::x86_64::*;
8750 let n = w.len();
8751 let flip = _mm512_set1_epi8(-128); #[inline(always)]
8753 unsafe fn step(
8754 w: *const u8,
8755 x: *const i8,
8756 flip: core::arch::x86_64::__m512i,
8757 acc: core::arch::x86_64::__m512i,
8758 ) -> core::arch::x86_64::__m512i {
8759 unsafe {
8760 use core::arch::x86_64::*;
8761 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8762 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8763 }
8764 }
8765 let (mut a0, mut a1, mut a2, mut a3) = (
8766 _mm512_setzero_si512(),
8767 _mm512_setzero_si512(),
8768 _mm512_setzero_si512(),
8769 _mm512_setzero_si512(),
8770 );
8771 let mut j = 0usize;
8772 while j + 256 <= n {
8773 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8774 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8775 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8776 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8777 j += 256;
8778 }
8779 while j + 64 <= n {
8780 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8781 j += 64;
8782 }
8783 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8784 _mm512_add_epi32(a0, a1),
8785 _mm512_add_epi32(a2, a3),
8786 ));
8787 while j < n {
8789 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8790 j += 1;
8791 }
8792 total
8793 }
8794}
8795
8796#[cfg(target_arch = "x86_64")]
8802#[target_feature(enable = "avx2")]
8803unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8804 unsafe {
8807 use core::arch::x86_64::*;
8808 let lomask = _mm_set1_epi8(0x0F);
8809 let eight = _mm256_set1_epi8(8);
8810 let ones = _mm256_set1_epi16(1);
8811 let mut acc = 0f32;
8812 for gi in 0..gpr {
8813 let g = g0 + gi;
8814 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8815 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8816 let lo = _mm_and_si128(b, lomask);
8817 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8818 let w = _mm256_sub_epi8(
8819 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8820 eight,
8821 );
8822 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8823 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8824 let d = _mm256_madd_epi16(p16, ones);
8825 let hi128 = _mm256_extracti128_si256::<1>(d);
8826 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8827 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8828 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8829 acc += _mm_cvtsi128_si32(s32) as f32 * s;
8830 }
8831 acc
8832 }
8833}
8834
8835#[cfg(target_arch = "x86_64")]
8838#[target_feature(enable = "avx2")]
8839unsafe fn dot_q4_row_avx2_2(
8840 packed: &[u8],
8841 scales: &[u8],
8842 g0: usize,
8843 gpr: usize,
8844 xq1: &[i8],
8845 xq2: &[i8],
8846) -> (f32, f32) {
8847 unsafe {
8849 use core::arch::x86_64::*;
8850 let lomask = _mm_set1_epi8(0x0F);
8851 let eight = _mm256_set1_epi8(8);
8852 let ones = _mm256_set1_epi16(1);
8853 let (mut acc1, mut acc2) = (0f32, 0f32);
8854 #[inline(always)]
8855 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8856 unsafe {
8857 use core::arch::x86_64::*;
8858 let hi128 = _mm256_extracti128_si256::<1>(d);
8859 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8860 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8861 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8862 _mm_cvtsi128_si32(s32)
8863 }
8864 }
8865 for gi in 0..gpr {
8866 let g = g0 + gi;
8867 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8868 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8869 let lo = _mm_and_si128(b, lomask);
8870 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8871 let w = _mm256_sub_epi8(
8872 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8873 eight,
8874 );
8875 let aw = _mm256_abs_epi8(w);
8876 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8877 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8878 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8879 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8880 acc1 += hsum(d1) as f32 * s;
8881 acc2 += hsum(d2) as f32 * s;
8882 }
8883 (acc1, acc2)
8884 }
8885}
8886
8887#[cfg(target_arch = "x86_64")]
8889fn q8_range_avx2(
8890 q: &[u8],
8891 row_scale: &[f32],
8892 act: &SplitAct,
8893 cols: usize,
8894 out_addr: SendMut,
8895 start: usize,
8896 end: usize,
8897) {
8898 for o in start..end {
8899 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8900 unsafe { *out_addr.at(o) = v };
8902 }
8903}
8904
8905#[cfg(target_arch = "x86_64")]
8907#[allow(clippy::too_many_arguments)]
8908fn q8_range2_avx2(
8909 q: &[u8],
8910 row_scale: &[f32],
8911 a1: &SplitAct,
8912 a2: &SplitAct,
8913 cols: usize,
8914 p1: SendMut,
8915 p2: SendMut,
8916 start: usize,
8917 end: usize,
8918) {
8919 for o in start..end {
8920 let row = &q[o * cols..(o + 1) * cols];
8921 unsafe {
8923 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8924 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8925 }
8926 }
8927}
8928
8929#[cfg(target_arch = "aarch64")]
8940fn i8mm_enabled() -> bool {
8941 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8942 *ON.get_or_init(|| {
8943 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8944 && std::arch::is_aarch64_feature_detected!("i8mm")
8945 })
8946}
8947
8948#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
8952fn sdot_enabled() -> bool {
8953 use std::sync::OnceLock;
8954 static ON: OnceLock<bool> = OnceLock::new();
8955 *ON.get_or_init(|| {
8956 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
8957 if !want {
8958 return false;
8959 }
8960
8961 #[cfg(target_arch = "aarch64")]
8962 {
8963 if std::arch::is_aarch64_feature_detected!("dotprod") {
8964 return true;
8965 }
8966 #[cfg(target_os = "android")]
8967 {
8968 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
8969 if cpuinfo.lines().any(|l| {
8970 (l.starts_with("Features") || l.starts_with("features"))
8971 && l.contains("asimddp")
8972 }) {
8973 return true;
8974 }
8975 }
8976 }
8977 false
8978 }
8979 #[cfg(not(target_arch = "aarch64"))]
8980 {
8981 false
8982 }
8983 })
8984}
8985
8986struct SplitAct {
8991 xq: Vec<i8>,
8992 sx: f32,
8993 outliers: Vec<(usize, f32)>,
8994 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
8997 xsum: i32,
8998}
8999
9000thread_local! {
9001 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9004 const { std::cell::RefCell::new(Vec::new()) };
9005}
9006
9007impl Drop for SplitAct {
9008 fn drop(&mut self) {
9009 let buf = std::mem::take(&mut self.xq);
9010 if buf.capacity() > 0 {
9011 XQ_FREE.with(|f| {
9012 let mut f = f.borrow_mut();
9013 if f.len() < 16 {
9014 f.push(buf);
9015 }
9016 });
9017 }
9018 }
9019}
9020
9021thread_local! {
9022 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9029}
9030
9031#[inline]
9034fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9035 KROW.with(|s| {
9036 let mut b = s.borrow_mut();
9037 if b.len() < n {
9038 b.resize(n, 0.0);
9039 }
9040 f(&mut b[..n])
9041 })
9042}
9043
9044fn split_act(x: &[f32]) -> SplitAct {
9045 let n = x.len();
9046 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9047 let thr = 8.0 * rms;
9048 let mut outliers: Vec<(usize, f32)> = Vec::new();
9052 let mut amax = 0f32;
9053 for (j, &v) in x.iter().enumerate() {
9054 let a = v.abs();
9055 if a > thr {
9056 outliers.push((j, v));
9057 } else if a > amax {
9058 amax = a;
9059 }
9060 }
9061 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9062 let inv = 1.0 / sx;
9063 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9064 xq.clear();
9065 xq.reserve(n);
9066 if outliers.is_empty() {
9067 xq.extend(
9068 x.iter()
9069 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9070 );
9071 } else {
9072 xq.extend(x.iter().map(|&v| {
9074 if v.abs() > thr {
9075 0
9076 } else {
9077 (v * inv).round().clamp(-127.0, 127.0) as i8
9078 }
9079 }));
9080 }
9081 let xsum = xq.iter().map(|&v| v as i32).sum();
9082 SplitAct {
9083 xq,
9084 sx,
9085 outliers,
9086 xsum,
9087 }
9088}
9089
9090fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9091 let n = x.len();
9092 let rms = (x
9093 .iter()
9094 .zip(col)
9095 .map(|(&a, &c)| {
9096 let v = a * c;
9097 (v * v) as f64
9098 })
9099 .sum::<f64>()
9100 / n.max(1) as f64)
9101 .sqrt() as f32;
9102 let thr = 8.0 * rms;
9103
9104 let mut outliers = Vec::new();
9105 let mut amax = 0f32;
9106 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9107 let v = a * c;
9108 let s = v.abs();
9109 if s > thr {
9110 outliers.push((j, v));
9111 } else if s > amax {
9112 amax = s;
9113 }
9114 }
9115
9116 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9117 let inv = 1.0 / sx;
9118 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9119 xq.clear();
9120 xq.reserve(n);
9121 if outliers.is_empty() {
9122 xq.extend(
9123 x.iter()
9124 .zip(col)
9125 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9126 );
9127 } else {
9128 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9129 let v = a * c;
9130 if v.abs() > thr {
9131 0
9132 } else {
9133 (v * inv).round().clamp(-127.0, 127.0) as i8
9134 }
9135 }));
9136 }
9137 let xsum = xq.iter().map(|&v| v as i32).sum();
9138 SplitAct {
9139 xq,
9140 sx,
9141 outliers,
9142 xsum,
9143 }
9144}
9145
9146#[cfg(target_arch = "aarch64")]
9149#[target_feature(enable = "neon,dotprod")]
9150unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9151 unsafe {
9153 use core::arch::aarch64::*;
9154 use core::arch::asm;
9155 let wp = w.as_ptr() as *const i8;
9156 let n = w.len();
9157 let (mut a0, mut a1, mut a2, mut a3) = (
9158 vdupq_n_s32(0),
9159 vdupq_n_s32(0),
9160 vdupq_n_s32(0),
9161 vdupq_n_s32(0),
9162 );
9163 let mut i = 0;
9164 while i + 64 <= n {
9165 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9166 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9167 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9168 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9169 asm!(
9170 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9171 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9172 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9173 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9174 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9175 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9176 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9177 options(pure, nomem, nostack),
9178 );
9179 i += 64;
9180 }
9181 while i + 16 <= n {
9182 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9183 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9184 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9185 i += 16;
9186 }
9187 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9188 while i < n {
9189 s += (*wp.add(i)) as i32 * xq[i] as i32;
9190 i += 1;
9191 }
9192 s
9193 }
9194}
9195
9196#[cfg(target_arch = "aarch64")]
9200#[target_feature(enable = "neon,dotprod")]
9201unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9202 unsafe {
9204 use core::arch::aarch64::*;
9205 use core::arch::asm;
9206 let n = xq.len();
9207 let px = xq.as_ptr();
9208 let (p0, p1, p2, p3) = (
9209 w0.as_ptr() as *const i8,
9210 w1.as_ptr() as *const i8,
9211 w2.as_ptr() as *const i8,
9212 w3.as_ptr() as *const i8,
9213 );
9214 let (mut a0, mut a1, mut a2, mut a3) = (
9215 vdupq_n_s32(0),
9216 vdupq_n_s32(0),
9217 vdupq_n_s32(0),
9218 vdupq_n_s32(0),
9219 );
9220 let mut i = 0;
9221 while i + 16 <= n {
9222 let x = vld1q_s8(px.add(i));
9223 let v0 = vld1q_s8(p0.add(i));
9224 let v1 = vld1q_s8(p1.add(i));
9225 let v2 = vld1q_s8(p2.add(i));
9226 let v3 = vld1q_s8(p3.add(i));
9227 asm!(
9228 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9229 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9230 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9231 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9232 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9233 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9234 options(pure, nomem, nostack),
9235 );
9236 i += 16;
9237 }
9238 let mut r = [
9239 vaddvq_s32(a0),
9240 vaddvq_s32(a1),
9241 vaddvq_s32(a2),
9242 vaddvq_s32(a3),
9243 ];
9244 while i < n {
9245 let xi = *px.add(i) as i32;
9246 r[0] += (*p0.add(i)) as i32 * xi;
9247 r[1] += (*p1.add(i)) as i32 * xi;
9248 r[2] += (*p2.add(i)) as i32 * xi;
9249 r[3] += (*p3.add(i)) as i32 * xi;
9250 i += 1;
9251 }
9252 r
9253 }
9254}
9255
9256#[cfg(target_arch = "aarch64")]
9263#[target_feature(enable = "neon,dotprod")]
9264unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9265 unsafe {
9268 use core::arch::aarch64::*;
9269 use core::arch::asm;
9270 let n = xq.len();
9271 let px = xq.as_ptr();
9272 let pg = g.as_ptr() as *const i8;
9273 let (mut a0, mut a1, mut a2, mut a3) = (
9274 vdupq_n_s32(0),
9275 vdupq_n_s32(0),
9276 vdupq_n_s32(0),
9277 vdupq_n_s32(0),
9278 );
9279 let mut i = 0;
9280 while i + 16 <= n {
9281 let x = vld1q_s8(px.add(i));
9282 let base = pg.add(4 * i);
9283 let v0 = vld1q_s8(base);
9284 let v1 = vld1q_s8(base.add(16));
9285 let v2 = vld1q_s8(base.add(32));
9286 let v3 = vld1q_s8(base.add(48));
9287 asm!(
9288 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9289 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9290 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9291 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9292 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9293 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9294 options(pure, nomem, nostack),
9295 );
9296 i += 16;
9297 }
9298 [
9299 vaddvq_s32(a0),
9300 vaddvq_s32(a1),
9301 vaddvq_s32(a2),
9302 vaddvq_s32(a3),
9303 ]
9304 }
9305}
9306
9307#[cfg(target_arch = "aarch64")]
9313fn q8_range_sdot(
9314 q: &[u8],
9315 rep: &[u8],
9316 row_scale: &[f32],
9317 act: &SplitAct,
9318 cols: usize,
9319 out_addr: SendMut,
9320 start: usize,
9321 end: usize,
9322) {
9323 let mut o = start;
9324 if !rep.is_empty() {
9327 while o < end && o % 4 != 0 {
9328 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9329 unsafe { *out_addr.at(o) = v };
9330 o += 1;
9331 }
9332 }
9333 while o + 4 <= end {
9334 let r = if rep.is_empty() {
9335 unsafe {
9336 dot_i8_sdot_4rows(
9337 &q[o * cols..(o + 1) * cols],
9338 &q[(o + 1) * cols..(o + 2) * cols],
9339 &q[(o + 2) * cols..(o + 3) * cols],
9340 &q[(o + 3) * cols..(o + 4) * cols],
9341 &act.xq,
9342 )
9343 }
9344 } else {
9345 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9346 };
9347 for k in 0..4 {
9348 let mut acc = r[k] as f32 * act.sx;
9349 for &(j, xv) in &act.outliers {
9350 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9351 }
9352 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9354 }
9355 o += 4;
9356 }
9357 while o < end {
9358 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9359 unsafe { *out_addr.at(o) = v };
9360 o += 1;
9361 }
9362}
9363
9364#[cfg(target_arch = "aarch64")]
9367#[allow(clippy::too_many_arguments)]
9368fn q8_range2_sdot(
9369 q: &[u8],
9370 row_scale: &[f32],
9371 a1: &SplitAct,
9372 a2: &SplitAct,
9373 cols: usize,
9374 p1: SendMut,
9375 p2: SendMut,
9376 start: usize,
9377 end: usize,
9378) {
9379 for o in start..end {
9380 let row = &q[o * cols..(o + 1) * cols];
9381 unsafe {
9383 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9384 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9385 }
9386 }
9387}
9388
9389#[allow(clippy::too_many_arguments)]
9391fn q8_range2_f32(
9392 q: &[u8],
9393 row_scale: &[f32],
9394 x1: &[f32],
9395 x2: &[f32],
9396 cols: usize,
9397 p1: SendMut,
9398 p2: SendMut,
9399 start: usize,
9400 end: usize,
9401) {
9402 for o in start..end {
9403 let row = &q[o * cols..(o + 1) * cols];
9404 unsafe {
9406 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9407 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9408 }
9409 }
9410}
9411
9412fn q8_range_f32(
9414 q: &[u8],
9415 row_scale: &[f32],
9416 xs: &[f32],
9417 cols: usize,
9418 out_addr: SendMut,
9419 start: usize,
9420 end: usize,
9421) {
9422 for o in start..end {
9423 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9424 unsafe { *out_addr.at(o) = v };
9426 }
9427}
9428
9429#[inline]
9433fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9434 #[cfg(target_arch = "aarch64")]
9435 return row_dot_sdot(row, act);
9436 #[cfg(target_arch = "x86_64")]
9437 return row_dot_avx2(row, act);
9438 #[allow(unreachable_code)]
9439 q8_row_dot_scalar(row, act)
9440}
9441
9442#[allow(dead_code)]
9443fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9444 let mut acc = 0i32;
9445 for (k, &b) in row.iter().enumerate() {
9446 acc += (b as i8) as i32 * act.xq[k] as i32;
9447 }
9448 let mut acc = acc as f32 * act.sx;
9449 for &(j, xv) in &act.outliers {
9450 acc += (row[j] as i8) as f32 * xv;
9451 }
9452 acc
9453}
9454
9455#[cfg(target_arch = "aarch64")]
9458#[inline]
9459fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9460 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9461 for &(j, xv) in &act.outliers {
9462 acc += (row[j] as i8) as f32 * xv;
9463 }
9464 acc
9465}
9466
9467#[cfg(target_arch = "aarch64")]
9475#[target_feature(enable = "neon,dotprod")]
9476unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9477 unsafe {
9480 use core::arch::aarch64::*;
9481 use core::arch::asm;
9482 let lomask = vdupq_n_u8(0x0F);
9483 let eight = vdupq_n_s8(8);
9484 let mut acc = 0f32;
9485 for gi in 0..gpr {
9486 let g = g0 + gi;
9487 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9488 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9489 let lo = vandq_u8(b, lomask);
9490 let hi = vshrq_n_u8::<4>(b);
9491 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9492 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9493 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9494 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9495 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9496 asm!(
9497 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9498 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9499 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9500 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9501 options(pure, nomem, nostack),
9502 );
9503 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9504 }
9505 acc
9506 }
9507}
9508
9509#[cfg(target_arch = "aarch64")]
9514#[target_feature(enable = "neon,dotprod")]
9515unsafe fn dot_q4_row_sdot2(
9516 packed: &[u8],
9517 scales: &[u8],
9518 g0: usize,
9519 gpr: usize,
9520 xq1: &[i8],
9521 xq2: &[i8],
9522) -> (f32, f32) {
9523 unsafe {
9526 use core::arch::aarch64::*;
9527 use core::arch::asm;
9528 let lomask = vdupq_n_u8(0x0F);
9529 let eight = vdupq_n_s8(8);
9530 let (mut acc1, mut acc2) = (0f32, 0f32);
9531 for gi in 0..gpr {
9532 let g = g0 + gi;
9533 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9534 let b = vld1q_u8(packed.as_ptr().add(g * 16));
9535 let lo = vandq_u8(b, lomask);
9536 let hi = vshrq_n_u8::<4>(b);
9537 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9538 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9539 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9540 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9541 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9542 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9543 let (mut a0, mut a1, mut b0, mut b1) = (
9544 vdupq_n_s32(0),
9545 vdupq_n_s32(0),
9546 vdupq_n_s32(0),
9547 vdupq_n_s32(0),
9548 );
9549 asm!(
9550 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9551 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9552 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9553 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9554 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9555 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9556 e0 = in(vreg) e0, e1 = in(vreg) e1,
9557 x10 = in(vreg) x10, x11 = in(vreg) x11,
9558 x20 = in(vreg) x20, x21 = in(vreg) x21,
9559 options(pure, nomem, nostack),
9560 );
9561 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9562 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9563 }
9564 (acc1, acc2)
9565 }
9566}
9567
9568#[inline]
9573pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9574 #[cfg(target_arch = "aarch64")]
9575 unsafe {
9576 return axpy_i8_f32_neon(acc, row, w);
9577 }
9578 #[cfg(target_arch = "x86_64")]
9579 if avx2_enabled() {
9580 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9581 }
9582 #[allow(unreachable_code)]
9583 {
9584 for (a, &b) in acc.iter_mut().zip(row) {
9585 *a += w * b as f32;
9586 }
9587 }
9588}
9589
9590#[cfg(target_arch = "x86_64")]
9592#[target_feature(enable = "avx2,fma")]
9593unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9594 unsafe {
9596 use core::arch::x86_64::*;
9597 let n = acc.len().min(row.len());
9598 let ap = acc.as_mut_ptr();
9599 let rp = row.as_ptr();
9600 let wv = _mm256_set1_ps(w);
9601 let mut j = 0usize;
9602 while j + 16 <= n {
9603 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9604 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9605 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9606 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9607 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9608 _mm256_storeu_ps(ap.add(j), v0);
9609 _mm256_storeu_ps(ap.add(j + 8), v1);
9610 j += 16;
9611 }
9612 while j < n {
9613 *ap.add(j) += w * (*rp.add(j)) as f32;
9614 j += 1;
9615 }
9616 }
9617}
9618
9619#[cfg(target_arch = "aarch64")]
9620#[target_feature(enable = "neon")]
9621unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9622 unsafe {
9624 use core::arch::aarch64::*;
9625 let n = acc.len().min(row.len());
9626 let ap = acc.as_mut_ptr();
9627 let rp = row.as_ptr();
9628 let wv = vdupq_n_f32(w);
9629 let mut j = 0usize;
9630 while j + 16 <= n {
9631 let rb = vld1q_s8(rp.add(j));
9632 let lo = vmovl_s8(vget_low_s8(rb));
9633 let hi = vmovl_s8(vget_high_s8(rb));
9634 for (off, half) in [(0, lo), (8, hi)] {
9635 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9636 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9637 let o = j + off;
9638 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9639 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9640 }
9641 j += 16;
9642 }
9643 while j < n {
9644 *ap.add(j) += w * (*rp.add(j)) as f32;
9645 j += 1;
9646 }
9647 }
9648}
9649
9650#[inline]
9653pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9654 #[cfg(target_arch = "aarch64")]
9655 unsafe {
9656 return dot_i8_f32_neon(w, x);
9657 }
9658 #[cfg(target_arch = "x86_64")]
9659 if avx2_enabled() {
9660 return unsafe { dot_i8_f32_avx2(w, x) };
9661 }
9662 #[allow(unreachable_code)]
9663 {
9664 let mut sum = 0.0f32;
9665 for (j, &b) in w.iter().enumerate() {
9666 sum += (b as i8) as f32 * x[j];
9667 }
9668 sum
9669 }
9670}
9671
9672#[inline]
9676fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9677 #[cfg(target_arch = "aarch64")]
9678 unsafe {
9679 return dot_i8_col_f32_neon(w, x, col);
9680 }
9681 #[allow(unreachable_code)]
9682 {
9683 let mut sum = 0.0f32;
9684 for (j, &b) in w.iter().enumerate() {
9685 sum += (b as i8) as f32 * x[j] * col[j];
9686 }
9687 sum
9688 }
9689}
9690
9691#[cfg(target_arch = "aarch64")]
9692#[target_feature(enable = "neon")]
9693unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9694 unsafe {
9696 use core::arch::aarch64::*;
9697 let n = x.len();
9698 let wp = w.as_ptr() as *const i8;
9699 let xp = x.as_ptr();
9700 let cp = col.as_ptr();
9701 let (mut a0, mut a1, mut a2, mut a3) = (
9702 vdupq_n_f32(0.0),
9703 vdupq_n_f32(0.0),
9704 vdupq_n_f32(0.0),
9705 vdupq_n_f32(0.0),
9706 );
9707 let mut j = 0usize;
9708 while j + 16 <= n {
9709 let wb = vld1q_s8(wp.add(j));
9710 let lo = vmovl_s8(vget_low_s8(wb));
9711 let hi = vmovl_s8(vget_high_s8(wb));
9712 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9713 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9714 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9715 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9716 a0 = vfmaq_f32(
9717 a0,
9718 w0,
9719 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9720 );
9721 a1 = vfmaq_f32(
9722 a1,
9723 w1,
9724 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9725 );
9726 a2 = vfmaq_f32(
9727 a2,
9728 w2,
9729 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9730 );
9731 a3 = vfmaq_f32(
9732 a3,
9733 w3,
9734 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9735 );
9736 j += 16;
9737 }
9738 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9739 while j < n {
9740 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9741 j += 1;
9742 }
9743 sum
9744 }
9745}
9746
9747#[cfg(target_arch = "aarch64")]
9748#[target_feature(enable = "neon")]
9749unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9750 unsafe {
9752 use core::arch::aarch64::*;
9753 let n = x.len();
9754 let wp = w.as_ptr() as *const i8;
9755 let xp = x.as_ptr();
9756 let (mut a0, mut a1, mut a2, mut a3) = (
9757 vdupq_n_f32(0.0),
9758 vdupq_n_f32(0.0),
9759 vdupq_n_f32(0.0),
9760 vdupq_n_f32(0.0),
9761 );
9762 let mut j = 0usize;
9763 while j + 16 <= n {
9764 let wb = vld1q_s8(wp.add(j));
9765 let lo = vmovl_s8(vget_low_s8(wb));
9766 let hi = vmovl_s8(vget_high_s8(wb));
9767 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9768 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9769 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9770 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9771 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9772 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9773 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9774 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9775 j += 16;
9776 }
9777 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9778 while j < n {
9779 sum += (*wp.add(j)) as f32 * *xp.add(j);
9780 j += 1;
9781 }
9782 sum
9783 }
9784}
9785
9786#[allow(clippy::too_many_arguments)]
9787fn qmatvec(
9788 q: &[u8],
9789 rep: &[u8],
9790 row_scale: &[f32],
9791 x: &[f32],
9792 col_field: &[f32],
9793 dtype: TensorDtype,
9794 rows: usize,
9795 cols: usize,
9796 out: &mut [f32],
9797 pool: Option<&Pool>,
9798) {
9799 debug_assert_eq!(out.len(), rows);
9800 #[cfg(not(target_arch = "aarch64"))]
9801 let _ = rep;
9802
9803 #[cfg(target_arch = "aarch64")]
9804 if sdot_enabled() {
9805 let act = if dtype == TensorDtype::Q8_2f {
9806 split_act_q8_2f(x, col_field)
9807 } else {
9808 split_act(x)
9809 };
9810 let out_addr = SendMut(out.as_mut_ptr());
9811 let run_range = |start: usize, end: usize| {
9812 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9813 };
9814 match pool {
9815 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9816 _ => run_range(0, rows),
9817 }
9818 return;
9819 }
9820 #[cfg(target_arch = "x86_64")]
9823 if avx2_a8w8_enabled() {
9824 let act = if dtype == TensorDtype::Q8_2f {
9825 split_act_q8_2f(x, col_field)
9826 } else {
9827 split_act(x)
9828 };
9829 let out_addr = SendMut(out.as_mut_ptr());
9830 let run_range = |start: usize, end: usize| {
9831 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9832 };
9833 match pool {
9834 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9835 _ => run_range(0, rows),
9836 }
9837 return;
9838 }
9839
9840 prescale_with(x, col_field, dtype, 1, |xs| {
9841 let out_addr = SendMut(out.as_mut_ptr());
9842 let run_range = move |start: usize, end: usize| {
9843 for o in start..end {
9844 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9845 unsafe { *out_addr.at(o) = v };
9847 }
9848 };
9849 match pool {
9850 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9851 _ => run_range(0, rows),
9852 }
9853 });
9854}
9855
9856#[allow(clippy::too_many_arguments)]
9857fn qmatvec2(
9858 q: &[u8],
9859 row_scale: &[f32],
9860 x1: &[f32],
9861 x2: &[f32],
9862 col_field: &[f32],
9863 dtype: TensorDtype,
9864 rows: usize,
9865 cols: usize,
9866 o1: &mut [f32],
9867 o2: &mut [f32],
9868 pool: Option<&Pool>,
9869) {
9870 #[cfg(target_arch = "aarch64")]
9871 if sdot_enabled() {
9872 let a1s = if dtype == TensorDtype::Q8_2f {
9873 split_act_q8_2f(x1, col_field)
9874 } else {
9875 split_act(x1)
9876 };
9877 let a2s = if dtype == TensorDtype::Q8_2f {
9878 split_act_q8_2f(x2, col_field)
9879 } else {
9880 split_act(x2)
9881 };
9882 let p1 = SendMut(o1.as_mut_ptr());
9883 let p2 = SendMut(o2.as_mut_ptr());
9884 let run_range = |start: usize, end: usize| {
9885 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9886 };
9887 match pool {
9888 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9889 _ => run_range(0, rows),
9890 }
9891 return;
9892 }
9893 #[cfg(target_arch = "x86_64")]
9894 if avx2_a8w8_enabled() {
9895 let a1s = if dtype == TensorDtype::Q8_2f {
9896 split_act_q8_2f(x1, col_field)
9897 } else {
9898 split_act(x1)
9899 };
9900 let a2s = if dtype == TensorDtype::Q8_2f {
9901 split_act_q8_2f(x2, col_field)
9902 } else {
9903 split_act(x2)
9904 };
9905 let p1 = SendMut(o1.as_mut_ptr());
9906 let p2 = SendMut(o2.as_mut_ptr());
9907 let run_range = |start: usize, end: usize| {
9908 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9909 };
9910 match pool {
9911 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9912 _ => run_range(0, rows),
9913 }
9914 return;
9915 }
9916
9917 prescale_with(x1, col_field, dtype, 1, |x1s| {
9918 prescale_with(x2, col_field, dtype, 2, |x2s| {
9919 let p1 = SendMut(o1.as_mut_ptr());
9920 let p2 = SendMut(o2.as_mut_ptr());
9921 let run_range = move |start: usize, end: usize| {
9922 for o in start..end {
9923 let row = &q[o * cols..(o + 1) * cols];
9924 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
9925 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
9926 unsafe {
9928 *p1.at(o) = s1;
9929 *p2.at(o) = s2;
9930 }
9931 }
9932 };
9933 match pool {
9934 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9935 _ => run_range(0, rows),
9936 }
9937 });
9938 });
9939}
9940
9941#[derive(Clone, Copy)]
9942struct SendMut(*mut f32);
9943unsafe impl Send for SendMut {}
9944unsafe impl Sync for SendMut {}
9945
9946impl SendMut {
9947 #[inline]
9948 fn at(self, i: usize) -> *mut f32 {
9949 unsafe { self.0.add(i) }
9950 }
9951}
9952
9953#[cfg(test)]
9954mod tests {
9955 use super::*;
9956
9957 #[test]
9958 fn q2tp_i8_dot_matches_exact_on_grid() {
9959 let (rows, cols) = (5, 64);
9963 let gpr = cols / GROUP_SIZE;
9964 let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
9968 .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
9969 .collect();
9970 let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
9971 let x: Vec<f32> = (0..cols).map(|i| if i % 3 == 0 { -1.0 } else { 1.0 }).collect();
9972 let act = split_act(&x);
9973 assert!(act.outliers.is_empty(), "on-grid input must have no outliers");
9974 let gsum = q1_group_sums(&act.xq, gpr);
9975 for r in 0..rows {
9976 let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
9977 let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
9978 assert!(
9979 (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
9980 "row {r}: exact {exact} vs i8 {fast}"
9981 );
9982 }
9983 }
9984
9985 #[test]
9986 fn q8_row_dot_fast_matches_scalar() {
9987 let cols = 96;
9990 let row: Vec<u8> = (0..cols)
9991 .map(|i| ((i as i32 * 37 % 251) - 125) as i8 as u8)
9992 .collect();
9993 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
9994 let act = split_act(&x);
9995 let fast = q8_row_dot(&row, &act);
9996 let scalar = q8_row_dot_scalar(&row, &act);
9997 assert!(
9998 (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
9999 "fast {fast} vs scalar {scalar}"
10000 );
10001 }
10002
10003 #[test]
10004 fn f32_matvec_matches_matvec_rows_bitexact() {
10005 let (rows, cols) = (300, 40);
10006 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10007 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10008 let qt = QTensor::from_f32(w.clone(), rows, cols);
10009
10010 let mut a = vec![0.0f32; rows];
10011 matvec_rows(None, &w, &x, &mut a);
10012 let mut b = vec![0.0f32; rows];
10013 qt.matvec(&x, &mut b, None);
10014 assert_eq!(a, b);
10015 }
10016
10017 #[test]
10018 fn sdot_kernel_exact_on_grid() {
10019 eprintln!("sdot_enabled = {}", sdot_enabled());
10024 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
10026 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10027 .collect();
10028 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10029 let x: Vec<f32> = (0..cols)
10030 .map(|i| match i % 3 {
10031 0 => 1.0,
10032 1 => -1.0,
10033 _ => 0.0,
10034 })
10035 .collect();
10036 let mut a = vec![0.0f32; rows];
10037 qmatvec(
10038 &w,
10039 &[],
10040 &scales,
10041 &x,
10042 &[],
10043 TensorDtype::Q8Row,
10044 rows,
10045 cols,
10046 &mut a,
10047 None,
10048 );
10049 for o in 0..rows {
10050 let mut acc = 0.0f32;
10051 for j in 0..cols {
10052 acc += (w[o * cols + j] as i8) as f32 * x[j];
10053 }
10054 let expect = acc * scales[o];
10055 assert!(
10056 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10057 "row {o}: {} vs {expect}",
10058 a[o]
10059 );
10060 }
10061 }
10062
10063 #[test]
10064 fn q1_tbl_fast_path_matches_reference() {
10065 let (rows, cols) = (5, 256);
10070 let gpr = cols / GROUP_SIZE;
10071 let mut bytes = Vec::new();
10072 for t in 0..rows * gpr {
10073 let s = 0.007 + (t % 11) as f32 * 0.004;
10074 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10075 for j in 0..4 {
10076 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10077 }
10078 }
10079 let x: Vec<f32> = (0..cols)
10080 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10081 .collect();
10082 let mut w = vec![0.0f32; rows * cols];
10083 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10084 let mut got = vec![0.0f32; rows];
10085 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10086 for o in 0..rows {
10087 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10088 assert!(
10089 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10090 "row {o}: {} vs {expect}",
10091 got[o]
10092 );
10093 }
10094 let b = 5usize;
10097 let mut xs_all = Vec::new();
10098 for bi in 0..b {
10099 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10100 }
10101 let mut mm = vec![0.0f32; b * rows];
10102 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10103 for bi in 0..b {
10104 let mut single = vec![0.0f32; rows];
10105 q1_matvec(
10106 &bytes,
10107 &xs_all[bi * cols..(bi + 1) * cols],
10108 rows,
10109 cols,
10110 &mut single,
10111 None,
10112 );
10113 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10114 }
10115 }
10116
10117 #[test]
10118 fn q1_kernels_match_exact_reference() {
10119 let (rows, cols) = (7, 96);
10121 let gpr = cols / GROUP_SIZE;
10122 let mut bytes = Vec::new();
10123 for t in 0..rows * gpr {
10124 let s = 0.01 + (t % 13) as f32 * 0.003;
10125 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10126 for j in 0..4 {
10127 bytes.push(((t * 31 + j * 97) % 251) as u8);
10128 }
10129 }
10130 let x: Vec<f32> = (0..cols)
10132 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10133 .collect();
10134 let mut w = vec![0.0f32; rows * cols];
10136 cortiq_core::quant::dequant_q1(&bytes, &mut w);
10137 let mut expect = vec![0.0f32; rows];
10138 for o in 0..rows {
10139 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10140 }
10141 let mut got = vec![0.0f32; rows];
10142 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10143 for o in 0..rows {
10144 assert!(
10145 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10146 "row {o}: {} vs {}",
10147 got[o],
10148 expect[o]
10149 );
10150 }
10151 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10153 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10154 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10155 assert_eq!(a1, got);
10156 let mut xs = x.clone();
10157 xs.extend_from_slice(&x2);
10158 let mut mm = vec![0.0f32; 2 * rows];
10159 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10160 assert_eq!(&mm[..rows], got.as_slice());
10161 assert_eq!(&mm[rows..], a2.as_slice());
10162 }
10163
10164 #[test]
10165 fn repack_is_bit_identical() {
10166 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
10172 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10173 .collect();
10174 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10175 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10176 let rep = q8_repack_layout(&w, rows, cols);
10177 for g in 0..rows / 4 {
10179 for c in 0..cols / 16 {
10180 for lane in 0..4 {
10181 assert_eq!(
10182 &rep[g * 4 * cols + c * 64 + lane * 16
10183 ..g * 4 * cols + c * 64 + lane * 16 + 16],
10184 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10185 );
10186 }
10187 }
10188 }
10189 let mut a = vec![0.0f32; rows];
10190 qmatvec(
10191 &w,
10192 &[],
10193 &scales,
10194 &x,
10195 &[],
10196 TensorDtype::Q8Row,
10197 rows,
10198 cols,
10199 &mut a,
10200 None,
10201 );
10202 let mut b = vec![0.0f32; rows];
10203 qmatvec(
10204 &w,
10205 &rep,
10206 &scales,
10207 &x,
10208 &[],
10209 TensorDtype::Q8Row,
10210 rows,
10211 cols,
10212 &mut b,
10213 None,
10214 );
10215 assert_eq!(a, b, "full-range repack output diverged");
10216
10217 #[cfg(target_arch = "aarch64")]
10218 if sdot_enabled() {
10219 let act = split_act(&x);
10221 let mut c1 = vec![0.0f32; rows];
10222 let mut c2 = vec![0.0f32; rows];
10223 q8_range_sdot(
10224 &w,
10225 &[],
10226 &scales,
10227 &act,
10228 cols,
10229 SendMut(c1.as_mut_ptr()),
10230 3,
10231 rows - 2,
10232 );
10233 q8_range_sdot(
10234 &w,
10235 &rep,
10236 &scales,
10237 &act,
10238 cols,
10239 SendMut(c2.as_mut_ptr()),
10240 3,
10241 rows - 2,
10242 );
10243 assert_eq!(c1, c2, "unaligned-range repack output diverged");
10244 }
10245 }
10246
10247 #[test]
10248 fn sdot_a8w8_noise_is_bounded() {
10249 let (rows, cols) = (16, 512);
10253 let w: Vec<u8> = (0..rows * cols)
10254 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10255 .collect();
10256 let scales = vec![0.01f32; rows];
10257 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10258 let mut a = vec![0.0f32; rows];
10259 qmatvec(
10260 &w,
10261 &[],
10262 &scales,
10263 &x,
10264 &[],
10265 TensorDtype::Q8Row,
10266 rows,
10267 cols,
10268 &mut a,
10269 None,
10270 );
10271 let (mut num, mut den) = (0f64, 0f64);
10272 for o in 0..rows {
10273 let mut acc = 0.0f32;
10274 for j in 0..cols {
10275 acc += (w[o * cols + j] as i8) as f32 * x[j];
10276 }
10277 let expect = acc * scales[o];
10278 num += ((a[o] - expect) as f64).powi(2);
10279 den += (expect as f64).powi(2);
10280 }
10281 let rel = (num / den.max(1e-12)).sqrt();
10282 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10283 }
10284
10285 #[test]
10286 fn i8_dot_neon_matches_scalar() {
10287 let n = 100;
10288 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10289 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10290 let mut scalar = 0.0f32;
10291 for j in 0..n {
10292 scalar += (w[j] as i8) as f32 * x[j];
10293 }
10294 let fast = dot_i8_f32(&w, &x);
10295 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10296 }
10297
10298 #[test]
10300 fn vbitmatvec_matches_full_dequant() {
10301 let (rows, cols) = (6, 64);
10302 let ng = cols / GROUP_SIZE;
10303 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10305 let mut bytes = bits.clone();
10306 for g in 0..rows * ng {
10307 let s = 0.02 + 0.001 * g as f32;
10308 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10309 }
10310 for r in 0..rows {
10311 let b = bits[r] as usize;
10312 let (mut acc, mut nb) = (0u64, 0usize);
10313 let mut rowbytes = Vec::new();
10314 for i in 0..cols {
10315 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10316 acc = (acc << b) | v;
10317 nb += b;
10318 while nb >= 8 {
10319 nb -= 8;
10320 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10321 }
10322 }
10323 if nb > 0 {
10324 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10325 }
10326 bytes.extend_from_slice(&rowbytes);
10327 }
10328 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10329
10330 let mut reference = vec![0f32; rows * cols];
10331 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10332 let mut expect = vec![0f32; rows];
10333 for r in 0..rows {
10334 expect[r] = reference[r * cols..(r + 1) * cols]
10335 .iter()
10336 .zip(&x)
10337 .map(|(w, xv)| w * xv)
10338 .sum();
10339 }
10340 let mut got = vec![0f32; rows];
10341 let offsets = vbit_row_offsets(&bytes, rows, cols);
10342 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10343 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10347 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10348 for r in 0..rows {
10349 assert!(
10350 (got[r] - expect[r]).abs() < tol * scale,
10351 "row {r}: {} vs {}",
10352 got[r],
10353 expect[r]
10354 );
10355 }
10356 }
10357
10358 #[test]
10363 #[cfg(target_arch = "x86_64")]
10364 fn vbit_matmat_blocked_matches_per_row() {
10365 let (rows, cols, b) = (64usize, 128usize, 9usize);
10366 let ng = cols / GROUP_SIZE;
10367 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10368 let mut bytes = bits.clone();
10369 for g in 0..rows * ng {
10370 let sc = 0.02 + 0.0005 * g as f32;
10371 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10372 }
10373 for r in 0..rows {
10374 let bw = bits[r] as usize;
10375 let (mut acc, mut nb) = (0u64, 0usize);
10376 let mut rowbytes = Vec::new();
10377 for i in 0..cols {
10378 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10379 acc = (acc << bw) | v;
10380 nb += bw;
10381 while nb >= 8 {
10382 nb -= 8;
10383 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10384 }
10385 }
10386 if nb > 0 {
10387 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10388 }
10389 bytes.extend_from_slice(&rowbytes);
10390 }
10391 let x: Vec<f32> = (0..b * cols)
10392 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10393 .collect();
10394 let offsets = vbit_row_offsets(&bytes, rows, cols);
10395 let mut y_a = vec![0f32; b * rows];
10396 let mut y_b = vec![0f32; b * rows];
10397 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10398 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10399 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10400 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10401 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10402 let max_d = y_a
10403 .iter()
10404 .zip(&y_b)
10405 .map(|(p, q)| (p - q).abs())
10406 .fold(0.0f32, f32::max);
10407 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10408 }
10409
10410 #[test]
10418 fn q4t_matmat_blocked_matches_per_row() {
10419 let (rows, cols, b) = (16usize, 64usize, 9usize);
10420 let gpr = cols / GROUP_SIZE;
10421 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10422 for r in 0..rows {
10423 for g in 0..gpr {
10424 let t = (r * gpr + g) * Q4_TILE;
10425 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10426 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10427 for k in 0..16 {
10428 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10429 }
10430 }
10431 }
10432 let x: Vec<f32> = (0..b * cols)
10433 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10434 .collect();
10435 let mut y_blk = vec![0f32; b * rows];
10436 let mut y_row = vec![0f32; b * rows];
10437 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10438 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10439 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10440 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10441 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10442 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10443 }
10444
10445 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10452 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10453 let gpr = cols / GROUP_SIZE;
10454 let stride = q4tp_code_stride(gpr);
10455 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10456 let mut b = vec![0u8; codes_off + rows * stride];
10457 for r in 0..rows {
10458 for g in 0..gpr {
10459 let t = (r * gpr + g) * Q4TP_NIB;
10460 for k in 0..16 {
10461 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10462 }
10463 }
10464 let lo = -6.0 - 0.03 * (r % 17) as f32;
10465 let step = 0.01 + 0.004 * (r % 11) as f32;
10466 let p = params_off + r * 4;
10467 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10468 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10469 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10470 for g in 0..gpr {
10471 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10472 }
10473 }
10474 b
10475 }
10476
10477 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10481 let gpr = cols / GROUP_SIZE;
10482 let v = Q4tpView::new(bytes, rows, cols);
10483 let mut out = vec![0u8; rows * gpr * Q4_TILE];
10484 let mut sc = vec![0f32; gpr];
10485 for r in 0..rows {
10486 v.scales_into(r, gpr, &mut sc);
10487 for g in 0..gpr {
10488 let t = (r * gpr + g) * Q4_TILE;
10489 let s = sc[g];
10490 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10491 let src = (r * gpr + g) * Q4TP_NIB;
10492 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10493 }
10494 }
10495 out
10496 }
10497
10498 #[test]
10504 fn q4tp_exact_path_matches_dequant_reference() {
10505 let (rows, cols) = (256usize, 512usize);
10506 let gpr = cols / GROUP_SIZE;
10507 let bytes = synth_q4tp(rows, cols);
10508 let mut w = vec![0f32; rows * cols];
10509 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10510
10511 let x: Vec<f32> = (0..cols)
10512 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10513 .collect();
10514 let v = Q4tpView::new(&bytes, rows, cols);
10515 let mut sc = vec![0f32; gpr];
10516 for r in 0..rows {
10517 v.scales_into(r, gpr, &mut sc);
10518 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10519 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10520 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10524 assert!(
10525 (got - want).abs() <= 1e-5 * mag,
10526 "row {r}: kernel {got} vs dequant {want}"
10527 );
10528 }
10529 }
10530
10531 #[test]
10537 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10538 let (rows, cols) = (256usize, 512usize);
10539 let bytes = synth_q4tp(rows, cols);
10540 let twin = q4tp_as_q4t(&bytes, rows, cols);
10541 let x: Vec<f32> = (0..cols)
10542 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10543 .collect();
10544
10545 let mut got = vec![0f32; rows];
10546 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10547 let mut want = vec![0f32; rows];
10548 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10549
10550 let mut w = vec![0f32; rows * cols];
10553 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10554 for r in 0..rows {
10555 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10556 assert!(
10557 (got[r] - want[r]).abs() <= 1e-3 * mag,
10558 "row {r}: q4tp {} vs q4t {}",
10559 got[r],
10560 want[r]
10561 );
10562 }
10563 }
10564
10565 #[test]
10570 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10571 let (rows, cols, b) = (256usize, 512usize, 5usize);
10572 let bytes = synth_q4tp(rows, cols);
10573 let twin = q4tp_as_q4t(&bytes, rows, cols);
10574 let xs: Vec<f32> = (0..b * cols)
10575 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10576 .collect();
10577
10578 let mut got = vec![0f32; b * rows];
10579 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10580 let mut want = vec![0f32; b * rows];
10581 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10582
10583 let mut w = vec![0f32; rows * cols];
10584 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10585 for t in 0..b {
10586 for r in 0..rows {
10587 let mag: f32 = (0..cols)
10588 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10589 .sum();
10590 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10591 assert!(
10592 (g - wa).abs() <= 1e-3 * mag,
10593 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10594 );
10595 }
10596 }
10597 }
10598
10599 #[test]
10600 fn q4tp_matvec2_matches_the_single_stream_kernel() {
10601 let (rows, cols) = (128usize, 256usize);
10602 let gpr = cols / GROUP_SIZE;
10603 let bytes = synth_q4tp(rows, cols);
10604 let xs: Vec<f32> = (0..2 * cols)
10605 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10606 .collect();
10607
10608 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10609 q4tp_matvec2(
10610 &bytes,
10611 &xs[..cols],
10612 &xs[cols..],
10613 rows,
10614 cols,
10615 &mut o1,
10616 &mut o2,
10617 None,
10618 );
10619
10620 let v = Q4tpView::new(&bytes, rows, cols);
10623 let mut sc = vec![0f32; gpr];
10624 for r in 0..rows {
10625 v.scales_into(r, gpr, &mut sc);
10626 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10627 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10628 }
10629 }
10630
10631 #[test]
10638 fn q4tp_matvec_keeps_pace_with_q4t() {
10639 let (rows, cols) = (4096usize, 3072usize);
10640 let bytes = synth_q4tp(rows, cols);
10641 let twin = q4tp_as_q4t(&bytes, rows, cols);
10642 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10643 let mut o = vec![0f32; rows];
10644 let n = 12;
10645 let mut best = (f64::MAX, f64::MAX);
10646 for _ in 0..3 {
10649 let t0 = std::time::Instant::now();
10650 for _ in 0..n {
10651 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10652 }
10653 best.0 = best.0.min(t0.elapsed().as_secs_f64());
10654 let t0 = std::time::Instant::now();
10655 for _ in 0..n {
10656 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10657 }
10658 best.1 = best.1.min(t0.elapsed().as_secs_f64());
10659 }
10660 let ratio = best.1 / best.0;
10661 println!(
10662 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10663 best.0 * 1e3 / n as f64,
10664 best.1 * 1e3 / n as f64
10665 );
10666 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10667 }
10668
10669 #[cfg(target_os = "macos")]
10670 #[test]
10671 fn q4t_matmat_accel_matches_dequant_reference() {
10672 if !accel_gemm_enabled() {
10673 return; }
10675 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
10677 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10678 for r in 0..rows {
10679 for g in 0..gpr {
10680 let t = (r * gpr + g) * Q4_TILE;
10681 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10682 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10683 for k in 0..16 {
10684 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10685 }
10686 }
10687 }
10688 let x: Vec<f32> = (0..b * cols)
10689 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10690 .collect();
10691 let mut got = vec![0f32; b * rows];
10692 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10693 let mut w = vec![0f32; rows * cols];
10695 for r in 0..rows {
10696 for g in 0..gpr {
10697 let t = (r * gpr + g) * Q4_TILE;
10698 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10699 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10700 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10701 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10702 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10703 }
10704 }
10705 }
10706 for bi in 0..b {
10707 for r in 0..rows {
10708 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10709 let d = (got[bi * rows + r] - want).abs();
10710 assert!(
10711 d <= want.abs().max(1.0) * 1e-4,
10712 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10713 got[bi * rows + r]
10714 );
10715 }
10716 }
10717 }
10718
10719 #[test]
10720 fn q4matvec_matches_full_dequant() {
10721 let (rows, cols) = (8, 64);
10722 let groups = rows * cols / GROUP_SIZE;
10723 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10725 for i in 0..groups * 16 {
10726 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10727 }
10728 for g in 0..groups {
10729 let s = 0.01 + 0.003 * g as f32;
10730 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10731 }
10732 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10733
10734 let mut reference = vec![0.0f32; rows * cols];
10735 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10736 let mut expect = vec![0.0f32; rows];
10737 for r in 0..rows {
10738 expect[r] = reference[r * cols..(r + 1) * cols]
10739 .iter()
10740 .zip(&x)
10741 .map(|(w, xv)| w * xv)
10742 .sum();
10743 }
10744
10745 let mut got = vec![0.0f32; rows];
10746 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10747 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10751 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10752 for r in 0..rows {
10753 assert!(
10754 (got[r] - expect[r]).abs() < tol * scale,
10755 "row {r}: {} vs {}",
10756 got[r],
10757 expect[r]
10758 );
10759 }
10760 }
10761
10762 #[test]
10765 fn vbitmatvec2_equals_two_singles() {
10766 let (rows, cols) = (6, 64);
10767 let ng = cols / GROUP_SIZE;
10768 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10769 let mut bytes = bits.clone();
10770 for g in 0..rows * ng {
10771 let s = 0.02 + 0.001 * g as f32;
10772 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10773 }
10774 for r in 0..rows {
10775 let b = bits[r] as usize;
10776 let (mut acc, mut nb) = (0u64, 0usize);
10777 let mut rowbytes = Vec::new();
10778 for i in 0..cols {
10779 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10780 acc = (acc << b) | v;
10781 nb += b;
10782 while nb >= 8 {
10783 nb -= 8;
10784 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10785 }
10786 }
10787 if nb > 0 {
10788 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10789 }
10790 bytes.extend_from_slice(&rowbytes);
10791 }
10792 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10793 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10794 let offsets = vbit_row_offsets(&bytes, rows, cols);
10795
10796 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10797 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10798 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10799 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10800 vbitmatvec2(
10801 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10802 );
10803 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10804 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10805 }
10806
10807 #[test]
10809 fn q4matvec2_equals_two_singles() {
10810 let (rows, cols) = (8, 128);
10811 let groups = rows * cols / GROUP_SIZE;
10812 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10813 for i in 0..groups * 16 {
10814 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10815 }
10816 for g in 0..groups {
10817 let s = 0.01 + 0.003 * g as f32;
10818 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10819 }
10820 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10823 x1[9] = 250.0;
10824 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10825
10826 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10827 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10828 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10829 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10830 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10831 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10832 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10833 }
10834
10835 #[test]
10838 fn matvec_many_equals_separate_matvecs() {
10839 use crate::pool::Pool;
10840 let (r1, r2, cols) = (300, 200, 64);
10841 let mk = |salt: usize, rows: usize| {
10842 QTensor::from_f32(
10843 (0..rows * cols)
10844 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10845 .collect(),
10846 rows,
10847 cols,
10848 )
10849 };
10850 let (a, b) = (mk(1, r1), mk(5, r2));
10851 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10852 let pool = Pool::new(3);
10853
10854 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10855 a.matvec(&x, &mut ea, Some(&pool));
10856 b.matvec(&x, &mut eb, Some(&pool));
10857 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10858 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10859 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10860 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10861 }
10862
10863 #[test]
10866 fn batched_matmat_equals_per_position_matvec() {
10867 let (rows, cols, b) = (8, 64, 5);
10868 let groups = rows * cols / GROUP_SIZE;
10870 let mut q4 = Vec::new();
10871 for i in 0..groups * 16 {
10872 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10873 }
10874 for g in 0..groups {
10875 q4.extend_from_slice(
10876 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10877 );
10878 }
10879 let ng = cols / GROUP_SIZE;
10881 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10882 let mut vb = bits.clone();
10883 for g in 0..rows * ng {
10884 vb.extend_from_slice(
10885 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10886 );
10887 }
10888 for r in 0..rows {
10889 let bw = bits[r] as usize;
10890 let (mut acc, mut nb) = (0u64, 0usize);
10891 let mut rowbytes = Vec::new();
10892 for i in 0..cols {
10893 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10894 acc = (acc << bw) | v;
10895 nb += bw;
10896 while nb >= 8 {
10897 nb -= 8;
10898 rowbytes.push(((acc >> nb) & 0xFF) as u8);
10899 }
10900 }
10901 if nb > 0 {
10902 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10903 }
10904 vb.extend_from_slice(&rowbytes);
10905 }
10906 let offsets = vbit_row_offsets(&vb, rows, cols);
10907
10908 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10909
10910 let mut got = vec![0f32; b * rows];
10912 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
10913 for bi in 0..b {
10914 let mut expect = vec![0f32; rows];
10915 q4matvec(
10916 &q4,
10917 &xs[bi * cols..(bi + 1) * cols],
10918 rows,
10919 cols,
10920 &mut expect,
10921 None,
10922 );
10923 assert_eq!(
10924 &got[bi * rows..(bi + 1) * rows],
10925 &expect[..],
10926 "q4 batch pos {bi}"
10927 );
10928 }
10929
10930 let mut got = vec![0f32; b * rows];
10932 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
10933 for bi in 0..b {
10934 let mut expect = vec![0f32; rows];
10935 vbitmatvec(
10936 &vb,
10937 &offsets,
10938 &xs[bi * cols..(bi + 1) * cols],
10939 rows,
10940 cols,
10941 &mut expect,
10942 None,
10943 );
10944 assert_eq!(
10945 &got[bi * rows..(bi + 1) * rows],
10946 &expect[..],
10947 "vbit batch pos {bi}"
10948 );
10949 }
10950 }
10951
10952 #[test]
10956 fn q4_tiled_matches_q4_block_bitexact() {
10957 let (rows, cols, b) = (8usize, 128usize, 3usize);
10958 let groups = rows * cols / GROUP_SIZE;
10959 let mut split = Vec::with_capacity(groups * 18);
10960 for i in 0..groups * 16 {
10961 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10962 }
10963 for g in 0..groups {
10964 split.extend_from_slice(
10965 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10966 );
10967 }
10968 let (packed, scales) = split.split_at(groups * 16);
10970 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
10971 for g in 0..groups {
10972 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
10973 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
10974 }
10975
10976 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10977 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10979
10980 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
10981 q4matvec(&split, &x1, rows, cols, &mut a, None);
10982 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
10983 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
10984
10985 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10986 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
10987 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
10988 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
10989 assert_eq!(a1, t1);
10990 assert_eq!(a2, t2);
10991
10992 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10993 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
10994 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
10995 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
10996 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
10997 }
10998
10999 #[test]
11006 fn q4matvec_sdot_outlier_exact() {
11007 let (rows, cols) = (4, 128);
11008 let groups = rows * cols / GROUP_SIZE;
11009 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11010 for i in 0..groups * 16 {
11011 bytes.push(((i * 11 + 5) % 256) as u8);
11012 }
11013 for g in 0..groups {
11014 let s = 0.02 + 0.002 * g as f32;
11015 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11016 }
11017 let mut x: Vec<f32> = (0..cols)
11018 .map(|i| match i % 3 {
11019 0 => 1.0,
11020 1 => -1.0,
11021 _ => 0.0,
11022 })
11023 .collect();
11024 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
11027 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11028 let mut expect = vec![0.0f32; rows];
11029 for r in 0..rows {
11030 expect[r] = reference[r * cols..(r + 1) * cols]
11031 .iter()
11032 .zip(&x)
11033 .map(|(w, xv)| w * xv)
11034 .sum();
11035 }
11036 let mut got = vec![0.0f32; rows];
11037 q4matvec(&bytes, &x, rows, cols, &mut got, None);
11038 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11039 for r in 0..rows {
11040 assert!(
11041 (got[r] - expect[r]).abs() < 2e-3 * scale,
11042 "row {r}: {} vs {} (outlier term must be exact)",
11043 got[r],
11044 expect[r]
11045 );
11046 }
11047 }
11048
11049 #[test]
11053 fn q1t_matvec_matches_reference() {
11054 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11055 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
11057 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11058 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11060 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11061 let mut bytes = Vec::new();
11062 for r in 0..rows {
11063 for g in 0..gpr {
11064 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11065 let mut c = [0u8; 7];
11066 for k in 0..GROUP_SIZE {
11067 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11069 0
11070 } else {
11071 ((k + r * 3 + g) % 3) as u8 };
11073 cortiq_core::quant::q1t_pack(&mut c, k, code);
11074 }
11075 bytes.extend_from_slice(&c);
11076 }
11077 }
11078 let mut row_ptr = vec![0u32; rows + 1];
11081 for &(idx, _) in &outliers {
11082 row_ptr[idx as usize / cols + 1] += 1;
11083 }
11084 for r in 0..rows {
11085 row_ptr[r + 1] += row_ptr[r];
11086 }
11087 for &p in &row_ptr {
11088 bytes.extend_from_slice(&p.to_le_bytes());
11089 }
11090 for &(idx, v) in &outliers {
11091 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11092 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11093 }
11094
11095 let mut refw = vec![0f32; rows * cols];
11096 dequant_q1t(&bytes, rows, cols, &mut refw);
11097 let x: Vec<f32> = (0..cols)
11100 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11101 .collect();
11102 let mut expect = vec![0f32; rows];
11103 for r in 0..rows {
11104 let mut a = 0.0f32;
11105 for j in 0..cols {
11106 a += refw[r * cols + j] * x[j];
11107 }
11108 expect[r] = a;
11109 }
11110 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11111 let mut got = vec![0f32; rows];
11112 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11113 for r in 0..rows {
11114 assert!(
11115 (got[r] - expect[r]).abs() < tol(expect[r]),
11116 "row {r}: {} vs {}",
11117 got[r],
11118 expect[r]
11119 );
11120 }
11121 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
11123 let mut gm = vec![0f32; 2 * rows];
11124 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11125 for r in 0..rows {
11126 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11127 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11128 }
11129 let xb: Vec<f32> = (0..cols)
11133 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11134 .collect();
11135 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11136 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11137 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11138 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11139 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11140 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11141 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11142 }
11143
11144 #[test]
11147 fn q1t_matvec2_odd_gpr_matches_singles() {
11148 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11149 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
11151 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11152 for r in 0..rows {
11153 for g in 0..gpr {
11154 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11155 let mut c = [0u8; 7];
11156 for k in 0..GROUP_SIZE {
11157 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11158 }
11159 bytes.extend_from_slice(&c);
11160 }
11161 }
11162 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11163 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11164 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11165 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11166 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11167 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11168 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11169 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11170 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11171 }
11172
11173 #[test]
11177 #[ignore]
11178 fn q1t_matvec2_speed() {
11179 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11180 use std::time::Instant;
11181 let (rows, cols) = (8192usize, 4096usize);
11182 let gpr = cols / GROUP_SIZE;
11183 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11184 for r in 0..rows {
11185 for g in 0..gpr {
11186 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11187 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11188 let mut c = [0u8; 7];
11189 for k in 0..GROUP_SIZE {
11190 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11191 }
11192 bytes.extend_from_slice(&c);
11193 }
11194 }
11195 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11196 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11197 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11198 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11199 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11201 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11202 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11203 for _ in 0..8 {
11204 let t0 = Instant::now();
11205 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11206 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11207 let t1 = Instant::now();
11208 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11209 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11210 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11211 }
11212 assert_eq!(p1, s1);
11213 assert_eq!(p2, s2);
11214 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11215 }
11216
11217 #[test]
11221 #[ignore]
11222 fn q1t_matvec_speed() {
11223 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11224 use std::time::Instant;
11225 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
11227 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11228 for r in 0..rows {
11229 for g in 0..gpr {
11230 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11231 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11232 let mut c = [0u8; 7];
11233 for k in 0..GROUP_SIZE {
11234 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11235 }
11236 bytes.extend_from_slice(&c);
11237 }
11238 }
11239 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
11241 let mut idx = 0usize;
11242 while idx < n {
11243 row_ptr[idx / cols + 1] += 1;
11244 idx += stride;
11245 }
11246 for r in 0..rows {
11247 row_ptr[r + 1] += row_ptr[r];
11248 }
11249 for &p in &row_ptr {
11250 bytes.extend_from_slice(&p.to_le_bytes());
11251 }
11252 let mut idx = 0usize;
11253 while idx < n {
11254 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11255 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11256 idx += stride;
11257 }
11258 let x: Vec<f32> = (0..cols)
11261 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11262 .collect();
11263 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11264
11265 let slow = |out: &mut [f32]| {
11267 let mut buf = vec![0f32; cols];
11268 for r in 0..rows {
11269 for g in 0..gpr {
11270 let off = (r * gpr + g) * Q1T_TILE;
11271 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11272 let codes = &bytes[off + 2..off + Q1T_TILE];
11273 for k in 0..GROUP_SIZE {
11274 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11275 1 => s,
11276 2 => -s,
11277 _ => 0.0,
11278 };
11279 }
11280 }
11281 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11282 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11283 }
11284 };
11285 let iters = 5;
11286 let mut a = vec![0f32; rows];
11287 slow(&mut a); let t = Instant::now();
11289 for _ in 0..iters {
11290 slow(&mut a);
11291 }
11292 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11293
11294 let mut b = vec![0f32; rows];
11295 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
11297 for _ in 0..iters {
11298 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11299 }
11300 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11301
11302 for r in 0..rows {
11303 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11304 }
11305 println!(
11306 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
11307 slow_ms / fast_ms
11308 );
11309 }
11310}
11311
11312
11313#[cfg(test)]
11314mod gemm_bench {
11315 #[test]
11326 #[ignore]
11327 fn q4tp_matmat_throughput() {
11328 let b: usize = std::env::var("CMF_BENCH_B")
11332 .ok()
11333 .and_then(|v| v.parse().ok())
11334 .unwrap_or(296);
11335 let (rows, cols) = (9216usize, 2304usize);
11336 let (_, _, _) = (rows, cols, b);
11337 let total = cortiq_core::quant::expected_nbytes(
11338 cortiq_core::TensorDtype::Q4TiledP,
11339 &[rows, cols],
11340 )
11341 .unwrap();
11342 let (params_off, codes_off, _) =
11347 cortiq_core::quant::q4tp_sections(rows, cols);
11348 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11349 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11350 let step = cortiq_core::quant::f32_to_f16(0.1);
11351 for r in 0..rows {
11352 let o = params_off + r * 4;
11353 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11354 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11355 }
11356 let _ = codes_off;
11357 let xs: Vec<f32> = (0..b * cols)
11358 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11359 .collect();
11360 let mut out = vec![0f32; b * rows];
11361 let pool = crate::pool::Pool::from_env();
11362 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11368 let reps: usize = std::env::var("CMF_BENCH_REPS")
11369 .ok()
11370 .and_then(|v| v.parse().ok())
11371 .unwrap_or(10);
11372 let mut best = [f64::MAX; 2];
11373 let mut sums = [0f32; 2];
11374 for _ in 0..reps {
11375 for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11376 super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11377 let t = std::time::Instant::now();
11378 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11379 best[k] = best[k].min(t.elapsed().as_secs_f64());
11380 sums[k] = out.iter().take(64).sum::<f32>();
11381 }
11382 }
11383 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11384 for (k, name) in ["previous", "tuned "].iter().enumerate() {
11385 println!(
11386 "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11387 best[k] * 1e3,
11388 flops / best[k] / 1e9,
11389 sums[k]
11390 );
11391 }
11392 assert!(
11393 (sums[0] - sums[1]).abs() < 1e-2,
11394 "the tuned kernel changed the result: {} vs {}",
11395 sums[0],
11396 sums[1]
11397 );
11398 }
11399
11400 #[test]
11406 fn q4tp_matmat_blocked_matches_scalar() {
11407 use std::sync::atomic::Ordering::Relaxed;
11408 for &(rows, cols, b) in &[
11416 (64usize, 128usize, 7usize),
11417 (33, 96, 4),
11418 (16, 256, 9),
11419 (192, 2304, 37),
11420 ] {
11421 let total = cortiq_core::quant::expected_nbytes(
11422 cortiq_core::TensorDtype::Q4TiledP,
11423 &[rows, cols],
11424 )
11425 .unwrap();
11426 let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11427 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11428 let lo = cortiq_core::quant::f32_to_f16(-4.0);
11429 let step = cortiq_core::quant::f32_to_f16(0.1);
11430 for r in 0..rows {
11431 let o = params_off + r * 4;
11432 bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11433 bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11434 }
11435 let xs: Vec<f32> = (0..b * cols)
11436 .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11437 .collect();
11438 let mut got = vec![0f32; b * rows];
11439 let mut want = vec![0f32; b * rows];
11440 let gpr = cols / 32;
11441 let view = super::Q4tpView::new(&bytes, rows, cols);
11442 let pool = crate::pool::Pool::from_env();
11443 super::Q4TP_ALT.store(2, Relaxed);
11444 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11445 super::Q4TP_ALT.store(1, Relaxed);
11446 super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11447 super::Q4TP_ALT.store(0, Relaxed);
11448 let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11455 let (mut worst, mut at) = (0f32, 0usize);
11456 for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11457 if (g - w).abs() > worst {
11458 worst = (g - w).abs();
11459 at = i;
11460 }
11461 }
11462 assert!(
11463 worst <= 1e-4 * scale,
11464 "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11465 (scale {scale:.3e}) at cell {at}: {} vs {}",
11466 got[at],
11467 want[at]
11468 );
11469
11470 let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11477 for bi in 0..b {
11478 let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11479 for r in 0..rows {
11480 let mut sc = vec![0f32; gpr];
11481 view.scales_into(r, gpr, &mut sc);
11482 let mut exact = 0f64;
11483 for j in 0..cols {
11484 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11485 exact += w as f64 * sq as f64 * act.xq[j] as f64;
11486 }
11487 exact *= act.sx as f64;
11488 for &(j, xv) in &act.outliers {
11489 let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11490 exact += w as f64 * sq as f64 * xv as f64;
11491 }
11492 let i = bi * rows + r;
11493 e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11494 e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11495 }
11496 }
11497 println!(
11498 "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11499 per-column {e_scalar:.3e}"
11500 );
11501 assert!(
11506 e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11507 "{rows}x{cols} b={b}: error against f64 too large — blocked \
11508 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11509 );
11510 }
11511 }
11512
11513 #[test]
11517 #[ignore]
11518 fn q4t_matmat_throughput() {
11519 let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11520 let total = cortiq_core::quant::expected_nbytes(
11521 cortiq_core::TensorDtype::Q4Tiled,
11522 &[rows, cols],
11523 )
11524 .unwrap();
11525 let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11528 let sc = cortiq_core::quant::f32_to_f16(0.02);
11529 for t in bytes.chunks_mut(super::Q4_TILE) {
11530 t[..2].copy_from_slice(&sc.to_le_bytes());
11531 }
11532 let xs: Vec<f32> = (0..b * cols)
11533 .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11534 .collect();
11535 let mut out = vec![0f32; b * rows];
11536 let pool = crate::pool::Pool::from_env();
11537 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11538 let reps: usize = std::env::var("CMF_BENCH_REPS")
11539 .ok()
11540 .and_then(|v| v.parse().ok())
11541 .unwrap_or(10);
11542 let mut best = f64::MAX;
11543 for _ in 0..reps {
11544 let t = std::time::Instant::now();
11545 super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11546 best = best.min(t.elapsed().as_secs_f64());
11547 }
11548 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11549 println!(
11550 "q4t matmat {rows}x{cols} b={b}: {:.1} ms {:.1} GFLOP/s (checksum {:.3})",
11551 best * 1e3,
11552 flops / best / 1e9,
11553 out.iter().take(64).sum::<f32>()
11554 );
11555 }
11556
11557}