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 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
123 *ON.get_or_init(|| std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true))
124}
125
126fn gpu_lmhead_enabled() -> bool {
127 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128 *ON.get_or_init(|| std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true))
129}
130
131fn gpu_split_frac() -> f32 {
132 static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
133 *F.get_or_init(|| {
134 std::env::var("CMF_GPU_SPLIT")
135 .ok()
136 .and_then(|v| v.parse::<f32>().ok())
137 .unwrap_or(0.5)
138 .clamp(0.0, 1.0)
139 })
140}
141
142impl QTensor {
143 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
144 debug_assert_eq!(data.len(), rows * cols);
145 Self::F32 { data, rows, cols }
146 }
147
148 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
151 let idx = model
154 .tensor_index(name)
155 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
156 let entry = &model.tensors[idx];
157 if entry.shape.len() != 2 {
158 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
159 }
160 let (rows, cols) = (entry.shape[0], entry.shape[1]);
161 let bytes = model.entry_bytes(entry);
162
163 match entry.dtype {
164 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
165 let n = rows * cols;
166 let scales_off = n;
167 let row_scale: Vec<f32> = (0..rows)
168 .map(|o| {
169 f16_to_f32(u16::from_le_bytes([
170 bytes[scales_off + o * 2],
171 bytes[scales_off + o * 2 + 1],
172 ]))
173 })
174 .collect();
175 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
176 let col_off = n + rows * 2;
177 (0..cols)
178 .map(|i| {
179 f16_to_f32(u16::from_le_bytes([
180 bytes[col_off + i * 2],
181 bytes[col_off + i * 2 + 1],
182 ]))
183 })
184 .collect()
185 } else {
186 Vec::new()
187 };
188 Ok(Self::Mapped {
189 model: model.clone(),
190 idx,
191 dtype: entry.dtype,
192 rows,
193 cols,
194 row_scale,
195 col_field,
196 vbit_offsets: Vec::new(),
197 repack: q8_repack(bytes, rows, cols),
198 })
199 }
200 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
202 model: model.clone(),
203 idx,
204 dtype: entry.dtype,
205 rows,
206 cols,
207 row_scale: Vec::new(),
208 col_field: Vec::new(),
209 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
210 repack: Vec::new(),
211 }),
212 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
216 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
217 let offsets: Vec<usize> = (0..=rows)
218 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
219 .collect();
220 Ok(Self::Mapped {
221 model: model.clone(),
222 idx,
223 dtype: entry.dtype,
224 rows,
225 cols,
226 row_scale: Vec::new(),
227 col_field: Vec::new(),
228 vbit_offsets: offsets,
229 repack: Vec::new(),
230 })
231 }
232 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
238 model: model.clone(),
239 idx,
240 dtype: entry.dtype,
241 rows,
242 cols,
243 row_scale: Vec::new(),
244 col_field: Vec::new(),
245 vbit_offsets: Vec::new(),
246 repack: Vec::new(),
247 }),
248 TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
251 model: model.clone(),
252 idx,
253 dtype: entry.dtype,
254 rows,
255 cols,
256 row_scale: Vec::new(),
257 col_field: Vec::new(),
258 vbit_offsets: Vec::new(),
259 repack: Vec::new(),
260 }),
261 TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
263 model: model.clone(),
264 idx,
265 dtype: entry.dtype,
266 rows,
267 cols,
268 row_scale: Vec::new(),
269 col_field: Vec::new(),
270 vbit_offsets: Vec::new(),
271 repack: Vec::new(),
272 }),
273 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
274 model: model.clone(),
275 idx,
276 dtype: entry.dtype,
277 rows,
278 cols,
279 row_scale: Vec::new(),
280 col_field: Vec::new(),
281 vbit_offsets: Vec::new(),
282 repack: Vec::new(),
283 }),
284 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
286 model: model.clone(),
287 idx,
288 dtype: entry.dtype,
289 rows,
290 cols,
291 row_scale: Vec::new(),
292 col_field: Vec::new(),
293 vbit_offsets: Vec::new(),
294 repack: Vec::new(),
295 }),
296 TensorDtype::Q1T 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 _ => {
312 let mut data = vec![0.0f32; rows * cols];
313 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
314 Ok(Self::from_f32(data, rows, cols))
315 }
316 }
317 }
318
319 pub(crate) fn is_q1(&self) -> bool {
322 matches!(
323 self,
324 Self::Mapped {
325 dtype: TensorDtype::Q1,
326 ..
327 }
328 )
329 }
330
331 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
334 match self {
335 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
336 _ => None,
337 }
338 }
339
340 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
347 match self {
348 #[cfg(target_os = "macos")]
349 Self::Mapped {
350 dtype: TensorDtype::Q1T,
351 ..
352 } if !crate::gpu::metal_q1t_enabled() => None,
353 Self::Mapped {
354 idx,
355 dtype:
356 TensorDtype::Q1
357 | TensorDtype::Q1T
358 | TensorDtype::Q4Block
359 | TensorDtype::Q4Tiled
360 | TensorDtype::Q4TiledP
361 | TensorDtype::Q2TiledP
362 | TensorDtype::Q8Row
363 | TensorDtype::Q8_2f,
364 rows,
365 cols,
366 ..
367 } => Some((*idx, *rows, *cols)),
368 _ => None,
369 }
370 }
371
372 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
378 match self {
379 Self::Mapped {
380 idx,
381 dtype: TensorDtype::Q4Tiled,
382 rows,
383 cols,
384 ..
385 } => Some((*idx, *rows, *cols)),
386 _ => None,
387 }
388 }
389
390 pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
394 match self {
395 Self::Mapped {
396 idx,
397 dtype: TensorDtype::Q4TiledP,
398 rows,
399 cols,
400 ..
401 } => Some((*idx, *rows, *cols)),
402 _ => None,
403 }
404 }
405
406 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
411 match self {
412 Self::Mapped {
413 idx,
414 dtype: TensorDtype::Q8Row,
415 rows,
416 cols,
417 row_scale,
418 col_field,
419 ..
420 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
421 _ => None,
422 }
423 }
424
425 pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
429 match self {
430 Self::Mapped { dtype, .. } => Some(*dtype),
431 _ => None,
432 }
433 }
434
435 pub fn model_idx(&self) -> Option<usize> {
440 match self {
441 Self::Mapped { idx, .. } => Some(*idx),
442 _ => None,
443 }
444 }
445
446 pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
451 match self {
452 Self::Mapped { model, .. } => Some(model.clone()),
453 _ => None,
454 }
455 }
456
457 pub fn rows(&self) -> usize {
458 match self {
459 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
460 }
461 }
462
463 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
466 match self {
467 Self::Mapped {
468 model,
469 idx,
470 dtype: TensorDtype::Q4Tiled,
471 ..
472 } => Some((model, *idx)),
473 _ => None,
474 }
475 }
476
477 pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
480 match self {
481 Self::Mapped {
482 model,
483 idx,
484 dtype: TensorDtype::Q4TiledP,
485 ..
486 } => Some((model, *idx)),
487 _ => None,
488 }
489 }
490
491 pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
494 match self {
495 Self::Mapped {
496 model,
497 idx,
498 dtype: TensorDtype::Q2TiledP,
499 ..
500 } => Some((model, *idx)),
501 _ => None,
502 }
503 }
504
505 pub fn cols(&self) -> usize {
506 match self {
507 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
508 }
509 }
510
511 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
514 match self {
515 Self::Mapped {
516 model,
517 idx,
518 dtype: TensorDtype::Q1,
519 ..
520 } => Some((model, *idx)),
521 _ => None,
522 }
523 }
524
525 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
529 match self {
530 Self::Mapped {
531 model,
532 idx,
533 dtype: TensorDtype::Q8Row,
534 row_scale,
535 ..
536 } => Some((model, *idx, 0, row_scale.as_slice())),
537 Self::Mapped {
538 model,
539 idx,
540 dtype: TensorDtype::Q1,
541 ..
542 } => Some((model, *idx, 1, &[])),
543 Self::Mapped {
548 model,
549 idx,
550 dtype: TensorDtype::Q4Tiled,
551 ..
552 } => Some((model, *idx, 5, &[])),
553 Self::Mapped {
557 model,
558 idx,
559 dtype: TensorDtype::Q4TiledP,
560 ..
561 } => Some((model, *idx, 6, &[])),
562 Self::Mapped {
563 model,
564 idx,
565 dtype: TensorDtype::Q4Block,
566 ..
567 } => Some((model, *idx, 2, &[])),
568 Self::Mapped {
569 model,
570 idx,
571 dtype: TensorDtype::Q1T,
572 ..
573 } => Some((model, *idx, 3, &[])),
574 _ => None,
575 }
576 }
577
578 pub fn as_f32(&self) -> Option<&[f32]> {
581 match self {
582 Self::F32 { data, .. } => Some(data),
583 Self::Mapped { .. } => None,
584 }
585 }
586
587 fn quant_bytes(&self) -> &[u8] {
588 match self {
589 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
590 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
591 }
592 }
593
594 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
596 let cols = self.cols();
597 debug_assert_eq!(dst.len(), cols);
598 match self {
599 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
600 Self::Mapped {
601 dtype,
602 row_scale,
603 col_field,
604 vbit_offsets,
605 ..
606 } => {
607 if *dtype == TensorDtype::Q4Tiled {
608 let bytes = self.quant_bytes();
609 let gpr = cols / GROUP_SIZE;
610 for gi in 0..gpr {
611 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
612 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
613 for (k, &b) in tile[2..].iter().enumerate() {
614 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
615 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
616 }
617 }
618 return;
619 }
620 if *dtype == TensorDtype::Q4TiledP {
621 let bytes = self.quant_bytes();
622 let gpr = cols / GROUP_SIZE;
623 let v = Q4tpView::new(bytes, self.rows(), cols);
624 let mut sc = vec![0f32; gpr];
625 v.scales_into(r, gpr, &mut sc);
626 for gi in 0..gpr {
627 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
628 let s = sc[gi];
629 for (k, &b) in tile.iter().enumerate() {
630 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
631 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
632 }
633 }
634 return;
635 }
636 if *dtype == TensorDtype::Q2TiledP {
637 let bytes = self.quant_bytes();
638 let gpr = cols / GROUP_SIZE;
639 let v = Q4tpView::new_q2(bytes, self.rows(), cols);
640 let mut sc = vec![0f32; gpr];
641 v.scales_into(r, gpr, &mut sc);
642 for gi in 0..gpr {
643 let ch =
644 &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
645 let s = sc[gi];
646 for (k, &b) in ch.iter().enumerate() {
647 for j in 0..4 {
648 dst[gi * GROUP_SIZE + k * 4 + j] =
649 (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
650 }
651 }
652 }
653 return;
654 }
655 if *dtype == TensorDtype::Q4Block {
656 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
657 let gpr = cols / GROUP_SIZE;
658 for gi in 0..gpr {
659 let g = r * gpr + gi;
660 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
661 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
662 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
663 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
664 }
665 }
666 return;
667 }
668 if *dtype == TensorDtype::Q1 {
669 let bytes = self.quant_bytes();
670 let gpr = cols / GROUP_SIZE;
671 for gi in 0..gpr {
672 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
673 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
674 for (j, &b) in tile[2..].iter().enumerate() {
675 for k in 0..8 {
676 dst[gi * GROUP_SIZE + j * 8 + k] =
677 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
678 }
679 }
680 }
681 return;
682 }
683 if *dtype == TensorDtype::Q1T {
684 let bytes = self.quant_bytes();
685 let gpr = cols / GROUP_SIZE;
686 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
687 for gi in 0..gpr {
688 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
689 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
690 bytes[off],
691 bytes[off + 1],
692 ]));
693 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
694 for k in 0..GROUP_SIZE {
695 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
696 {
697 1 => s,
698 2 => -s,
699 _ => 0.0,
700 };
701 }
702 }
703 let rows = self.rows();
705 let entries = base_len + (rows + 1) * 4;
706 if entries <= bytes.len() {
707 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
708 let r0 = u32::from_le_bytes([
709 ptrs[r * 4],
710 ptrs[r * 4 + 1],
711 ptrs[r * 4 + 2],
712 ptrs[r * 4 + 3],
713 ]) as usize;
714 let r1 = u32::from_le_bytes([
715 ptrs[(r + 1) * 4],
716 ptrs[(r + 1) * 4 + 1],
717 ptrs[(r + 1) * 4 + 2],
718 ptrs[(r + 1) * 4 + 3],
719 ]) as usize;
720 let off = entries + r0 * 4;
721 for i in 0..r1 - r0 {
722 let item = &bytes[off + i * 4..off + i * 4 + 4];
723 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
724 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
725 item[2], item[3],
726 ]));
727 if c < cols {
728 dst[c] = v;
729 }
730 }
731 }
732 return;
733 }
734 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
735 let bytes = self.quant_bytes();
736 let rows = self.rows();
737 let ng = cols / GROUP_SIZE;
738 let bits = &bytes[..rows];
739 let sc_off = rows;
740 let off = vbit_offsets[r];
743 let b = bits[r] as usize;
744 let l = ((1usize << (b - 1)) - 1) as f32;
745 let data = &bytes[off..];
746 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
747 for (i, d) in dst.iter_mut().enumerate() {
748 while nbits < b {
749 acc = (acc << 8) | data[idx] as u64;
750 idx += 1;
751 nbits += 8;
752 }
753 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
754 nbits -= b;
755 let so = (r * ng + i / GROUP_SIZE) * 2;
756 let sv = f16_to_f32(u16::from_le_bytes([
757 bytes[sc_off + so],
758 bytes[sc_off + so + 1],
759 ]));
760 *d = (u - l) * sv;
761 }
762 return;
763 }
764 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
765 let s = row_scale[r];
766 match dtype {
767 TensorDtype::Q8Row => {
768 for (d, &b) in dst.iter_mut().zip(q) {
769 *d = (b as i8) as f32 * s;
770 }
771 }
772 TensorDtype::Q8_2f => {
773 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
774 *d = (b as i8) as f32 * s * col_field[i];
775 }
776 }
777 _ => unreachable!(),
778 }
779 }
780 }
781 }
782
783 pub fn sparse_col_ok(&self) -> bool {
788 match self {
789 Self::F32 { .. } => true,
790 Self::Mapped { dtype, .. } => {
791 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
792 }
793 }
794 }
795
796 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
800 let inter = self.cols();
801 let hidden = self.rows();
802 debug_assert_eq!(out.len(), hidden);
803 match self {
804 Self::F32 { data, .. } => {
805 for (k, o) in out.iter_mut().enumerate() {
806 *o += w * data[k * inter + c];
807 }
808 }
809 Self::Mapped {
810 dtype,
811 row_scale,
812 col_field,
813 ..
814 } => {
815 let q = self.quant_bytes();
816 let colf = if *dtype == TensorDtype::Q8_2f {
817 col_field[c]
818 } else {
819 1.0
820 };
821 let wc = w * colf;
822 for (k, o) in out.iter_mut().enumerate() {
823 let b = q[k * inter + c] as i8 as f32;
824 *o += wc * b * row_scale[k];
825 }
826 }
827 }
828 }
829
830 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
834 let cols = self.cols();
835 match self {
836 Self::F32 { data, .. } => {
837 let row = &data[r * cols..(r + 1) * cols];
838 row.iter().zip(x).map(|(w, v)| w * v).sum()
839 }
840 Self::Mapped {
841 dtype,
842 row_scale,
843 col_field,
844 ..
845 } => match dtype {
846 TensorDtype::Q8Row => {
847 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
848 dot_i8_f32(q, x) * row_scale[r]
849 }
850 TensorDtype::Q8_2f => {
851 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
852 dot_i8_col_f32(q, x, col_field) * row_scale[r]
853 }
854 _ => {
855 self.row_f32(r, scratch);
856 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
857 }
858 },
859 }
860 }
861
862 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
865 match self {
866 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
870 Self::Mapped {
871 model,
872 idx,
873 dtype,
874 rows,
875 cols,
876 row_scale,
877 col_field,
878 vbit_offsets,
879 repack,
880 } => {
881 let _ = (model, idx);
882 assert!(
892 out.len() >= *rows && x.len() >= *cols,
893 "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
894 out.len(),
895 x.len(),
896 );
897 if *dtype == TensorDtype::Q4Block {
898 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
902 let t0 = std::time::Instant::now();
903 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
904 crate::gpu::ProbeArm::Gpu => {
905 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
906 crate::gpu::probe_record(
907 crate::gpu::OpClass::Matvec,
908 true,
909 t0.elapsed(),
910 );
911 return;
912 }
913 }
914 crate::gpu::ProbeArm::CpuTimed => {
915 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
916 crate::gpu::probe_record(
917 crate::gpu::OpClass::Matvec,
918 false,
919 t0.elapsed(),
920 );
921 return;
922 }
923 crate::gpu::ProbeArm::Cpu => {}
924 }
925 }
926 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
927 return;
928 }
929 if *dtype == TensorDtype::Q4Tiled {
930 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
931 return;
932 }
933 if *dtype == TensorDtype::Q4TiledP {
934 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
940 let t0 = std::time::Instant::now();
941 let cls = crate::gpu::matvec_class(*rows, *cols);
942 match crate::gpu::probe_arm(cls) {
943 crate::gpu::ProbeArm::Gpu => {
944 if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
945 crate::gpu::probe_record(cls, true, t0.elapsed());
946 return;
947 }
948 }
949 crate::gpu::ProbeArm::CpuTimed => {
950 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
951 crate::gpu::probe_record(cls, false, t0.elapsed());
952 return;
953 }
954 crate::gpu::ProbeArm::Cpu => {}
955 }
956 }
957 q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
958 return;
959 }
960 if *dtype == TensorDtype::Q2TiledP {
961 q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
962 return;
963 }
964 if *dtype == TensorDtype::Q1 {
965 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
970 let t0 = std::time::Instant::now();
971 let arm = if crate::gpu::q1_force() {
972 crate::gpu::ProbeArm::Gpu
973 } else {
974 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
975 };
976 match arm {
977 crate::gpu::ProbeArm::Gpu => {
978 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
979 crate::gpu::probe_record(
980 crate::gpu::OpClass::Matvec,
981 true,
982 t0.elapsed(),
983 );
984 return;
985 }
986 }
987 crate::gpu::ProbeArm::CpuTimed => {
988 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
989 crate::gpu::probe_record(
990 crate::gpu::OpClass::Matvec,
991 false,
992 t0.elapsed(),
993 );
994 return;
995 }
996 crate::gpu::ProbeArm::Cpu => {}
997 }
998 }
999 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1000 return;
1001 }
1002 if *dtype == TensorDtype::Q1T {
1003 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1007 let t0 = std::time::Instant::now();
1008 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1009 crate::gpu::ProbeArm::Gpu => {
1010 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1011 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1012 crate::gpu::probe_record(
1013 crate::gpu::OpClass::Matvec,
1014 true,
1015 t0.elapsed(),
1016 );
1017 return;
1018 }
1019 }
1020 crate::gpu::ProbeArm::CpuTimed => {
1021 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1022 crate::gpu::probe_record(
1023 crate::gpu::OpClass::Matvec,
1024 false,
1025 t0.elapsed(),
1026 );
1027 return;
1028 }
1029 crate::gpu::ProbeArm::Cpu => {}
1030 }
1031 }
1032 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1033 return;
1034 }
1035 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1036 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1037 return;
1038 }
1039 let xs = prescale(x, col_field, *dtype);
1040 if *rows >= crate::gpu::min_rows()
1045 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1046 && gpu_lmhead_enabled()
1047 && crate::gpu::enabled_here()
1048 {
1049 let t0 = std::time::Instant::now();
1052 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1053 crate::gpu::ProbeArm::Gpu => {}
1054 crate::gpu::ProbeArm::CpuTimed => {
1055 qmatvec(
1056 self.quant_bytes(),
1057 repack,
1058 row_scale,
1059 x,
1060 col_field,
1061 *dtype,
1062 *rows,
1063 *cols,
1064 out,
1065 pool,
1066 );
1067 crate::gpu::probe_record(
1068 crate::gpu::OpClass::Matvec,
1069 false,
1070 t0.elapsed(),
1071 );
1072 return;
1073 }
1074 crate::gpu::ProbeArm::Cpu => {
1075 qmatvec(
1076 self.quant_bytes(),
1077 repack,
1078 row_scale,
1079 x,
1080 col_field,
1081 *dtype,
1082 *rows,
1083 *cols,
1084 out,
1085 pool,
1086 );
1087 return;
1088 }
1089 }
1090 let frac = gpu_split_frac();
1091 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1092 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1093 let bytes = self.quant_bytes();
1094 let ok = std::thread::scope(|sc| {
1095 let g = sc.spawn(|| {
1096 crate::gpu::q8_matvec_range(
1097 model,
1098 *idx,
1099 cpu_rows,
1100 &row_scale[cpu_rows..],
1101 &xs,
1102 *rows - cpu_rows,
1103 *cols,
1104 out_gpu,
1105 )
1106 });
1107 if cpu_rows > 0 {
1108 let rep_cpu = if repack.is_empty() {
1111 &[][..]
1112 } else {
1113 &repack[..(cpu_rows / 4) * 4 * *cols]
1114 };
1115 qmatvec(
1116 &bytes[..cpu_rows * *cols],
1117 rep_cpu,
1118 &row_scale[..cpu_rows],
1119 x,
1120 col_field,
1121 *dtype,
1122 cpu_rows,
1123 *cols,
1124 out_cpu,
1125 pool,
1126 );
1127 }
1128 g.join().unwrap_or(false)
1129 });
1130 if ok {
1131 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1132 return;
1133 }
1134 qmatvec(
1137 &bytes[cpu_rows * *cols..(*rows) * *cols],
1138 &[],
1139 &row_scale[cpu_rows..],
1140 x,
1141 col_field,
1142 *dtype,
1143 *rows - cpu_rows,
1144 *cols,
1145 out_gpu,
1146 pool,
1147 );
1148 return;
1149 }
1150 qmatvec(
1151 self.quant_bytes(),
1152 repack,
1153 row_scale,
1154 x,
1155 col_field,
1156 *dtype,
1157 *rows,
1158 *cols,
1159 out,
1160 pool,
1161 );
1162 }
1163 }
1164 }
1165
1166 pub fn matvec2(
1168 &self,
1169 x1: &[f32],
1170 x2: &[f32],
1171 o1: &mut [f32],
1172 o2: &mut [f32],
1173 pool: Option<&Pool>,
1174 ) {
1175 match self {
1176 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1177 Self::Mapped {
1178 dtype,
1179 rows,
1180 cols,
1181 row_scale,
1182 col_field,
1183 vbit_offsets,
1184 ..
1185 } => {
1186 if *dtype == TensorDtype::Q4Block {
1187 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1188 return;
1189 }
1190 if *dtype == TensorDtype::Q4Tiled {
1191 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1192 return;
1193 }
1194 if *dtype == TensorDtype::Q4TiledP {
1195 q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1196 return;
1197 }
1198 if *dtype == TensorDtype::Q2TiledP {
1199 q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1200 return;
1201 }
1202 if *dtype == TensorDtype::Q1 {
1203 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1204 return;
1205 }
1206 if *dtype == TensorDtype::Q1T {
1207 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1213 return;
1214 }
1215 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1216 vbitmatvec2(
1217 self.quant_bytes(),
1218 vbit_offsets,
1219 x1,
1220 x2,
1221 *rows,
1222 *cols,
1223 o1,
1224 o2,
1225 pool,
1226 );
1227 return;
1228 }
1229 qmatvec2(
1230 self.quant_bytes(),
1231 row_scale,
1232 x1,
1233 x2,
1234 col_field,
1235 *dtype,
1236 *rows,
1237 *cols,
1238 o1,
1239 o2,
1240 pool,
1241 );
1242 }
1243 }
1244 }
1245}
1246
1247impl QTensor {
1248 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1253 let cols = self.cols();
1254 let rows = self.rows();
1255 debug_assert_eq!(xs_all.len(), b * cols);
1256 debug_assert_eq!(out.len(), b * rows);
1257 if crate::gptq_capture::capturing() {
1261 if let Self::Mapped { model, idx, .. } = self {
1262 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1263 }
1264 }
1265 match self {
1266 Self::F32 { data, .. } => {
1267 let out_addr = SendMut(out.as_mut_ptr());
1268 let run = |start: usize, end: usize| {
1269 for o in start..end {
1270 let row = &data[o * cols..(o + 1) * cols];
1271 for bi in 0..b {
1272 let x = &xs_all[bi * cols..(bi + 1) * cols];
1273 let mut acc = 0f32;
1274 for j in 0..cols {
1275 acc += row[j] * x[j];
1276 }
1277 unsafe { *out_addr.at(bi * rows + o) = acc };
1278 }
1279 }
1280 };
1281 dispatch_rows(pool, rows, &run);
1282 }
1283 Self::Mapped {
1284 dtype,
1285 row_scale,
1286 col_field,
1287 vbit_offsets,
1288 ..
1289 } => {
1290 if *dtype == TensorDtype::Q4Block {
1291 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1292 return;
1293 }
1294 if *dtype == TensorDtype::Q4TiledP {
1295 if b >= 32
1307 && b * rows * cols >= 128_000_000
1308 && cols % 32 == 0
1309 && !crate::gpu::mm_killed()
1310 && crate::gpu::enabled_here()
1311 {
1312 let class = if b >= 128 {
1313 crate::gpu::OpClass::MatmatWide
1314 } else {
1315 crate::gpu::OpClass::Matmat
1316 };
1317 if let Self::Mapped { model, idx, .. } = self {
1318 let t0 = std::time::Instant::now();
1319 match crate::gpu::probe_arm(class) {
1320 crate::gpu::ProbeArm::Gpu => {
1321 if crate::gpu::q4tp_matmat(
1322 model, *idx, xs_all, b, rows, cols, out,
1323 ) {
1324 let el = t0.elapsed();
1325 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1335 let budget = std::time::Duration::from_secs_f64(
1336 flops / 1.5e12 * 8.0 + 0.020,
1337 );
1338 if el > budget && !crate::gpu::probe_was_cold() {
1339 tracing::warn!(
1340 "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1341 device contended, CPU for the rest of the process"
1342 );
1343 crate::gpu::mm_kill();
1344 }
1345 crate::gpu::probe_record(class, true, el);
1346 return;
1347 }
1348 }
1349 crate::gpu::ProbeArm::CpuTimed => {
1350 q4tp_matmat(
1351 self.quant_bytes(),
1352 xs_all,
1353 b,
1354 rows,
1355 cols,
1356 out,
1357 pool,
1358 );
1359 crate::gpu::probe_record(class, false, t0.elapsed());
1360 return;
1361 }
1362 crate::gpu::ProbeArm::Cpu => {}
1363 }
1364 }
1365 }
1366 q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1367 return;
1368 }
1369 if *dtype == TensorDtype::Q2TiledP {
1370 q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1375 return;
1376 }
1377 if *dtype == TensorDtype::Q4Tiled {
1378 if b >= 32
1390 && b * rows * cols >= 128_000_000
1391 && cols % 32 == 0
1392 && !crate::gpu::mm_killed()
1393 && crate::gpu::enabled_here()
1394 {
1395 let class = if b >= 128 {
1396 crate::gpu::OpClass::MatmatWide
1397 } else {
1398 crate::gpu::OpClass::Matmat
1399 };
1400 if let Self::Mapped { model, idx, .. } = self {
1401 let t0 = std::time::Instant::now();
1402 match crate::gpu::probe_arm(class) {
1403 crate::gpu::ProbeArm::Gpu => {
1404 if crate::gpu::q4t_matmat(
1405 model, *idx, xs_all, b, rows, cols, out,
1406 ) {
1407 let el = t0.elapsed();
1408 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1418 let budget = std::time::Duration::from_secs_f64(
1419 flops / 1.5e12 * 8.0 + 0.020,
1420 );
1421 if el > budget && !crate::gpu::probe_was_cold() {
1422 tracing::warn!(
1423 "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1424 device contended, CPU for the rest of the process"
1425 );
1426 crate::gpu::mm_kill();
1427 }
1428 crate::gpu::probe_record(class, true, el);
1429 return;
1430 }
1431 }
1432 crate::gpu::ProbeArm::CpuTimed => {
1433 q4t_matmat(
1434 self.quant_bytes(),
1435 xs_all,
1436 b,
1437 rows,
1438 cols,
1439 out,
1440 pool,
1441 );
1442 crate::gpu::probe_record(class, false, t0.elapsed());
1443 return;
1444 }
1445 crate::gpu::ProbeArm::Cpu => {}
1446 }
1447 }
1448 }
1449 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1450 return;
1451 }
1452 if *dtype == TensorDtype::Q1 {
1453 if b >= 32
1456 && b * rows * cols >= 128_000_000
1457 && cols % 64 == 0
1458 && crate::gpu::enabled_here()
1459 {
1460 if let Self::Mapped { model, idx, .. } = self {
1461 let t0 = std::time::Instant::now();
1462 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1463 crate::gpu::ProbeArm::Gpu => {
1464 if crate::gpu::q1_matmat(
1465 model, *idx, xs_all, b, rows, cols, out,
1466 ) {
1467 crate::gpu::probe_record(
1468 crate::gpu::OpClass::Matmat,
1469 true,
1470 t0.elapsed(),
1471 );
1472 return;
1473 }
1474 }
1475 crate::gpu::ProbeArm::CpuTimed => {
1476 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1477 crate::gpu::probe_record(
1478 crate::gpu::OpClass::Matmat,
1479 false,
1480 t0.elapsed(),
1481 );
1482 return;
1483 }
1484 crate::gpu::ProbeArm::Cpu => {}
1485 }
1486 }
1487 }
1488 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1489 return;
1490 }
1491 if *dtype == TensorDtype::Q1T {
1492 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1495 if let Self::Mapped { model, idx, .. } = self {
1496 let t0 = std::time::Instant::now();
1497 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1498 crate::gpu::ProbeArm::Gpu => {
1499 if crate::gpu::q1t_matmat(
1500 model, *idx, xs_all, b, rows, cols, out,
1501 ) {
1502 crate::gpu::probe_record(
1503 crate::gpu::OpClass::Matmat,
1504 true,
1505 t0.elapsed(),
1506 );
1507 return;
1508 }
1509 }
1510 crate::gpu::ProbeArm::CpuTimed => {
1511 q1t_matmat(
1512 self.quant_bytes(),
1513 xs_all,
1514 b,
1515 rows,
1516 cols,
1517 out,
1518 pool,
1519 );
1520 crate::gpu::probe_record(
1521 crate::gpu::OpClass::Matmat,
1522 false,
1523 t0.elapsed(),
1524 );
1525 return;
1526 }
1527 crate::gpu::ProbeArm::Cpu => {}
1528 }
1529 }
1530 }
1531 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1532 return;
1533 }
1534 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1535 vbitmatmat(
1536 self.quant_bytes(),
1537 vbit_offsets,
1538 xs_all,
1539 b,
1540 rows,
1541 cols,
1542 out,
1543 pool,
1544 );
1545 return;
1546 }
1547 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1548 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1549 .collect();
1550 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1556 if let Self::Mapped { model, idx, .. } = self {
1557 let t0 = std::time::Instant::now();
1558 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1559 crate::gpu::ProbeArm::Gpu
1560 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1561 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1562 {
1563 let q = self.quant_bytes();
1567 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1568 return;
1569 }
1570 crate::gpu::ProbeArm::Gpu => {
1571 let flat: Vec<f32> =
1572 pre.iter().flat_map(|v| v.iter().copied()).collect();
1573 if crate::gpu::q8_matmat(
1574 model, *idx, row_scale, &flat, b, rows, cols, out,
1575 ) {
1576 crate::gpu::probe_record(
1577 crate::gpu::OpClass::Matmat,
1578 true,
1579 t0.elapsed(),
1580 );
1581 return;
1582 }
1583 }
1584 crate::gpu::ProbeArm::CpuTimed => {
1585 let q = self.quant_bytes();
1586 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1587 crate::gpu::probe_record(
1588 crate::gpu::OpClass::Matmat,
1589 false,
1590 t0.elapsed(),
1591 );
1592 return;
1593 }
1594 crate::gpu::ProbeArm::Cpu => {}
1595 }
1596 }
1597 }
1598 let q = self.quant_bytes();
1599 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1600 }
1601 }
1602 }
1603}
1604
1605impl QTensor {
1606 pub fn matvec_many<const N: usize>(
1613 ts: [&QTensor; N],
1614 x: &[f32],
1615 mut outs: [&mut [f32]; N],
1616 pool: Option<&Pool>,
1617 ) {
1618 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1619 let uniform_q8 = ts.iter().all(|t| {
1620 matches!(
1621 t,
1622 Self::Mapped {
1623 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1624 ..
1625 }
1626 )
1627 });
1628 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1629 let uniform_q4 = ts.iter().all(|t| {
1630 matches!(
1631 t,
1632 Self::Mapped {
1633 dtype: TensorDtype::Q4Block,
1634 ..
1635 }
1636 )
1637 });
1638 let uniform_vbit = ts.iter().all(|t| {
1639 matches!(
1640 t,
1641 Self::Mapped {
1642 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1643 ..
1644 }
1645 )
1646 });
1647 let uniform_q1 = ts.iter().all(|t| {
1648 matches!(
1649 t,
1650 Self::Mapped {
1651 dtype: TensorDtype::Q1,
1652 ..
1653 }
1654 )
1655 });
1656 let uniform_q1t = ts.iter().all(|t| {
1657 matches!(
1658 t,
1659 Self::Mapped {
1660 dtype: TensorDtype::Q1T,
1661 ..
1662 }
1663 )
1664 });
1665 let uniform_q4tp = ts.iter().all(|t| {
1670 matches!(
1671 t,
1672 Self::Mapped {
1673 dtype: TensorDtype::Q4TiledP,
1674 ..
1675 }
1676 )
1677 }) && ts.iter().all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1678 let Some(pool) = pool else {
1679 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1680 t.matvec(x, o, None);
1681 }
1682 return;
1683 };
1684 if total_rows < 256
1685 || !(uniform_q8
1686 || uniform_f32
1687 || uniform_q4
1688 || uniform_vbit
1689 || uniform_q1
1690 || uniform_q1t
1691 || uniform_q4tp)
1692 {
1693 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1694 t.matvec(x, o, Some(pool));
1695 }
1696 return;
1697 }
1698
1699 if uniform_q4tp {
1700 let cols = ts[0].cols();
1706 let gpr = cols / GROUP_SIZE;
1707 let views: Vec<Q4tpView> = ts
1708 .iter()
1709 .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1710 .collect();
1711 let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1712 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1713 let locate = |flat: usize| -> (usize, usize) {
1715 let mut acc = 0;
1716 for (i, &r) in rows_of.iter().enumerate() {
1717 if flat < acc + r {
1718 return (i, flat - acc);
1719 }
1720 acc += r;
1721 }
1722 (rows_of.len() - 1, 0)
1723 };
1724 let (views, outs_addr) = (&views, &outs_addr);
1725 if a8w8_enabled() {
1726 let act = split_act(x);
1727 let act = &act;
1728 let run = |start: usize, end: usize| {
1729 let mut sc = vec![0f32; gpr];
1730 for flat in start..end {
1731 let (t, r) = locate(flat);
1732 let v = &views[t];
1733 v.scales_into(r, gpr, &mut sc);
1734 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1735 for &(j, xv) in &act.outliers {
1736 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1737 acc += w * s * xv;
1738 }
1739 unsafe { *outs_addr[t].at(r) = acc };
1741 }
1742 };
1743 pool.run_rows(total_rows, &run);
1744 } else {
1745 let run = |start: usize, end: usize| {
1746 let mut sc = vec![0f32; gpr];
1747 for flat in start..end {
1748 let (t, r) = locate(flat);
1749 let v = &views[t];
1750 v.scales_into(r, gpr, &mut sc);
1751 unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1753 }
1754 };
1755 pool.run_rows(total_rows, &run);
1756 }
1757 return;
1758 }
1759
1760 if uniform_q1 {
1761 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1764 if a8w8_enabled() {
1765 let act = split_act(x);
1766 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1767 let (act, gsum) = (&act, &gsum);
1768 let closures: [_; N] = std::array::from_fn(|i| {
1769 let (bytes, gpr, out) =
1770 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1771 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1772 });
1773 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1774 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1775 pool.run_many(&parts);
1776 } else {
1777 let closures: [_; N] = std::array::from_fn(|i| {
1778 let (bytes, gpr, out) =
1779 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1780 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1781 });
1782 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1783 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1784 pool.run_many(&parts);
1785 }
1786 return;
1787 }
1788
1789 if uniform_q1t {
1790 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1794 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1795 if a8w8_enabled() {
1796 let act = split_act(x);
1797 let act = &act;
1798 let x_ref = x;
1799 let closures: [_; N] = std::array::from_fn(|i| {
1800 let bytes = ts[i].quant_bytes();
1801 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1802 let gpr = cols / GROUP_SIZE;
1803 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1804 let out = outs_addr[i];
1805 move |s: usize, e: usize| {
1806 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1807 }
1808 });
1809 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1810 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1811 pool.run_many(&parts);
1812 } else {
1813 let x_ref = x;
1814 let closures: [_; N] = std::array::from_fn(|i| {
1815 let bytes = ts[i].quant_bytes();
1816 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1817 let gpr = cols / GROUP_SIZE;
1818 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1819 let out = outs_addr[i];
1820 move |s: usize, e: usize| {
1821 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1822 }
1823 });
1824 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1825 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1826 pool.run_many(&parts);
1827 }
1828 return;
1829 }
1830
1831 if uniform_q4 || uniform_vbit {
1832 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1833 if a8w8_enabled() {
1835 let act = split_act(x);
1836 let act = &act;
1837 if uniform_q4 {
1838 let closures: [_; N] = std::array::from_fn(|i| {
1839 let (packed, scales) =
1840 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1841 let (gpr, cols, out) =
1842 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1843 move |s: usize, e: usize| {
1844 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1845 }
1846 });
1847 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1848 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1849 pool.run_many(&parts);
1850 } else {
1851 let closures: [_; N] = std::array::from_fn(|i| {
1852 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1853 unreachable!()
1854 };
1855 let (bytes, rows, cols, out) = (
1856 ts[i].quant_bytes(),
1857 ts[i].rows(),
1858 ts[i].cols(),
1859 outs_addr[i],
1860 );
1861 move |s: usize, e: usize| {
1862 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1863 }
1864 });
1865 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1866 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1867 pool.run_many(&parts);
1868 }
1869 return;
1870 }
1871 if uniform_q4 {
1872 let closures: [_; N] = std::array::from_fn(|i| {
1873 let (packed, scales) =
1874 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1875 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1876 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1877 });
1878 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1879 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1880 pool.run_many(&parts);
1881 } else {
1882 let closures: [_; N] = std::array::from_fn(|i| {
1883 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1884 unreachable!()
1885 };
1886 let (bytes, rows, cols, out) = (
1887 ts[i].quant_bytes(),
1888 ts[i].rows(),
1889 ts[i].cols(),
1890 outs_addr[i],
1891 );
1892 move |s: usize, e: usize| {
1893 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1894 }
1895 });
1896 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1897 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1898 pool.run_many(&parts);
1899 }
1900 return;
1901 }
1902
1903 if uniform_f32 {
1904 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1905 let closures: [_; N] = std::array::from_fn(|i| {
1906 let Self::F32 { data, cols, .. } = ts[i] else {
1907 unreachable!()
1908 };
1909 let out = outs_addr[i];
1910 move |start: usize, end: usize| {
1911 for o in start..end {
1912 let row = &data[o * cols..(o + 1) * cols];
1913 let mut sum = 0.0f32;
1914 for j in 0..*cols {
1915 sum += row[j] * x[j];
1916 }
1917 unsafe { *out.at(o) = sum };
1919 }
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 return;
1926 }
1927
1928 struct Ctx<'a> {
1931 bytes: &'a [u8],
1932 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1933 rep: &'a [u8],
1934 row_scale: &'a [f32],
1935 cols: usize,
1936 xs: std::borrow::Cow<'a, [f32]>,
1937 }
1938 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1939 let Self::Mapped {
1940 dtype,
1941 cols,
1942 row_scale,
1943 col_field,
1944 repack,
1945 ..
1946 } = ts[i]
1947 else {
1948 unreachable!()
1949 };
1950 Ctx {
1951 bytes: ts[i].quant_bytes(),
1952 rep: repack,
1953 row_scale,
1954 cols: *cols,
1955 xs: prescale(x, col_field, *dtype),
1956 }
1957 });
1958 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1959 #[cfg(target_arch = "aarch64")]
1960 if sdot_enabled() {
1961 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1962 let closures: [_; N] = std::array::from_fn(|i| {
1963 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1964 move |start: usize, end: usize| {
1965 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1966 }
1967 });
1968 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1969 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1970 pool.run_many(&parts);
1971 return;
1972 }
1973 #[cfg(target_arch = "x86_64")]
1974 if avx2_a8w8_enabled() {
1975 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1976 let closures: [_; N] = std::array::from_fn(|i| {
1977 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1978 move |start: usize, end: usize| {
1979 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1980 }
1981 });
1982 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1983 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1984 pool.run_many(&parts);
1985 return;
1986 }
1987 let closures: [_; N] = std::array::from_fn(|i| {
1988 let (c, out) = (&ctxs[i], outs_addr[i]);
1989 move |start: usize, end: usize| {
1990 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1991 }
1992 });
1993 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1994 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1995 pool.run_many(&parts);
1996 }
1997}
1998
1999impl QTensor {
2000 #[allow(clippy::needless_range_loop)]
2005 pub fn matvec2_many<const N: usize>(
2006 ts: [&QTensor; N],
2007 x1: &[f32],
2008 x2: &[f32],
2009 mut o1s: [&mut [f32]; N],
2010 mut o2s: [&mut [f32]; N],
2011 pool: Option<&Pool>,
2012 ) {
2013 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2014 let uniform_q8 = ts.iter().all(|t| {
2015 matches!(
2016 t,
2017 Self::Mapped {
2018 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2019 ..
2020 }
2021 )
2022 });
2023 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2024 let uniform_q4 = ts.iter().all(|t| {
2025 matches!(
2026 t,
2027 Self::Mapped {
2028 dtype: TensorDtype::Q4Block,
2029 ..
2030 }
2031 )
2032 });
2033 let uniform_vbit = ts.iter().all(|t| {
2034 matches!(
2035 t,
2036 Self::Mapped {
2037 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2038 ..
2039 }
2040 )
2041 });
2042 let fusable = pool.is_some()
2043 && total_rows >= 256
2044 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2045 if !fusable {
2046 for i in 0..N {
2047 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2048 }
2049 return;
2050 }
2051 let pool = pool.unwrap();
2052
2053 if uniform_q4 || uniform_vbit {
2054 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2055 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2056 if a8w8_enabled() {
2058 let a1 = split_act(x1);
2059 let a2 = split_act(x2);
2060 let (a1, a2) = (&a1, &a2);
2061 if uniform_q4 {
2062 let closures: [_; N] = std::array::from_fn(|i| {
2063 let (packed, scales) =
2064 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2065 let (gpr, cols, o1, o2) =
2066 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2067 move |s: usize, e: usize| {
2068 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2069 }
2070 });
2071 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2072 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2073 pool.run_many(&parts);
2074 } else {
2075 let closures: [_; N] = std::array::from_fn(|i| {
2076 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2077 unreachable!()
2078 };
2079 let (bytes, rows, cols, o1, o2) = (
2080 ts[i].quant_bytes(),
2081 ts[i].rows(),
2082 ts[i].cols(),
2083 p1[i],
2084 p2[i],
2085 );
2086 move |s: usize, e: usize| {
2087 vbit_range2_a8w8(
2088 bytes,
2089 vbit_offsets,
2090 x1,
2091 x2,
2092 a1,
2093 a2,
2094 rows,
2095 cols,
2096 o1,
2097 o2,
2098 s,
2099 e,
2100 )
2101 }
2102 });
2103 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2104 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2105 pool.run_many(&parts);
2106 }
2107 return;
2108 }
2109 if uniform_q4 {
2110 let closures: [_; N] = std::array::from_fn(|i| {
2111 let (packed, scales) =
2112 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2113 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2114 move |s: usize, e: usize| {
2115 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2116 }
2117 });
2118 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2119 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2120 pool.run_many(&parts);
2121 } else {
2122 let closures: [_; N] = std::array::from_fn(|i| {
2123 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2124 unreachable!()
2125 };
2126 let (bytes, rows, cols, o1, o2) = (
2127 ts[i].quant_bytes(),
2128 ts[i].rows(),
2129 ts[i].cols(),
2130 p1[i],
2131 p2[i],
2132 );
2133 move |s: usize, e: usize| {
2134 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2135 }
2136 });
2137 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2138 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2139 pool.run_many(&parts);
2140 }
2141 return;
2142 }
2143
2144 if uniform_f32 {
2145 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2146 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2147 let closures: [_; N] = std::array::from_fn(|i| {
2148 let Self::F32 { data, cols, .. } = ts[i] else {
2149 unreachable!()
2150 };
2151 let (o1, o2) = (p1[i], p2[i]);
2152 move |start: usize, end: usize| {
2153 for o in start..end {
2154 let row = &data[o * cols..(o + 1) * cols];
2155 let (mut s1, mut s2) = (0.0f32, 0.0f32);
2156 for j in 0..*cols {
2157 s1 += row[j] * x1[j];
2158 s2 += row[j] * x2[j];
2159 }
2160 unsafe {
2162 *o1.at(o) = s1;
2163 *o2.at(o) = s2;
2164 }
2165 }
2166 }
2167 });
2168 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2169 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2170 pool.run_many(&parts);
2171 return;
2172 }
2173
2174 struct Ctx<'a> {
2175 bytes: &'a [u8],
2176 row_scale: &'a [f32],
2177 cols: usize,
2178 xs1: std::borrow::Cow<'a, [f32]>,
2179 xs2: std::borrow::Cow<'a, [f32]>,
2180 }
2181 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2182 let Self::Mapped {
2183 dtype,
2184 cols,
2185 row_scale,
2186 col_field,
2187 ..
2188 } = ts[i]
2189 else {
2190 unreachable!()
2191 };
2192 Ctx {
2193 bytes: ts[i].quant_bytes(),
2194 row_scale,
2195 cols: *cols,
2196 xs1: prescale(x1, col_field, *dtype),
2197 xs2: prescale(x2, col_field, *dtype),
2198 }
2199 });
2200 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2201 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2202 #[cfg(target_arch = "aarch64")]
2203 if sdot_enabled() {
2204 let acts: [(SplitAct, SplitAct); N] =
2205 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2206 let closures: [_; N] = std::array::from_fn(|i| {
2207 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2208 move |start: usize, end: usize| {
2209 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2210 }
2211 });
2212 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2213 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2214 pool.run_many(&parts);
2215 return;
2216 }
2217 #[cfg(target_arch = "x86_64")]
2218 if avx2_a8w8_enabled() {
2219 let acts: [(SplitAct, SplitAct); N] =
2220 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2221 let closures: [_; N] = std::array::from_fn(|i| {
2222 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2223 move |start: usize, end: usize| {
2224 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2225 }
2226 });
2227 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2228 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2229 pool.run_many(&parts);
2230 return;
2231 }
2232 let closures: [_; N] = std::array::from_fn(|i| {
2233 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2234 move |start: usize, end: usize| {
2235 q8_range2_f32(
2236 c.bytes,
2237 c.row_scale,
2238 &c.xs1,
2239 &c.xs2,
2240 c.cols,
2241 o1,
2242 o2,
2243 start,
2244 end,
2245 )
2246 }
2247 });
2248 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2249 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2250 pool.run_many(&parts);
2251 }
2252
2253 pub fn matvec_silu_mul(
2258 gate: &QTensor,
2259 up: &QTensor,
2260 x: &[f32],
2261 out: &mut [f32],
2262 pool: Option<&Pool>,
2263 ) -> bool {
2264 let inter = gate.rows();
2265 debug_assert_eq!(up.rows(), inter);
2266 debug_assert_eq!(out.len(), inter);
2267 debug_assert_eq!(gate.cols(), up.cols());
2268 if !a8w8_enabled() {
2269 return false;
2270 }
2271 let act = split_act(x);
2272 let act = &act;
2273 let x_ref = x;
2274 let out_addr = SendMut(out.as_mut_ptr());
2275
2276 match (gate, up) {
2277 (
2279 Self::Mapped {
2280 dtype: TensorDtype::Q4Block,
2281 ..
2282 },
2283 Self::Mapped {
2284 dtype: TensorDtype::Q4Block,
2285 ..
2286 },
2287 ) => {
2288 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2289 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2290 let gpr = gate.cols() / GROUP_SIZE;
2291 let cols = gate.cols();
2292 let run = move |start: usize, end: usize| {
2293 for r in start..end {
2294 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2295 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2296 for &(j, xv) in &act.outliers {
2297 let flat = r * cols + j;
2298 let gb = gp[flat / 2];
2299 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2300 let gsc = f16_to_f32(u16::from_le_bytes([
2301 gs[(flat / GROUP_SIZE) * 2],
2302 gs[(flat / GROUP_SIZE) * 2 + 1],
2303 ]));
2304 gv += ((gn as i32 - 8) as f32) * gsc * xv;
2305 let ub = up_p[flat / 2];
2306 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2307 let usc = f16_to_f32(u16::from_le_bytes([
2308 up_s[(flat / GROUP_SIZE) * 2],
2309 up_s[(flat / GROUP_SIZE) * 2 + 1],
2310 ]));
2311 uv += ((un as i32 - 8) as f32) * usc * xv;
2312 }
2313 let silu_g = gv / (1.0 + (-gv).exp());
2314 unsafe { *out_addr.at(r) = silu_g * uv };
2316 }
2317 };
2318 dispatch_rows(pool, inter, &run);
2319 true
2320 }
2321 (
2325 Self::Mapped {
2326 dtype: TensorDtype::Q4Tiled,
2327 ..
2328 },
2329 Self::Mapped {
2330 dtype: TensorDtype::Q4Tiled,
2331 ..
2332 },
2333 ) => {
2334 let g_bytes = gate.quant_bytes();
2335 let u_bytes = up.quant_bytes();
2336 let gpr = gate.cols() / GROUP_SIZE;
2337 let run = move |start: usize, end: usize| {
2338 for r in start..end {
2339 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2340 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2341 for &(j, xv) in &act.outliers {
2342 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2343 gv += w * s * xv;
2344 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2345 uv += w * s * xv;
2346 }
2347 let silu_g = gv / (1.0 + (-gv).exp());
2348 unsafe { *out_addr.at(r) = silu_g * uv };
2350 }
2351 };
2352 dispatch_rows(pool, inter, &run);
2353 true
2354 }
2355 (
2358 Self::Mapped {
2359 dtype: TensorDtype::Q4TiledP,
2360 ..
2361 },
2362 Self::Mapped {
2363 dtype: TensorDtype::Q4TiledP,
2364 ..
2365 },
2366 ) => {
2367 let cols = gate.cols();
2368 let gpr = cols / GROUP_SIZE;
2369 let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2370 let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2371 let run = |start: usize, end: usize| {
2372 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2373 for r in start..end {
2374 gv_view.scales_into(r, gpr, &mut gsc);
2375 uv_view.scales_into(r, gpr, &mut usc);
2376 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2377 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2378 for &(j, xv) in &act.outliers {
2379 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2380 gv += w * s * xv;
2381 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2382 uv += w * s * xv;
2383 }
2384 let silu_g = gv / (1.0 + (-gv).exp());
2385 unsafe { *out_addr.at(r) = silu_g * uv };
2387 }
2388 };
2389 dispatch_rows(pool, inter, &run);
2390 true
2391 }
2392 (
2394 Self::Mapped {
2395 dtype: TensorDtype::Q1T,
2396 ..
2397 },
2398 Self::Mapped {
2399 dtype: TensorDtype::Q1T,
2400 ..
2401 },
2402 ) => {
2403 const TILE: usize = cortiq_core::quant::Q1T_TILE;
2404 let g_bytes = gate.quant_bytes();
2405 let u_bytes = up.quant_bytes();
2406 let gpr = gate.cols() / GROUP_SIZE;
2407 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2408 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2409 let run = move |start: usize, end: usize| {
2410 for r in start..end {
2411 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2412 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2413 for &(j, xv) in &act.outliers {
2414 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2415 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2416 }
2417 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2418 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2419 let silu_g = gv / (1.0 + (-gv).exp());
2420 unsafe { *out_addr.at(r) = silu_g * uv };
2422 }
2423 };
2424 dispatch_rows(pool, inter, &run);
2425 true
2426 }
2427 _ => false,
2428 }
2429 }
2430
2431 pub fn moe_gate_up_many(
2445 pairs: &[(&QTensor, &QTensor)],
2446 x: &[f32],
2447 outs: &mut [Vec<f32>],
2448 pool: Option<&Pool>,
2449 ) -> bool {
2450 if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2451 return false;
2452 }
2453 let inter = pairs[0].0.rows();
2454 let cols = pairs[0].0.cols();
2455 if cols % GROUP_SIZE != 0 {
2456 return false;
2457 }
2458 let gpr = cols / GROUP_SIZE;
2459 let mut views = Vec::with_capacity(pairs.len() * 2);
2460 for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2461 let both_q4tp = matches!(
2462 g,
2463 Self::Mapped {
2464 dtype: TensorDtype::Q4TiledP,
2465 ..
2466 }
2467 ) && matches!(
2468 u,
2469 Self::Mapped {
2470 dtype: TensorDtype::Q4TiledP,
2471 ..
2472 }
2473 );
2474 if !both_q4tp
2475 || g.rows() != inter
2476 || u.rows() != inter
2477 || g.cols() != cols
2478 || u.cols() != cols
2479 || o.len() != inter
2480 {
2481 return false;
2482 }
2483 views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
2484 views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
2485 }
2486 let act = split_act(x);
2487 let act = &act;
2488 let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2489 let (views, ptrs) = (&views, &ptrs);
2490 let run = |start: usize, end: usize| {
2491 let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2492 for flat in start..end {
2493 let (e, r) = (flat / inter, flat % inter);
2494 let gv_view = &views[e * 2];
2495 let uv_view = &views[e * 2 + 1];
2496 gv_view.scales_into(r, gpr, &mut gsc);
2497 uv_view.scales_into(r, gpr, &mut usc);
2498 let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2499 let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2500 for &(j, xv) in &act.outliers {
2501 let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2502 gv += w * s * xv;
2503 let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2504 uv += w * s * xv;
2505 }
2506 let silu_g = gv / (1.0 + (-gv).exp());
2507 unsafe { *ptrs[e].at(r) = silu_g * uv };
2509 }
2510 };
2511 dispatch_rows(pool, pairs.len() * inter, &run);
2512 true
2513 }
2514
2515 pub fn moe_down_many(
2524 downs: &[&QTensor],
2525 gs: &[Vec<f32>],
2526 weights: &[f32],
2527 out: &mut [f32],
2528 pool: Option<&Pool>,
2529 ) -> bool {
2530 if downs.is_empty()
2531 || downs.len() != gs.len()
2532 || downs.len() != weights.len()
2533 || !a8w8_enabled()
2534 {
2535 return false;
2536 }
2537 let rows = out.len();
2538 let cols = downs[0].cols();
2539 if cols % GROUP_SIZE != 0 {
2540 return false;
2541 }
2542 let gpr = cols / GROUP_SIZE;
2543 let mut views = Vec::with_capacity(downs.len());
2544 for (d, g) in downs.iter().zip(gs.iter()) {
2545 if !matches!(
2546 d,
2547 Self::Mapped {
2548 dtype: TensorDtype::Q4TiledP,
2549 ..
2550 }
2551 ) || d.rows() != rows
2552 || d.cols() != cols
2553 || g.len() != cols
2554 {
2555 return false;
2556 }
2557 views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2558 }
2559 let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2561 let out_addr = SendMut(out.as_mut_ptr());
2569 let (views, acts, weights) = (&views, &acts, &weights);
2570 let run = |start: usize, end: usize| {
2571 let mut sc = vec![0f32; gpr];
2572 for r in start..end {
2573 let mut acc = 0f32;
2574 for (e, v) in views.iter().enumerate() {
2575 v.scales_into(r, gpr, &mut sc);
2576 let a = &acts[e];
2577 let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2578 for &(j, xv) in &a.outliers {
2579 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2580 d += w * s * xv;
2581 }
2582 acc += weights[e] * d;
2583 }
2584 unsafe { *out_addr.at(r) = acc };
2586 }
2587 };
2588 dispatch_rows(pool, rows, &run);
2589 true
2590 }
2591}
2592
2593#[cfg(target_os = "macos")]
2598mod accel_blas {
2599 #[link(name = "Accelerate", kind = "framework")]
2600 unsafe extern "C" {
2601 pub fn cblas_sgemm(
2602 order: i32,
2603 trans_a: i32,
2604 trans_b: i32,
2605 m: i32,
2606 n: i32,
2607 k: i32,
2608 alpha: f32,
2609 a: *const f32,
2610 lda: i32,
2611 b: *const f32,
2612 ldb: i32,
2613 beta: f32,
2614 c: *mut f32,
2615 ldc: i32,
2616 );
2617 }
2618}
2619
2620#[cfg(target_os = "macos")]
2621pub(crate) fn accel_gemm_enabled() -> bool {
2622 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2623 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2624}
2625
2626#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2629pub(crate) fn accel_gemm_enabled() -> bool {
2630 true
2631}
2632
2633#[cfg(target_arch = "aarch64")]
2640#[allow(clippy::too_many_arguments)]
2641pub(crate) fn neon_gemm_rm(
2642 m: usize,
2643 n: usize,
2644 k: usize,
2645 alpha: f32,
2646 a: &[f32],
2647 lda: usize,
2648 b_mat: &[f32],
2649 ldb: usize,
2650 b_rows_are_n: bool,
2651 c: &mut [f32],
2652 ldc: usize,
2653) {
2654 debug_assert!(a.len() >= (m - 1) * lda + k);
2655 debug_assert!(c.len() >= (m - 1) * ldc + n);
2656 unsafe {
2658 use core::arch::aarch64::*;
2659 let mut i = 0usize;
2660 while i < m {
2661 let mi = (m - i).min(4);
2662 let mut j = 0usize;
2663 while j < n {
2664 let nj = (n - j).min(8);
2665 if mi == 4 && nj == 8 {
2666 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2667 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2668 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2669 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2670 for p in 0..k {
2671 let (b0, b1) = if b_rows_are_n {
2672 let base = b_mat.as_ptr().add(j * ldb + p);
2675 let g = |o: usize| *base.add(o * ldb);
2676 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2677 } else {
2678 let base = b_mat.as_ptr().add(p * ldb + j);
2679 (
2680 [*base, *base.add(1), *base.add(2), *base.add(3)],
2681 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2682 )
2683 };
2684 let bv0 = vld1q_f32(b0.as_ptr());
2685 let bv1 = vld1q_f32(b1.as_ptr());
2686 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2687 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2688 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2689 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2690 c0a = vfmaq_f32(c0a, a0, bv0);
2691 c0b = vfmaq_f32(c0b, a0, bv1);
2692 c1a = vfmaq_f32(c1a, a1, bv0);
2693 c1b = vfmaq_f32(c1b, a1, bv1);
2694 c2a = vfmaq_f32(c2a, a2, bv0);
2695 c2b = vfmaq_f32(c2b, a2, bv1);
2696 c3a = vfmaq_f32(c3a, a3, bv0);
2697 c3b = vfmaq_f32(c3b, a3, bv1);
2698 }
2699 let al = vdupq_n_f32(alpha);
2700 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2701 .iter()
2702 .enumerate()
2703 {
2704 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2705 vst1q_f32(dst, vmulq_f32(*ca, al));
2706 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2707 }
2708 } else {
2709 for r in 0..mi {
2710 for q in 0..nj {
2711 let mut acc = 0f32;
2712 for p in 0..k {
2713 let bv = if b_rows_are_n {
2714 b_mat[(j + q) * ldb + p]
2715 } else {
2716 b_mat[p * ldb + j + q]
2717 };
2718 acc += a[(i + r) * lda + p] * bv;
2719 }
2720 c[(i + r) * ldc + j + q] = acc * alpha;
2721 }
2722 }
2723 }
2724 j += nj;
2725 }
2726 i += mi;
2727 }
2728 }
2729}
2730
2731#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2733#[allow(clippy::too_many_arguments)]
2734pub(crate) fn sgemm_rm(
2735 m: usize,
2736 n: usize,
2737 k: usize,
2738 alpha: f32,
2739 a: &[f32],
2740 lda: usize,
2741 b_mat: &[f32],
2742 ldb: usize,
2743 b_rows_are_n: bool,
2744 c: &mut [f32],
2745 ldc: usize,
2746) {
2747 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2748}
2749
2750#[allow(clippy::too_many_arguments)]
2754pub fn sgemm_public(
2755 m: usize,
2756 n: usize,
2757 k: usize,
2758 alpha: f32,
2759 a: &[f32],
2760 lda: usize,
2761 b_mat: &[f32],
2762 ldb: usize,
2763 b_rows_are_n: bool,
2764 c: &mut [f32],
2765 ldc: usize,
2766) {
2767 #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
2768 {
2769 sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2770 }
2771 #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
2776 {
2777 for i in 0..m {
2778 for j in 0..n {
2779 let mut acc = 0f32;
2780 for p in 0..k {
2781 let bv = if b_rows_are_n {
2782 b_mat[j * ldb + p]
2783 } else {
2784 b_mat[p * ldb + j]
2785 };
2786 acc += a[i * lda + p] * bv;
2787 }
2788 c[i * ldc + j] = alpha * acc;
2789 }
2790 }
2791 }
2792}
2793
2794#[cfg(target_os = "macos")]
2797#[allow(clippy::too_many_arguments)]
2798pub(crate) fn sgemm_rm(
2799 m: usize,
2800 n: usize,
2801 k: usize,
2802 alpha: f32,
2803 a: &[f32],
2804 lda: usize,
2805 b_mat: &[f32],
2806 ldb: usize,
2807 b_rows_are_n: bool,
2808 c: &mut [f32],
2809 ldc: usize,
2810) {
2811 debug_assert!(a.len() >= (m - 1) * lda + k);
2812 debug_assert!(c.len() >= (m - 1) * ldc + n);
2813 #[cfg(target_arch = "aarch64")]
2818 if std::env::var("CMF_FORCE_NEON_GEMM")
2819 .map(|v| v == "1")
2820 .unwrap_or(false)
2821 {
2822 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2823 }
2824 unsafe {
2825 accel_blas::cblas_sgemm(
2826 101, 111, if b_rows_are_n { 112 } else { 111 },
2829 m as i32,
2830 n as i32,
2831 k as i32,
2832 alpha,
2833 a.as_ptr(),
2834 lda as i32,
2835 b_mat.as_ptr(),
2836 ldb as i32,
2837 0.0,
2838 c.as_mut_ptr(),
2839 ldc as i32,
2840 );
2841 }
2842}
2843
2844#[cfg(target_os = "macos")]
2851fn qmatmat_accel(
2852 q: &[u8],
2853 row_scale: &[f32],
2854 pre: &[std::borrow::Cow<'_, [f32]>],
2855 rows: usize,
2856 cols: usize,
2857 out: &mut [f32],
2858 pool: Option<&Pool>,
2859) {
2860 const TR: usize = 2048;
2865 let b = pre.len();
2866 thread_local! {
2867 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2868 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2869 }
2870 XPANEL.with(|xp| {
2871 WTILE.with(|wt| {
2872 let mut xpanel = xp.borrow_mut();
2873 xpanel.clear();
2874 for x in pre {
2875 xpanel.extend_from_slice(x);
2876 }
2877 let mut wtile = wt.borrow_mut();
2878 wtile.resize(TR * cols, 0.0);
2879 let mut r0 = 0usize;
2880 while r0 < rows {
2881 let tr = TR.min(rows - r0);
2882 let wt_addr = SendMut(wtile.as_mut_ptr());
2884 let run = |start: usize, end: usize| {
2885 for r in start..end {
2886 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2887 let s = row_scale[r0 + r];
2888 let dst =
2890 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2891 for (d, &v) in dst.iter_mut().zip(row) {
2892 *d = (v as i8) as f32 * s;
2893 }
2894 }
2895 };
2896 dispatch_rows(pool, tr, &run);
2897 unsafe {
2899 accel_blas::cblas_sgemm(
2900 101, 111, 112, b as i32,
2904 tr as i32,
2905 cols as i32,
2906 1.0,
2907 xpanel.as_ptr(),
2908 cols as i32,
2909 wtile.as_ptr(),
2910 cols as i32,
2911 0.0,
2912 out.as_mut_ptr().add(r0),
2913 rows as i32,
2914 );
2915 }
2916 r0 += tr;
2917 }
2918 })
2919 });
2920}
2921
2922fn qmatmat(
2923 q: &[u8],
2924 row_scale: &[f32],
2925 pre: &[std::borrow::Cow<'_, [f32]>],
2926 rows: usize,
2927 cols: usize,
2928 out: &mut [f32],
2929 pool: Option<&Pool>,
2930) {
2931 let b = pre.len();
2932 debug_assert_eq!(out.len(), b * rows);
2933 #[cfg(target_os = "macos")]
2938 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
2939 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
2940 return;
2941 }
2942 #[cfg(target_arch = "aarch64")]
2943 if sdot_enabled() {
2944 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2945 let out_addr = SendMut(out.as_mut_ptr());
2946 let blocked_ok = blocked_enabled();
2949 let use_i8mm = i8mm_enabled();
2950 if blocked_ok {
2951 let run = |start: usize, end: usize| {
2952 let mut o = start;
2953 while o < end {
2954 if o + 2 <= end {
2955 let r0 = &q[o * cols..(o + 1) * cols];
2956 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2957 let mut bi = 0usize;
2958 while bi + 4 <= acts.len() {
2959 let xs = [
2960 acts[bi].xq.as_slice(),
2961 acts[bi + 1].xq.as_slice(),
2962 acts[bi + 2].xq.as_slice(),
2963 acts[bi + 3].xq.as_slice(),
2964 ];
2965 let d = if use_i8mm {
2966 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
2967 } else {
2968 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
2969 };
2970 for (r, row) in [r0, r1].into_iter().enumerate() {
2971 for k in 0..4 {
2972 let act = &acts[bi + k];
2973 let mut v = d[r][k] as f32 * act.sx;
2974 for &(j, xv) in &act.outliers {
2975 v += (row[j] as i8) as f32 * xv;
2976 }
2977 unsafe {
2978 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2979 };
2980 }
2981 }
2982 bi += 4;
2983 }
2984 while bi < acts.len() {
2985 for (r, row) in [r0, r1].into_iter().enumerate() {
2986 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
2987 unsafe { *out_addr.at(bi * rows + o + r) = v };
2988 }
2989 bi += 1;
2990 }
2991 o += 2;
2992 } else {
2993 let row = &q[o * cols..(o + 1) * cols];
2994 for (bi, act) in acts.iter().enumerate() {
2995 let v = row_dot_sdot(row, act) * row_scale[o];
2996 unsafe { *out_addr.at(bi * rows + o) = v };
2997 }
2998 o += 1;
2999 }
3000 }
3001 };
3002 dispatch_rows(pool, rows, &run);
3003 return;
3004 }
3005 let run = |start: usize, end: usize| {
3006 for o in start..end {
3007 let row = &q[o * cols..(o + 1) * cols];
3008 for (bi, act) in acts.iter().enumerate() {
3009 let v = row_dot_sdot(row, act) * row_scale[o];
3010 unsafe { *out_addr.at(bi * rows + o) = v };
3011 }
3012 }
3013 };
3014 dispatch_rows(pool, rows, &run);
3015 return;
3016 }
3017 #[cfg(target_arch = "x86_64")]
3022 if avx2_a8w8_enabled() {
3023 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3024 let out_addr = SendMut(out.as_mut_ptr());
3025 let blocked_ok = blocked_enabled();
3028 if !avx512vnni_enabled() && blocked_ok {
3029 let run = |start: usize, end: usize| {
3030 let mut o = start;
3031 while o < end {
3032 if o + 2 <= end {
3033 let r0 = &q[o * cols..(o + 1) * cols];
3034 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3035 let mut bi = 0usize;
3036 while bi + 4 <= acts.len() {
3037 let xs = [
3038 acts[bi].xq.as_slice(),
3039 acts[bi + 1].xq.as_slice(),
3040 acts[bi + 2].xq.as_slice(),
3041 acts[bi + 3].xq.as_slice(),
3042 ];
3043 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3044 for (r, row) in [r0, r1].into_iter().enumerate() {
3045 for k in 0..4 {
3046 let act = &acts[bi + k];
3047 let mut v = d[r][k] as f32 * act.sx;
3048 for &(j, xv) in &act.outliers {
3049 v += (row[j] as i8) as f32 * xv;
3050 }
3051 unsafe {
3052 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3053 };
3054 }
3055 }
3056 bi += 4;
3057 }
3058 while bi < acts.len() {
3059 for (r, row) in [r0, r1].into_iter().enumerate() {
3060 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3061 unsafe { *out_addr.at(bi * rows + o + r) = v };
3062 }
3063 bi += 1;
3064 }
3065 o += 2;
3066 } else {
3067 let row = &q[o * cols..(o + 1) * cols];
3068 for (bi, act) in acts.iter().enumerate() {
3069 let v = row_dot_avx2(row, act) * row_scale[o];
3070 unsafe { *out_addr.at(bi * rows + o) = v };
3071 }
3072 o += 1;
3073 }
3074 }
3075 };
3076 dispatch_rows(pool, rows, &run);
3077 return;
3078 }
3079 let run = |start: usize, end: usize| {
3080 for o in start..end {
3081 let row = &q[o * cols..(o + 1) * cols];
3082 for (bi, act) in acts.iter().enumerate() {
3083 let v = row_dot_avx2(row, act) * row_scale[o];
3084 unsafe { *out_addr.at(bi * rows + o) = v };
3085 }
3086 }
3087 };
3088 dispatch_rows(pool, rows, &run);
3089 return;
3090 }
3091 let out_addr = SendMut(out.as_mut_ptr());
3092 let run = |start: usize, end: usize| {
3093 for o in start..end {
3094 let row = &q[o * cols..(o + 1) * cols];
3095 for (bi, x) in pre.iter().enumerate() {
3096 let mut acc = 0f32;
3097 for j in 0..cols {
3098 acc += (row[j] as i8) as f32 * x[j];
3099 }
3100 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3101 }
3102 }
3103 };
3104 dispatch_rows(pool, rows, &run);
3105}
3106
3107fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3110 match pool {
3111 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3112 _ => run(0, rows),
3113 }
3114}
3115
3116fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3118 let groups = rows * cols / GROUP_SIZE;
3119 bytes.split_at(groups * 16)
3120}
3121
3122#[inline]
3127fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3128 #[cfg(target_arch = "aarch64")]
3129 unsafe {
3130 return vbit_fill4_neon(data, buf);
3131 }
3132 #[cfg(target_arch = "x86_64")]
3133 if avx2_enabled() {
3134 return unsafe { vbit_fill4_avx2(data, buf) };
3135 }
3136 #[allow(unreachable_code)]
3137 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3138 let u = unpack8::<4>(&data[blk * 4..]);
3139 for k in 0..8 {
3140 chunk[k] = (u[k] - 7) as i8 as u8;
3141 }
3142 }
3143}
3144
3145#[cfg(target_arch = "aarch64")]
3146#[target_feature(enable = "neon")]
3147unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3148 unsafe {
3151 use core::arch::aarch64::*;
3152 let n = buf.len();
3153 let mask = vdupq_n_u8(0x0F);
3154 let seven = vdupq_n_s8(7);
3155 let mut g = 0usize;
3156 while g * 32 + 32 <= n {
3157 let b = vld1q_u8(data.as_ptr().add(g * 16));
3158 let hi = vshrq_n_u8::<4>(b);
3159 let lo = vandq_u8(b, mask);
3160 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3161 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3162 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3163 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3164 g += 1;
3165 }
3166 }
3167}
3168
3169#[cfg(target_arch = "x86_64")]
3170#[target_feature(enable = "avx2")]
3171unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3172 unsafe {
3174 use core::arch::x86_64::*;
3175 let n = buf.len();
3176 let mask = _mm_set1_epi8(0x0F);
3177 let seven = _mm256_set1_epi8(7);
3178 let mut g = 0usize;
3179 while g * 32 + 32 <= n {
3180 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3181 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3182 let lo = _mm_and_si128(b, mask);
3183 let z = _mm256_sub_epi8(
3184 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3185 seven,
3186 );
3187 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3188 g += 1;
3189 }
3190 }
3191}
3192
3193#[inline(always)]
3198fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3199 let mut acc = 0u64;
3200 for i in 0..B {
3201 acc = (acc << 8) | data[i] as u64;
3202 }
3203 let mask = (1u64 << B) - 1;
3204 let mut out = [0i32; 8];
3205 for (k, o) in out.iter_mut().enumerate() {
3206 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3207 }
3208 out
3209}
3210
3211#[allow(clippy::too_many_arguments)]
3217fn vbitmatvec(
3218 bytes: &[u8],
3219 offsets: &[usize],
3220 x: &[f32],
3221 rows: usize,
3222 cols: usize,
3223 out: &mut [f32],
3224 pool: Option<&Pool>,
3225) {
3226 debug_assert_eq!(out.len(), rows);
3227 debug_assert_eq!(offsets.len(), rows + 1);
3228
3229 if a8w8_enabled() {
3233 let act = split_act(x);
3234 let out_addr = SendMut(out.as_mut_ptr());
3235 let run = move |start: usize, end: usize| {
3236 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3237 };
3238 dispatch_rows(pool, rows, &run);
3239 return;
3240 }
3241
3242 let out_addr = SendMut(out.as_mut_ptr());
3243 let run = move |start: usize, end: usize| {
3244 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3245 };
3246 dispatch_rows(pool, rows, &run);
3247}
3248
3249#[allow(clippy::too_many_arguments)]
3253fn vbit_range_a8w8(
3254 bytes: &[u8],
3255 offsets: &[usize],
3256 x: &[f32],
3257 act: &SplitAct,
3258 rows: usize,
3259 cols: usize,
3260 out: SendMut,
3261 start: usize,
3262 end: usize,
3263) {
3264 let ng = cols / GROUP_SIZE;
3265 let bits = &bytes[..rows];
3266 let sc_off = rows;
3267 let row_dot = |r: usize| -> f32 {
3268 let b = bits[r] as usize;
3269 let l = (1i32 << (b - 1)) - 1;
3270 let mask = (1u64 << b) - 1;
3271 let data = &bytes[offsets[r]..offsets[r + 1]];
3272 if b == 8 {
3273 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3275 let mut dot = 0f32;
3276 for g in 0..ng {
3277 let so = (r * ng + g) * 2;
3278 let sgf = f16_to_f32(u16::from_le_bytes([
3279 bytes[sc_off + so],
3280 bytes[sc_off + so + 1],
3281 ]));
3282 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3283 let mut gd = 0f32;
3284 for &xv in xg.iter() {
3285 if nbits < 8 {
3286 acc = (acc << 8) | data[idx] as u64;
3287 idx += 1;
3288 nbits += 8;
3289 }
3290 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3291 nbits -= 8;
3292 gd += (u - l) as f32 * xv;
3293 }
3294 dot += gd * sgf;
3295 }
3296 return dot;
3297 }
3298 thread_local! {
3302 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3303 const { std::cell::RefCell::new(Vec::new()) };
3304 }
3305 #[inline(always)]
3306 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3307 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3308 let u = unpack8::<B>(&data[blk * B..]);
3309 for k in 0..8 {
3310 chunk[k] = (u[k] - l) as i8 as u8;
3311 }
3312 }
3313 }
3314 let _ = mask;
3315 VBIT_SCRATCH.with(|scratch| {
3316 let mut buf = scratch.borrow_mut();
3317 buf.resize(cols, 0);
3318 match b {
3319 3 => fill::<3>(data, l, &mut buf),
3320 4 => vbit_fill4(data, &mut buf),
3321 5 => fill::<5>(data, l, &mut buf),
3322 6 => fill::<6>(data, l, &mut buf),
3323 _ => unreachable!(),
3324 }
3325 let mut dot = 0f32;
3326 for g in 0..ng {
3327 let so = (r * ng + g) * 2;
3328 let s = f16_to_f32(u16::from_le_bytes([
3329 bytes[sc_off + so],
3330 bytes[sc_off + so + 1],
3331 ]));
3332 let d = dot_i8_i8(
3333 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3334 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3335 ) as f32
3336 * act.sx;
3337 dot += d * s;
3338 }
3339 for &(j, xv) in &act.outliers {
3340 let so = (r * ng + j / GROUP_SIZE) * 2;
3341 let s = f16_to_f32(u16::from_le_bytes([
3342 bytes[sc_off + so],
3343 bytes[sc_off + so + 1],
3344 ]));
3345 dot += (buf[j] as i8) as f32 * s * xv;
3347 }
3348 dot
3349 })
3350 };
3351 for r in start..end {
3352 unsafe { *out.at(r) = row_dot(r) };
3354 }
3355}
3356
3357#[allow(clippy::too_many_arguments)]
3359fn vbit_range_f32(
3360 bytes: &[u8],
3361 offsets: &[usize],
3362 x: &[f32],
3363 rows: usize,
3364 cols: usize,
3365 out: SendMut,
3366 start: usize,
3367 end: usize,
3368) {
3369 let ng = cols / GROUP_SIZE;
3370 let bits = &bytes[..rows];
3371 let sc_off = rows;
3372 #[inline(always)]
3376 fn dot_row<const B: usize>(
3377 data: &[u8],
3378 bytes: &[u8],
3379 sc_off: usize,
3380 r: usize,
3381 ng: usize,
3382 x: &[f32],
3383 ) -> f32 {
3384 let l = ((1i32 << (B - 1)) - 1) as f32;
3385 let gbytes = GROUP_SIZE * B / 8;
3386 let mut dot = 0f32;
3387 for g in 0..ng {
3388 let so = (r * ng + g) * 2;
3389 let s = f16_to_f32(u16::from_le_bytes([
3390 bytes[sc_off + so],
3391 bytes[sc_off + so + 1],
3392 ]));
3393 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3394 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3395 let mut gd = 0f32;
3396 for blk in 0..GROUP_SIZE / 8 {
3397 let u = unpack8::<B>(&gd0[blk * B..]);
3398 let xb = &xg[blk * 8..blk * 8 + 8];
3399 for k in 0..8 {
3400 gd += (u[k] as f32 - l) * xb[k];
3401 }
3402 }
3403 dot += gd * s;
3404 }
3405 dot
3406 }
3407 for r in start..end {
3408 let data = &bytes[offsets[r]..offsets[r + 1]];
3409 let v = match bits[r] {
3410 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3411 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3412 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3413 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3414 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3415 b => unreachable!("vbit bit-width {b} (validated at load)"),
3416 };
3417 unsafe { *out.at(r) = v };
3419 }
3420}
3421
3422#[allow(clippy::too_many_arguments)]
3427fn vbitmatvec2(
3428 bytes: &[u8],
3429 offsets: &[usize],
3430 x1: &[f32],
3431 x2: &[f32],
3432 rows: usize,
3433 cols: usize,
3434 o1: &mut [f32],
3435 o2: &mut [f32],
3436 pool: Option<&Pool>,
3437) {
3438 debug_assert_eq!(o1.len(), rows);
3439 debug_assert_eq!(o2.len(), rows);
3440
3441 if a8w8_enabled() {
3442 let a1 = split_act(x1);
3443 let a2 = split_act(x2);
3444 let p1 = SendMut(o1.as_mut_ptr());
3445 let p2 = SendMut(o2.as_mut_ptr());
3446 let run = move |start: usize, end: usize| {
3447 vbit_range2_a8w8(
3448 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3449 )
3450 };
3451 dispatch_rows(pool, rows, &run);
3452 return;
3453 }
3454
3455 let p1 = SendMut(o1.as_mut_ptr());
3456 let p2 = SendMut(o2.as_mut_ptr());
3457 let run = move |start: usize, end: usize| {
3458 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3459 };
3460 dispatch_rows(pool, rows, &run);
3461}
3462
3463#[allow(clippy::too_many_arguments)]
3467fn vbit_range2_a8w8(
3468 bytes: &[u8],
3469 offsets: &[usize],
3470 x1: &[f32],
3471 x2: &[f32],
3472 a1: &SplitAct,
3473 a2: &SplitAct,
3474 rows: usize,
3475 cols: usize,
3476 p1: SendMut,
3477 p2: SendMut,
3478 start: usize,
3479 end: usize,
3480) {
3481 let ng = cols / GROUP_SIZE;
3482 let bits = &bytes[..rows];
3483 let sc_off = rows;
3484 let row_dots = |r: usize| -> (f32, f32) {
3485 let b = bits[r] as usize;
3486 let l = (1i32 << (b - 1)) - 1;
3487 let data = &bytes[offsets[r]..offsets[r + 1]];
3488 if b == 8 {
3489 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3492 let (mut d1, mut d2) = (0f32, 0f32);
3493 for g in 0..ng {
3494 let so = (r * ng + g) * 2;
3495 let sgf = f16_to_f32(u16::from_le_bytes([
3496 bytes[sc_off + so],
3497 bytes[sc_off + so + 1],
3498 ]));
3499 let (mut g1, mut g2) = (0f32, 0f32);
3500 for k in 0..GROUP_SIZE {
3501 if nbits < 8 {
3502 acc = (acc << 8) | data[idx] as u64;
3503 idx += 1;
3504 nbits += 8;
3505 }
3506 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3507 nbits -= 8;
3508 let w = (u - l) as f32;
3509 g1 += w * x1[g * GROUP_SIZE + k];
3510 g2 += w * x2[g * GROUP_SIZE + k];
3511 }
3512 d1 += g1 * sgf;
3513 d2 += g2 * sgf;
3514 }
3515 return (d1, d2);
3516 }
3517 thread_local! {
3518 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3519 const { std::cell::RefCell::new(Vec::new()) };
3520 }
3521 #[inline(always)]
3522 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3523 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3524 let u = unpack8::<B>(&data[blk * B..]);
3525 for k in 0..8 {
3526 chunk[k] = (u[k] - l) as i8 as u8;
3527 }
3528 }
3529 }
3530 VBIT_SCRATCH2.with(|scratch| {
3531 let mut buf = scratch.borrow_mut();
3532 buf.resize(cols, 0);
3533 match b {
3534 3 => fill::<3>(data, l, &mut buf),
3535 4 => vbit_fill4(data, &mut buf),
3536 5 => fill::<5>(data, l, &mut buf),
3537 6 => fill::<6>(data, l, &mut buf),
3538 _ => unreachable!(),
3539 }
3540 let (mut d1, mut d2) = (0f32, 0f32);
3541 for g in 0..ng {
3542 let so = (r * ng + g) * 2;
3543 let s = f16_to_f32(u16::from_le_bytes([
3544 bytes[sc_off + so],
3545 bytes[sc_off + so + 1],
3546 ]));
3547 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3548 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3549 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3550 d1 += v1 * s;
3551 d2 += v2 * s;
3552 }
3553 for &(j, xv) in &a1.outliers {
3554 let so = (r * ng + j / GROUP_SIZE) * 2;
3555 let s = f16_to_f32(u16::from_le_bytes([
3556 bytes[sc_off + so],
3557 bytes[sc_off + so + 1],
3558 ]));
3559 d1 += (buf[j] as i8) as f32 * s * xv;
3560 }
3561 for &(j, xv) in &a2.outliers {
3562 let so = (r * ng + j / GROUP_SIZE) * 2;
3563 let s = f16_to_f32(u16::from_le_bytes([
3564 bytes[sc_off + so],
3565 bytes[sc_off + so + 1],
3566 ]));
3567 d2 += (buf[j] as i8) as f32 * s * xv;
3568 }
3569 (d1, d2)
3570 })
3571 };
3572 for r in start..end {
3573 let (v1, v2) = row_dots(r);
3574 unsafe {
3576 *p1.at(r) = v1;
3577 *p2.at(r) = v2;
3578 }
3579 }
3580}
3581
3582#[allow(clippy::too_many_arguments)]
3586fn vbit_range2_f32(
3587 bytes: &[u8],
3588 offsets: &[usize],
3589 x1: &[f32],
3590 x2: &[f32],
3591 rows: usize,
3592 cols: usize,
3593 p1: SendMut,
3594 p2: SendMut,
3595 start: usize,
3596 end: usize,
3597) {
3598 let ng = cols / GROUP_SIZE;
3599 let bits = &bytes[..rows];
3600 let sc_off = rows;
3601 #[inline(always)]
3602 #[allow(clippy::too_many_arguments)]
3603 fn dot_row2<const B: usize>(
3604 data: &[u8],
3605 bytes: &[u8],
3606 sc_off: usize,
3607 r: usize,
3608 ng: usize,
3609 x1: &[f32],
3610 x2: &[f32],
3611 ) -> (f32, f32) {
3612 let l = ((1i32 << (B - 1)) - 1) as f32;
3613 let gbytes = GROUP_SIZE * B / 8;
3614 let (mut d1, mut d2) = (0f32, 0f32);
3615 for g in 0..ng {
3616 let so = (r * ng + g) * 2;
3617 let s = f16_to_f32(u16::from_le_bytes([
3618 bytes[sc_off + so],
3619 bytes[sc_off + so + 1],
3620 ]));
3621 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3622 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3623 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3624 let (mut g1, mut g2) = (0f32, 0f32);
3625 for blk in 0..GROUP_SIZE / 8 {
3626 let u = unpack8::<B>(&gd0[blk * B..]);
3627 for k in 0..8 {
3628 let w = u[k] as f32 - l;
3629 g1 += w * x1g[blk * 8 + k];
3630 g2 += w * x2g[blk * 8 + k];
3631 }
3632 }
3633 d1 += g1 * s;
3634 d2 += g2 * s;
3635 }
3636 (d1, d2)
3637 }
3638 for r in start..end {
3639 let data = &bytes[offsets[r]..offsets[r + 1]];
3640 let (v1, v2) = match bits[r] {
3641 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3642 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3643 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3644 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3645 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3646 b => unreachable!("vbit bit-width {b} (validated at load)"),
3647 };
3648 unsafe {
3650 *p1.at(r) = v1;
3651 *p2.at(r) = v2;
3652 }
3653 }
3654}
3655
3656#[inline]
3663#[allow(unreachable_code)]
3664fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3665 #[cfg(target_arch = "aarch64")]
3666 unsafe {
3667 return dot_q4t_row_sdot(bytes, r, gpr, xq);
3668 }
3669 #[cfg(target_arch = "x86_64")]
3670 unsafe {
3671 if vnni_tiles_enabled() {
3672 return dot_q4t_row_vnni(bytes, r, gpr, xq);
3673 }
3674 return dot_q4t_row_avx2(bytes, r, gpr, xq);
3675 }
3676 let mut acc = 0f32;
3677 for gi in 0..gpr {
3678 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3679 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3680 let mut d = 0i32;
3681 for (k, &b) in tile[2..].iter().enumerate() {
3682 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3683 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3684 }
3685 acc += d as f32 * s;
3686 }
3687 acc
3688}
3689
3690#[cfg(target_arch = "aarch64")]
3691#[target_feature(enable = "neon,dotprod")]
3692unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3693 unsafe {
3696 use core::arch::aarch64::*;
3697 use core::arch::asm;
3698 let lomask = vdupq_n_u8(0x0F);
3699 let eight = vdupq_n_s8(8);
3700 let mut acc = 0f32;
3701 for gi in 0..gpr {
3702 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3703 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3704 let b = vld1q_u8(t.add(2));
3705 let lo = vandq_u8(b, lomask);
3706 let hi = vshrq_n_u8::<4>(b);
3707 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3708 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3709 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3710 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3711 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3712 asm!(
3713 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3714 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3715 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3716 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3717 options(pure, nomem, nostack),
3718 );
3719 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3720 }
3721 acc
3722 }
3723}
3724
3725#[cfg(target_arch = "x86_64")]
3726#[target_feature(enable = "avx2")]
3727unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3728 unsafe {
3730 use core::arch::x86_64::*;
3731 let lomask = _mm_set1_epi8(0x0F);
3732 let eight = _mm256_set1_epi8(8);
3733 let ones = _mm256_set1_epi16(1);
3734 let mut acc = 0f32;
3735 for gi in 0..gpr {
3736 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3737 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3738 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3739 let lo = _mm_and_si128(b, lomask);
3740 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3741 let w = _mm256_sub_epi8(
3742 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3743 eight,
3744 );
3745 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3746 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3747 let d = _mm256_madd_epi16(p16, ones);
3748 let hi128 = _mm256_extracti128_si256::<1>(d);
3749 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3750 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3751 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3752 acc += _mm_cvtsi128_si32(s32) as f32 * s;
3753 }
3754 acc
3755 }
3756}
3757
3758#[cfg(target_arch = "x86_64")]
3762#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3763unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3764 unsafe {
3766 use core::arch::x86_64::*;
3767 let lomask = _mm_set1_epi8(0x0F);
3768 let eight = _mm256_set1_epi8(8);
3769 let mut acc = 0f32;
3770 for gi in 0..gpr {
3771 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3772 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3773 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3774 let lo = _mm_and_si128(b, lomask);
3775 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3776 let w = _mm256_sub_epi8(
3777 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3778 eight,
3779 );
3780 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3781 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3782 acc += d as f32 * s;
3783 }
3784 acc
3785 }
3786}
3787
3788#[cfg(target_arch = "x86_64")]
3793#[target_feature(enable = "avx2,fma")]
3798unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3799 unsafe {
3801 use core::arch::x86_64::*;
3802 let lomask = _mm_set1_epi8(0x0F);
3803 let eight = _mm256_set1_epi8(8);
3804 let ones = _mm256_set1_epi16(1);
3805 let mut f0 = _mm256_setzero_ps();
3818 let mut f1 = _mm256_setzero_ps();
3819 let mut f2 = _mm256_setzero_ps();
3820 let mut f3 = _mm256_setzero_ps();
3821 for gi in 0..gpr {
3822 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3823 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3824 let sv = _mm256_set1_ps(s);
3825 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3826 let lo = _mm_and_si128(bb, lomask);
3827 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3828 let w = _mm256_sub_epi8(
3829 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3830 eight,
3831 );
3832 let aw = _mm256_abs_epi8(w);
3833 let off = gi * GROUP_SIZE;
3834 let dot = |xq: &[i8]| {
3835 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3836 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3837 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
3838 };
3839 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3840 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3841 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3842 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3843 }
3844 [
3845 hsum256_ps(f0),
3846 hsum256_ps(f1),
3847 hsum256_ps(f2),
3848 hsum256_ps(f3),
3849 ]
3850 }
3851}
3852
3853#[cfg(target_arch = "x86_64")]
3856#[target_feature(enable = "avx2")]
3857#[inline]
3858unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
3859 unsafe {
3861 use core::arch::x86_64::*;
3862 let hi = _mm256_extractf128_ps::<1>(v);
3863 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
3864 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
3865 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
3866 _mm_cvtss_f32(s)
3867 }
3868}
3869
3870#[cfg(target_arch = "x86_64")]
3872#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
3873unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3874 unsafe {
3876 use core::arch::x86_64::*;
3877 let lomask = _mm_set1_epi8(0x0F);
3878 let eight = _mm256_set1_epi8(8);
3879 let mut f0 = _mm256_setzero_ps();
3882 let mut f1 = _mm256_setzero_ps();
3883 let mut f2 = _mm256_setzero_ps();
3884 let mut f3 = _mm256_setzero_ps();
3885 for gi in 0..gpr {
3886 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3887 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3888 let sv = _mm256_set1_ps(s);
3889 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3890 let lo = _mm_and_si128(bb, lomask);
3891 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3892 let w = _mm256_sub_epi8(
3893 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3894 eight,
3895 );
3896 let aw = _mm256_abs_epi8(w);
3897 let off = gi * GROUP_SIZE;
3898 let dot = |xq: &[i8]| {
3899 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3900 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
3901 _mm256_setzero_si256(),
3902 aw,
3903 _mm256_sign_epi8(x, w),
3904 ))
3905 };
3906 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3907 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3908 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3909 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3910 }
3911 let acc = [
3912 hsum256_ps(f0),
3913 hsum256_ps(f1),
3914 hsum256_ps(f2),
3915 hsum256_ps(f3),
3916 ];
3917 acc
3918 }
3919}
3920
3921#[cfg(target_arch = "aarch64")]
3926#[target_feature(enable = "neon,dotprod")]
3927unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3928 unsafe {
3930 use core::arch::aarch64::*;
3931 use core::arch::asm;
3932 let lomask = vdupq_n_u8(0x0F);
3933 let eight = vdupq_n_s8(8);
3934 let mut acc = [0f32; 4];
3935 for gi in 0..gpr {
3936 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3937 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3938 let b = vld1q_u8(t.add(2));
3939 let lo = vandq_u8(b, lomask);
3940 let hi = vshrq_n_u8::<4>(b);
3941 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3942 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3943 for (k, xq) in xs.iter().enumerate() {
3944 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3945 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3946 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3947 asm!(
3948 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3949 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3950 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3951 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3952 options(pure, nomem, nostack),
3953 );
3954 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3955 }
3956 }
3957 acc
3958 }
3959}
3960
3961#[inline]
3963fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3964 let gi = j / GROUP_SIZE;
3965 let k = j % GROUP_SIZE;
3966 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3967 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3968 let byte = tile[2 + k / 2];
3969 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3970 ((nib as i32 - 8) as f32, s)
3971}
3972
3973#[inline]
3976fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3977 let mut acc = 0f32;
3978 for gi in 0..gpr {
3979 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3980 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3981 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3982 let mut ga = 0f32;
3983 for (k, &b) in tile[2..].iter().enumerate() {
3984 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3985 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3986 }
3987 acc += ga * s;
3988 }
3989 acc
3990}
3991
3992struct Q4tpView<'a> {
3996 nib: &'a [u8],
3997 params: &'a [u8],
3998 codes: &'a [u8],
3999 stride: usize,
4000 zero_rung: bool,
4002}
4003
4004impl<'a> Q4tpView<'a> {
4005 fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4006 let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4007 Self {
4008 nib: &bytes[..params_off],
4009 params: &bytes[params_off..codes_off],
4010 codes: &bytes[codes_off..],
4011 stride,
4012 zero_rung: false,
4013 }
4014 }
4015
4016 fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4018 let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4019 Self {
4020 nib: &bytes[..params_off],
4021 params: &bytes[params_off..codes_off],
4022 codes: &bytes[codes_off..],
4023 stride,
4024 zero_rung: true,
4025 }
4026 }
4027
4028 #[inline]
4044 fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4045 let tab = if self.zero_rung {
4046 q2tp_ladder(self.params, r)
4047 } else {
4048 q4tp_ladder(self.params, r)
4049 };
4050 let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4051 let out = &mut out[..gpr];
4052 let mut chunks = out.chunks_exact_mut(8);
4053 let mut ci = 0usize;
4054 for c in &mut chunks {
4055 let w = u64::from(codes[ci])
4056 | u64::from(codes[ci + 1]) << 8
4057 | u64::from(codes[ci + 2]) << 16
4058 | u64::from(codes[ci + 3]) << 24
4059 | u64::from(codes[ci + 4]) << 32;
4060 for (k, o) in c.iter_mut().enumerate() {
4061 *o = tab[((w >> (5 * k)) & 31) as usize];
4062 }
4063 ci += 5;
4064 }
4065 let tail = &codes[ci..];
4068 for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4069 *o = tab[q4tp_code(tail, k)];
4070 }
4071 }
4072}
4073
4074#[inline]
4075fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4076 #[cfg(target_arch = "aarch64")]
4077 unsafe {
4078 return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4079 }
4080 #[cfg(target_arch = "x86_64")]
4081 unsafe {
4082 if vnni_tiles_enabled() {
4083 return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4084 }
4085 return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4086 }
4087 #[allow(unreachable_code)]
4088 {
4089 let mut acc = 0f32;
4090 for gi in 0..gpr {
4091 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4092 let s = scales[gi];
4093 let mut d = 0i32;
4094 for (k, &b) in tile.iter().enumerate() {
4095 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4096 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4097 }
4098 acc += d as f32 * s;
4099 }
4100 acc
4101 }
4102}
4103
4104#[cfg(target_arch = "aarch64")]
4107#[target_feature(enable = "neon,dotprod")]
4108unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4109 unsafe {
4112 use core::arch::aarch64::*;
4113 use core::arch::asm;
4114 let lomask = vdupq_n_u8(0x0F);
4115 let eight = vdupq_n_s8(8);
4116 let mut acc = 0f32;
4117 for gi in 0..gpr {
4118 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4119 let s = *scales.get_unchecked(gi);
4120 let b = vld1q_u8(t);
4121 let lo = vandq_u8(b, lomask);
4122 let hi = vshrq_n_u8::<4>(b);
4123 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4124 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4125 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4126 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4127 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4128 asm!(
4129 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4130 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4131 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4132 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4133 options(pure, nomem, nostack),
4134 );
4135 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4136 }
4137 acc
4138 }
4139}
4140
4141#[cfg(target_arch = "x86_64")]
4142#[target_feature(enable = "avx2")]
4143unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4144 unsafe {
4146 use core::arch::x86_64::*;
4147 let lomask = _mm_set1_epi8(0x0F);
4148 let eight = _mm256_set1_epi8(8);
4149 let ones = _mm256_set1_epi16(1);
4150 let mut acc = 0f32;
4151 for gi in 0..gpr {
4152 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4153 let s = *scales.get_unchecked(gi);
4154 let b = _mm_loadu_si128(t as *const __m128i);
4155 let lo = _mm_and_si128(b, lomask);
4156 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4157 let w = _mm256_sub_epi8(
4158 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4159 eight,
4160 );
4161 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4162 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4163 let d = _mm256_madd_epi16(p16, ones);
4164 let hi128 = _mm256_extracti128_si256::<1>(d);
4165 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4166 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4167 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4168 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4169 }
4170 acc
4171 }
4172}
4173
4174#[cfg(target_arch = "x86_64")]
4177#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4178unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4179 unsafe {
4181 use core::arch::x86_64::*;
4182 let lomask = _mm_set1_epi8(0x0F);
4183 let eight = _mm256_set1_epi8(8);
4184 let mut acc = 0f32;
4185 for gi in 0..gpr {
4186 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4187 let s = *scales.get_unchecked(gi);
4188 let b = _mm_loadu_si128(t as *const __m128i);
4189 let lo = _mm_and_si128(b, lomask);
4190 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4191 let w = _mm256_sub_epi8(
4192 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4193 eight,
4194 );
4195 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4196 acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4197 }
4198 acc
4199 }
4200}
4201
4202#[inline]
4205fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4206 let mut acc = 0f32;
4207 for gi in 0..gpr {
4208 let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4209 let s = scales[gi];
4210 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4211 let mut ga = 0f32;
4212 for (k, &b) in tile.iter().enumerate() {
4213 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4214 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4215 }
4216 acc += ga * s;
4217 }
4218 acc
4219}
4220
4221#[inline]
4224fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4225 let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4226 let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4227 let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4228 ((n as i32 - 8) as f32, scales[gi])
4229}
4230
4231fn q4tp_matvec(
4233 bytes: &[u8],
4234 x: &[f32],
4235 rows: usize,
4236 cols: usize,
4237 out: &mut [f32],
4238 pool: Option<&Pool>,
4239) {
4240 debug_assert_eq!(out.len(), rows);
4241 let gpr = cols / GROUP_SIZE;
4242 let v = Q4tpView::new(bytes, rows, cols);
4243 let out_addr = SendMut(out.as_mut_ptr());
4244 if a8w8_enabled() {
4245 let act = split_act(x);
4246 let run = |start: usize, end: usize| {
4247 with_krow(gpr, |sc| {
4249 for r in start..end {
4250 v.scales_into(r, gpr, sc);
4251 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4252 for &(j, xv) in &act.outliers {
4253 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4254 acc += w * s * xv;
4255 }
4256 unsafe { *out_addr.at(r) = acc };
4258 }
4259 })
4260 };
4261 dispatch_rows(pool, rows, &run);
4262 return;
4263 }
4264 let run = |start: usize, end: usize| {
4265 with_krow(gpr, |sc| {
4266 for r in start..end {
4267 v.scales_into(r, gpr, sc);
4268 unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4270 }
4271 })
4272 };
4273 dispatch_rows(pool, rows, &run);
4274}
4275
4276#[allow(clippy::too_many_arguments)]
4279fn q4tp_matvec2(
4280 bytes: &[u8],
4281 x1: &[f32],
4282 x2: &[f32],
4283 rows: usize,
4284 cols: usize,
4285 o1: &mut [f32],
4286 o2: &mut [f32],
4287 pool: Option<&Pool>,
4288) {
4289 let gpr = cols / GROUP_SIZE;
4290 let v = Q4tpView::new(bytes, rows, cols);
4291 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4292 let run = |start: usize, end: usize| {
4293 let mut sc = vec![0f32; gpr];
4294 for r in start..end {
4295 v.scales_into(r, gpr, &mut sc);
4296 unsafe {
4298 *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4299 *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4300 }
4301 }
4302 };
4303 dispatch_rows(pool, rows, &run);
4304}
4305
4306fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4310 let mut acc = 0f32;
4311 for gi in 0..gpr {
4312 let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4313 let s = scales[gi];
4314 let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4315 let mut g = 0f32;
4316 for (k, &b) in ch.iter().enumerate() {
4317 g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4318 + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4319 + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4320 + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4321 }
4322 acc += s * g;
4323 }
4324 acc
4325}
4326
4327fn q2tp_matvec(
4328 bytes: &[u8],
4329 x: &[f32],
4330 rows: usize,
4331 cols: usize,
4332 out: &mut [f32],
4333 pool: Option<&Pool>,
4334) {
4335 debug_assert_eq!(out.len(), rows);
4336 let gpr = cols / GROUP_SIZE;
4337 let v = Q4tpView::new_q2(bytes, rows, cols);
4338 let out_addr = SendMut(out.as_mut_ptr());
4339 let run = |start: usize, end: usize| {
4340 with_krow(gpr, |sc| {
4341 for r in start..end {
4342 v.scales_into(r, gpr, sc);
4343 unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4345 }
4346 })
4347 };
4348 dispatch_rows(pool, rows, &run);
4349}
4350
4351#[allow(clippy::too_many_arguments)]
4353fn q2tp_matvec2(
4354 bytes: &[u8],
4355 x1: &[f32],
4356 x2: &[f32],
4357 rows: usize,
4358 cols: usize,
4359 o1: &mut [f32],
4360 o2: &mut [f32],
4361 pool: Option<&Pool>,
4362) {
4363 let gpr = cols / GROUP_SIZE;
4364 let v = Q4tpView::new_q2(bytes, rows, cols);
4365 let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4366 let run = |start: usize, end: usize| {
4367 let mut sc = vec![0f32; gpr];
4368 for r in start..end {
4369 v.scales_into(r, gpr, &mut sc);
4370 unsafe {
4372 *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4373 *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4374 }
4375 }
4376 };
4377 dispatch_rows(pool, rows, &run);
4378}
4379
4380pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4387 q2tp_matvec(bytes, x, rows, cols, out, None);
4388}
4389
4390pub fn q2tp_matmat_for_test(
4391 bytes: &[u8],
4392 xs_all: &[f32],
4393 b: usize,
4394 rows: usize,
4395 cols: usize,
4396 out: &mut [f32],
4397) {
4398 q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4399}
4400
4401fn q2tp_matmat(
4402 bytes: &[u8],
4403 xs_all: &[f32],
4404 b: usize,
4405 rows: usize,
4406 cols: usize,
4407 out: &mut [f32],
4408 pool: Option<&Pool>,
4409) {
4410 debug_assert_eq!(out.len(), b * rows);
4411 let gpr = cols / GROUP_SIZE;
4412 let v = Q4tpView::new_q2(bytes, rows, cols);
4413 let out_addr = SendMut(out.as_mut_ptr());
4414 let run = |start: usize, end: usize| {
4415 let mut sc = vec![0f32; gpr];
4416 for r in start..end {
4417 v.scales_into(r, gpr, &mut sc);
4418 for bi in 0..b {
4419 let x = &xs_all[bi * cols..(bi + 1) * cols];
4420 unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4422 }
4423 }
4424 };
4425 dispatch_rows(pool, rows, &run);
4426}
4427
4428#[cfg(target_arch = "aarch64")]
4432#[target_feature(enable = "neon,dotprod")]
4433unsafe fn dot_q4tp_row_1x4_sdot(
4434 nib: &[u8],
4435 r: usize,
4436 gpr: usize,
4437 xs: [&[i8]; 4],
4438 scales: &[f32],
4439) -> [f32; 4] {
4440 unsafe {
4442 use core::arch::aarch64::*;
4443 use core::arch::asm;
4444 let lomask = vdupq_n_u8(0x0F);
4445 let eight = vdupq_n_s8(8);
4446 let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4450 for gi in 0..gpr {
4451 let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4452 let s = *scales.get_unchecked(gi);
4453 let bb = vld1q_u8(t);
4454 let lo = vandq_u8(bb, lomask);
4455 let hi = vshrq_n_u8::<4>(bb);
4456 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4457 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4458 let mut d = [0f32; 4];
4459 for (k, dk) in d.iter_mut().enumerate() {
4460 let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4461 let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4462 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4463 asm!(
4464 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4465 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4466 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4467 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4468 options(pure, nomem, nostack),
4469 );
4470 *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4471 }
4472 f0 += d[0];
4473 f1 += d[1];
4474 f2 += d[2];
4475 f3 += d[3];
4476 }
4477 [f0, f1, f2, f3]
4478 }
4479}
4480
4481fn q4tp_matmat(
4485 bytes: &[u8],
4486 xs_all: &[f32],
4487 b: usize,
4488 rows: usize,
4489 cols: usize,
4490 out: &mut [f32],
4491 pool: Option<&Pool>,
4492) {
4493 debug_assert_eq!(out.len(), b * rows);
4494 let gpr = cols / GROUP_SIZE;
4495 let v = Q4tpView::new(bytes, rows, cols);
4496
4497 #[cfg(target_os = "macos")]
4499 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4500 dequant_matmat_accel(
4501 &|r, dst| {
4502 let mut sc = [0f32; 32];
4503 let mut scv;
4504 let s: &[f32] = if gpr <= 32 {
4505 v.scales_into(r, gpr, &mut sc);
4506 &sc[..gpr]
4507 } else {
4508 scv = vec![0f32; gpr];
4509 v.scales_into(r, gpr, &mut scv);
4510 &scv
4511 };
4512 for gi in 0..gpr {
4513 let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4514 for (k, &bb) in tile.iter().enumerate() {
4515 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
4516 dst[gi * GROUP_SIZE + k * 2 + 1] =
4517 (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
4518 }
4519 }
4520 },
4521 xs_all,
4522 b,
4523 rows,
4524 cols,
4525 out,
4526 pool,
4527 );
4528 return;
4529 }
4530
4531 let out_addr = SendMut(out.as_mut_ptr());
4532 if a8w8_enabled() {
4533 let acts: Vec<SplitAct> = (0..b)
4534 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4535 .collect();
4536 let acts = &acts;
4537 #[cfg(target_arch = "aarch64")]
4538 let blocked_ok = sdot_enabled() && blocked_enabled();
4539 #[cfg(not(target_arch = "aarch64"))]
4540 let blocked_ok = false;
4541 let run = |start: usize, end: usize| {
4542 let mut sc = vec![0f32; gpr];
4543 for r in start..end {
4544 v.scales_into(r, gpr, &mut sc);
4545 let mut bi = 0usize;
4546 #[cfg(target_arch = "aarch64")]
4547 if blocked_ok {
4548 while bi + 4 <= acts.len() {
4549 let xs = [
4550 acts[bi].xq.as_slice(),
4551 acts[bi + 1].xq.as_slice(),
4552 acts[bi + 2].xq.as_slice(),
4553 acts[bi + 3].xq.as_slice(),
4554 ];
4555 let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc) };
4556 for k in 0..4 {
4557 let act = &acts[bi + k];
4558 let mut acc = d[k] * act.sx;
4559 for &(j, xv) in &act.outliers {
4560 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4561 acc += w * s * xv;
4562 }
4563 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4565 }
4566 bi += 4;
4567 }
4568 }
4569 let _ = blocked_ok;
4570 while bi < acts.len() {
4571 let act = &acts[bi];
4572 let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
4573 for &(j, xv) in &act.outliers {
4574 let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4575 acc += w * s * xv;
4576 }
4577 unsafe { *out_addr.at(bi * rows + r) = acc };
4579 bi += 1;
4580 }
4581 }
4582 };
4583 dispatch_rows(pool, rows, &run);
4584 return;
4585 }
4586
4587 let run = |start: usize, end: usize| {
4588 let mut sc = vec![0f32; gpr];
4589 for r in start..end {
4590 v.scales_into(r, gpr, &mut sc);
4591 for bi in 0..b {
4592 let x = &xs_all[bi * cols..(bi + 1) * cols];
4593 unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
4595 }
4596 }
4597 };
4598 dispatch_rows(pool, rows, &run);
4599}
4600
4601fn q4t_matvec(
4603 bytes: &[u8],
4604 x: &[f32],
4605 rows: usize,
4606 cols: usize,
4607 out: &mut [f32],
4608 pool: Option<&Pool>,
4609) {
4610 debug_assert_eq!(out.len(), rows);
4611 let gpr = cols / GROUP_SIZE;
4612 let out_addr = SendMut(out.as_mut_ptr());
4613 if a8w8_enabled() {
4614 let act = split_act(x);
4615 let run = move |start: usize, end: usize| {
4616 for r in start..end {
4617 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4618 for &(j, xv) in &act.outliers {
4619 let (w, s) = q4t_outlier(bytes, r, gpr, j);
4620 acc += w * s * xv;
4621 }
4622 unsafe { *out_addr.at(r) = acc };
4624 }
4625 };
4626 dispatch_rows(pool, rows, &run);
4627 return;
4628 }
4629 let run = move |start: usize, end: usize| {
4630 for r in start..end {
4631 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
4633 }
4634 };
4635 dispatch_rows(pool, rows, &run);
4636}
4637
4638#[allow(clippy::too_many_arguments)]
4640fn q4t_matvec2(
4641 bytes: &[u8],
4642 x1: &[f32],
4643 x2: &[f32],
4644 rows: usize,
4645 cols: usize,
4646 o1: &mut [f32],
4647 o2: &mut [f32],
4648 pool: Option<&Pool>,
4649) {
4650 let gpr = cols / GROUP_SIZE;
4651 let p1 = SendMut(o1.as_mut_ptr());
4652 let p2 = SendMut(o2.as_mut_ptr());
4653 if a8w8_enabled() {
4654 let a1 = split_act(x1);
4655 let a2 = split_act(x2);
4656 let run = move |start: usize, end: usize| {
4657 for r in start..end {
4658 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
4659 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
4660 for &(j, xv) in &a1.outliers {
4661 let (w, s) = q4t_outlier(bytes, r, gpr, j);
4662 v1 += w * s * xv;
4663 }
4664 for &(j, xv) in &a2.outliers {
4665 let (w, s) = q4t_outlier(bytes, r, gpr, j);
4666 v2 += w * s * xv;
4667 }
4668 unsafe {
4670 *p1.at(r) = v1;
4671 *p2.at(r) = v2;
4672 }
4673 }
4674 };
4675 dispatch_rows(pool, rows, &run);
4676 return;
4677 }
4678 let run = move |start: usize, end: usize| {
4679 for r in start..end {
4680 unsafe {
4682 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
4683 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
4684 }
4685 }
4686 };
4687 dispatch_rows(pool, rows, &run);
4688}
4689
4690#[allow(clippy::too_many_arguments)]
4692#[cfg(target_os = "macos")]
4698fn dequant_matmat_accel(
4699 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
4700 xs_all: &[f32],
4701 b: usize,
4702 rows: usize,
4703 cols: usize,
4704 out: &mut [f32],
4705 pool: Option<&Pool>,
4706) {
4707 const TR: usize = 2048;
4708 thread_local! {
4709 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4710 }
4711 WTILE.with(|wt| {
4712 let mut wtile = wt.borrow_mut();
4713 wtile.resize(TR * cols, 0.0);
4714 let mut r0 = 0usize;
4715 while r0 < rows {
4716 let tr = TR.min(rows - r0);
4717 let wt_addr = SendMut(wtile.as_mut_ptr());
4718 let run = |start: usize, end: usize| {
4719 for r in start..end {
4720 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4722 dequant_row(r0 + r, dst);
4723 }
4724 };
4725 dispatch_rows(pool, tr, &run);
4726 unsafe {
4727 accel_blas::cblas_sgemm(
4728 101, 111, 112, b as i32,
4732 tr as i32,
4733 cols as i32,
4734 1.0,
4735 xs_all.as_ptr(),
4736 cols as i32,
4737 wtile.as_ptr(),
4738 cols as i32,
4739 0.0,
4740 out.as_mut_ptr().add(r0),
4741 rows as i32,
4742 );
4743 }
4744 r0 += tr;
4745 }
4746 });
4747}
4748
4749fn q4t_matmat(
4750 bytes: &[u8],
4751 xs_all: &[f32],
4752 b: usize,
4753 rows: usize,
4754 cols: usize,
4755 out: &mut [f32],
4756 pool: Option<&Pool>,
4757) {
4758 debug_assert_eq!(out.len(), b * rows);
4759 let gpr = cols / GROUP_SIZE;
4760 #[cfg(target_os = "macos")]
4764 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4765 dequant_matmat_accel(
4766 &|r, dst| {
4767 for gi in 0..gpr {
4768 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4769 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4770 for (k, &bb) in tile[2..].iter().enumerate() {
4771 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
4772 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
4773 }
4774 }
4775 },
4776 xs_all,
4777 b,
4778 rows,
4779 cols,
4780 out,
4781 pool,
4782 );
4783 return;
4784 }
4785 let out_addr = SendMut(out.as_mut_ptr());
4786 if a8w8_enabled() {
4787 let acts: Vec<SplitAct> = (0..b)
4788 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4789 .collect();
4790 let acts = &acts;
4791 #[cfg(target_arch = "x86_64")]
4792 let blocked_ok = avx2_enabled()
4793 && blocked_enabled();
4794 #[cfg(target_arch = "aarch64")]
4795 let blocked_ok = sdot_enabled()
4796 && blocked_enabled();
4797 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
4798 let blocked_ok = false;
4799 let run = move |start: usize, end: usize| {
4800 for r in start..end {
4801 let mut bi = 0usize;
4802 #[cfg(target_arch = "aarch64")]
4803 if blocked_ok {
4804 while bi + 4 <= acts.len() {
4805 let xs = [
4806 acts[bi].xq.as_slice(),
4807 acts[bi + 1].xq.as_slice(),
4808 acts[bi + 2].xq.as_slice(),
4809 acts[bi + 3].xq.as_slice(),
4810 ];
4811 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
4812 for k in 0..4 {
4813 let act = &acts[bi + k];
4814 let mut acc = d[k] * act.sx;
4815 for &(j, xv) in &act.outliers {
4816 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4817 acc += w * sc * xv;
4818 }
4819 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4821 }
4822 bi += 4;
4823 }
4824 }
4825 #[cfg(target_arch = "x86_64")]
4826 if blocked_ok {
4827 while bi + 4 <= acts.len() {
4828 let xs = [
4829 acts[bi].xq.as_slice(),
4830 acts[bi + 1].xq.as_slice(),
4831 acts[bi + 2].xq.as_slice(),
4832 acts[bi + 3].xq.as_slice(),
4833 ];
4834 let d = unsafe {
4835 if vnni_tiles_enabled() {
4836 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
4837 } else {
4838 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
4839 }
4840 };
4841 for k in 0..4 {
4842 let act = &acts[bi + k];
4843 let mut acc = d[k] * act.sx;
4844 for &(j, xv) in &act.outliers {
4845 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4846 acc += w * sc * xv;
4847 }
4848 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4850 }
4851 bi += 4;
4852 }
4853 }
4854 let _ = blocked_ok;
4855 while bi < acts.len() {
4856 let act = &acts[bi];
4857 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4858 for &(j, xv) in &act.outliers {
4859 let (w, s) = q4t_outlier(bytes, r, gpr, j);
4860 acc += w * s * xv;
4861 }
4862 unsafe { *out_addr.at(bi * rows + r) = acc };
4864 bi += 1;
4865 }
4866 }
4867 };
4868 dispatch_rows(pool, rows, &run);
4869 return;
4870 }
4871 let run = move |start: usize, end: usize| {
4872 for r in start..end {
4873 for bi in 0..b {
4874 let x = &xs_all[bi * cols..(bi + 1) * cols];
4875 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
4877 }
4878 }
4879 };
4880 dispatch_rows(pool, rows, &run);
4881}
4882
4883fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
4892 (0..gpr)
4893 .map(|gi| {
4894 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
4895 .iter()
4896 .map(|&v| v as i32)
4897 .sum()
4898 })
4899 .collect()
4900}
4901
4902#[inline]
4906#[allow(unreachable_code)]
4907#[cfg(target_arch = "x86_64")]
4912#[target_feature(enable = "avx2")]
4913unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4914 unsafe {
4916 use core::arch::x86_64::*;
4917 let expand = _mm256_setr_epi8(
4919 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,
4920 3, 3, 3,
4921 );
4922 let bitsel = _mm256_setr_epi8(
4923 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4924 -128, 1, 2, 4, 8, 16, 32, 64, -128,
4925 );
4926 let ones8 = _mm256_set1_epi8(1);
4927 let ones16 = _mm256_set1_epi16(1);
4928 let mut acc = 0f32;
4929 for gi in 0..gpr {
4930 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4931 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4932 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4933 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4934 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4935 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4936 let sel = _mm256_and_si256(x, mask);
4937 let p16 = _mm256_maddubs_epi16(ones8, sel);
4939 let d32 = _mm256_madd_epi16(p16, ones16);
4940 let hi128 = _mm256_extracti128_si256::<1>(d32);
4941 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4942 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4943 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4944 let msum = _mm_cvtsi128_si32(s32);
4945 let d = 2 * msum - gsum[gi];
4948 acc += d as f32 * s;
4949 }
4950 acc
4951 }
4952}
4953
4954#[cfg(target_arch = "x86_64")]
4957#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4958unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4959 unsafe {
4961 use core::arch::x86_64::*;
4962 let expand = _mm256_setr_epi8(
4963 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,
4964 3, 3, 3,
4965 );
4966 let bitsel = _mm256_setr_epi8(
4967 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4968 -128, 1, 2, 4, 8, 16, 32, 64, -128,
4969 );
4970 let ones8 = _mm256_set1_epi8(1);
4971 let mut acc = 0f32;
4972 for gi in 0..gpr {
4973 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4974 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4975 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4976 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4977 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4978 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4979 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4980 let d = 2 * msum - gsum[gi];
4981 acc += d as f32 * s;
4982 }
4983 acc
4984 }
4985}
4986
4987#[cfg(target_arch = "x86_64")]
4989#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4990unsafe fn dot_q1_row_1x4_vnni(
4991 bytes: &[u8],
4992 r: usize,
4993 gpr: usize,
4994 xs: [&[i8]; 4],
4995 gsums: [&[i32]; 4],
4996) -> [f32; 4] {
4997 unsafe {
4999 use core::arch::x86_64::*;
5000 let expand = _mm256_setr_epi8(
5001 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,
5002 3, 3, 3,
5003 );
5004 let bitsel = _mm256_setr_epi8(
5005 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5006 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5007 );
5008 let ones8 = _mm256_set1_epi8(1);
5009 let mut acc = [0f32; 4];
5010 for gi in 0..gpr {
5011 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5012 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5013 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5014 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5015 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5016 for (k, xq) in xs.iter().enumerate() {
5017 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5018 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5019 let d = 2 * msum - gsums[k][gi];
5020 acc[k] += d as f32 * s;
5021 }
5022 }
5023 acc
5024 }
5025}
5026
5027#[cfg(target_arch = "x86_64")]
5030#[target_feature(enable = "avx2")]
5031unsafe fn dot_q1_row_1x4_avx2(
5032 bytes: &[u8],
5033 r: usize,
5034 gpr: usize,
5035 xs: [&[i8]; 4],
5036 gsums: [&[i32]; 4],
5037) -> [f32; 4] {
5038 unsafe {
5040 use core::arch::x86_64::*;
5041 let expand = _mm256_setr_epi8(
5042 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,
5043 3, 3, 3,
5044 );
5045 let bitsel = _mm256_setr_epi8(
5046 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5047 -128, 1, 2, 4, 8, 16, 32, 64, -128,
5048 );
5049 let ones8 = _mm256_set1_epi8(1);
5050 let ones16 = _mm256_set1_epi16(1);
5051 let mut acc = [0f32; 4];
5052 for gi in 0..gpr {
5053 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5054 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5055 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5056 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5057 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5058 for (k, xq) in xs.iter().enumerate() {
5059 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5060 let sel = _mm256_and_si256(x, mask);
5061 let p16 = _mm256_maddubs_epi16(ones8, sel);
5062 let d32 = _mm256_madd_epi16(p16, ones16);
5063 let hi128 = _mm256_extracti128_si256::<1>(d32);
5064 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5065 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5066 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5067 let msum = _mm_cvtsi128_si32(s32);
5068 let d = 2 * msum - gsums[k][gi];
5069 acc[k] += d as f32 * s;
5070 }
5071 }
5072 acc
5073 }
5074}
5075
5076#[allow(unreachable_code)]
5077fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5078 #[cfg(target_arch = "aarch64")]
5079 unsafe {
5080 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
5081 }
5082 #[cfg(target_arch = "x86_64")]
5083 if avx2_enabled() {
5084 unsafe {
5085 if vnni_tiles_enabled() {
5086 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
5087 }
5088 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
5089 }
5090 }
5091 let _ = gsum;
5092 let mut acc = 0f32;
5093 for gi in 0..gpr {
5094 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5095 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5096 let mut d = 0i32;
5097 for (j, &b) in tile[2..].iter().enumerate() {
5098 for k in 0..8 {
5099 let w = ((b >> k) & 1) as i32 * 2 - 1;
5100 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
5101 }
5102 }
5103 acc += d as f32 * s;
5104 }
5105 acc
5106}
5107
5108#[cfg(target_arch = "aarch64")]
5117#[target_feature(enable = "neon,dotprod")]
5118unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5119 unsafe {
5122 use core::arch::aarch64::*;
5123 use core::arch::asm;
5124 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
5125 let m = vld1q_u8(MASKS.as_ptr());
5126 macro_rules! tile_dot {
5128 ($t:expr, $x:expr) => {{
5129 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
5130 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
5131 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
5132 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
5133 let x0 = vld1q_s8($x);
5134 let x1 = vld1q_s8($x.add(16));
5135 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5136 asm!(
5137 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5138 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5139 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5140 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5141 options(pure, nomem, nostack),
5142 );
5143 vaddq_s32(a0, a1)
5144 }};
5145 }
5146 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
5155 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
5156 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
5157 const IW11: [u8; 16] = [
5158 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
5159 ];
5160 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
5161 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
5162 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
5163 let isc = vld1_u8(ISC.as_ptr());
5164 macro_rules! tile_dot_tbl {
5166 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
5167 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
5168 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
5169 let x0 = vld1q_s8($x);
5170 let x1 = vld1q_s8($x.add(16));
5171 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5172 asm!(
5173 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5174 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5175 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5176 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5177 options(pure, nomem, nostack),
5178 );
5179 vaddq_s32(a0, a1)
5180 }};
5181 }
5182 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
5183 let row_base = r * gpr * Q1_TILE;
5184 let abs_end = bytes.len();
5185 let xp = xq.as_ptr();
5186 let gp = gsum.as_ptr();
5187 let mut accv = vdupq_n_f32(0.0);
5188 let mut gi = 0;
5189 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
5192 let t0 = base.add(gi * Q1_TILE);
5193 let ld_a = vld1q_u8(t0);
5194 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
5195 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
5196 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
5197 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
5198 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
5199 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
5201 let g = vld1q_s32(gp.add(gi));
5202 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
5203 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
5204 let scf: float32x4_t;
5205 asm!(
5206 "fcvtl {o:v}.4s, {i:v}.4h",
5207 o = out(vreg) scf, i = in(vreg) sc16,
5208 options(pure, nomem, nostack),
5209 );
5210 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
5211 gi += 4;
5212 }
5213 let mut acc = vaddvq_f32(accv);
5214 while gi < gpr {
5215 let t = base.add(gi * Q1_TILE);
5216 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5217 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
5218 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
5219 gi += 1;
5220 }
5221 acc
5222 }
5223}
5224
5225#[cfg(target_arch = "aarch64")]
5230#[target_feature(enable = "neon,dotprod")]
5231unsafe fn dot_q1_row_1x4_sdot(
5232 bytes: &[u8],
5233 r: usize,
5234 gpr: usize,
5235 xs: [&[i8]; 4],
5236 gs: [&[i32]; 4],
5237) -> [f32; 4] {
5238 unsafe {
5240 use core::arch::aarch64::*;
5241 use core::arch::asm;
5242 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
5243 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
5244 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
5245 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
5246 const IW11: [u8; 16] = [
5247 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
5248 ];
5249 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
5250 let m = vld1q_u8(MASKS.as_ptr());
5251 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
5252 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
5253 let isc = vld1_u8(ISC.as_ptr());
5254 macro_rules! sdot2 {
5255 ($w0:expr, $w1:expr, $x:expr) => {{
5256 let x0 = vld1q_s8($x);
5257 let x1 = vld1q_s8($x.add(16));
5258 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5259 asm!(
5260 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5261 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5262 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5263 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
5264 options(pure, nomem, nostack),
5265 );
5266 vaddq_s32(a0, a1)
5267 }};
5268 }
5269 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
5270 let row_base = r * gpr * Q1_TILE;
5271 let abs_end = bytes.len();
5272 let mut accv = [vdupq_n_f32(0.0); 4];
5273 let mut gi = 0;
5274 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
5275 let t0 = base.add(gi * Q1_TILE);
5276 let ld_a = vld1q_u8(t0);
5277 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
5278 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
5280 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
5281 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
5282 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
5283 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
5284 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
5285 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
5286 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
5287 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
5288 let scf: float32x4_t;
5289 asm!(
5290 "fcvtl {o:v}.4s, {i:v}.4h",
5291 o = out(vreg) scf, i = in(vreg) sc16,
5292 options(pure, nomem, nostack),
5293 );
5294 for k in 0..4 {
5295 let xp = xs[k].as_ptr();
5296 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
5297 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
5298 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
5299 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
5300 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
5301 let g = vld1q_s32(gs[k].as_ptr().add(gi));
5302 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
5303 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
5304 }
5305 gi += 4;
5306 }
5307 let mut acc = [
5308 vaddvq_f32(accv[0]),
5309 vaddvq_f32(accv[1]),
5310 vaddvq_f32(accv[2]),
5311 vaddvq_f32(accv[3]),
5312 ];
5313 while gi < gpr {
5314 let t = base.add(gi * Q1_TILE);
5315 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5316 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
5317 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
5318 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
5319 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
5320 for k in 0..4 {
5321 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
5322 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
5323 }
5324 gi += 1;
5325 }
5326 acc
5327 }
5328}
5329
5330#[inline]
5332fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5333 let gi = j / GROUP_SIZE;
5334 let k = j % GROUP_SIZE;
5335 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5336 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5337 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
5338 ((bit as i32 * 2 - 1) as f32, s)
5339}
5340
5341#[inline]
5343fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5344 let mut acc = 0f32;
5345 for gi in 0..gpr {
5346 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5347 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5348 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5349 let mut ga = 0f32;
5350 for (j, &b) in tile[2..].iter().enumerate() {
5351 for k in 0..8 {
5352 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
5353 }
5354 }
5355 acc += ga * s;
5356 }
5357 acc
5358}
5359
5360#[allow(clippy::too_many_arguments)]
5363fn q1_range_a8w8(
5364 bytes: &[u8],
5365 gpr: usize,
5366 act: &SplitAct,
5367 gsum: &[i32],
5368 out: SendMut,
5369 start: usize,
5370 end: usize,
5371) {
5372 for r in start..end {
5373 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5374 for &(j, xv) in &act.outliers {
5375 let (w, s) = q1_outlier(bytes, r, gpr, j);
5376 acc += w * s * xv;
5377 }
5378 unsafe { *out.at(r) = acc };
5380 }
5381}
5382
5383fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
5385 for r in start..end {
5386 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
5388 }
5389}
5390
5391fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
5396 let entries = base_len + (rows + 1) * 4;
5397 (base_len, entries, entries <= bytes.len())
5398}
5399
5400#[inline]
5402fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
5403 let o = rp_off + r * 4;
5404 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
5405}
5406
5407const SIGN5: [[f32; 5]; 256] = {
5411 let mut lut = [[0.0f32; 5]; 256];
5412 let pow3 = [1u16, 3, 9, 27, 81];
5413 let mut byte = 0usize;
5414 while byte < 256 {
5415 let mut i = 0usize;
5416 while i < 5 {
5417 let code = (byte as u16 / pow3[i]) % 3;
5418 lut[byte][i] = if code == 1 {
5419 1.0
5420 } else if code == 2 {
5421 -1.0
5422 } else {
5423 0.0
5424 };
5425 i += 1;
5426 }
5427 byte += 1;
5428 }
5429 lut
5430};
5431
5432const SIGN5_I8: [[i8; 5]; 256] = {
5434 let mut lut = [[0i8; 5]; 256];
5435 let pow3 = [1u16, 3, 9, 27, 81];
5436 let mut byte = 0usize;
5437 while byte < 256 {
5438 let mut i = 0usize;
5439 while i < 5 {
5440 let code = (byte as u16 / pow3[i]) % 3;
5441 lut[byte][i] = if code == 1 {
5442 1
5443 } else if code == 2 {
5444 -1
5445 } else {
5446 0
5447 };
5448 i += 1;
5449 }
5450 byte += 1;
5451 }
5452 lut
5453};
5454
5455const SIGN5_U64: [u64; 256] = {
5461 let mut lut = [0u64; 256];
5462 let pow3 = [1u16, 3, 9, 27, 81];
5463 let mut byte = 0usize;
5464 while byte < 256 {
5465 let mut v = 0u64;
5466 let mut i = 0usize;
5467 while i < 5 {
5468 let code = (byte as u16 / pow3[i]) % 3;
5469 let s: u8 = if code == 1 {
5470 1
5471 } else if code == 2 {
5472 0xFF
5473 } else {
5474 0
5475 };
5476 v |= (s as u64) << (i * 8);
5477 i += 1;
5478 }
5479 lut[byte] = v;
5480 byte += 1;
5481 }
5482 lut
5483};
5484
5485#[inline]
5490fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
5491 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5492 let off = (r * gpr + j / GROUP_SIZE) * TILE;
5493 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5494 let within = j % GROUP_SIZE;
5495 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
5496}
5497
5498#[cfg(target_arch = "aarch64")]
5501#[target_feature(enable = "neon,dotprod")]
5502#[inline]
5503unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
5504 unsafe {
5506 use core::arch::aarch64::*;
5507 use core::arch::asm;
5508 let w0 = vld1q_s8(w);
5509 let w1 = vld1q_s8(w.add(16));
5510 let x0 = vld1q_s8(x);
5511 let x1 = vld1q_s8(x.add(16));
5512 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5513 asm!(
5514 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5515 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5516 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5517 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5518 options(pure, nomem, nostack),
5519 );
5520 vaddvq_s32(vaddq_s32(a0, a1))
5521 }
5522}
5523
5524#[cfg(target_arch = "x86_64")]
5527#[target_feature(enable = "avx2")]
5528#[inline]
5529unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
5530 unsafe {
5532 use core::arch::x86_64::*;
5533 let wv = _mm256_loadu_si256(w as *const __m256i);
5534 let xv = _mm256_loadu_si256(x as *const __m256i);
5535 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5536 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
5537 let hi128 = _mm256_extracti128_si256::<1>(d);
5538 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5539 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5540 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5541 _mm_cvtsi128_si32(s32)
5542 }
5543}
5544
5545#[inline]
5550fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
5551 debug_assert!(dst.len() >= 40);
5552 unsafe {
5555 let p = dst.as_mut_ptr();
5556 for bi in 0..7 {
5557 core::ptr::write_unaligned(
5558 p.add(bi * 5) as *mut u64,
5559 SIGN5_U64[*codes.add(bi) as usize],
5560 );
5561 }
5562 }
5563}
5564
5565#[inline]
5570fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
5571 #[cfg(target_arch = "aarch64")]
5572 unsafe {
5573 return sdot32_i8(w, x);
5574 }
5575 #[cfg(target_arch = "x86_64")]
5576 unsafe {
5577 return i8dot32_avx2(w, x);
5578 }
5579 #[allow(unreachable_code)]
5580 unsafe {
5581 let mut s = 0i32;
5582 for k in 0..GROUP_SIZE {
5583 s += *w.add(k) as i32 * *x.add(k) as i32;
5584 }
5585 s
5586 }
5587}
5588
5589#[inline]
5590unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
5591 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
5592 (
5593 SIGN5_U64[*codes as usize],
5594 SIGN5_U64[*codes.add(1) as usize],
5595 SIGN5_U64[*codes.add(2) as usize],
5596 SIGN5_U64[*codes.add(3) as usize],
5597 SIGN5_U64[*codes.add(4) as usize],
5598 SIGN5_U64[*codes.add(5) as usize],
5599 SIGN5_U64[*codes.add(6) as usize],
5600 )
5601 };
5602
5603 let u0 = s0 | (s1 << 40);
5604 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
5605 let u2 = (s3 >> 8) | (s4 << 32);
5606 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
5607
5608 (u0, u1, u2, u3)
5609}
5610
5611#[cfg(target_arch = "aarch64")]
5615#[target_feature(enable = "neon,dotprod")]
5616unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5617 use core::arch::aarch64::*;
5618 use core::arch::asm;
5619 unsafe {
5620 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5621 let mut acc = 0f32;
5622 let bytes_ptr = bytes.as_ptr();
5623 let xq_ptr = xq.as_ptr();
5624 let row_off = r * gpr * TILE;
5625
5626 let gpr2 = gpr & !1;
5627 let mut gi = 0;
5628 while gi < gpr2 {
5629 let off0 = row_off + gi * TILE;
5630 let off1 = off0 + TILE;
5631 let s0 = f16_to_f32(u16::from_le_bytes([
5632 *bytes_ptr.add(off0),
5633 *bytes_ptr.add(off0 + 1),
5634 ]));
5635 let s1 = f16_to_f32(u16::from_le_bytes([
5636 *bytes_ptr.add(off1),
5637 *bytes_ptr.add(off1 + 1),
5638 ]));
5639
5640 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5641 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5642
5643 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5644 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5645 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5646 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5647
5648 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5649 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5650 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
5651 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
5652
5653 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
5654 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5655 asm!(
5656 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
5657 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
5658 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
5659 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
5660 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
5661 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
5662 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
5663 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
5664 options(pure, nomem, nostack),
5665 );
5666 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
5667 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
5668 acc += d0 as f32 * s0 + d1 as f32 * s1;
5669 gi += 2;
5670 }
5671
5672 if gi < gpr {
5673 let off = row_off + gi * TILE;
5674 let s = f16_to_f32(u16::from_le_bytes([
5675 *bytes_ptr.add(off),
5676 *bytes_ptr.add(off + 1),
5677 ]));
5678 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5679 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5680 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5681 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5682 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5683 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5684 asm!(
5685 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5686 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5687 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5688 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5689 options(pure, nomem, nostack),
5690 );
5691 let d = vaddvq_s32(vaddq_s32(a0, a1));
5692 acc += d as f32 * s;
5693 }
5694 acc
5695 }
5696}
5697
5698#[cfg(target_arch = "x86_64")]
5700#[target_feature(enable = "avx2")]
5701unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5702 use core::arch::x86_64::*;
5703 unsafe {
5704 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5705 let mut acc = 0f32;
5706 let bytes_ptr = bytes.as_ptr();
5707 let xq_ptr = xq.as_ptr();
5708 let row_off = r * gpr * TILE;
5709
5710 let ones = _mm256_set1_epi16(1);
5711 for gi in 0..gpr {
5712 let off = row_off + gi * TILE;
5713 let s = f16_to_f32(u16::from_le_bytes([
5714 *bytes_ptr.add(off),
5715 *bytes_ptr.add(off + 1),
5716 ]));
5717 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5718 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5719 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5720 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5721 let d256 = _mm256_madd_epi16(p16, ones);
5722 let d128 = _mm_add_epi32(
5723 _mm256_castsi256_si128(d256),
5724 _mm256_extracti128_si256(d256, 1),
5725 );
5726 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
5727 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
5728 acc += d32 as f32 * s;
5729 }
5730 acc
5731 }
5732}
5733
5734#[cfg(target_arch = "x86_64")]
5736#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5737unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5738 use core::arch::x86_64::*;
5739 unsafe {
5741 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5742 let mut acc = 0f32;
5743 let bytes_ptr = bytes.as_ptr();
5744 let xq_ptr = xq.as_ptr();
5745 let row_off = r * gpr * TILE;
5746 for gi in 0..gpr {
5747 let off = row_off + gi * TILE;
5748 let s = f16_to_f32(u16::from_le_bytes([
5749 *bytes_ptr.add(off),
5750 *bytes_ptr.add(off + 1),
5751 ]));
5752 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5753 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5754 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5755 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5756 acc += d as f32 * s;
5757 }
5758 acc
5759 }
5760}
5761
5762#[inline]
5766fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5767 #[cfg(target_arch = "aarch64")]
5768 unsafe {
5769 return q1t_dot_row_sdot(bytes, r, gpr, xq);
5770 }
5771 #[cfg(target_arch = "x86_64")]
5772 unsafe {
5773 if vnni_tiles_enabled() {
5774 return q1t_dot_row_vnni(bytes, r, gpr, xq);
5775 }
5776 return q1t_dot_row_avx2(bytes, r, gpr, xq);
5777 }
5778 #[allow(unreachable_code)]
5779 {
5780 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5781 let mut acc = 0f32;
5782 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
5784 let off = (r * gpr + gi) * TILE;
5785 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5786 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
5787 let mut d = 0i32;
5788 for k in 0..GROUP_SIZE {
5789 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
5790 }
5791 acc += d as f32 * s;
5792 }
5793 acc
5794 }
5795}
5796
5797fn q1t_row_outlier_correction(
5804 bytes: &[u8],
5805 r: usize,
5806 rp_off: usize,
5807 entries_off: usize,
5808 has_ov: bool,
5809 x: &[f32],
5810) -> f32 {
5811 if !has_ov {
5812 return 0.0;
5813 }
5814 let (c0, c1) = (
5815 q1t_rowptr(bytes, rp_off, r),
5816 q1t_rowptr(bytes, rp_off, r + 1),
5817 );
5818 let mut corr = 0f32;
5819 for p in c0..c1 {
5820 let e = entries_off + p * 4;
5821 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5822 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5823 corr += val * x[col];
5824 }
5825 corr
5826}
5827
5828fn q1t_dequant_row(
5832 bytes: &[u8],
5833 r: usize,
5834 gpr: usize,
5835 rp_off: usize,
5836 entries_off: usize,
5837 has_ov: bool,
5838 buf: &mut [f32],
5839) {
5840 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5841 for g in 0..gpr {
5842 let off = (r * gpr + g) * TILE;
5843 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5844 let codes = &bytes[off + 2..off + TILE];
5845 let bc = g * GROUP_SIZE;
5846 for bi in 0..6 {
5848 let lut = &SIGN5[codes[bi] as usize];
5849 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
5850 for i in 0..5 {
5851 d[i] = lut[i] * s;
5852 }
5853 }
5854 let lut = &SIGN5[codes[6] as usize];
5855 buf[bc + 30] = lut[0] * s;
5856 buf[bc + 31] = lut[1] * s;
5857 }
5858 if !has_ov {
5859 return;
5860 }
5861 let (c0, c1) = (
5862 q1t_rowptr(bytes, rp_off, r),
5863 q1t_rowptr(bytes, rp_off, r + 1),
5864 );
5865 for p in c0..c1 {
5866 let e = entries_off + p * 4;
5867 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5868 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5869 }
5870}
5871
5872fn q1t_add_overlay(
5876 bytes: &[u8],
5877 x: &[f32],
5878 rows: usize,
5879 cols: usize,
5880 out: &mut [f32],
5881 pool: Option<&Pool>,
5882) {
5883 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5884 let gpr = cols / GROUP_SIZE;
5885 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5886 if !has_ov {
5887 return;
5888 }
5889 let out_addr = SendMut(out.as_mut_ptr());
5890 let run = move |start: usize, end: usize| {
5891 for r in start..end {
5892 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5893 unsafe { *out_addr.at(r) += corr };
5895 }
5896 };
5897 dispatch_rows(pool, rows, &run);
5898}
5899
5900#[allow(clippy::too_many_arguments)]
5903fn q1t_range_a8w8(
5904 bytes: &[u8],
5905 gpr: usize,
5906 rp_off: usize,
5907 ent_off: usize,
5908 has_ov: bool,
5909 act: &SplitAct,
5910 x: &[f32],
5911 out: SendMut,
5912 start: usize,
5913 end: usize,
5914) {
5915 for r in start..end {
5916 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5917 for &(j, xv) in &act.outliers {
5918 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5919 }
5920 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5921 unsafe { *out.at(r) = acc };
5923 }
5924}
5925
5926#[allow(clippy::too_many_arguments)]
5929fn q1t_range_f32_batch(
5930 bytes: &[u8],
5931 gpr: usize,
5932 rp_off: usize,
5933 ent_off: usize,
5934 has_ov: bool,
5935 x: &[f32],
5936 out: SendMut,
5937 start: usize,
5938 end: usize,
5939) {
5940 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5941 let mut sg = [0f32; GROUP_SIZE];
5942 for r in start..end {
5943 let mut acc = 0f32;
5944 for g in 0..gpr {
5945 let off = (r * gpr + g) * TILE;
5946 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5947 let codes = &bytes[off + 2..off + TILE];
5948 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5949 for bi in 0..6 {
5950 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5951 }
5952 let lut = &SIGN5[codes[6] as usize];
5953 sg[30] = lut[0];
5954 sg[31] = lut[1];
5955 let mut gsum = 0f32;
5956 for k in 0..GROUP_SIZE {
5957 gsum += sg[k] * xg[k];
5958 }
5959 acc += s * gsum;
5960 }
5961 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5962 unsafe { *out.at(r) = acc };
5964 }
5965}
5966
5967fn q1t_matvec(
5971 bytes: &[u8],
5972 x: &[f32],
5973 rows: usize,
5974 cols: usize,
5975 out: &mut [f32],
5976 pool: Option<&Pool>,
5977) {
5978 debug_assert_eq!(out.len(), rows);
5979 const TILE: usize = cortiq_core::quant::Q1T_TILE;
5980 let gpr = cols / GROUP_SIZE;
5981 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5982 let out_addr = SendMut(out.as_mut_ptr());
5983 if a8w8_enabled() {
5987 let act = split_act(x);
5988 let act = &act;
5989 let run = move |start: usize, end: usize| {
5990 for r in start..end {
5991 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5992 for &(j, xv) in &act.outliers {
5993 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5994 }
5995 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5996 unsafe { *out_addr.at(r) = acc };
5998 }
5999 };
6000 dispatch_rows(pool, rows, &run);
6001 return;
6002 }
6003 let run = move |start: usize, end: usize| {
6004 let mut sg = [0f32; GROUP_SIZE];
6008 for r in start..end {
6009 let mut acc = 0f32;
6010 for g in 0..gpr {
6011 let off = (r * gpr + g) * TILE;
6012 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6013 let codes = &bytes[off + 2..off + TILE];
6014 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6015 for bi in 0..6 {
6016 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6017 }
6018 let lut = &SIGN5[codes[6] as usize];
6019 sg[30] = lut[0];
6020 sg[31] = lut[1];
6021 let mut gsum = 0f32;
6022 for k in 0..GROUP_SIZE {
6023 gsum += sg[k] * xg[k];
6024 }
6025 acc += s * gsum;
6026 }
6027 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6028 unsafe { *out_addr.at(r) = acc };
6029 }
6030 };
6031 dispatch_rows(pool, rows, &run);
6032}
6033
6034#[cfg(target_arch = "aarch64")]
6040#[target_feature(enable = "neon,dotprod")]
6041unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6042 use core::arch::aarch64::*;
6043 use core::arch::asm;
6044 unsafe {
6046 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6047 let bytes_ptr = bytes.as_ptr();
6048 let row_off = r * gpr * TILE;
6049 let xp = [xa.as_ptr(), xb.as_ptr()];
6050 let mut acc = [0f32; 2];
6051 macro_rules! sdot2 {
6052 ($w0:expr, $w1:expr, $x:expr) => {{
6053 let x0 = vld1q_s8($x);
6054 let x1 = vld1q_s8($x.add(16));
6055 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6056 asm!(
6057 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6058 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6059 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6060 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6061 options(pure, nomem, nostack),
6062 );
6063 vaddvq_s32(vaddq_s32(a0, a1))
6064 }};
6065 }
6066 let gpr2 = gpr & !1;
6067 let mut gi = 0;
6068 while gi < gpr2 {
6069 let off0 = row_off + gi * TILE;
6070 let off1 = off0 + TILE;
6071 let s0 = f16_to_f32(u16::from_le_bytes([
6072 *bytes_ptr.add(off0),
6073 *bytes_ptr.add(off0 + 1),
6074 ]));
6075 let s1 = f16_to_f32(u16::from_le_bytes([
6076 *bytes_ptr.add(off1),
6077 *bytes_ptr.add(off1 + 1),
6078 ]));
6079 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6080 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6081 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6082 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6083 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6084 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6085 for k in 0..2 {
6086 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
6087 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
6088 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
6089 }
6090 gi += 2;
6091 }
6092 if gi < gpr {
6093 let off = row_off + gi * TILE;
6094 let s = f16_to_f32(u16::from_le_bytes([
6095 *bytes_ptr.add(off),
6096 *bytes_ptr.add(off + 1),
6097 ]));
6098 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6099 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6100 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6101 for k in 0..2 {
6102 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
6103 acc[k] += d as f32 * s;
6104 }
6105 }
6106 acc
6107 }
6108}
6109
6110fn q1t_matvec2(
6116 bytes: &[u8],
6117 x1: &[f32],
6118 x2: &[f32],
6119 rows: usize,
6120 cols: usize,
6121 o1: &mut [f32],
6122 o2: &mut [f32],
6123 pool: Option<&Pool>,
6124) {
6125 debug_assert_eq!(o1.len(), rows);
6126 debug_assert_eq!(o2.len(), rows);
6127 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6128 let gpr = cols / GROUP_SIZE;
6129 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6130 let out1 = SendMut(o1.as_mut_ptr());
6131 let out2 = SendMut(o2.as_mut_ptr());
6132 if a8w8_enabled() {
6133 let a1 = split_act(x1);
6134 let a2 = split_act(x2);
6135 let (a1, a2) = (&a1, &a2);
6136 let run = move |start: usize, end: usize| {
6137 for r in start..end {
6138 #[cfg(target_arch = "aarch64")]
6139 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
6142 #[cfg(not(target_arch = "aarch64"))]
6143 let ds = [
6144 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
6145 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
6146 ];
6147 let mut acc1 = ds[0] * a1.sx;
6148 for &(j, xv) in &a1.outliers {
6149 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
6150 }
6151 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
6152 let mut acc2 = ds[1] * a2.sx;
6153 for &(j, xv) in &a2.outliers {
6154 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
6155 }
6156 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
6157 unsafe {
6159 *out1.at(r) = acc1;
6160 *out2.at(r) = acc2;
6161 }
6162 }
6163 };
6164 dispatch_rows(pool, rows, &run);
6165 return;
6166 }
6167 let run = move |start: usize, end: usize| {
6168 let mut sg = [0f32; GROUP_SIZE];
6171 for r in start..end {
6172 let mut acc1 = 0f32;
6173 let mut acc2 = 0f32;
6174 for g in 0..gpr {
6175 let off = (r * gpr + g) * TILE;
6176 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6177 let codes = &bytes[off + 2..off + TILE];
6178 for bi in 0..6 {
6179 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6180 }
6181 let lut = &SIGN5[codes[6] as usize];
6182 sg[30] = lut[0];
6183 sg[31] = lut[1];
6184 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6185 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6186 let mut gsum1 = 0f32;
6187 for k in 0..GROUP_SIZE {
6188 gsum1 += sg[k] * xg1[k];
6189 }
6190 acc1 += s * gsum1;
6191 let mut gsum2 = 0f32;
6192 for k in 0..GROUP_SIZE {
6193 gsum2 += sg[k] * xg2[k];
6194 }
6195 acc2 += s * gsum2;
6196 }
6197 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
6198 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
6199 unsafe {
6201 *out1.at(r) = acc1;
6202 *out2.at(r) = acc2;
6203 }
6204 }
6205 };
6206 dispatch_rows(pool, rows, &run);
6207}
6208
6209fn q1t_matmat(
6212 bytes: &[u8],
6213 xs: &[f32],
6214 b: usize,
6215 rows: usize,
6216 cols: usize,
6217 out: &mut [f32],
6218 pool: Option<&Pool>,
6219) {
6220 debug_assert_eq!(out.len(), b * rows);
6221 const TILE: usize = cortiq_core::quant::Q1T_TILE;
6222 let gpr = cols / GROUP_SIZE;
6223 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6224 let out_addr = SendMut(out.as_mut_ptr());
6225 if a8w8_enabled() {
6229 let acts: Vec<SplitAct> = (0..b)
6230 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
6231 .collect();
6232 let acts = &acts;
6233 let run = move |start: usize, end: usize| {
6234 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
6238 for g in 0..gpr {
6239 let off = (r * gpr + g) * TILE;
6240 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6241 q1t_unpack_group_i8(
6242 bytes.as_ptr().wrapping_add(off + 2),
6243 &mut sg[g * GROUP_SIZE..],
6244 );
6245 }
6246 for bi in 0..b {
6247 let act = &acts[bi];
6248 let mut isum = 0f32;
6249 for g in 0..gpr {
6250 let d = q1t_i8dot32(
6251 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
6252 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
6253 );
6254 isum += d as f32 * sc[g];
6255 }
6256 let mut acc = isum * act.sx;
6257 for &(j, xv) in &act.outliers {
6258 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6259 }
6260 accs[bi] = acc;
6261 }
6262 if has_ov {
6266 let (c0, c1) = (
6267 q1t_rowptr(bytes, rp_off, r),
6268 q1t_rowptr(bytes, rp_off, r + 1),
6269 );
6270 for p in c0..c1 {
6271 let e = ent_off + p * 4;
6272 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6273 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6274 for bi in 0..b {
6275 accs[bi] += val * xs[bi * cols + col];
6276 }
6277 }
6278 }
6279 for bi in 0..b {
6280 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
6281 }
6282 }
6283 };
6284 dispatch_rows(pool, rows, &run);
6285 return;
6286 }
6287 let run = move |start: usize, end: usize| {
6288 let mut buf = vec![0f32; cols];
6289 for r in start..end {
6290 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
6291 for bi in 0..b {
6292 let xr = &xs[bi * cols..(bi + 1) * cols];
6293 let mut acc = 0f32;
6294 for j in 0..cols {
6295 acc += buf[j] * xr[j];
6296 }
6297 unsafe { *out_addr.at(bi * rows + r) = acc };
6298 }
6299 }
6300 };
6301 dispatch_rows(pool, rows, &run);
6302}
6303
6304fn q1_matvec(
6305 bytes: &[u8],
6306 x: &[f32],
6307 rows: usize,
6308 cols: usize,
6309 out: &mut [f32],
6310 pool: Option<&Pool>,
6311) {
6312 debug_assert_eq!(out.len(), rows);
6313 let gpr = cols / GROUP_SIZE;
6314 let out_addr = SendMut(out.as_mut_ptr());
6315 if a8w8_enabled() {
6316 let act = split_act(x);
6317 let gsum = q1_group_sums(&act.xq, gpr);
6318 let (act, gsum) = (&act, &gsum);
6319 let run = move |start: usize, end: usize| {
6320 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
6321 };
6322 dispatch_rows(pool, rows, &run);
6323 return;
6324 }
6325 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
6326 dispatch_rows(pool, rows, &run);
6327}
6328
6329#[allow(clippy::too_many_arguments)]
6331fn q1_matvec2(
6332 bytes: &[u8],
6333 x1: &[f32],
6334 x2: &[f32],
6335 rows: usize,
6336 cols: usize,
6337 o1: &mut [f32],
6338 o2: &mut [f32],
6339 pool: Option<&Pool>,
6340) {
6341 let gpr = cols / GROUP_SIZE;
6342 let p1 = SendMut(o1.as_mut_ptr());
6343 let p2 = SendMut(o2.as_mut_ptr());
6344 if a8w8_enabled() {
6345 let a1 = split_act(x1);
6346 let a2 = split_act(x2);
6347 let g1 = q1_group_sums(&a1.xq, gpr);
6348 let g2 = q1_group_sums(&a2.xq, gpr);
6349 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
6350 let run = move |start: usize, end: usize| {
6351 for r in start..end {
6352 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
6353 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
6354 for &(j, xv) in &a1.outliers {
6355 let (w, s) = q1_outlier(bytes, r, gpr, j);
6356 v1 += w * s * xv;
6357 }
6358 for &(j, xv) in &a2.outliers {
6359 let (w, s) = q1_outlier(bytes, r, gpr, j);
6360 v2 += w * s * xv;
6361 }
6362 unsafe {
6364 *p1.at(r) = v1;
6365 *p2.at(r) = v2;
6366 }
6367 }
6368 };
6369 dispatch_rows(pool, rows, &run);
6370 return;
6371 }
6372 let run = move |start: usize, end: usize| {
6373 for r in start..end {
6374 unsafe {
6376 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
6377 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
6378 }
6379 }
6380 };
6381 dispatch_rows(pool, rows, &run);
6382}
6383
6384#[allow(clippy::too_many_arguments)]
6386fn q1_matmat(
6387 bytes: &[u8],
6388 xs_all: &[f32],
6389 b: usize,
6390 rows: usize,
6391 cols: usize,
6392 out: &mut [f32],
6393 pool: Option<&Pool>,
6394) {
6395 debug_assert_eq!(out.len(), b * rows);
6396 let gpr = cols / GROUP_SIZE;
6397 let out_addr = SendMut(out.as_mut_ptr());
6398 if a8w8_enabled() {
6399 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
6400 .map(|bi| {
6401 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
6402 let gsum = q1_group_sums(&act.xq, gpr);
6403 (act, gsum)
6404 })
6405 .collect();
6406 let acts = &acts;
6407 #[cfg(target_arch = "x86_64")]
6408 let blocked_ok = avx2_enabled()
6409 && blocked_enabled();
6410 #[cfg(target_arch = "aarch64")]
6411 let blocked_ok = sdot_enabled()
6412 && blocked_enabled();
6413 let run = move |start: usize, end: usize| {
6414 for r in start..end {
6415 let mut bi = 0usize;
6416 #[cfg(target_arch = "aarch64")]
6419 if blocked_ok {
6420 while bi + 4 <= acts.len() {
6421 let xs = [
6422 acts[bi].0.xq.as_slice(),
6423 acts[bi + 1].0.xq.as_slice(),
6424 acts[bi + 2].0.xq.as_slice(),
6425 acts[bi + 3].0.xq.as_slice(),
6426 ];
6427 let gs = [
6428 acts[bi].1.as_slice(),
6429 acts[bi + 1].1.as_slice(),
6430 acts[bi + 2].1.as_slice(),
6431 acts[bi + 3].1.as_slice(),
6432 ];
6433 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
6434 for k in 0..4 {
6435 let (act, _) = &acts[bi + k];
6436 let mut acc = d[k] * act.sx;
6437 for &(j, xv) in &act.outliers {
6438 let (w, sc) = q1_outlier(bytes, r, gpr, j);
6439 acc += w * sc * xv;
6440 }
6441 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6443 }
6444 bi += 4;
6445 }
6446 }
6447 #[cfg(target_arch = "x86_64")]
6448 if blocked_ok {
6449 while bi + 4 <= acts.len() {
6450 let xs = [
6451 acts[bi].0.xq.as_slice(),
6452 acts[bi + 1].0.xq.as_slice(),
6453 acts[bi + 2].0.xq.as_slice(),
6454 acts[bi + 3].0.xq.as_slice(),
6455 ];
6456 let gs = [
6457 acts[bi].1.as_slice(),
6458 acts[bi + 1].1.as_slice(),
6459 acts[bi + 2].1.as_slice(),
6460 acts[bi + 3].1.as_slice(),
6461 ];
6462 let d = unsafe {
6463 if vnni_tiles_enabled() {
6464 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
6465 } else {
6466 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
6467 }
6468 };
6469 for k in 0..4 {
6470 let (act, _) = &acts[bi + k];
6471 let mut acc = d[k] * act.sx;
6472 for &(j, xv) in &act.outliers {
6473 let (w, sc) = q1_outlier(bytes, r, gpr, j);
6474 acc += w * sc * xv;
6475 }
6476 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6478 }
6479 bi += 4;
6480 }
6481 }
6482 while bi < acts.len() {
6483 let (act, gsum) = &acts[bi];
6484 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6485 for &(j, xv) in &act.outliers {
6486 let (w, s) = q1_outlier(bytes, r, gpr, j);
6487 acc += w * s * xv;
6488 }
6489 unsafe { *out_addr.at(bi * rows + r) = acc };
6491 bi += 1;
6492 }
6493 }
6494 };
6495 dispatch_rows(pool, rows, &run);
6496 return;
6497 }
6498 let run = move |start: usize, end: usize| {
6499 for r in start..end {
6500 for bi in 0..b {
6501 let x = &xs_all[bi * cols..(bi + 1) * cols];
6502 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
6504 }
6505 }
6506 };
6507 dispatch_rows(pool, rows, &run);
6508}
6509
6510fn q4matvec(
6516 bytes: &[u8],
6517 x: &[f32],
6518 rows: usize,
6519 cols: usize,
6520 out: &mut [f32],
6521 pool: Option<&Pool>,
6522) {
6523 debug_assert_eq!(out.len(), rows);
6524 let (packed, scales) = q4_split(bytes, rows, cols);
6525 let gpr = cols / GROUP_SIZE;
6526 let out_addr = SendMut(out.as_mut_ptr());
6527
6528 if a8w8_enabled() {
6529 let act = split_act(x);
6530 let run = move |start: usize, end: usize| {
6531 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
6532 };
6533 dispatch_rows(pool, rows, &run);
6534 return;
6535 }
6536
6537 let run =
6538 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
6539 dispatch_rows(pool, rows, &run);
6540}
6541
6542#[inline]
6545#[allow(unreachable_code)]
6546#[cfg(target_arch = "x86_64")]
6551#[target_feature(enable = "avx2")]
6552unsafe fn dot_q4b_row_1x4_avx2(
6553 buf: &[u8],
6554 scales: &[u8],
6555 g0: usize,
6556 gpr: usize,
6557 xs: [&[i8]; 4],
6558) -> [f32; 4] {
6559 unsafe {
6561 use core::arch::x86_64::*;
6562 let ones = _mm256_set1_epi16(1);
6563 let mut acc = [0f32; 4];
6564 for gi in 0..gpr {
6565 let s = f16_to_f32(u16::from_le_bytes([
6566 scales[(g0 + gi) * 2],
6567 scales[(g0 + gi) * 2 + 1],
6568 ]));
6569 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6570 let aw = _mm256_abs_epi8(w);
6571 for (k, xq) in xs.iter().enumerate() {
6572 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6573 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6574 let d = _mm256_madd_epi16(p16, ones);
6575 let hi128 = _mm256_extracti128_si256::<1>(d);
6576 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6577 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6578 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6579 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
6580 }
6581 }
6582 acc
6583 }
6584}
6585
6586#[cfg(target_arch = "x86_64")]
6588#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6589unsafe fn dot_q4b_row_1x4_vnni(
6590 buf: &[u8],
6591 scales: &[u8],
6592 g0: usize,
6593 gpr: usize,
6594 xs: [&[i8]; 4],
6595) -> [f32; 4] {
6596 unsafe {
6598 use core::arch::x86_64::*;
6599 let mut acc = [0f32; 4];
6600 for gi in 0..gpr {
6601 let s = f16_to_f32(u16::from_le_bytes([
6602 scales[(g0 + gi) * 2],
6603 scales[(g0 + gi) * 2 + 1],
6604 ]));
6605 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6606 let aw = _mm256_abs_epi8(w);
6607 for (k, xq) in xs.iter().enumerate() {
6608 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6609 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6610 acc[k] += d as f32 * s;
6611 }
6612 }
6613 acc
6614 }
6615}
6616
6617#[cfg(target_arch = "x86_64")]
6623#[target_feature(enable = "avx2")]
6624unsafe fn dot_q4b_row_1x4_sx_avx2(
6625 buf: &[u8],
6626 scales: &[u8],
6627 g0: usize,
6628 gpr: usize,
6629 xs: [&[i8]; 4],
6630 sxs: [f32; 4],
6631) -> [f32; 4] {
6632 unsafe {
6634 use core::arch::x86_64::*;
6635 let ones = _mm256_set1_epi16(1);
6636 let mut acc = [0f32; 4];
6637 for gi in 0..gpr {
6638 let s = f16_to_f32(u16::from_le_bytes([
6639 scales[(g0 + gi) * 2],
6640 scales[(g0 + gi) * 2 + 1],
6641 ]));
6642 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6643 let aw = _mm256_abs_epi8(w);
6644 for (k, xq) in xs.iter().enumerate() {
6645 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6646 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6647 let d = _mm256_madd_epi16(p16, ones);
6648 let hi128 = _mm256_extracti128_si256::<1>(d);
6649 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6650 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6651 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6652 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
6653 }
6654 }
6655 acc
6656 }
6657}
6658
6659#[cfg(target_arch = "x86_64")]
6662#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6663unsafe fn dot_q4b_row_1x4_sx_vnni(
6664 buf: &[u8],
6665 scales: &[u8],
6666 g0: usize,
6667 gpr: usize,
6668 xs: [&[i8]; 4],
6669 sxs: [f32; 4],
6670) -> [f32; 4] {
6671 unsafe {
6673 use core::arch::x86_64::*;
6674 let mut acc = [0f32; 4];
6675 for gi in 0..gpr {
6676 let s = f16_to_f32(u16::from_le_bytes([
6677 scales[(g0 + gi) * 2],
6678 scales[(g0 + gi) * 2 + 1],
6679 ]));
6680 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6681 let aw = _mm256_abs_epi8(w);
6682 for (k, xq) in xs.iter().enumerate() {
6683 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6684 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6685 acc[k] += (d as f32 * sxs[k]) * s;
6686 }
6687 }
6688 acc
6689 }
6690}
6691
6692#[allow(unreachable_code)]
6693fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6694 #[cfg(target_arch = "aarch64")]
6695 unsafe {
6696 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
6697 }
6698 #[cfg(target_arch = "x86_64")]
6699 unsafe {
6700 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
6701 }
6702 let mut acc = 0f32;
6703 for gi in 0..gpr {
6704 let g = g0 + gi;
6705 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6706 let mut d = 0i32;
6707 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6708 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6709 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6710 }
6711 acc += d as f32 * s;
6712 }
6713 acc
6714}
6715
6716#[inline]
6718#[allow(unreachable_code)]
6719fn dot_q4_row_i8_2(
6720 packed: &[u8],
6721 scales: &[u8],
6722 g0: usize,
6723 gpr: usize,
6724 xq1: &[i8],
6725 xq2: &[i8],
6726) -> (f32, f32) {
6727 #[cfg(target_arch = "aarch64")]
6728 unsafe {
6729 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
6730 }
6731 #[cfg(target_arch = "x86_64")]
6732 unsafe {
6733 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
6734 }
6735 (
6736 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
6737 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
6738 )
6739}
6740
6741#[allow(clippy::too_many_arguments)]
6744fn q4_range_a8w8(
6745 packed: &[u8],
6746 scales: &[u8],
6747 gpr: usize,
6748 cols: usize,
6749 act: &SplitAct,
6750 out: SendMut,
6751 start: usize,
6752 end: usize,
6753) {
6754 for r in start..end {
6755 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
6756 for &(j, xv) in &act.outliers {
6758 let flat = r * cols + j;
6759 let byte = packed[flat / 2];
6760 let nib = if flat & 1 == 0 {
6761 byte & 0x0F
6762 } else {
6763 byte >> 4
6764 };
6765 let s = f16_to_f32(u16::from_le_bytes([
6766 scales[(flat / GROUP_SIZE) * 2],
6767 scales[(flat / GROUP_SIZE) * 2 + 1],
6768 ]));
6769 acc += ((nib as i32 - 8) as f32) * s * xv;
6770 }
6771 unsafe { *out.at(r) = acc };
6773 }
6774}
6775
6776#[allow(clippy::too_many_arguments)]
6779fn q4_range2_a8w8(
6780 packed: &[u8],
6781 scales: &[u8],
6782 gpr: usize,
6783 cols: usize,
6784 a1: &SplitAct,
6785 a2: &SplitAct,
6786 p1: SendMut,
6787 p2: SendMut,
6788 start: usize,
6789 end: usize,
6790) {
6791 for r in start..end {
6792 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
6793 let mut acc1 = s1 * a1.sx;
6794 let mut acc2 = s2 * a2.sx;
6795 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
6797 for &(j, xv) in outliers {
6798 let flat = r * cols + j;
6799 let byte = packed[flat / 2];
6800 let nib = if flat & 1 == 0 {
6801 byte & 0x0F
6802 } else {
6803 byte >> 4
6804 };
6805 let s = f16_to_f32(u16::from_le_bytes([
6806 scales[(flat / GROUP_SIZE) * 2],
6807 scales[(flat / GROUP_SIZE) * 2 + 1],
6808 ]));
6809 *acc += ((nib as i32 - 8) as f32) * s * xv;
6810 }
6811 };
6812 fix(&a1.outliers, &mut acc1);
6813 fix(&a2.outliers, &mut acc2);
6814 unsafe {
6816 *p1.at(r) = acc1;
6817 *p2.at(r) = acc2;
6818 }
6819 }
6820}
6821
6822fn q4_range_f32(
6824 packed: &[u8],
6825 scales: &[u8],
6826 gpr: usize,
6827 x: &[f32],
6828 out: SendMut,
6829 start: usize,
6830 end: usize,
6831) {
6832 for r in start..end {
6833 let mut acc = 0f32;
6834 for gi in 0..gpr {
6835 let g = r * gpr + gi;
6836 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6837 let pk = &packed[g * 16..(g + 1) * 16];
6838 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6839 let mut ga = 0f32;
6840 for (k, &b) in pk.iter().enumerate() {
6841 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6842 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6843 }
6844 acc += ga * s;
6845 }
6846 unsafe { *out.at(r) = acc };
6848 }
6849}
6850
6851#[allow(clippy::too_many_arguments)]
6855fn q4matvec2(
6856 bytes: &[u8],
6857 x1: &[f32],
6858 x2: &[f32],
6859 rows: usize,
6860 cols: usize,
6861 o1: &mut [f32],
6862 o2: &mut [f32],
6863 pool: Option<&Pool>,
6864) {
6865 debug_assert_eq!(o1.len(), rows);
6866 debug_assert_eq!(o2.len(), rows);
6867 let (packed, scales) = q4_split(bytes, rows, cols);
6868 let gpr = cols / GROUP_SIZE;
6869
6870 if a8w8_enabled() {
6871 let a1 = split_act(x1);
6872 let a2 = split_act(x2);
6873 let p1 = SendMut(o1.as_mut_ptr());
6874 let p2 = SendMut(o2.as_mut_ptr());
6875 let run = move |start: usize, end: usize| {
6876 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
6877 };
6878 dispatch_rows(pool, rows, &run);
6879 return;
6880 }
6881
6882 let p1 = SendMut(o1.as_mut_ptr());
6883 let p2 = SendMut(o2.as_mut_ptr());
6884 let run = move |start: usize, end: usize| {
6885 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
6886 };
6887 dispatch_rows(pool, rows, &run);
6888}
6889
6890#[allow(clippy::too_many_arguments)]
6892fn q4_range2_f32(
6893 packed: &[u8],
6894 scales: &[u8],
6895 gpr: usize,
6896 x1: &[f32],
6897 x2: &[f32],
6898 p1: SendMut,
6899 p2: SendMut,
6900 start: usize,
6901 end: usize,
6902) {
6903 for r in start..end {
6904 let (mut acc1, mut acc2) = (0f32, 0f32);
6905 for gi in 0..gpr {
6906 let g = r * gpr + gi;
6907 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6908 let pk = &packed[g * 16..(g + 1) * 16];
6909 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6910 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6911 let (mut g1, mut g2) = (0f32, 0f32);
6912 for (k, &b) in pk.iter().enumerate() {
6913 let wl = (b & 0x0F) as f32 - 8.0;
6914 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
6915 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
6916 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
6917 }
6918 acc1 += g1 * s;
6919 acc2 += g2 * s;
6920 }
6921 unsafe {
6923 *p1.at(r) = acc1;
6924 *p2.at(r) = acc2;
6925 }
6926 }
6927}
6928
6929thread_local! {
6930 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
6933 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6934}
6935
6936#[allow(clippy::too_many_arguments)]
6942fn q4matmat(
6943 bytes: &[u8],
6944 xs_all: &[f32],
6945 b: usize,
6946 rows: usize,
6947 cols: usize,
6948 out: &mut [f32],
6949 pool: Option<&Pool>,
6950) {
6951 debug_assert_eq!(xs_all.len(), b * cols);
6952 debug_assert_eq!(out.len(), b * rows);
6953 let (packed, scales) = q4_split(bytes, rows, cols);
6954 let gpr = cols / GROUP_SIZE;
6955 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6956
6957 if a8w8_enabled() {
6958 let acts: Vec<SplitAct> = (0..b)
6959 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6960 .collect();
6961 let acts = &acts;
6962 let out_addr = SendMut(out.as_mut_ptr());
6963 let run = move |start: usize, end: usize| {
6964 ROW_I8.with(|rb| {
6965 let mut buf = rb.borrow_mut();
6966 buf.resize(cols, 0);
6967 for r in start..end {
6968 for gi in 0..gpr {
6972 let g = r * gpr + gi;
6973 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6974 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
6975 buf[gi * GROUP_SIZE + k * 2 + 1] =
6976 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
6977 }
6978 }
6979 let mut bi = 0usize;
6980 #[cfg(target_arch = "x86_64")]
6981 if avx2_enabled()
6982 && blocked_enabled()
6983 {
6984 while bi + 4 <= acts.len() {
6985 let xs = [
6986 acts[bi].xq.as_slice(),
6987 acts[bi + 1].xq.as_slice(),
6988 acts[bi + 2].xq.as_slice(),
6989 acts[bi + 3].xq.as_slice(),
6990 ];
6991 let d = unsafe {
6992 if vnni_tiles_enabled() {
6993 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
6994 } else {
6995 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
6996 }
6997 };
6998 for k in 0..4 {
6999 let act = &acts[bi + k];
7000 let mut acc = d[k] * act.sx;
7001 for &(j, xv) in &act.outliers {
7002 acc += (buf[j] as i8) as f32
7003 * gscale((r * cols + j) / GROUP_SIZE)
7004 * xv;
7005 }
7006 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7008 }
7009 bi += 4;
7010 }
7011 }
7012 while bi < acts.len() {
7013 let act = &acts[bi];
7014 let mut acc = 0f32;
7015 for gi in 0..gpr {
7016 let d = dot_i8_i8(
7017 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7018 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7019 );
7020 acc += d as f32 * gscale(r * gpr + gi);
7021 }
7022 acc *= act.sx;
7023 for &(j, xv) in &act.outliers {
7025 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7026 }
7027 unsafe { *out_addr.at(bi * rows + r) = acc };
7029 bi += 1;
7030 }
7031 }
7032 })
7033 };
7034 dispatch_rows(pool, rows, &run);
7035 return;
7036 }
7037
7038 let out_addr = SendMut(out.as_mut_ptr());
7039 let run = move |start: usize, end: usize| {
7040 ROW_F32.with(|rb| {
7041 let mut buf = rb.borrow_mut();
7042 buf.resize(cols, 0.0);
7043 for r in start..end {
7044 for gi in 0..gpr {
7047 let g = r * gpr + gi;
7048 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7049 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7050 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7051 }
7052 }
7053 for bi in 0..b {
7054 let x = &xs_all[bi * cols..(bi + 1) * cols];
7055 let mut acc = 0f32;
7056 for gi in 0..gpr {
7057 let mut ga = 0f32;
7058 for k in 0..GROUP_SIZE / 2 {
7063 let e = gi * GROUP_SIZE + k * 2;
7064 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
7065 }
7066 acc += ga * gscale(r * gpr + gi);
7067 }
7068 unsafe { *out_addr.at(bi * rows + r) = acc };
7070 }
7071 }
7072 })
7073 };
7074 dispatch_rows(pool, rows, &run);
7075}
7076
7077#[allow(clippy::too_many_arguments)]
7082fn vbitmatmat(
7083 bytes: &[u8],
7084 offsets: &[usize],
7085 xs_all: &[f32],
7086 b: usize,
7087 rows: usize,
7088 cols: usize,
7089 out: &mut [f32],
7090 pool: Option<&Pool>,
7091) {
7092 debug_assert_eq!(xs_all.len(), b * cols);
7093 debug_assert_eq!(out.len(), b * rows);
7094 debug_assert_eq!(offsets.len(), rows + 1);
7095 let ng = cols / GROUP_SIZE;
7096 let bits = &bytes[..rows];
7097 let sc_off = rows;
7098 let gscale = |r: usize, g: usize| {
7099 let so = (r * ng + g) * 2;
7100 f16_to_f32(u16::from_le_bytes([
7101 bytes[sc_off + so],
7102 bytes[sc_off + so + 1],
7103 ]))
7104 };
7105
7106 let decode_f32 = |r: usize, dst: &mut [f32]| {
7108 let bw = bits[r] as usize;
7109 let l = ((1i32 << (bw - 1)) - 1) as f32;
7110 let data = &bytes[offsets[r]..offsets[r + 1]];
7111 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
7112 for d in dst.iter_mut() {
7113 while nbits < bw {
7114 acc = (acc << 8) | data[idx] as u64;
7115 idx += 1;
7116 nbits += 8;
7117 }
7118 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
7119 nbits -= bw;
7120 *d = u - l;
7121 }
7122 };
7123
7124 if a8w8_enabled() {
7125 let acts: Vec<SplitAct> = (0..b)
7126 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7127 .collect();
7128 let acts = &acts;
7129 let out_addr = SendMut(out.as_mut_ptr());
7130 let run = move |start: usize, end: usize| {
7131 for r in start..end {
7132 let bw = bits[r] as usize;
7133 if bw == 8 {
7134 ROW_F32.with(|rb| {
7137 let mut buf = rb.borrow_mut();
7138 buf.resize(cols, 0.0);
7139 decode_f32(r, &mut buf);
7140 for bi in 0..b {
7141 let x = &xs_all[bi * cols..(bi + 1) * cols];
7142 let mut dot = 0f32;
7143 for g in 0..ng {
7144 let mut gd = 0f32;
7145 for k in 0..GROUP_SIZE {
7146 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
7147 }
7148 dot += gd * gscale(r, g);
7149 }
7150 unsafe { *out_addr.at(bi * rows + r) = dot };
7152 }
7153 });
7154 continue;
7155 }
7156 let l = (1i32 << (bw - 1)) - 1;
7157 let data = &bytes[offsets[r]..offsets[r + 1]];
7158 ROW_I8.with(|rb| {
7159 let mut buf = rb.borrow_mut();
7160 buf.resize(cols, 0);
7161 #[inline(always)]
7162 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
7163 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
7164 let u = unpack8::<B>(&data[blk * B..]);
7165 for k in 0..8 {
7166 chunk[k] = (u[k] - l) as i8 as u8;
7167 }
7168 }
7169 }
7170 match bw {
7171 3 => fill::<3>(data, l, &mut buf),
7172 4 => vbit_fill4(data, &mut buf),
7173 5 => fill::<5>(data, l, &mut buf),
7174 6 => fill::<6>(data, l, &mut buf),
7175 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
7176 }
7177 let mut bi = 0usize;
7178 #[cfg(target_arch = "x86_64")]
7182 if avx2_enabled()
7183 && blocked_enabled()
7184 {
7185 while bi + 4 <= acts.len() {
7186 let xs = [
7187 acts[bi].xq.as_slice(),
7188 acts[bi + 1].xq.as_slice(),
7189 acts[bi + 2].xq.as_slice(),
7190 acts[bi + 3].xq.as_slice(),
7191 ];
7192 let sxs = [
7193 acts[bi].sx,
7194 acts[bi + 1].sx,
7195 acts[bi + 2].sx,
7196 acts[bi + 3].sx,
7197 ];
7198 let d = unsafe {
7199 if vnni_tiles_enabled() {
7200 dot_q4b_row_1x4_sx_vnni(
7201 &buf,
7202 &bytes[sc_off..],
7203 r * ng,
7204 ng,
7205 xs,
7206 sxs,
7207 )
7208 } else {
7209 dot_q4b_row_1x4_sx_avx2(
7210 &buf,
7211 &bytes[sc_off..],
7212 r * ng,
7213 ng,
7214 xs,
7215 sxs,
7216 )
7217 }
7218 };
7219 for k in 0..4 {
7220 let act = &acts[bi + k];
7221 let mut dot = d[k];
7222 for &(j, xv) in &act.outliers {
7223 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
7224 }
7225 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
7227 }
7228 bi += 4;
7229 }
7230 }
7231 while bi < acts.len() {
7232 let act = &acts[bi];
7233 let mut dot = 0f32;
7234 for g in 0..ng {
7235 let d = dot_i8_i8(
7236 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
7237 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
7238 ) as f32
7239 * act.sx;
7240 dot += d * gscale(r, g);
7241 }
7242 for &(j, xv) in &act.outliers {
7243 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
7244 }
7245 unsafe { *out_addr.at(bi * rows + r) = dot };
7247 bi += 1;
7248 }
7249 });
7250 }
7251 };
7252 dispatch_rows(pool, rows, &run);
7253 return;
7254 }
7255
7256 let out_addr = SendMut(out.as_mut_ptr());
7257 let run = move |start: usize, end: usize| {
7258 ROW_F32.with(|rb| {
7259 let mut buf = rb.borrow_mut();
7260 buf.resize(cols, 0.0);
7261 for r in start..end {
7262 decode_f32(r, &mut buf);
7263 for bi in 0..b {
7264 let x = &xs_all[bi * cols..(bi + 1) * cols];
7265 let mut dot = 0f32;
7266 for g in 0..ng {
7267 let mut gd = 0f32;
7268 for k in 0..GROUP_SIZE {
7269 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
7270 }
7271 dot += gd * gscale(r, g);
7272 }
7273 unsafe { *out_addr.at(bi * rows + r) = dot };
7275 }
7276 }
7277 })
7278 };
7279 dispatch_rows(pool, rows, &run);
7280}
7281
7282pub(crate) fn gpu_batch_job<'a>(
7286 t: &'a QTensor,
7287 x: &[f32],
7288) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
7289 match t {
7290 QTensor::Mapped {
7291 model,
7292 idx,
7293 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
7294 rows,
7295 cols,
7296 row_scale,
7297 col_field,
7298 ..
7299 } => Some((
7300 model.clone(),
7301 crate::gpu::BatchJob {
7302 idx: *idx,
7303 rows: *rows,
7304 cols: *cols,
7305 row_scale,
7306 xs: prescale(x, col_field, *dt).into_owned(),
7307 layout: crate::gpu::BatchLayout::Q8,
7308 },
7309 )),
7310 QTensor::Mapped {
7312 model,
7313 idx,
7314 dtype: TensorDtype::Q1,
7315 rows,
7316 cols,
7317 ..
7318 } => Some((
7319 model.clone(),
7320 crate::gpu::BatchJob {
7321 idx: *idx,
7322 rows: *rows,
7323 cols: *cols,
7324 row_scale: &[],
7325 xs: x.to_vec(),
7326 layout: crate::gpu::BatchLayout::Q1,
7327 },
7328 )),
7329 _ => None,
7330 }
7331}
7332
7333thread_local! {
7334 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7335 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7336}
7337
7338pub(crate) fn prescale<'a>(
7339 x: &'a [f32],
7340 col_field: &[f32],
7341 dtype: TensorDtype,
7342) -> std::borrow::Cow<'a, [f32]> {
7343 if dtype == TensorDtype::Q8_2f {
7344 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
7345 } else {
7346 std::borrow::Cow::Borrowed(x)
7347 }
7348}
7349
7350pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
7353 x: &[f32],
7354 col_field: &[f32],
7355 dtype: TensorDtype,
7356 buf_id: u8,
7357 f: F,
7358) -> R {
7359 if dtype == TensorDtype::Q8_2f {
7360 if buf_id == 1 {
7361 PRESCALE_BUF1.with(|b| {
7362 let mut buf = b.borrow_mut();
7363 buf.clear();
7364 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7365 f(&buf)
7366 })
7367 } else {
7368 PRESCALE_BUF2.with(|b| {
7369 let mut buf = b.borrow_mut();
7370 buf.clear();
7371 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7372 f(&buf)
7373 })
7374 }
7375 } else {
7376 f(x)
7377 }
7378}
7379
7380#[cfg(target_arch = "x86_64")]
7385pub(crate) fn avx2_enabled() -> bool {
7386 use std::sync::OnceLock;
7387 static ON: OnceLock<bool> = OnceLock::new();
7388 *ON.get_or_init(|| {
7389 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
7390 && std::arch::is_x86_feature_detected!("avx2")
7391 && std::arch::is_x86_feature_detected!("fma")
7392 })
7393}
7394
7395#[cfg(target_arch = "x86_64")]
7400fn avx2_a8w8_enabled() -> bool {
7401 use std::sync::OnceLock;
7402 static ON: OnceLock<bool> = OnceLock::new();
7403 *ON.get_or_init(|| {
7404 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
7405 })
7406}
7407
7408#[inline]
7412pub(crate) fn a8w8_enabled() -> bool {
7413 #[cfg(target_arch = "aarch64")]
7414 {
7415 sdot_enabled()
7416 }
7417 #[cfg(target_arch = "x86_64")]
7418 {
7419 avx2_a8w8_enabled()
7420 }
7421 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7422 {
7423 false
7424 }
7425}
7426
7427#[inline]
7430#[allow(unreachable_code)]
7431fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
7432 #[cfg(target_arch = "aarch64")]
7433 unsafe {
7434 return dot_i8_sdot(w, xq);
7435 }
7436 #[cfg(target_arch = "x86_64")]
7437 unsafe {
7438 if avx512vnni_enabled() {
7439 return dot_i8_i8_vnni(w, xq);
7440 }
7441 return dot_i8_i8_avx2(w, xq);
7442 }
7443 w.iter()
7444 .zip(xq)
7445 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
7446 .sum()
7447}
7448
7449#[cfg(target_arch = "x86_64")]
7453fn avx512vnni_enabled() -> bool {
7454 use std::sync::OnceLock;
7455 static ON: OnceLock<bool> = OnceLock::new();
7456 *ON.get_or_init(|| {
7457 std::env::var("CMF_AVX512")
7458 .map(|v| v != "0")
7459 .unwrap_or(true)
7460 && std::arch::is_x86_feature_detected!("avx512f")
7461 && std::arch::is_x86_feature_detected!("avx512bw")
7462 && std::arch::is_x86_feature_detected!("avx512vl")
7463 && std::arch::is_x86_feature_detected!("avx512vnni")
7464 })
7465}
7466
7467#[cfg(target_arch = "x86_64")]
7475fn vnni_tiles_enabled() -> bool {
7476 use std::sync::OnceLock;
7477 static ON: OnceLock<bool> = OnceLock::new();
7478 *ON.get_or_init(|| {
7479 std::env::var("CMF_VNNI_TILES")
7480 .map(|v| v != "0")
7481 .unwrap_or(true)
7482 && avx512vnni_enabled()
7483 })
7484}
7485
7486#[cfg(target_arch = "x86_64")]
7491#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7492#[inline]
7493unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
7494 unsafe {
7496 use core::arch::x86_64::*;
7497 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
7498 let hi128 = _mm256_extracti128_si256::<1>(d);
7499 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7500 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7501 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7502 _mm_cvtsi128_si32(s32)
7503 }
7504}
7505
7506#[cfg(target_arch = "x86_64")]
7511#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7512unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7513 unsafe {
7515 use core::arch::x86_64::*;
7516 let n = w.len();
7517 let mut j = 0usize;
7518 let mut total: i32;
7519 {
7524 #[inline(always)]
7525 unsafe fn step(
7526 w: *const u8,
7527 x: *const i8,
7528 acc: core::arch::x86_64::__m512i,
7529 ) -> core::arch::x86_64::__m512i {
7530 unsafe {
7531 use core::arch::x86_64::*;
7532 let wv = _mm512_loadu_si512(w as *const _);
7533 let xv = _mm512_loadu_si512(x as *const _);
7534 let aw = _mm512_abs_epi8(wv);
7535 let neg = _mm512_movepi8_mask(wv);
7536 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
7537 _mm512_dpbusd_epi32(acc, aw, sx)
7538 }
7539 }
7540 let (mut a0, mut a1, mut a2, mut a3) = (
7541 _mm512_setzero_si512(),
7542 _mm512_setzero_si512(),
7543 _mm512_setzero_si512(),
7544 _mm512_setzero_si512(),
7545 );
7546 while j + 256 <= n {
7547 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7548 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
7549 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
7550 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
7551 j += 256;
7552 }
7553 while j + 64 <= n {
7554 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7555 j += 64;
7556 }
7557 let s01 = _mm512_add_epi32(a0, a1);
7558 let s23 = _mm512_add_epi32(a2, a3);
7559 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
7560 }
7561 if j + 32 <= n {
7563 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7564 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7565 let d = _mm256_dpbusd_epi32(
7566 _mm256_setzero_si256(),
7567 _mm256_abs_epi8(wv),
7568 _mm256_sign_epi8(xv, wv),
7569 );
7570 let hi128 = _mm256_extracti128_si256::<1>(d);
7571 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7572 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7573 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7574 total += _mm_cvtsi128_si32(s32);
7575 j += 32;
7576 }
7577 while j < n {
7578 total += (w[j] as i8) as i32 * xq[j] as i32;
7579 j += 1;
7580 }
7581 total
7582 }
7583}
7584
7585#[cfg(target_arch = "x86_64")]
7587#[target_feature(enable = "avx2,fma")]
7588unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
7589 unsafe {
7591 use core::arch::x86_64::*;
7592 let n = x.len();
7593 let wp = w.as_ptr();
7594 let xp = x.as_ptr();
7595 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
7596 let mut j = 0usize;
7597 while j + 16 <= n {
7598 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
7599 let lo = _mm256_cvtepi8_epi32(wb);
7600 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
7601 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
7602 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
7603 j += 16;
7604 }
7605 let acc = _mm256_add_ps(a0, a1);
7606 let hi128 = _mm256_extractf128_ps::<1>(acc);
7607 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
7608 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
7609 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
7610 let mut sum = _mm_cvtss_f32(s32);
7611 while j < n {
7612 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
7613 j += 1;
7614 }
7615 sum
7616 }
7617}
7618
7619#[cfg(target_arch = "x86_64")]
7624#[target_feature(enable = "avx2")]
7625unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
7626 unsafe {
7628 use core::arch::x86_64::*;
7629 let n = w.len();
7630 let ones = _mm256_set1_epi16(1);
7631 let mut acc = _mm256_setzero_si256();
7632 let mut j = 0usize;
7633 while j + 32 <= n {
7634 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7635 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7636 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7637 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
7638 j += 32;
7639 }
7640 let hi128 = _mm256_extracti128_si256::<1>(acc);
7641 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
7642 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7643 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7644 let mut s = _mm_cvtsi128_si32(s32);
7645 while j < n {
7646 s += (w[j] as i8) as i32 * xq[j] as i32;
7647 j += 1;
7648 }
7649 s
7650 }
7651}
7652
7653#[cfg(target_arch = "aarch64")]
7657#[target_feature(enable = "neon,i8mm")]
7658unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7659 unsafe {
7661 use core::arch::aarch64::*;
7662 use core::arch::asm;
7663 let n = w0.len();
7664 let w0p = w0.as_ptr() as *const i8;
7665 let w1p = w1.as_ptr() as *const i8;
7666 let mut acc01 = vdupq_n_s32(0);
7669 let mut acc23 = vdupq_n_s32(0);
7670 let mut i = 0usize;
7671 while i + 8 <= n {
7672 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
7673 let xb01 = vcombine_s8(
7674 vld1_s8(xs[0].as_ptr().add(i)),
7675 vld1_s8(xs[1].as_ptr().add(i)),
7676 );
7677 let xb23 = vcombine_s8(
7678 vld1_s8(xs[2].as_ptr().add(i)),
7679 vld1_s8(xs[3].as_ptr().add(i)),
7680 );
7681 asm!(
7682 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
7683 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
7684 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
7685 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
7686 options(pure, nomem, nostack),
7687 );
7688 i += 8;
7689 }
7690 let mut out = [[0i32; 4]; 2];
7691 let a01: [i32; 4] = core::mem::transmute(acc01);
7692 let a23: [i32; 4] = core::mem::transmute(acc23);
7693 out[0][0] = a01[0];
7694 out[0][1] = a01[1];
7695 out[1][0] = a01[2];
7696 out[1][1] = a01[3];
7697 out[0][2] = a23[0];
7698 out[0][3] = a23[1];
7699 out[1][2] = a23[2];
7700 out[1][3] = a23[3];
7701 if i < n {
7702 for (k, x) in xs.iter().enumerate() {
7703 for j in i..n {
7704 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7705 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7706 }
7707 }
7708 }
7709 out
7710 }
7711}
7712
7713#[cfg(target_arch = "aarch64")]
7717#[target_feature(enable = "neon,dotprod")]
7718unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7719 unsafe {
7721 use core::arch::aarch64::*;
7722 use core::arch::asm;
7723 let n = w0.len();
7724 let w0p = w0.as_ptr() as *const i8;
7725 let w1p = w1.as_ptr() as *const i8;
7726 let mut acc = [[vdupq_n_s32(0); 4]; 2];
7727 let mut i = 0usize;
7728 while i + 16 <= n {
7729 let wv0 = vld1q_s8(w0p.add(i));
7730 let wv1 = vld1q_s8(w1p.add(i));
7731 for (k, x) in xs.iter().enumerate() {
7732 let xv = vld1q_s8(x.as_ptr().add(i));
7733 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
7734 asm!(
7735 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
7736 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
7737 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7738 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
7739 options(pure, nomem, nostack),
7740 );
7741 acc[0][k] = a0;
7742 acc[1][k] = a1;
7743 }
7744 i += 16;
7745 }
7746 let mut out = [[0i32; 4]; 2];
7747 for r in 0..2 {
7748 for k in 0..4 {
7749 out[r][k] = vaddvq_s32(acc[r][k]);
7750 }
7751 }
7752 if i < n {
7753 for (k, x) in xs.iter().enumerate() {
7754 for j in i..n {
7755 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7756 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7757 }
7758 }
7759 }
7760 out
7761 }
7762}
7763
7764#[cfg(target_arch = "x86_64")]
7770#[target_feature(enable = "avx2")]
7771unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7772 unsafe {
7774 use core::arch::x86_64::*;
7775 let n = w0.len();
7776 let ones = _mm256_set1_epi16(1);
7777 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
7778 let mut j = 0usize;
7779 while j + 32 <= n {
7780 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
7781 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
7782 let aw0 = _mm256_abs_epi8(wv0);
7783 let aw1 = _mm256_abs_epi8(wv1);
7784 for (k, x) in xs.iter().enumerate() {
7785 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
7786 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
7787 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
7788 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
7789 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
7790 }
7791 j += 32;
7792 }
7793 let mut out = [[0i32; 4]; 2];
7794 for r in 0..2 {
7795 for k in 0..4 {
7796 let a = acc[r][k];
7797 let hi128 = _mm256_extracti128_si256::<1>(a);
7798 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
7799 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7800 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7801 out[r][k] = _mm_cvtsi128_si32(s32);
7802 }
7803 }
7804 if j < n {
7805 for (k, x) in xs.iter().enumerate() {
7806 for i in j..n {
7807 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
7808 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
7809 }
7810 }
7811 }
7812 out
7813 }
7814}
7815
7816#[cfg(target_arch = "x86_64")]
7821#[inline]
7822fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
7823 let dot = if avx512vnni_enabled() && row.len() >= 64 {
7824 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
7825 } else {
7826 unsafe { dot_i8_i8_avx2(row, &act.xq) }
7827 };
7828 let mut acc = dot as f32 * act.sx;
7829 for &(j, xv) in &act.outliers {
7830 acc += (row[j] as i8) as f32 * xv;
7831 }
7832 acc
7833}
7834
7835#[cfg(target_arch = "x86_64")]
7839#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7840unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7841 unsafe {
7843 use core::arch::x86_64::*;
7844 let n = w.len();
7845 let flip = _mm512_set1_epi8(-128); #[inline(always)]
7847 unsafe fn step(
7848 w: *const u8,
7849 x: *const i8,
7850 flip: core::arch::x86_64::__m512i,
7851 acc: core::arch::x86_64::__m512i,
7852 ) -> core::arch::x86_64::__m512i {
7853 unsafe {
7854 use core::arch::x86_64::*;
7855 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
7856 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
7857 }
7858 }
7859 let (mut a0, mut a1, mut a2, mut a3) = (
7860 _mm512_setzero_si512(),
7861 _mm512_setzero_si512(),
7862 _mm512_setzero_si512(),
7863 _mm512_setzero_si512(),
7864 );
7865 let mut j = 0usize;
7866 while j + 256 <= n {
7867 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7868 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
7869 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
7870 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
7871 j += 256;
7872 }
7873 while j + 64 <= n {
7874 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7875 j += 64;
7876 }
7877 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
7878 _mm512_add_epi32(a0, a1),
7879 _mm512_add_epi32(a2, a3),
7880 ));
7881 while j < n {
7883 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
7884 j += 1;
7885 }
7886 total
7887 }
7888}
7889
7890#[cfg(target_arch = "x86_64")]
7896#[target_feature(enable = "avx2")]
7897unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7898 unsafe {
7901 use core::arch::x86_64::*;
7902 let lomask = _mm_set1_epi8(0x0F);
7903 let eight = _mm256_set1_epi8(8);
7904 let ones = _mm256_set1_epi16(1);
7905 let mut acc = 0f32;
7906 for gi in 0..gpr {
7907 let g = g0 + gi;
7908 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7909 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7910 let lo = _mm_and_si128(b, lomask);
7911 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7912 let w = _mm256_sub_epi8(
7913 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7914 eight,
7915 );
7916 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7917 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
7918 let d = _mm256_madd_epi16(p16, ones);
7919 let hi128 = _mm256_extracti128_si256::<1>(d);
7920 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7921 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7922 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7923 acc += _mm_cvtsi128_si32(s32) as f32 * s;
7924 }
7925 acc
7926 }
7927}
7928
7929#[cfg(target_arch = "x86_64")]
7932#[target_feature(enable = "avx2")]
7933unsafe fn dot_q4_row_avx2_2(
7934 packed: &[u8],
7935 scales: &[u8],
7936 g0: usize,
7937 gpr: usize,
7938 xq1: &[i8],
7939 xq2: &[i8],
7940) -> (f32, f32) {
7941 unsafe {
7943 use core::arch::x86_64::*;
7944 let lomask = _mm_set1_epi8(0x0F);
7945 let eight = _mm256_set1_epi8(8);
7946 let ones = _mm256_set1_epi16(1);
7947 let (mut acc1, mut acc2) = (0f32, 0f32);
7948 #[inline(always)]
7949 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
7950 unsafe {
7951 use core::arch::x86_64::*;
7952 let hi128 = _mm256_extracti128_si256::<1>(d);
7953 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7954 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7955 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7956 _mm_cvtsi128_si32(s32)
7957 }
7958 }
7959 for gi in 0..gpr {
7960 let g = g0 + gi;
7961 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7962 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7963 let lo = _mm_and_si128(b, lomask);
7964 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7965 let w = _mm256_sub_epi8(
7966 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7967 eight,
7968 );
7969 let aw = _mm256_abs_epi8(w);
7970 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7971 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7972 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
7973 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
7974 acc1 += hsum(d1) as f32 * s;
7975 acc2 += hsum(d2) as f32 * s;
7976 }
7977 (acc1, acc2)
7978 }
7979}
7980
7981#[cfg(target_arch = "x86_64")]
7983fn q8_range_avx2(
7984 q: &[u8],
7985 row_scale: &[f32],
7986 act: &SplitAct,
7987 cols: usize,
7988 out_addr: SendMut,
7989 start: usize,
7990 end: usize,
7991) {
7992 for o in start..end {
7993 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7994 unsafe { *out_addr.at(o) = v };
7996 }
7997}
7998
7999#[cfg(target_arch = "x86_64")]
8001#[allow(clippy::too_many_arguments)]
8002fn q8_range2_avx2(
8003 q: &[u8],
8004 row_scale: &[f32],
8005 a1: &SplitAct,
8006 a2: &SplitAct,
8007 cols: usize,
8008 p1: SendMut,
8009 p2: SendMut,
8010 start: usize,
8011 end: usize,
8012) {
8013 for o in start..end {
8014 let row = &q[o * cols..(o + 1) * cols];
8015 unsafe {
8017 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8018 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8019 }
8020 }
8021}
8022
8023#[cfg(target_arch = "aarch64")]
8034fn i8mm_enabled() -> bool {
8035 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8036 *ON.get_or_init(|| {
8037 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8038 && std::arch::is_aarch64_feature_detected!("i8mm")
8039 })
8040}
8041
8042#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
8046fn sdot_enabled() -> bool {
8047 use std::sync::OnceLock;
8048 static ON: OnceLock<bool> = OnceLock::new();
8049 *ON.get_or_init(|| {
8050 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
8051 if !want {
8052 return false;
8053 }
8054
8055 #[cfg(target_arch = "aarch64")]
8056 {
8057 if std::arch::is_aarch64_feature_detected!("dotprod") {
8058 return true;
8059 }
8060 #[cfg(target_os = "android")]
8061 {
8062 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
8063 if cpuinfo.lines().any(|l| {
8064 (l.starts_with("Features") || l.starts_with("features"))
8065 && l.contains("asimddp")
8066 }) {
8067 return true;
8068 }
8069 }
8070 }
8071 false
8072 }
8073 #[cfg(not(target_arch = "aarch64"))]
8074 {
8075 false
8076 }
8077 })
8078}
8079
8080struct SplitAct {
8085 xq: Vec<i8>,
8086 sx: f32,
8087 outliers: Vec<(usize, f32)>,
8088 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
8091 xsum: i32,
8092}
8093
8094thread_local! {
8095 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
8098 const { std::cell::RefCell::new(Vec::new()) };
8099}
8100
8101impl Drop for SplitAct {
8102 fn drop(&mut self) {
8103 let buf = std::mem::take(&mut self.xq);
8104 if buf.capacity() > 0 {
8105 XQ_FREE.with(|f| {
8106 let mut f = f.borrow_mut();
8107 if f.len() < 16 {
8108 f.push(buf);
8109 }
8110 });
8111 }
8112 }
8113}
8114
8115thread_local! {
8116 static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8123}
8124
8125#[inline]
8128fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
8129 KROW.with(|s| {
8130 let mut b = s.borrow_mut();
8131 if b.len() < n {
8132 b.resize(n, 0.0);
8133 }
8134 f(&mut b[..n])
8135 })
8136}
8137
8138fn split_act(x: &[f32]) -> SplitAct {
8139 let n = x.len();
8140 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
8141 let thr = 8.0 * rms;
8142 let mut outliers: Vec<(usize, f32)> = Vec::new();
8146 let mut amax = 0f32;
8147 for (j, &v) in x.iter().enumerate() {
8148 let a = v.abs();
8149 if a > thr {
8150 outliers.push((j, v));
8151 } else if a > amax {
8152 amax = a;
8153 }
8154 }
8155 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
8156 let inv = 1.0 / sx;
8157 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
8158 xq.clear();
8159 xq.reserve(n);
8160 if outliers.is_empty() {
8161 xq.extend(
8162 x.iter()
8163 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
8164 );
8165 } else {
8166 xq.extend(x.iter().map(|&v| {
8168 if v.abs() > thr {
8169 0
8170 } else {
8171 (v * inv).round().clamp(-127.0, 127.0) as i8
8172 }
8173 }));
8174 }
8175 let xsum = xq.iter().map(|&v| v as i32).sum();
8176 SplitAct {
8177 xq,
8178 sx,
8179 outliers,
8180 xsum,
8181 }
8182}
8183
8184fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
8185 let n = x.len();
8186 let rms = (x
8187 .iter()
8188 .zip(col)
8189 .map(|(&a, &c)| {
8190 let v = a * c;
8191 (v * v) as f64
8192 })
8193 .sum::<f64>()
8194 / n.max(1) as f64)
8195 .sqrt() as f32;
8196 let thr = 8.0 * rms;
8197
8198 let mut outliers = Vec::new();
8199 let mut amax = 0f32;
8200 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
8201 let v = a * c;
8202 let s = v.abs();
8203 if s > thr {
8204 outliers.push((j, v));
8205 } else if s > amax {
8206 amax = s;
8207 }
8208 }
8209
8210 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
8211 let inv = 1.0 / sx;
8212 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
8213 xq.clear();
8214 xq.reserve(n);
8215 if outliers.is_empty() {
8216 xq.extend(
8217 x.iter()
8218 .zip(col)
8219 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
8220 );
8221 } else {
8222 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
8223 let v = a * c;
8224 if v.abs() > thr {
8225 0
8226 } else {
8227 (v * inv).round().clamp(-127.0, 127.0) as i8
8228 }
8229 }));
8230 }
8231 let xsum = xq.iter().map(|&v| v as i32).sum();
8232 SplitAct {
8233 xq,
8234 sx,
8235 outliers,
8236 xsum,
8237 }
8238}
8239
8240#[cfg(target_arch = "aarch64")]
8243#[target_feature(enable = "neon,dotprod")]
8244unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
8245 unsafe {
8247 use core::arch::aarch64::*;
8248 use core::arch::asm;
8249 let wp = w.as_ptr() as *const i8;
8250 let n = w.len();
8251 let (mut a0, mut a1, mut a2, mut a3) = (
8252 vdupq_n_s32(0),
8253 vdupq_n_s32(0),
8254 vdupq_n_s32(0),
8255 vdupq_n_s32(0),
8256 );
8257 let mut i = 0;
8258 while i + 64 <= n {
8259 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
8260 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
8261 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
8262 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
8263 asm!(
8264 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8265 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8266 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
8267 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
8268 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8269 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8270 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
8271 options(pure, nomem, nostack),
8272 );
8273 i += 64;
8274 }
8275 while i + 16 <= n {
8276 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
8277 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
8278 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
8279 i += 16;
8280 }
8281 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
8282 while i < n {
8283 s += (*wp.add(i)) as i32 * xq[i] as i32;
8284 i += 1;
8285 }
8286 s
8287 }
8288}
8289
8290#[cfg(target_arch = "aarch64")]
8294#[target_feature(enable = "neon,dotprod")]
8295unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
8296 unsafe {
8298 use core::arch::aarch64::*;
8299 use core::arch::asm;
8300 let n = xq.len();
8301 let px = xq.as_ptr();
8302 let (p0, p1, p2, p3) = (
8303 w0.as_ptr() as *const i8,
8304 w1.as_ptr() as *const i8,
8305 w2.as_ptr() as *const i8,
8306 w3.as_ptr() as *const i8,
8307 );
8308 let (mut a0, mut a1, mut a2, mut a3) = (
8309 vdupq_n_s32(0),
8310 vdupq_n_s32(0),
8311 vdupq_n_s32(0),
8312 vdupq_n_s32(0),
8313 );
8314 let mut i = 0;
8315 while i + 16 <= n {
8316 let x = vld1q_s8(px.add(i));
8317 let v0 = vld1q_s8(p0.add(i));
8318 let v1 = vld1q_s8(p1.add(i));
8319 let v2 = vld1q_s8(p2.add(i));
8320 let v3 = vld1q_s8(p3.add(i));
8321 asm!(
8322 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8323 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8324 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8325 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8326 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8327 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8328 options(pure, nomem, nostack),
8329 );
8330 i += 16;
8331 }
8332 let mut r = [
8333 vaddvq_s32(a0),
8334 vaddvq_s32(a1),
8335 vaddvq_s32(a2),
8336 vaddvq_s32(a3),
8337 ];
8338 while i < n {
8339 let xi = *px.add(i) as i32;
8340 r[0] += (*p0.add(i)) as i32 * xi;
8341 r[1] += (*p1.add(i)) as i32 * xi;
8342 r[2] += (*p2.add(i)) as i32 * xi;
8343 r[3] += (*p3.add(i)) as i32 * xi;
8344 i += 1;
8345 }
8346 r
8347 }
8348}
8349
8350#[cfg(target_arch = "aarch64")]
8357#[target_feature(enable = "neon,dotprod")]
8358unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
8359 unsafe {
8362 use core::arch::aarch64::*;
8363 use core::arch::asm;
8364 let n = xq.len();
8365 let px = xq.as_ptr();
8366 let pg = g.as_ptr() as *const i8;
8367 let (mut a0, mut a1, mut a2, mut a3) = (
8368 vdupq_n_s32(0),
8369 vdupq_n_s32(0),
8370 vdupq_n_s32(0),
8371 vdupq_n_s32(0),
8372 );
8373 let mut i = 0;
8374 while i + 16 <= n {
8375 let x = vld1q_s8(px.add(i));
8376 let base = pg.add(4 * i);
8377 let v0 = vld1q_s8(base);
8378 let v1 = vld1q_s8(base.add(16));
8379 let v2 = vld1q_s8(base.add(32));
8380 let v3 = vld1q_s8(base.add(48));
8381 asm!(
8382 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8383 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8384 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8385 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8386 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8387 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8388 options(pure, nomem, nostack),
8389 );
8390 i += 16;
8391 }
8392 [
8393 vaddvq_s32(a0),
8394 vaddvq_s32(a1),
8395 vaddvq_s32(a2),
8396 vaddvq_s32(a3),
8397 ]
8398 }
8399}
8400
8401#[cfg(target_arch = "aarch64")]
8407fn q8_range_sdot(
8408 q: &[u8],
8409 rep: &[u8],
8410 row_scale: &[f32],
8411 act: &SplitAct,
8412 cols: usize,
8413 out_addr: SendMut,
8414 start: usize,
8415 end: usize,
8416) {
8417 let mut o = start;
8418 if !rep.is_empty() {
8421 while o < end && o % 4 != 0 {
8422 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8423 unsafe { *out_addr.at(o) = v };
8424 o += 1;
8425 }
8426 }
8427 while o + 4 <= end {
8428 let r = if rep.is_empty() {
8429 unsafe {
8430 dot_i8_sdot_4rows(
8431 &q[o * cols..(o + 1) * cols],
8432 &q[(o + 1) * cols..(o + 2) * cols],
8433 &q[(o + 2) * cols..(o + 3) * cols],
8434 &q[(o + 3) * cols..(o + 4) * cols],
8435 &act.xq,
8436 )
8437 }
8438 } else {
8439 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
8440 };
8441 for k in 0..4 {
8442 let mut acc = r[k] as f32 * act.sx;
8443 for &(j, xv) in &act.outliers {
8444 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
8445 }
8446 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
8448 }
8449 o += 4;
8450 }
8451 while o < end {
8452 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8453 unsafe { *out_addr.at(o) = v };
8454 o += 1;
8455 }
8456}
8457
8458#[cfg(target_arch = "aarch64")]
8461#[allow(clippy::too_many_arguments)]
8462fn q8_range2_sdot(
8463 q: &[u8],
8464 row_scale: &[f32],
8465 a1: &SplitAct,
8466 a2: &SplitAct,
8467 cols: usize,
8468 p1: SendMut,
8469 p2: SendMut,
8470 start: usize,
8471 end: usize,
8472) {
8473 for o in start..end {
8474 let row = &q[o * cols..(o + 1) * cols];
8475 unsafe {
8477 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
8478 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
8479 }
8480 }
8481}
8482
8483#[allow(clippy::too_many_arguments)]
8485fn q8_range2_f32(
8486 q: &[u8],
8487 row_scale: &[f32],
8488 x1: &[f32],
8489 x2: &[f32],
8490 cols: usize,
8491 p1: SendMut,
8492 p2: SendMut,
8493 start: usize,
8494 end: usize,
8495) {
8496 for o in start..end {
8497 let row = &q[o * cols..(o + 1) * cols];
8498 unsafe {
8500 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
8501 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
8502 }
8503 }
8504}
8505
8506fn q8_range_f32(
8508 q: &[u8],
8509 row_scale: &[f32],
8510 xs: &[f32],
8511 cols: usize,
8512 out_addr: SendMut,
8513 start: usize,
8514 end: usize,
8515) {
8516 for o in start..end {
8517 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8518 unsafe { *out_addr.at(o) = v };
8520 }
8521}
8522
8523#[cfg(target_arch = "aarch64")]
8526#[inline]
8527fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
8528 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
8529 for &(j, xv) in &act.outliers {
8530 acc += (row[j] as i8) as f32 * xv;
8531 }
8532 acc
8533}
8534
8535#[cfg(target_arch = "aarch64")]
8543#[target_feature(enable = "neon,dotprod")]
8544unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8545 unsafe {
8548 use core::arch::aarch64::*;
8549 use core::arch::asm;
8550 let lomask = vdupq_n_u8(0x0F);
8551 let eight = vdupq_n_s8(8);
8552 let mut acc = 0f32;
8553 for gi in 0..gpr {
8554 let g = g0 + gi;
8555 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8556 let b = vld1q_u8(packed.as_ptr().add(g * 16));
8557 let lo = vandq_u8(b, lomask);
8558 let hi = vshrq_n_u8::<4>(b);
8559 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8560 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8561 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
8562 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
8563 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8564 asm!(
8565 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8566 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8567 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8568 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8569 options(pure, nomem, nostack),
8570 );
8571 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8572 }
8573 acc
8574 }
8575}
8576
8577#[cfg(target_arch = "aarch64")]
8582#[target_feature(enable = "neon,dotprod")]
8583unsafe fn dot_q4_row_sdot2(
8584 packed: &[u8],
8585 scales: &[u8],
8586 g0: usize,
8587 gpr: usize,
8588 xq1: &[i8],
8589 xq2: &[i8],
8590) -> (f32, f32) {
8591 unsafe {
8594 use core::arch::aarch64::*;
8595 use core::arch::asm;
8596 let lomask = vdupq_n_u8(0x0F);
8597 let eight = vdupq_n_s8(8);
8598 let (mut acc1, mut acc2) = (0f32, 0f32);
8599 for gi in 0..gpr {
8600 let g = g0 + gi;
8601 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8602 let b = vld1q_u8(packed.as_ptr().add(g * 16));
8603 let lo = vandq_u8(b, lomask);
8604 let hi = vshrq_n_u8::<4>(b);
8605 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8606 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8607 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
8608 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
8609 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
8610 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
8611 let (mut a0, mut a1, mut b0, mut b1) = (
8612 vdupq_n_s32(0),
8613 vdupq_n_s32(0),
8614 vdupq_n_s32(0),
8615 vdupq_n_s32(0),
8616 );
8617 asm!(
8618 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
8619 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
8620 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
8621 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
8622 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8623 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
8624 e0 = in(vreg) e0, e1 = in(vreg) e1,
8625 x10 = in(vreg) x10, x11 = in(vreg) x11,
8626 x20 = in(vreg) x20, x21 = in(vreg) x21,
8627 options(pure, nomem, nostack),
8628 );
8629 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8630 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
8631 }
8632 (acc1, acc2)
8633 }
8634}
8635
8636#[inline]
8641pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
8642 #[cfg(target_arch = "aarch64")]
8643 unsafe {
8644 return axpy_i8_f32_neon(acc, row, w);
8645 }
8646 #[cfg(target_arch = "x86_64")]
8647 if avx2_enabled() {
8648 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
8649 }
8650 #[allow(unreachable_code)]
8651 {
8652 for (a, &b) in acc.iter_mut().zip(row) {
8653 *a += w * b as f32;
8654 }
8655 }
8656}
8657
8658#[cfg(target_arch = "x86_64")]
8660#[target_feature(enable = "avx2,fma")]
8661unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
8662 unsafe {
8664 use core::arch::x86_64::*;
8665 let n = acc.len().min(row.len());
8666 let ap = acc.as_mut_ptr();
8667 let rp = row.as_ptr();
8668 let wv = _mm256_set1_ps(w);
8669 let mut j = 0usize;
8670 while j + 16 <= n {
8671 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
8672 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
8673 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
8674 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
8675 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
8676 _mm256_storeu_ps(ap.add(j), v0);
8677 _mm256_storeu_ps(ap.add(j + 8), v1);
8678 j += 16;
8679 }
8680 while j < n {
8681 *ap.add(j) += w * (*rp.add(j)) as f32;
8682 j += 1;
8683 }
8684 }
8685}
8686
8687#[cfg(target_arch = "aarch64")]
8688#[target_feature(enable = "neon")]
8689unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
8690 unsafe {
8692 use core::arch::aarch64::*;
8693 let n = acc.len().min(row.len());
8694 let ap = acc.as_mut_ptr();
8695 let rp = row.as_ptr();
8696 let wv = vdupq_n_f32(w);
8697 let mut j = 0usize;
8698 while j + 16 <= n {
8699 let rb = vld1q_s8(rp.add(j));
8700 let lo = vmovl_s8(vget_low_s8(rb));
8701 let hi = vmovl_s8(vget_high_s8(rb));
8702 for (off, half) in [(0, lo), (8, hi)] {
8703 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
8704 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
8705 let o = j + off;
8706 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
8707 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
8708 }
8709 j += 16;
8710 }
8711 while j < n {
8712 *ap.add(j) += w * (*rp.add(j)) as f32;
8713 j += 1;
8714 }
8715 }
8716}
8717
8718#[inline]
8721pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
8722 #[cfg(target_arch = "aarch64")]
8723 unsafe {
8724 return dot_i8_f32_neon(w, x);
8725 }
8726 #[cfg(target_arch = "x86_64")]
8727 if avx2_enabled() {
8728 return unsafe { dot_i8_f32_avx2(w, x) };
8729 }
8730 #[allow(unreachable_code)]
8731 {
8732 let mut sum = 0.0f32;
8733 for (j, &b) in w.iter().enumerate() {
8734 sum += (b as i8) as f32 * x[j];
8735 }
8736 sum
8737 }
8738}
8739
8740#[inline]
8744fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8745 #[cfg(target_arch = "aarch64")]
8746 unsafe {
8747 return dot_i8_col_f32_neon(w, x, col);
8748 }
8749 #[allow(unreachable_code)]
8750 {
8751 let mut sum = 0.0f32;
8752 for (j, &b) in w.iter().enumerate() {
8753 sum += (b as i8) as f32 * x[j] * col[j];
8754 }
8755 sum
8756 }
8757}
8758
8759#[cfg(target_arch = "aarch64")]
8760#[target_feature(enable = "neon")]
8761unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8762 unsafe {
8764 use core::arch::aarch64::*;
8765 let n = x.len();
8766 let wp = w.as_ptr() as *const i8;
8767 let xp = x.as_ptr();
8768 let cp = col.as_ptr();
8769 let (mut a0, mut a1, mut a2, mut a3) = (
8770 vdupq_n_f32(0.0),
8771 vdupq_n_f32(0.0),
8772 vdupq_n_f32(0.0),
8773 vdupq_n_f32(0.0),
8774 );
8775 let mut j = 0usize;
8776 while j + 16 <= n {
8777 let wb = vld1q_s8(wp.add(j));
8778 let lo = vmovl_s8(vget_low_s8(wb));
8779 let hi = vmovl_s8(vget_high_s8(wb));
8780 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8781 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8782 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8783 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8784 a0 = vfmaq_f32(
8785 a0,
8786 w0,
8787 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
8788 );
8789 a1 = vfmaq_f32(
8790 a1,
8791 w1,
8792 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
8793 );
8794 a2 = vfmaq_f32(
8795 a2,
8796 w2,
8797 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
8798 );
8799 a3 = vfmaq_f32(
8800 a3,
8801 w3,
8802 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
8803 );
8804 j += 16;
8805 }
8806 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8807 while j < n {
8808 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
8809 j += 1;
8810 }
8811 sum
8812 }
8813}
8814
8815#[cfg(target_arch = "aarch64")]
8816#[target_feature(enable = "neon")]
8817unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
8818 unsafe {
8820 use core::arch::aarch64::*;
8821 let n = x.len();
8822 let wp = w.as_ptr() as *const i8;
8823 let xp = x.as_ptr();
8824 let (mut a0, mut a1, mut a2, mut a3) = (
8825 vdupq_n_f32(0.0),
8826 vdupq_n_f32(0.0),
8827 vdupq_n_f32(0.0),
8828 vdupq_n_f32(0.0),
8829 );
8830 let mut j = 0usize;
8831 while j + 16 <= n {
8832 let wb = vld1q_s8(wp.add(j));
8833 let lo = vmovl_s8(vget_low_s8(wb));
8834 let hi = vmovl_s8(vget_high_s8(wb));
8835 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8836 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8837 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8838 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8839 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
8840 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
8841 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
8842 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
8843 j += 16;
8844 }
8845 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8846 while j < n {
8847 sum += (*wp.add(j)) as f32 * *xp.add(j);
8848 j += 1;
8849 }
8850 sum
8851 }
8852}
8853
8854#[allow(clippy::too_many_arguments)]
8855fn qmatvec(
8856 q: &[u8],
8857 rep: &[u8],
8858 row_scale: &[f32],
8859 x: &[f32],
8860 col_field: &[f32],
8861 dtype: TensorDtype,
8862 rows: usize,
8863 cols: usize,
8864 out: &mut [f32],
8865 pool: Option<&Pool>,
8866) {
8867 debug_assert_eq!(out.len(), rows);
8868 #[cfg(not(target_arch = "aarch64"))]
8869 let _ = rep;
8870
8871 #[cfg(target_arch = "aarch64")]
8872 if sdot_enabled() {
8873 let act = if dtype == TensorDtype::Q8_2f {
8874 split_act_q8_2f(x, col_field)
8875 } else {
8876 split_act(x)
8877 };
8878 let out_addr = SendMut(out.as_mut_ptr());
8879 let run_range = |start: usize, end: usize| {
8880 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
8881 };
8882 match pool {
8883 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8884 _ => run_range(0, rows),
8885 }
8886 return;
8887 }
8888 #[cfg(target_arch = "x86_64")]
8891 if avx2_a8w8_enabled() {
8892 let act = if dtype == TensorDtype::Q8_2f {
8893 split_act_q8_2f(x, col_field)
8894 } else {
8895 split_act(x)
8896 };
8897 let out_addr = SendMut(out.as_mut_ptr());
8898 let run_range = |start: usize, end: usize| {
8899 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
8900 };
8901 match pool {
8902 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8903 _ => run_range(0, rows),
8904 }
8905 return;
8906 }
8907
8908 prescale_with(x, col_field, dtype, 1, |xs| {
8909 let out_addr = SendMut(out.as_mut_ptr());
8910 let run_range = move |start: usize, end: usize| {
8911 for o in start..end {
8912 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8913 unsafe { *out_addr.at(o) = v };
8915 }
8916 };
8917 match pool {
8918 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8919 _ => run_range(0, rows),
8920 }
8921 });
8922}
8923
8924#[allow(clippy::too_many_arguments)]
8925fn qmatvec2(
8926 q: &[u8],
8927 row_scale: &[f32],
8928 x1: &[f32],
8929 x2: &[f32],
8930 col_field: &[f32],
8931 dtype: TensorDtype,
8932 rows: usize,
8933 cols: usize,
8934 o1: &mut [f32],
8935 o2: &mut [f32],
8936 pool: Option<&Pool>,
8937) {
8938 #[cfg(target_arch = "aarch64")]
8939 if sdot_enabled() {
8940 let a1s = if dtype == TensorDtype::Q8_2f {
8941 split_act_q8_2f(x1, col_field)
8942 } else {
8943 split_act(x1)
8944 };
8945 let a2s = if dtype == TensorDtype::Q8_2f {
8946 split_act_q8_2f(x2, col_field)
8947 } else {
8948 split_act(x2)
8949 };
8950 let p1 = SendMut(o1.as_mut_ptr());
8951 let p2 = SendMut(o2.as_mut_ptr());
8952 let run_range = |start: usize, end: usize| {
8953 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8954 };
8955 match pool {
8956 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8957 _ => run_range(0, rows),
8958 }
8959 return;
8960 }
8961 #[cfg(target_arch = "x86_64")]
8962 if avx2_a8w8_enabled() {
8963 let a1s = if dtype == TensorDtype::Q8_2f {
8964 split_act_q8_2f(x1, col_field)
8965 } else {
8966 split_act(x1)
8967 };
8968 let a2s = if dtype == TensorDtype::Q8_2f {
8969 split_act_q8_2f(x2, col_field)
8970 } else {
8971 split_act(x2)
8972 };
8973 let p1 = SendMut(o1.as_mut_ptr());
8974 let p2 = SendMut(o2.as_mut_ptr());
8975 let run_range = |start: usize, end: usize| {
8976 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8977 };
8978 match pool {
8979 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8980 _ => run_range(0, rows),
8981 }
8982 return;
8983 }
8984
8985 prescale_with(x1, col_field, dtype, 1, |x1s| {
8986 prescale_with(x2, col_field, dtype, 2, |x2s| {
8987 let p1 = SendMut(o1.as_mut_ptr());
8988 let p2 = SendMut(o2.as_mut_ptr());
8989 let run_range = move |start: usize, end: usize| {
8990 for o in start..end {
8991 let row = &q[o * cols..(o + 1) * cols];
8992 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
8993 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
8994 unsafe {
8996 *p1.at(o) = s1;
8997 *p2.at(o) = s2;
8998 }
8999 }
9000 };
9001 match pool {
9002 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9003 _ => run_range(0, rows),
9004 }
9005 });
9006 });
9007}
9008
9009#[derive(Clone, Copy)]
9010struct SendMut(*mut f32);
9011unsafe impl Send for SendMut {}
9012unsafe impl Sync for SendMut {}
9013
9014impl SendMut {
9015 #[inline]
9016 fn at(self, i: usize) -> *mut f32 {
9017 unsafe { self.0.add(i) }
9018 }
9019}
9020
9021#[cfg(test)]
9022mod tests {
9023 use super::*;
9024
9025 #[test]
9026 fn f32_matvec_matches_matvec_rows_bitexact() {
9027 let (rows, cols) = (300, 40);
9028 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
9029 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
9030 let qt = QTensor::from_f32(w.clone(), rows, cols);
9031
9032 let mut a = vec![0.0f32; rows];
9033 matvec_rows(None, &w, &x, &mut a);
9034 let mut b = vec![0.0f32; rows];
9035 qt.matvec(&x, &mut b, None);
9036 assert_eq!(a, b);
9037 }
9038
9039 #[test]
9040 fn sdot_kernel_exact_on_grid() {
9041 eprintln!("sdot_enabled = {}", sdot_enabled());
9046 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
9048 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
9049 .collect();
9050 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
9051 let x: Vec<f32> = (0..cols)
9052 .map(|i| match i % 3 {
9053 0 => 1.0,
9054 1 => -1.0,
9055 _ => 0.0,
9056 })
9057 .collect();
9058 let mut a = vec![0.0f32; rows];
9059 qmatvec(
9060 &w,
9061 &[],
9062 &scales,
9063 &x,
9064 &[],
9065 TensorDtype::Q8Row,
9066 rows,
9067 cols,
9068 &mut a,
9069 None,
9070 );
9071 for o in 0..rows {
9072 let mut acc = 0.0f32;
9073 for j in 0..cols {
9074 acc += (w[o * cols + j] as i8) as f32 * x[j];
9075 }
9076 let expect = acc * scales[o];
9077 assert!(
9078 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
9079 "row {o}: {} vs {expect}",
9080 a[o]
9081 );
9082 }
9083 }
9084
9085 #[test]
9086 fn q1_tbl_fast_path_matches_reference() {
9087 let (rows, cols) = (5, 256);
9092 let gpr = cols / GROUP_SIZE;
9093 let mut bytes = Vec::new();
9094 for t in 0..rows * gpr {
9095 let s = 0.007 + (t % 11) as f32 * 0.004;
9096 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9097 for j in 0..4 {
9098 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
9099 }
9100 }
9101 let x: Vec<f32> = (0..cols)
9102 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
9103 .collect();
9104 let mut w = vec![0.0f32; rows * cols];
9105 cortiq_core::quant::dequant_q1(&bytes, &mut w);
9106 let mut got = vec![0.0f32; rows];
9107 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
9108 for o in 0..rows {
9109 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
9110 assert!(
9111 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
9112 "row {o}: {} vs {expect}",
9113 got[o]
9114 );
9115 }
9116 let b = 5usize;
9119 let mut xs_all = Vec::new();
9120 for bi in 0..b {
9121 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
9122 }
9123 let mut mm = vec![0.0f32; b * rows];
9124 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
9125 for bi in 0..b {
9126 let mut single = vec![0.0f32; rows];
9127 q1_matvec(
9128 &bytes,
9129 &xs_all[bi * cols..(bi + 1) * cols],
9130 rows,
9131 cols,
9132 &mut single,
9133 None,
9134 );
9135 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
9136 }
9137 }
9138
9139 #[test]
9140 fn q1_kernels_match_exact_reference() {
9141 let (rows, cols) = (7, 96);
9143 let gpr = cols / GROUP_SIZE;
9144 let mut bytes = Vec::new();
9145 for t in 0..rows * gpr {
9146 let s = 0.01 + (t % 13) as f32 * 0.003;
9147 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9148 for j in 0..4 {
9149 bytes.push(((t * 31 + j * 97) % 251) as u8);
9150 }
9151 }
9152 let x: Vec<f32> = (0..cols)
9154 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
9155 .collect();
9156 let mut w = vec![0.0f32; rows * cols];
9158 cortiq_core::quant::dequant_q1(&bytes, &mut w);
9159 let mut expect = vec![0.0f32; rows];
9160 for o in 0..rows {
9161 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
9162 }
9163 let mut got = vec![0.0f32; rows];
9164 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
9165 for o in 0..rows {
9166 assert!(
9167 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
9168 "row {o}: {} vs {}",
9169 got[o],
9170 expect[o]
9171 );
9172 }
9173 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
9175 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
9176 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
9177 assert_eq!(a1, got);
9178 let mut xs = x.clone();
9179 xs.extend_from_slice(&x2);
9180 let mut mm = vec![0.0f32; 2 * rows];
9181 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
9182 assert_eq!(&mm[..rows], got.as_slice());
9183 assert_eq!(&mm[rows..], a2.as_slice());
9184 }
9185
9186 #[test]
9187 fn repack_is_bit_identical() {
9188 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
9194 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
9195 .collect();
9196 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
9197 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
9198 let rep = q8_repack_layout(&w, rows, cols);
9199 for g in 0..rows / 4 {
9201 for c in 0..cols / 16 {
9202 for lane in 0..4 {
9203 assert_eq!(
9204 &rep[g * 4 * cols + c * 64 + lane * 16
9205 ..g * 4 * cols + c * 64 + lane * 16 + 16],
9206 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
9207 );
9208 }
9209 }
9210 }
9211 let mut a = vec![0.0f32; rows];
9212 qmatvec(
9213 &w,
9214 &[],
9215 &scales,
9216 &x,
9217 &[],
9218 TensorDtype::Q8Row,
9219 rows,
9220 cols,
9221 &mut a,
9222 None,
9223 );
9224 let mut b = vec![0.0f32; rows];
9225 qmatvec(
9226 &w,
9227 &rep,
9228 &scales,
9229 &x,
9230 &[],
9231 TensorDtype::Q8Row,
9232 rows,
9233 cols,
9234 &mut b,
9235 None,
9236 );
9237 assert_eq!(a, b, "full-range repack output diverged");
9238
9239 #[cfg(target_arch = "aarch64")]
9240 if sdot_enabled() {
9241 let act = split_act(&x);
9243 let mut c1 = vec![0.0f32; rows];
9244 let mut c2 = vec![0.0f32; rows];
9245 q8_range_sdot(
9246 &w,
9247 &[],
9248 &scales,
9249 &act,
9250 cols,
9251 SendMut(c1.as_mut_ptr()),
9252 3,
9253 rows - 2,
9254 );
9255 q8_range_sdot(
9256 &w,
9257 &rep,
9258 &scales,
9259 &act,
9260 cols,
9261 SendMut(c2.as_mut_ptr()),
9262 3,
9263 rows - 2,
9264 );
9265 assert_eq!(c1, c2, "unaligned-range repack output diverged");
9266 }
9267 }
9268
9269 #[test]
9270 fn sdot_a8w8_noise_is_bounded() {
9271 let (rows, cols) = (16, 512);
9275 let w: Vec<u8> = (0..rows * cols)
9276 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
9277 .collect();
9278 let scales = vec![0.01f32; rows];
9279 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
9280 let mut a = vec![0.0f32; rows];
9281 qmatvec(
9282 &w,
9283 &[],
9284 &scales,
9285 &x,
9286 &[],
9287 TensorDtype::Q8Row,
9288 rows,
9289 cols,
9290 &mut a,
9291 None,
9292 );
9293 let (mut num, mut den) = (0f64, 0f64);
9294 for o in 0..rows {
9295 let mut acc = 0.0f32;
9296 for j in 0..cols {
9297 acc += (w[o * cols + j] as i8) as f32 * x[j];
9298 }
9299 let expect = acc * scales[o];
9300 num += ((a[o] - expect) as f64).powi(2);
9301 den += (expect as f64).powi(2);
9302 }
9303 let rel = (num / den.max(1e-12)).sqrt();
9304 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
9305 }
9306
9307 #[test]
9308 fn i8_dot_neon_matches_scalar() {
9309 let n = 100;
9310 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
9311 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
9312 let mut scalar = 0.0f32;
9313 for j in 0..n {
9314 scalar += (w[j] as i8) as f32 * x[j];
9315 }
9316 let fast = dot_i8_f32(&w, &x);
9317 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
9318 }
9319
9320 #[test]
9322 fn vbitmatvec_matches_full_dequant() {
9323 let (rows, cols) = (6, 64);
9324 let ng = cols / GROUP_SIZE;
9325 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9327 let mut bytes = bits.clone();
9328 for g in 0..rows * ng {
9329 let s = 0.02 + 0.001 * g as f32;
9330 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9331 }
9332 for r in 0..rows {
9333 let b = bits[r] as usize;
9334 let (mut acc, mut nb) = (0u64, 0usize);
9335 let mut rowbytes = Vec::new();
9336 for i in 0..cols {
9337 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9338 acc = (acc << b) | v;
9339 nb += b;
9340 while nb >= 8 {
9341 nb -= 8;
9342 rowbytes.push(((acc >> nb) & 0xFF) as u8);
9343 }
9344 }
9345 if nb > 0 {
9346 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9347 }
9348 bytes.extend_from_slice(&rowbytes);
9349 }
9350 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9351
9352 let mut reference = vec![0f32; rows * cols];
9353 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
9354 let mut expect = vec![0f32; rows];
9355 for r in 0..rows {
9356 expect[r] = reference[r * cols..(r + 1) * cols]
9357 .iter()
9358 .zip(&x)
9359 .map(|(w, xv)| w * xv)
9360 .sum();
9361 }
9362 let mut got = vec![0f32; rows];
9363 let offsets = vbit_row_offsets(&bytes, rows, cols);
9364 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
9365 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9369 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
9370 for r in 0..rows {
9371 assert!(
9372 (got[r] - expect[r]).abs() < tol * scale,
9373 "row {r}: {} vs {}",
9374 got[r],
9375 expect[r]
9376 );
9377 }
9378 }
9379
9380 #[test]
9385 #[cfg(target_arch = "x86_64")]
9386 fn vbit_matmat_blocked_matches_per_row() {
9387 let (rows, cols, b) = (64usize, 128usize, 9usize);
9388 let ng = cols / GROUP_SIZE;
9389 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
9390 let mut bytes = bits.clone();
9391 for g in 0..rows * ng {
9392 let sc = 0.02 + 0.0005 * g as f32;
9393 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9394 }
9395 for r in 0..rows {
9396 let bw = bits[r] as usize;
9397 let (mut acc, mut nb) = (0u64, 0usize);
9398 let mut rowbytes = Vec::new();
9399 for i in 0..cols {
9400 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9401 acc = (acc << bw) | v;
9402 nb += bw;
9403 while nb >= 8 {
9404 nb -= 8;
9405 rowbytes.push(((acc >> nb) & 0xFF) as u8);
9406 }
9407 }
9408 if nb > 0 {
9409 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9410 }
9411 bytes.extend_from_slice(&rowbytes);
9412 }
9413 let x: Vec<f32> = (0..b * cols)
9414 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9415 .collect();
9416 let offsets = vbit_row_offsets(&bytes, rows, cols);
9417 let mut y_a = vec![0f32; b * rows];
9418 let mut y_b = vec![0f32; b * rows];
9419 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
9420 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
9421 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
9422 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
9423 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
9424 let max_d = y_a
9425 .iter()
9426 .zip(&y_b)
9427 .map(|(p, q)| (p - q).abs())
9428 .fold(0.0f32, f32::max);
9429 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
9430 }
9431
9432 #[test]
9440 fn q4t_matmat_blocked_matches_per_row() {
9441 let (rows, cols, b) = (16usize, 64usize, 9usize);
9442 let gpr = cols / GROUP_SIZE;
9443 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9444 for r in 0..rows {
9445 for g in 0..gpr {
9446 let t = (r * gpr + g) * Q4_TILE;
9447 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
9448 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9449 for k in 0..16 {
9450 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9451 }
9452 }
9453 }
9454 let x: Vec<f32> = (0..b * cols)
9455 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9456 .collect();
9457 let mut y_blk = vec![0f32; b * rows];
9458 let mut y_row = vec![0f32; b * rows];
9459 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
9460 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
9461 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
9462 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
9463 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
9464 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
9465 }
9466
9467 fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
9474 use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
9475 let gpr = cols / GROUP_SIZE;
9476 let stride = q4tp_code_stride(gpr);
9477 let (params_off, codes_off, _) = q4tp_sections(rows, cols);
9478 let mut b = vec![0u8; codes_off + rows * stride];
9479 for r in 0..rows {
9480 for g in 0..gpr {
9481 let t = (r * gpr + g) * Q4TP_NIB;
9482 for k in 0..16 {
9483 b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9484 }
9485 }
9486 let lo = -6.0 - 0.03 * (r % 17) as f32;
9487 let step = 0.01 + 0.004 * (r % 11) as f32;
9488 let p = params_off + r * 4;
9489 b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
9490 b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
9491 let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
9492 for g in 0..gpr {
9493 q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
9494 }
9495 }
9496 b
9497 }
9498
9499 fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
9503 let gpr = cols / GROUP_SIZE;
9504 let v = Q4tpView::new(bytes, rows, cols);
9505 let mut out = vec![0u8; rows * gpr * Q4_TILE];
9506 let mut sc = vec![0f32; gpr];
9507 for r in 0..rows {
9508 v.scales_into(r, gpr, &mut sc);
9509 for g in 0..gpr {
9510 let t = (r * gpr + g) * Q4_TILE;
9511 let s = sc[g];
9512 out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9513 let src = (r * gpr + g) * Q4TP_NIB;
9514 out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
9515 }
9516 }
9517 out
9518 }
9519
9520 #[test]
9526 fn q4tp_exact_path_matches_dequant_reference() {
9527 let (rows, cols) = (256usize, 512usize);
9528 let gpr = cols / GROUP_SIZE;
9529 let bytes = synth_q4tp(rows, cols);
9530 let mut w = vec![0f32; rows * cols];
9531 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9532
9533 let x: Vec<f32> = (0..cols)
9534 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9535 .collect();
9536 let v = Q4tpView::new(&bytes, rows, cols);
9537 let mut sc = vec![0f32; gpr];
9538 for r in 0..rows {
9539 v.scales_into(r, gpr, &mut sc);
9540 let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
9541 let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
9542 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9546 assert!(
9547 (got - want).abs() <= 1e-5 * mag,
9548 "row {r}: kernel {got} vs dequant {want}"
9549 );
9550 }
9551 }
9552
9553 #[test]
9559 fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
9560 let (rows, cols) = (256usize, 512usize);
9561 let bytes = synth_q4tp(rows, cols);
9562 let twin = q4tp_as_q4t(&bytes, rows, cols);
9563 let x: Vec<f32> = (0..cols)
9564 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9565 .collect();
9566
9567 let mut got = vec![0f32; rows];
9568 q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
9569 let mut want = vec![0f32; rows];
9570 q4t_matvec(&twin, &x, rows, cols, &mut want, None);
9571
9572 let mut w = vec![0f32; rows * cols];
9575 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9576 for r in 0..rows {
9577 let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9578 assert!(
9579 (got[r] - want[r]).abs() <= 1e-3 * mag,
9580 "row {r}: q4tp {} vs q4t {}",
9581 got[r],
9582 want[r]
9583 );
9584 }
9585 }
9586
9587 #[test]
9592 fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
9593 let (rows, cols, b) = (256usize, 512usize, 5usize);
9594 let bytes = synth_q4tp(rows, cols);
9595 let twin = q4tp_as_q4t(&bytes, rows, cols);
9596 let xs: Vec<f32> = (0..b * cols)
9597 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9598 .collect();
9599
9600 let mut got = vec![0f32; b * rows];
9601 q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
9602 let mut want = vec![0f32; b * rows];
9603 q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
9604
9605 let mut w = vec![0f32; rows * cols];
9606 cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9607 for t in 0..b {
9608 for r in 0..rows {
9609 let mag: f32 = (0..cols)
9610 .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
9611 .sum();
9612 let (g, wa) = (got[t * rows + r], want[t * rows + r]);
9613 assert!(
9614 (g - wa).abs() <= 1e-3 * mag,
9615 "batch {t} row {r}: q4tp {g} vs q4t {wa}"
9616 );
9617 }
9618 }
9619 }
9620
9621 #[test]
9622 fn q4tp_matvec2_matches_the_single_stream_kernel() {
9623 let (rows, cols) = (128usize, 256usize);
9624 let gpr = cols / GROUP_SIZE;
9625 let bytes = synth_q4tp(rows, cols);
9626 let xs: Vec<f32> = (0..2 * cols)
9627 .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9628 .collect();
9629
9630 let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
9631 q4tp_matvec2(
9632 &bytes,
9633 &xs[..cols],
9634 &xs[cols..],
9635 rows,
9636 cols,
9637 &mut o1,
9638 &mut o2,
9639 None,
9640 );
9641
9642 let v = Q4tpView::new(&bytes, rows, cols);
9645 let mut sc = vec![0f32; gpr];
9646 for r in 0..rows {
9647 v.scales_into(r, gpr, &mut sc);
9648 assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
9649 assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
9650 }
9651 }
9652
9653 #[test]
9660 fn q4tp_matvec_keeps_pace_with_q4t() {
9661 let (rows, cols) = (4096usize, 3072usize);
9662 let bytes = synth_q4tp(rows, cols);
9663 let twin = q4tp_as_q4t(&bytes, rows, cols);
9664 let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
9665 let mut o = vec![0f32; rows];
9666 let n = 12;
9667 let mut best = (f64::MAX, f64::MAX);
9668 for _ in 0..3 {
9671 let t0 = std::time::Instant::now();
9672 for _ in 0..n {
9673 q4t_matvec(&twin, &x, rows, cols, &mut o, None);
9674 }
9675 best.0 = best.0.min(t0.elapsed().as_secs_f64());
9676 let t0 = std::time::Instant::now();
9677 for _ in 0..n {
9678 q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
9679 }
9680 best.1 = best.1.min(t0.elapsed().as_secs_f64());
9681 }
9682 let ratio = best.1 / best.0;
9683 println!(
9684 "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
9685 best.0 * 1e3 / n as f64,
9686 best.1 * 1e3 / n as f64
9687 );
9688 assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
9689 }
9690
9691 #[cfg(target_os = "macos")]
9692 #[test]
9693 fn q4t_matmat_accel_matches_dequant_reference() {
9694 if !accel_gemm_enabled() {
9695 return; }
9697 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
9699 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9700 for r in 0..rows {
9701 for g in 0..gpr {
9702 let t = (r * gpr + g) * Q4_TILE;
9703 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
9704 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9705 for k in 0..16 {
9706 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9707 }
9708 }
9709 }
9710 let x: Vec<f32> = (0..b * cols)
9711 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9712 .collect();
9713 let mut got = vec![0f32; b * rows];
9714 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
9715 let mut w = vec![0f32; rows * cols];
9717 for r in 0..rows {
9718 for g in 0..gpr {
9719 let t = (r * gpr + g) * Q4_TILE;
9720 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
9721 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
9722 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
9723 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
9724 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
9725 }
9726 }
9727 }
9728 for bi in 0..b {
9729 for r in 0..rows {
9730 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
9731 let d = (got[bi * rows + r] - want).abs();
9732 assert!(
9733 d <= want.abs().max(1.0) * 1e-4,
9734 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
9735 got[bi * rows + r]
9736 );
9737 }
9738 }
9739 }
9740
9741 #[test]
9742 fn q4matvec_matches_full_dequant() {
9743 let (rows, cols) = (8, 64);
9744 let groups = rows * cols / GROUP_SIZE;
9745 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9747 for i in 0..groups * 16 {
9748 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9749 }
9750 for g in 0..groups {
9751 let s = 0.01 + 0.003 * g as f32;
9752 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9753 }
9754 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9755
9756 let mut reference = vec![0.0f32; rows * cols];
9757 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9758 let mut expect = vec![0.0f32; rows];
9759 for r in 0..rows {
9760 expect[r] = reference[r * cols..(r + 1) * cols]
9761 .iter()
9762 .zip(&x)
9763 .map(|(w, xv)| w * xv)
9764 .sum();
9765 }
9766
9767 let mut got = vec![0.0f32; rows];
9768 q4matvec(&bytes, &x, rows, cols, &mut got, None);
9769 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9773 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9774 for r in 0..rows {
9775 assert!(
9776 (got[r] - expect[r]).abs() < tol * scale,
9777 "row {r}: {} vs {}",
9778 got[r],
9779 expect[r]
9780 );
9781 }
9782 }
9783
9784 #[test]
9787 fn vbitmatvec2_equals_two_singles() {
9788 let (rows, cols) = (6, 64);
9789 let ng = cols / GROUP_SIZE;
9790 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9791 let mut bytes = bits.clone();
9792 for g in 0..rows * ng {
9793 let s = 0.02 + 0.001 * g as f32;
9794 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9795 }
9796 for r in 0..rows {
9797 let b = bits[r] as usize;
9798 let (mut acc, mut nb) = (0u64, 0usize);
9799 let mut rowbytes = Vec::new();
9800 for i in 0..cols {
9801 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9802 acc = (acc << b) | v;
9803 nb += b;
9804 while nb >= 8 {
9805 nb -= 8;
9806 rowbytes.push(((acc >> nb) & 0xFF) as u8);
9807 }
9808 }
9809 if nb > 0 {
9810 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9811 }
9812 bytes.extend_from_slice(&rowbytes);
9813 }
9814 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9815 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
9816 let offsets = vbit_row_offsets(&bytes, rows, cols);
9817
9818 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9819 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
9820 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
9821 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9822 vbitmatvec2(
9823 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
9824 );
9825 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
9826 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
9827 }
9828
9829 #[test]
9831 fn q4matvec2_equals_two_singles() {
9832 let (rows, cols) = (8, 128);
9833 let groups = rows * cols / GROUP_SIZE;
9834 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9835 for i in 0..groups * 16 {
9836 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9837 }
9838 for g in 0..groups {
9839 let s = 0.01 + 0.003 * g as f32;
9840 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9841 }
9842 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9845 x1[9] = 250.0;
9846 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9847
9848 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9849 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
9850 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
9851 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9852 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
9853 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
9854 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
9855 }
9856
9857 #[test]
9860 fn matvec_many_equals_separate_matvecs() {
9861 use crate::pool::Pool;
9862 let (r1, r2, cols) = (300, 200, 64);
9863 let mk = |salt: usize, rows: usize| {
9864 QTensor::from_f32(
9865 (0..rows * cols)
9866 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
9867 .collect(),
9868 rows,
9869 cols,
9870 )
9871 };
9872 let (a, b) = (mk(1, r1), mk(5, r2));
9873 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
9874 let pool = Pool::new(3);
9875
9876 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
9877 a.matvec(&x, &mut ea, Some(&pool));
9878 b.matvec(&x, &mut eb, Some(&pool));
9879 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
9880 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
9881 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
9882 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
9883 }
9884
9885 #[test]
9888 fn batched_matmat_equals_per_position_matvec() {
9889 let (rows, cols, b) = (8, 64, 5);
9890 let groups = rows * cols / GROUP_SIZE;
9892 let mut q4 = Vec::new();
9893 for i in 0..groups * 16 {
9894 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9895 }
9896 for g in 0..groups {
9897 q4.extend_from_slice(
9898 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9899 );
9900 }
9901 let ng = cols / GROUP_SIZE;
9903 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
9904 let mut vb = bits.clone();
9905 for g in 0..rows * ng {
9906 vb.extend_from_slice(
9907 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
9908 );
9909 }
9910 for r in 0..rows {
9911 let bw = bits[r] as usize;
9912 let (mut acc, mut nb) = (0u64, 0usize);
9913 let mut rowbytes = Vec::new();
9914 for i in 0..cols {
9915 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9916 acc = (acc << bw) | v;
9917 nb += bw;
9918 while nb >= 8 {
9919 nb -= 8;
9920 rowbytes.push(((acc >> nb) & 0xFF) as u8);
9921 }
9922 }
9923 if nb > 0 {
9924 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9925 }
9926 vb.extend_from_slice(&rowbytes);
9927 }
9928 let offsets = vbit_row_offsets(&vb, rows, cols);
9929
9930 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9931
9932 let mut got = vec![0f32; b * rows];
9934 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
9935 for bi in 0..b {
9936 let mut expect = vec![0f32; rows];
9937 q4matvec(
9938 &q4,
9939 &xs[bi * cols..(bi + 1) * cols],
9940 rows,
9941 cols,
9942 &mut expect,
9943 None,
9944 );
9945 assert_eq!(
9946 &got[bi * rows..(bi + 1) * rows],
9947 &expect[..],
9948 "q4 batch pos {bi}"
9949 );
9950 }
9951
9952 let mut got = vec![0f32; b * rows];
9954 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
9955 for bi in 0..b {
9956 let mut expect = vec![0f32; rows];
9957 vbitmatvec(
9958 &vb,
9959 &offsets,
9960 &xs[bi * cols..(bi + 1) * cols],
9961 rows,
9962 cols,
9963 &mut expect,
9964 None,
9965 );
9966 assert_eq!(
9967 &got[bi * rows..(bi + 1) * rows],
9968 &expect[..],
9969 "vbit batch pos {bi}"
9970 );
9971 }
9972 }
9973
9974 #[test]
9978 fn q4_tiled_matches_q4_block_bitexact() {
9979 let (rows, cols, b) = (8usize, 128usize, 3usize);
9980 let groups = rows * cols / GROUP_SIZE;
9981 let mut split = Vec::with_capacity(groups * 18);
9982 for i in 0..groups * 16 {
9983 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9984 }
9985 for g in 0..groups {
9986 split.extend_from_slice(
9987 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9988 );
9989 }
9990 let (packed, scales) = split.split_at(groups * 16);
9992 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
9993 for g in 0..groups {
9994 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
9995 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
9996 }
9997
9998 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9999 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10001
10002 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
10003 q4matvec(&split, &x1, rows, cols, &mut a, None);
10004 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
10005 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
10006
10007 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10008 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
10009 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
10010 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
10011 assert_eq!(a1, t1);
10012 assert_eq!(a2, t2);
10013
10014 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10015 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
10016 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
10017 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
10018 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
10019 }
10020
10021 #[test]
10028 fn q4matvec_sdot_outlier_exact() {
10029 let (rows, cols) = (4, 128);
10030 let groups = rows * cols / GROUP_SIZE;
10031 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10032 for i in 0..groups * 16 {
10033 bytes.push(((i * 11 + 5) % 256) as u8);
10034 }
10035 for g in 0..groups {
10036 let s = 0.02 + 0.002 * g as f32;
10037 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10038 }
10039 let mut x: Vec<f32> = (0..cols)
10040 .map(|i| match i % 3 {
10041 0 => 1.0,
10042 1 => -1.0,
10043 _ => 0.0,
10044 })
10045 .collect();
10046 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
10049 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10050 let mut expect = vec![0.0f32; rows];
10051 for r in 0..rows {
10052 expect[r] = reference[r * cols..(r + 1) * cols]
10053 .iter()
10054 .zip(&x)
10055 .map(|(w, xv)| w * xv)
10056 .sum();
10057 }
10058 let mut got = vec![0.0f32; rows];
10059 q4matvec(&bytes, &x, rows, cols, &mut got, None);
10060 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10061 for r in 0..rows {
10062 assert!(
10063 (got[r] - expect[r]).abs() < 2e-3 * scale,
10064 "row {r}: {} vs {} (outlier term must be exact)",
10065 got[r],
10066 expect[r]
10067 );
10068 }
10069 }
10070
10071 #[test]
10075 fn q1t_matvec_matches_reference() {
10076 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
10077 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
10079 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
10080 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
10082 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
10083 let mut bytes = Vec::new();
10084 for r in 0..rows {
10085 for g in 0..gpr {
10086 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
10087 let mut c = [0u8; 7];
10088 for k in 0..GROUP_SIZE {
10089 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
10091 0
10092 } else {
10093 ((k + r * 3 + g) % 3) as u8 };
10095 cortiq_core::quant::q1t_pack(&mut c, k, code);
10096 }
10097 bytes.extend_from_slice(&c);
10098 }
10099 }
10100 let mut row_ptr = vec![0u32; rows + 1];
10103 for &(idx, _) in &outliers {
10104 row_ptr[idx as usize / cols + 1] += 1;
10105 }
10106 for r in 0..rows {
10107 row_ptr[r + 1] += row_ptr[r];
10108 }
10109 for &p in &row_ptr {
10110 bytes.extend_from_slice(&p.to_le_bytes());
10111 }
10112 for &(idx, v) in &outliers {
10113 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
10114 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
10115 }
10116
10117 let mut refw = vec![0f32; rows * cols];
10118 dequant_q1t(&bytes, rows, cols, &mut refw);
10119 let x: Vec<f32> = (0..cols)
10122 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
10123 .collect();
10124 let mut expect = vec![0f32; rows];
10125 for r in 0..rows {
10126 let mut a = 0.0f32;
10127 for j in 0..cols {
10128 a += refw[r * cols + j] * x[j];
10129 }
10130 expect[r] = a;
10131 }
10132 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
10133 let mut got = vec![0f32; rows];
10134 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
10135 for r in 0..rows {
10136 assert!(
10137 (got[r] - expect[r]).abs() < tol(expect[r]),
10138 "row {r}: {} vs {}",
10139 got[r],
10140 expect[r]
10141 );
10142 }
10143 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
10145 let mut gm = vec![0f32; 2 * rows];
10146 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
10147 for r in 0..rows {
10148 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
10149 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
10150 }
10151 let xb: Vec<f32> = (0..cols)
10155 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
10156 .collect();
10157 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10158 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
10159 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
10160 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10161 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
10162 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
10163 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
10164 }
10165
10166 #[test]
10169 fn q1t_matvec2_odd_gpr_matches_singles() {
10170 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
10171 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
10173 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
10174 for r in 0..rows {
10175 for g in 0..gpr {
10176 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
10177 let mut c = [0u8; 7];
10178 for k in 0..GROUP_SIZE {
10179 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
10180 }
10181 bytes.extend_from_slice(&c);
10182 }
10183 }
10184 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
10185 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
10186 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10187 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10188 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
10189 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10190 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10191 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
10192 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
10193 }
10194
10195 #[test]
10199 #[ignore]
10200 fn q1t_matvec2_speed() {
10201 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
10202 use std::time::Instant;
10203 let (rows, cols) = (8192usize, 4096usize);
10204 let gpr = cols / GROUP_SIZE;
10205 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
10206 for r in 0..rows {
10207 for g in 0..gpr {
10208 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
10209 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
10210 let mut c = [0u8; 7];
10211 for k in 0..GROUP_SIZE {
10212 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
10213 }
10214 bytes.extend_from_slice(&c);
10215 }
10216 }
10217 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
10218 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
10219 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10220 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10221 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10223 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10224 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
10225 for _ in 0..8 {
10226 let t0 = Instant::now();
10227 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10228 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
10229 let t1 = Instant::now();
10230 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10231 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
10232 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
10233 }
10234 assert_eq!(p1, s1);
10235 assert_eq!(p2, s2);
10236 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
10237 }
10238
10239 #[test]
10243 #[ignore]
10244 fn q1t_matvec_speed() {
10245 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
10246 use std::time::Instant;
10247 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
10249 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
10250 for r in 0..rows {
10251 for g in 0..gpr {
10252 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
10253 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
10254 let mut c = [0u8; 7];
10255 for k in 0..GROUP_SIZE {
10256 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
10257 }
10258 bytes.extend_from_slice(&c);
10259 }
10260 }
10261 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
10263 let mut idx = 0usize;
10264 while idx < n {
10265 row_ptr[idx / cols + 1] += 1;
10266 idx += stride;
10267 }
10268 for r in 0..rows {
10269 row_ptr[r + 1] += row_ptr[r];
10270 }
10271 for &p in &row_ptr {
10272 bytes.extend_from_slice(&p.to_le_bytes());
10273 }
10274 let mut idx = 0usize;
10275 while idx < n {
10276 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
10277 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
10278 idx += stride;
10279 }
10280 let x: Vec<f32> = (0..cols)
10283 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
10284 .collect();
10285 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
10286
10287 let slow = |out: &mut [f32]| {
10289 let mut buf = vec![0f32; cols];
10290 for r in 0..rows {
10291 for g in 0..gpr {
10292 let off = (r * gpr + g) * Q1T_TILE;
10293 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10294 let codes = &bytes[off + 2..off + Q1T_TILE];
10295 for k in 0..GROUP_SIZE {
10296 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
10297 1 => s,
10298 2 => -s,
10299 _ => 0.0,
10300 };
10301 }
10302 }
10303 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
10304 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
10305 }
10306 };
10307 let iters = 5;
10308 let mut a = vec![0f32; rows];
10309 slow(&mut a); let t = Instant::now();
10311 for _ in 0..iters {
10312 slow(&mut a);
10313 }
10314 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
10315
10316 let mut b = vec![0f32; rows];
10317 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
10319 for _ in 0..iters {
10320 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
10321 }
10322 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
10323
10324 for r in 0..rows {
10325 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
10326 }
10327 println!(
10328 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
10329 slow_ms / fast_ms
10330 );
10331 }
10332}