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