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